diff --git a/bin/omarchy-agent b/bin/omarchy-agent index b69c809d612..ec632b98339 100755 --- a/bin/omarchy-agent +++ b/bin/omarchy-agent @@ -140,18 +140,99 @@ kimi) command=(kimi --yolo) fi ;; +qwen) + # Bind the value so a prompt starting with a dash stays literal. + command=(qwen --yolo) + [[ -n ${prompt:-} ]] && command+=("--prompt-interactive=$prompt") + ;; +dsh) + # The browser UI has no prompt option, and the headless profile would run the + # prompt with full access in a window that closes with its answer, so a + # prompt goes to the clipboard below instead. + command=(env DSH_PERMISSION_MODE=danger-full-access dsh web) + prompt_via_clipboard="DeepSeek Harness" + ;; +mcode) + # The TUI has no no-approval flag; permission options belong to exec mode. + command=(mcode) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; +kiro-cli) + # A positional input seeds the interactive chat; -- keeps a prompt starting + # with a dash literal. + command=(kiro-cli chat --trust-all-tools) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; +freebuff) + # Freebuff takes no starting prompt and asks no tool permissions, so a + # prompt goes to the clipboard below instead. + command=(freebuff) + prompt_via_clipboard="Freebuff" + ;; +kilo) + # Kilo is built on OpenCode and shares its flags. Bind the prompt so one + # starting with a dash stays literal. + command=(kilo --auto) + [[ -n ${prompt:-} ]] && command+=("--prompt=$prompt") + ;; +devin) + # A bare positional is a path that opens Devin Desktop, so the prompt + # always follows --, which starts an interactive session with it. + command=(devin --permission-mode dangerous --respect-workspace-trust false) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; +amp) + # Amp takes its first interactive message from piped input; a positional + # message is for --execute, which exits after one turn. The binary rides in + # as $0 so a user install below can replace it. + if [[ -n ${prompt:-} ]]; then + command=(bash -c 'printf "%s\n" "$1" | exec "$0" --dangerously-allow-all' amp "$prompt") + else + command=(amp --dangerously-allow-all) + fi + ;; +commandcode) + # --yolo skips permission prompts and implies --trust. A positional prompt + # seeds the interactive session; -- keeps one starting with a dash literal. + command=(commandcode --yolo) + [[ -n ${prompt:-} ]] && command+=(-- "$prompt") + ;; *) echo "Unsupported default agent: $agent" >&2 exit 1 ;; esac -# A user-installed Cursor or Muse still wins in the agent launcher even when -# mise's system lazy shim is earlier on PATH. -if [[ $agent == "cursor-agent" || $agent == "muse" ]] && - [[ -f $HOME/.local/bin/$agent && -x $HOME/.local/bin/$agent ]] && - { [[ -L $HOME/.local/bin/$agent ]] || ! grep -q '^mise use -g' "$HOME/.local/bin/$agent"; }; then - command[0]="$HOME/.local/bin/$agent" +# An agent with no way to take a starting prompt still opens, with the prompt +# on the clipboard to paste as its first message, so a prompted launch from the +# agents panel or a crash report is not lost. +if [[ -n ${prompt:-} && -n ${prompt_via_clipboard:-} ]]; then + printf '%s' "$prompt" | wl-copy + omarchy-notification-send -g 󰅇 "$(omarchy-i18n agents.promptOnClipboard "$prompt_via_clipboard")" \ + "$(omarchy-i18n agents.promptOnClipboardBody)" +fi + +# omarchy-mise-install's wrappers run `mise use` and `mise x`, current ones on +# lines of their own and older ones after an inline variable; either way +# running one installs the tool it wraps. A symlink is the user's own install. +mise_wrapper() { + [[ ! -L $1 ]] && head -c 4096 "$1" | grep -qaE 'mise (use|x) ' +} + +# The user's own install in ~/.local/bin, which omarchy-default-agent counts as +# installed, still wins over mise's lazy shim earlier on PATH, which would +# install and run mise's copy instead. Once mise has installed its own, a +# terminal runs that one, so the launcher does too. OpenClaw's command runs its +# launcher, so its prompt is the only argument that could match the name. +user_install="$HOME/.local/bin/$agent" +if [[ $agent != "openclaw" && -f $user_install && -x $user_install ]] && + ! mise_wrapper "$user_install" && ! MISE_OFFLINE=1 mise which "$agent" &>/dev/null; then + for i in "${!command[@]}"; do + if [[ ${command[i]} == "$agent" ]]; then + command[i]=$user_install + break + fi + done fi # A new session uses the current selection rather than its parent's account diff --git a/bin/omarchy-agent-account-state b/bin/omarchy-agent-account-state index afc65f9e0ee..b77402a8a1b 100755 --- a/bin/omarchy-agent-account-state +++ b/bin/omarchy-agent-account-state @@ -186,7 +186,7 @@ def grok_identity(home: Path) -> dict[str, str]: "accountId": str(login.get("user_id") or ""), "email": str(login.get("email") or ""), "org": "", - "plan": str(settings.get("subscription_tier_display") or ""), + "plan": str((settings if isinstance(settings, dict) else {}).get("subscription_tier_display") or ""), } diff --git a/bin/omarchy-agent-usage-agy b/bin/omarchy-agent-usage-agy new file mode 100755 index 00000000000..44c0d2bbabd --- /dev/null +++ b/bin/omarchy-agent-usage-agy @@ -0,0 +1,775 @@ +#!/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/, and from the +Antigravity answers pi, omp, OpenClaw, and opencode's Antigravity plugin +recorded; no other collector counts those. The plan and +limits come from the Google Cloud Code API (v1internal:loadCodeAssist and +v1internal:retrieveUserQuotaSummary), asked with the sign-in agy keeps in the +Secret Service keyring (or its antigravity-oauth-token file), or with the +Antigravity sign-in omp, pi, or opencode's Antigravity plugin keeps; a lapsed +omp sign-in is asked through `omp usage`, which renews it. Those stores are +only ever read. The agents panel only ever reads the JSON this prints. +""" + +import argparse +import json +import math +import os +import re +import shutil +import sqlite3 +import subprocess +import sys +import time +import urllib.request +from datetime import datetime, timezone +from pathlib import Path + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record, + cached_scan, + canonical_roots, + emit, + jwt_claims, + now_ms, + number, + open_json, + opencode_db, + pi_session_roots, + read_json, + scan_max_age, + scan_opencode, + scan_pi_sessions, +) + +AGENT_ID = "agy" +AGENT_NAME = "Antigravity" +AUTH_HELP = "Run `agy` to sign in, or sign in to Antigravity in omp, pi, or opencode." +# How each source of a sign-in is named when it needs attention. +SOURCE_NAMES = {"agy": "`agy`", "omp": "omp", "pi": "pi", "opencode": "opencode"} +DEFAULT_CLOUD_CODE_URL = "https://daily-cloudcode-pa.googleapis.com" + + +def app_root(): + return Path(os.environ.get("AGY_DIR") or (Path.home() / ".gemini" / "antigravity-cli")) + + +# Times on disk are untrusted: a boolean, an epoch at or below zero, one too +# large for a date, NaN, or text that isn't ISO is no time at all, and the +# row it dates stays undated rather than taking the record down. +def parse_time(value): + if value is None or value == "" or isinstance(value, bool): + return None + try: + if isinstance(value, (int, float)) or str(value).isdigit(): + stamp = float(value) + if not stamp > 0: + return None + parsed = datetime.fromtimestamp(stamp / 1000 if stamp > 1e12 else stamp, timezone.utc) + else: + parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00")) + if not parsed.tzinfo: + parsed = parsed.replace(tzinfo=timezone.utc) + except (ValueError, OverflowError, OSError): + return None + # Go's zero time marks an uninitialized summary, not dated activity. + return None if parsed == datetime.min.replace(tzinfo=timezone.utc) else parsed + + +def local_day(value): + parsed = parse_time(value) + if parsed is None: + return "" + try: + day = parsed.astimezone().date() + except (ValueError, OverflowError, OSError): + # Even a parsed date can cross years 1/9999 during local conversion. + return "" + # No agent wrote a row before 2000 or after 2100: a time out there is a + # damaged stamp (a stray 1 reads as 1970), not activity to date. + return day.isoformat() if 2000 <= day.year <= 2100 else "" + + +# The sign-in agy keeps in the keyring: its access token, when that lapses, +# whether a refresh token stands behind it, and which kind of account it is. +# Without a Secret Service, agy keeps the same sign-in in a file of its own. +def login(): + try: + raw = subprocess.check_output( + ["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): + data = read_json(app_root() / "antigravity-oauth-token") or 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 { + "source": "agy", + "accessToken": str(token.get("access_token") or data.get("access_token") or ""), + "expiresAt": expires.timestamp() if expires else 0, + "refresh": str(token.get("refresh_token") or data.get("refresh_token") or ""), + "authMethod": str(data.get("auth_method") or "consumer"), + "email": email_of(data.get("id_token")), + } + + +# omp, pi, and opencode's Antigravity plugin sign in with the same Google +# client agy uses, so the access tokens they keep reach the same plan and +# quota. Each refreshes its own token while it runs; these are only ever +# read, never refreshed or written back. +def harness_login(source, data): + if not isinstance(data, dict) or data.get("type", "oauth") != "oauth" or not isinstance(data.get("access"), str) or not data["access"]: + return None + expires = parse_time(data.get("expires")) + return { + "source": source, + "accessToken": data["access"], + "expiresAt": expires.timestamp() if expires else 0, + "refresh": str(data.get("refresh") or ""), + "authMethod": "consumer", + "email": str(data.get("email") or ""), + } + + +def omp_logins(): + database = Path.home() / ".omp" / "agent" / "agent.db" + if not database.is_file(): + return [] + try: + conn = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + rows = conn.execute( + "SELECT data FROM auth_credentials WHERE provider = 'google-antigravity' " + "AND credential_type = 'oauth' AND disabled_cause IS NULL" + ).fetchall() + finally: + conn.close() + except sqlite3.Error: + return [] + logins = [] + for (data,) in rows: + try: + entry = harness_login("omp", json.loads(data)) + except (TypeError, ValueError): + continue + if entry: + logins.append(entry) + return logins + + +def pi_agent_dir(): + return Path(os.environ.get("PI_CODING_AGENT_DIR") or (Path.home() / ".pi" / "agent")) + + +def pi_logins(): + auth = read_json(pi_agent_dir() / "auth.json") + entry = harness_login("pi", auth.get("google-antigravity")) if isinstance(auth, dict) else None + return [entry] if entry else [] + + +# opencode files the Antigravity plugin's sign-in under its "google" +# provider, which other Google sign-ins share, so it only counts alongside +# the plugin's own accounts file. +def opencode_logins(): + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + if not (config_home / "opencode" / "antigravity-accounts.json").is_file(): + return [] + data_home = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + auth = read_json(data_home / "opencode" / "auth.json") + entry = harness_login("opencode", auth.get("google")) if isinstance(auth, dict) else None + return [entry] if entry else [] + + +def logins(): + entry = login() + return ([entry] if entry else []) + omp_logins() + pi_logins() + opencode_logins() + + +def email_of(id_token): + return str(jwt_claims(id_token).get("email") or "") + + +# Which Google account a sign-in belongs to, so the limits kept for it are +# its own: the email when the sign-in names one, otherwise its refresh token, +# which outlives the access tokens it mints. Only a digest of either names +# the cache file. +def limits_cache(entry): + identity = entry.get("email") or entry.get("refresh") or entry.get("accessToken") + return LimitsCache(AGENT_ID, identity) if identity else None + + +def token_fresh(entry, now): + return bool(entry["accessToken"]) and (not entry["expiresAt"] or entry["expiresAt"] > now) + + +def normalize_tier_label(name, tier_id=""): + for word in ("ultra", "pro", "free"): + if word in tier_id.lower() or word in name.lower(): + 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(data, dict) else None + 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 "" + + +# How much of a bucket is spent, or None when its fraction isn't a number. +# A bucket's remainder is either a fraction or an amount, which gives no +# share of the whole to show unless it's zero, and either arrives even when +# it is zero. Proto3 JSON leaves zero values out of answers without that +# choice, so a bucket with a reset time and neither has none left: fully +# spent. A disabled bucket counts against nothing. +def used(bucket): + if bucket.get("disabled") is True: + return None + remaining = bucket.get("remainingFraction", bucket.get("remaining_fraction")) + amount = bucket.get("remainingAmount", bucket.get("remaining_amount")) + if remaining is None and amount is not None: + return 1.0 if amount in (0, "0") and not isinstance(amount, bool) else None + if remaining is None and reset_of(bucket): + return 1.0 + if remaining is None or isinstance(remaining, bool): + return None + try: + remaining = float(remaining) + except (TypeError, ValueError): + return None + if not math.isfinite(remaining): + return None + return round(max(0.0, min(1.0, 1.0 - remaining)), 4) + + +def reset_of(bucket): + parsed = parse_time(bucket.get("resetTime") or bucket.get("reset_time")) + return parsed.isoformat() if parsed else "" + + +# The quota summary is asked for the project the plan's answer names, as agy +# and omp ask it, its id given as a string or as {id}; an answer naming none +# leaves the body empty. +def project_body(data): + project = data.get("cloudaicompanionProject") if isinstance(data, dict) else None + project = project.get("id") if isinstance(project, dict) else project + return {"project": project.strip()} if isinstance(project, str) and project.strip() else {} + + +# Whether an answer's `key` is a list, or not there at all. +def listed(data, key): + return data.get(key) is None or isinstance(data[key], list) + + +# 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. An answer that lists its buckets without groups +# lists the Gemini ones. An answer of another shape (not an object, an error, +# groups or the Gemini buckets that aren't a list) is None, a failed check +# rather than a plan without quotas; another group's buckets of the wrong +# shape skip that group. +def extract_limits(data): + if not isinstance(data, dict) or "error" in data or not listed(data, "groups") or not listed(data, "buckets"): + return None + groups = [g for g in data.get("groups") or [] if isinstance(g, dict)] + if not groups and data.get("buckets"): + groups = [{"displayName": "Gemini", "buckets": data["buckets"]}] + if not groups: + return [] + gemini = next((g for g in groups if "gemini" in str(g.get("displayName") or "").lower()), groups[0]) + if not listed(gemini, "buckets"): + return None + + limits = [] + session = weekly = None + for bucket in gemini.get("buckets") or []: + if not isinstance(bucket, dict) or used(bucket) is None: + continue + window = str(bucket.get("window") or "").lower() + bucket_id = str(bucket.get("bucketId") or "").lower() + 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") + buckets = group.get("buckets") + for bucket in buckets if isinstance(buckets, list) else []: + if isinstance(bucket, dict) and (used(bucket) or 0) > 0: + title = f"{short} {'Weekly' if 'week' in str(bucket.get('window') or '').lower() else 'Session'}" + limits.append({"label": title, "title": title, "percent": used(bucket), "resetsAt": reset_of(bucket)}) + return limits + + +# The bearer token goes wherever CLOUD_CODE_URL points, so a destination +# named there is taken only over https, and anything else is Google's own. +def post(endpoint, token, body=None): + base = (os.environ.get("CLOUD_CODE_URL") or DEFAULT_CLOUD_CODE_URL).rstrip("/") + if not base.startswith("https://"): + base = DEFAULT_CLOUD_CODE_URL + request = urllib.request.Request( + f"{base}/v1internal:{endpoint}", data=json.dumps(body or {}).encode(), + headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json", "User-Agent": "antigravity"}, + ) + return open_json(request, timeout=8) + + +def kept_tier(entry): + cache = limits_cache(entry) + return str(cache.data.get("tierLabel") or "") if cache else "" + + +def lapsed_help(entry, kept): + name = SOURCE_NAMES[entry["source"]] + shown = " — showing the last known limits." if kept else "." + if entry["refresh"]: + return f"The Antigravity access token {name} keeps has lapsed{shown} Start {name} to refresh it; signing in again is not needed." + return f"The Antigravity sign-in {name} keeps expired{shown} Start {name} to sign in again." + + +def refused_help(entry): + name = SOURCE_NAMES[entry["source"]] + return f"Google refused the Antigravity sign-in {name} keeps. Start {name} to sign in again." + + +# The plan and limits, asked of Google with one sign-in after another until +# one answers. Each account keeps its own last good answer, which stands in, +# marked stale, when its check fails. +def collect_limits(entries, force): + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "live": False, "fetchedAtMs": 0, "retryAdvised": False} + now = time.time() + # Each harness refreshes its own token while it runs, so the one lapsing + # last was refreshed most recently. A token of unknown expiry comes after + # those, since only Google can say whether it still holds, and a lapsed + # one last. + ordered = sorted(entries, key=lambda e: (not token_fresh(e, now), -e["expiresAt"])) + fresh = [e for e in ordered if token_fresh(e, now)] + lead = next(iter(fresh or [e for e in ordered if e["accessToken"]] or ordered), None) + omp = shutil.which("omp") if any(e["source"] == "omp" for e in entries) else None + omp_entries = [e for e in ordered if e["source"] == "omp"] if omp else [] + + def keep(entry): + # A window kept from an earlier check whose reset has passed is back to + # untouched, not unknown, so it reads zero until Google names the next + # reset. + cache = limits_cache(entry) if entry else None + if cache: + result.update(limits=cache.limits, tierLabel=str(cache.data.get("tierLabel") or ""), fetchedAtMs=cache.fetched_at) + else: + result.update(limits=[], tierLabel="", fetchedAtMs=0) + + def answered(entry, tier, limits): + # A plan with no quota windows to report is healthy, with no meters. + cache = limits_cache(entry) + fetched = cache.save(limits, tierLabel=tier) if cache else now_ms() + result.update(limits=limits, tierLabel=tier, live=True, fetchedAtMs=fetched) + return result + + keep(lead) + if not lead or not lead["accessToken"]: + result.update(usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + return result + + # A panel that is opened and shut repeatedly must not turn into a request + # per flick; --force is a person asking for fresh numbers. + if not force: + for entry in fresh + [e for e in omp_entries if e not in fresh]: + # An answer with no windows is an answer too: a plan without quotas. + cache = limits_cache(entry) + if cache and cache.fresh(force, allow_empty=True): + keep(entry) + result["live"] = True + return result + + refused, failure = [], None + for entry in fresh: + try: + assist = post("loadCodeAssist", entry["accessToken"]) + tier = extract_tier(assist) + limits = extract_limits(post("retrieveUserQuotaSummary", entry["accessToken"], project_body(assist))) + if limits is None: + raise HttpError() + except HttpError as error: + # A refused sign-in gives way to the next one; any other status is + # Google answering for every sign-in alike. + if error.status in (401, 403): + refused.append(entry) + continue + if error.status: + failure = entry, f"Google's quota service answered with status {error.status}.", False + elif error.transport: + failure = entry, "Couldn't reach Google's quota service. Retrying shortly.", True + else: + failure = entry, "Google's quota service sent an answer that couldn't be read.", False + break + return answered(entry, tier, limits) + + if failure: + entry, help_text, transport = failure + keep(entry) + result.update(authHelpText=help_text, retryAdvised=transport) + if not result["limits"]: + result["usageStatusText"] = "Antigravity limits unavailable" + return result + + # No token in hand answered, but omp renews its own sign-in when asked + # for usage. + answer = omp_answer(omp, omp_entries) if omp_entries else None + if answer: + entry, tier, limits = answer + return answered(entry, tier or kept_tier(entry), limits) + + if refused: + keep(refused[0]) + result.update(usageStatusText="Sign-in expired", authHelpText=refused_help(refused[0])) + return result + + # Every harness refreshes its token while it runs; one left to lapse can't + # be used until that harness starts again, which needs no new sign-in + # while a refresh token stands behind it. + result.update(usageStatusText="Limits paused" if lead["refresh"] else "Sign-in expired", authHelpText=lapsed_help(lead, result["limits"])) + return result + + +# omp refreshes its own Antigravity sign-in when asked for usage, so its +# limits come from `omp usage` even after the token it keeps has lapsed. The +# report is folded back into the shape Google's quota summary has. +def omp_usage(omp): + try: + output = subprocess.run( + [omp, "usage", "--json", "--provider", "google-antigravity"], + stdin=subprocess.DEVNULL, capture_output=True, text=True, timeout=20, check=True, + ).stdout + reports = json.loads(output).get("reports") + except (OSError, ValueError, AttributeError, subprocess.SubprocessError): + return None + # A report of another shape is no answer, never one without limits. + reports = [r for r in reports if isinstance(r, dict) and r.get("provider") == "google-antigravity"] if isinstance(reports, list) else [] + if not reports or not listed(reports[0], "limits"): + return None + groups, seen = {}, set() + for limit in reports[0].get("limits") or []: + if not isinstance(limit, dict): + continue + window = limit.get("window") if isinstance(limit.get("window"), dict) else {} + amount = limit.get("amount") if isinstance(limit.get("amount"), dict) else {} + label = str(limit.get("label") or "") + # A shared allowance is listed once per model family that draws on it. An + # id of another shape (a list, an object) still names one window. + key = (label, str(window.get("id"))) + if key in seen: + continue + seen.add(key) + fraction = amount.get("usedFraction") + if isinstance(fraction, bool) or not isinstance(fraction, (int, float)): + continue + groups.setdefault(label, []).append({"window": str(window.get("id") or ""), "remainingFraction": 1.0 - fraction, "resetTime": window.get("resetsAt")}) + metadata = reports[0].get("metadata") if isinstance(reports[0].get("metadata"), dict) else {} + quota = {"groups": [{"displayName": name, "buckets": buckets} for name, buckets in groups.items()]} + return quota, str(metadata.get("email") or "") + + +# The sign-in `omp usage` answered for, its plan when the token omp just +# renewed can still ask Google for it, and its limits. +def omp_answer(omp, entries): + answer = omp_usage(omp) + if not answer: + return None + quota, email = answer + if email: + renewed = [e for e in omp_logins() if e["email"] == email] + entry = (renewed or [e for e in entries if e["email"] == email] or [dict(entries[0], email=email)])[0] + else: + # A report naming no account answers for omp's only sign-in. With several, + # there is no telling whose it is: it is shown, and kept for none of them. + signed_in = omp_logins() or entries + entry = signed_in[0] if len(signed_in) == 1 else dict(signed_in[0], email="", refresh="", accessToken="", expiresAt=0) + tier = "" + if token_fresh(entry, time.time()): + try: + tier = extract_tier(post("loadCodeAssist", entry["accessToken"])) + except HttpError: + pass + return entry, tier, extract_limits(quota) + + +# A transcript step's tokens as (input, output, cache read, cache write). +# Steps nest Gemini's counts under usage, where the prompt count includes +# the cached part and thinking sits beside the candidates; newer builds put +# input_tokens, output_tokens, cache_read_tokens and cache_write_tokens on the +# step itself. Every reader of those adds the four up, so input here excludes +# the cache reads, the way Anthropic's counts do. None when the step reports +# no usage. A damaged count (a boolean, negative, NaN, huge) is nothing. +def step_tokens(step): + usage = step.get("usage") + for source in ([usage] if isinstance(usage, dict) else []) + [step]: + if "prompt_token_count" in source or "candidates_token_count" in source: + prompt = number(source.get("prompt_token_count")) + cached = min(prompt, number(source.get("cached_content_token_count"))) + output = number(source.get("candidates_token_count")) + number(source.get("thoughts_token_count")) + return prompt - cached, output, cached, 0 + if any(key in source for key in ("input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens")): + return (number(source.get("input_tokens")), number(source.get("output_tokens")), + number(source.get("cache_read_tokens")), number(source.get("cache_write_tokens"))) + return None + + +# A model step names its model by id; a "Model Selection" step names it the +# way the picker shows it ("Gemini 3.8 Flash (High)"). Both fold to the same +# key, so a conversation's tokens don't split across two spellings. +MODEL_SELECTION = re.compile(r"Model Selection` (?:from .+? )?to (.+?)\.(?:\s|$|<)") + + +def model_key(name): + key = re.sub(r"[\s_]+", "-", re.sub(r"[()]", " ", name).strip().lower()) + return re.sub(r"-+", "-", key).strip("-") + + +def step_model(step): + for field in ("model", "model_name", "model_id"): + value = step.get(field) + if isinstance(value, str) and model_key(value): + return model_key(value) + return "" + + +def selected_model(step): + content = step.get("content") + if not isinstance(content, str) or "Model Selection" not in content: + return "" + match = MODEL_SELECTION.search(content) + return model_key(match.group(1)) if match else "" + + +def history_entries(root, problems): + try: + lines = (root / "history.jsonl").read_text(encoding="utf-8", errors="ignore").splitlines() + except FileNotFoundError: + return + except OSError: + problems.append("history") + return + for line in lines: + try: + entry = json.loads(line) + except ValueError: + continue + if isinstance(entry, dict): + yield entry + + +# Each conversation in the summaries with when it last changed and, where +# agy records it, the conversation it was spawned from. +def conversation_summaries(root, problems): + summaries = root / "conversation_summaries.db" + if not summaries.is_file(): + return [] + try: + conn = sqlite3.connect(summaries.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + columns = {row[1] for row in conn.execute("PRAGMA table_info(conversation_summaries)")} + parent = "parent_conversation_id" if "parent_conversation_id" in columns else "NULL" + return conn.execute(f"SELECT conversation_id, last_modified_time, {parent} FROM conversation_summaries").fetchall() + finally: + conn.close() + except sqlite3.Error: + problems.append("summaries") + return [] + + +# OpenClaw keeps pi's transcript format, one directory per agent, beside +# pi's and omp's own trees. +def session_roots(): + roots = pi_session_roots() + openclaw = Path(os.environ.get("OPENCLAW_STATE_DIR") or (Path.home() / ".openclaw")) + try: + # Sorted so the scan order doesn't depend on directory order. + roots.extend(sorted(child / "sessions" for child in (openclaw / "agents").iterdir() if child.is_dir())) + except OSError: + pass + # An alias is scanned under one spelling, and a root inside another is + # dropped rather than counted twice. + return canonical_roots(roots) + + +# pi, omp, and OpenClaw answers through their google-antigravity provider. +def pi_billed_here(provider, model, entry): + return provider == "google-antigravity" + + +# opencode's Antigravity plugin runs under a "google" provider, with model ids +# it prefixes with antigravity- to keep them apart from Gemini API ones. +def opencode_billed_here(provider, model, entry): + return provider.startswith("google") and model.startswith("antigravity-") + + +# Each answer pi, omp, OpenClaw, or opencode's Antigravity plugin recorded, +# counted once. Hermes reaches Antigravity only through plugins that run agy +# itself, so its sign-in and transcripts are agy's own, which local_stats +# reads already. +def harness_rows(problems): + # OpenClaw's trees are read along with pi's and omp's, so a message copied + # between any of them is still counted once. + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"google-antigravity"', pi_billed_here, session_roots()) + db_rows, db_complete = scan_opencode(["google"], opencode_billed_here) + if not pi_complete: + problems.append("sessions") + if not db_complete: + problems.append("opencode") + return pi_rows + db_rows + + +# Every prompt from history.jsonl, every conversation from it and from the +# conversation summaries, and the tokens each model step in a transcript +# reports. Steps that report no usage add no tokens; steps without a usable +# date still count toward the lifetime totals. Antigravity driven from +# another harness adds each answer it recorded as a prompt, the way the +# Codex collector counts them, with every token it reports. +def local_stats(root, problems=None): + problems = [] if problems is None else problems + stats = Stats() + + # A subagent's conversation is part of the one that spawned it, so it + # counts toward that conversation rather than as a session of its own. + summaries = conversation_summaries(root, problems) + parents = {str(conversation): str(parent) for conversation, _, parent in summaries if conversation and parent} + + def session_of(conversation): + conversation, seen = str(conversation), set() + while conversation in parents and conversation not in seen: + seen.add(conversation) + conversation = parents[conversation] + return conversation + + for entry in history_entries(root, problems): + day = local_day(entry.get("timestamp")) + if not day: + continue + conversation = entry.get("conversationId") or entry.get("conversation_id") or entry.get("sessionId") + stats.add_prompts(day, session_of(conversation) if conversation else None) + + # A conversation with no prompt in sight is a session all the same, and + # today's when it was last touched today. + for conversation, modified, _ in summaries: + if conversation: + stats.add_prompts(stats.today if local_day(modified) == stats.today else None, session_of(conversation), 0) + + for transcript in (root / "brain").glob("*/.system_generated/logs/transcript.jsonl"): + try: + lines = transcript.read_text(encoding="utf-8", errors="ignore").splitlines() + except OSError: + problems.append(str(transcript)) + continue + # The model a conversation runs on carries from step to step until a + # step names another one. + current = "" + for line in lines: + try: + step = json.loads(line) + except ValueError: + continue + if not isinstance(step, dict): + continue + current = step_model(step) or selected_model(step) or current + tokens = step_tokens(step) + if not tokens or not any(tokens): + continue + # Tokens with no prompt, and maybe no date. + stats.add(local_day(step.get("created_at") or step.get("timestamp")) or None, None, current or "unknown", *tokens, prompts=0) + + for row in harness_rows(problems): + stats.add_row(row, model_key(row.model) or "unknown") + + snapshot = stats.snapshot() + return {"hasLocalStats": snapshot["totalPrompts"] > 0 or bool(stats.sessions) or bool(snapshot["modelUsage"]), **snapshot} + + +# A scan is reused for max_age seconds, and only on the day it was made: one +# from before midnight would count yesterday's prompts as today's. One that +# couldn't read every source isn't kept. +def cached_local_stats(root, max_age): + def scan(): + problems = [] + stats = local_stats(root, problems) + return stats, not problems + + return cached_scan(AGENT_ID, [root, opencode_db(), *session_roots()], max_age, scan) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + stats = cached_local_stats(app_root(), scan_max_age(args.force, args.limits_only)) + + # Nobody signed in to Antigravity here and no local history: no limits and + # no plan, and the panel skips the record. + entries = logins() + ready = bool(entries) or bool(stats.get("hasLocalStats")) + limits = collect_limits(entries, args.force) if ready else { + "limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "live": False, "fetchedAtMs": 0, "retryAdvised": False, + } + tier = limits["tierLabel"] + agy = next((e for e in entries if e["source"] == "agy"), None) + if not tier and agy and agy["authMethod"] == "enterprise": + tier = "Enterprise" + elif not tier and agy and agy["authMethod"] in ("cloud", "vertex"): + tier = "Cloud" + + record = base_record( + AGENT_ID, AGENT_NAME, + ready=ready, + tierLabel=tier, + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + limits=limits["limits"], + # Stale says limits were kept past a failed check; with none kept there + # is nothing to be stale. + limitsStale=not limits["live"] and bool(limits["limits"]), + limitsFetchedAt=limits["fetchedAtMs"], + retryAdvised=limits["retryAdvised"], + ) + record.update(stats) + emit(record) + + +if __name__ == "__main__": + main() diff --git a/bin/omarchy-agent-usage-claude b/bin/omarchy-agent-usage-claude index e615c5836eb..0b81f76c9ba 100755 --- a/bin/omarchy-agent-usage-claude +++ b/bin/omarchy-agent-usage-claude @@ -16,25 +16,54 @@ from __future__ import annotations import argparse import datetime as dt -import fcntl -import hashlib import json import os import re -import sqlite3 import sys -import tempfile import time -import urllib.error import urllib.request from pathlib import Path from typing import Any +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + LIMITS_ONLY_REUSE_SECONDS, + SCAN_REUSE_SECONDS, + HttpError, + LimitsCache, + Stats, + account_entry, + base_record, + cache_root, + cached_scan, + claude_code_base_url, + claude_code_dir, + claude_code_index_path, + claude_code_owner, + claude_code_scan_match, + emit, + empty_bucket, + number, + obj, + open_json, + opencode_db, + pi_session_roots, + read_json, + recent_dates, + registered_accounts, + claude_code_usage, + scan_opencode, + scan_pi_sessions, + today as today_string, + warn, +) + AGENT_ID = "claude" AGENT_NAME = "Claude Code" AUTH_HELP = "Run `claude auth login` to restore authoritative usage." USAGE_ENDPOINT = "https://api.anthropic.com/api/oauth/usage" -PROBE_MIN_INTERVAL_SECONDS = 15 # Anthropic rate-limits its usage endpoint readily. A refused re-check of # numbers this recent still describes the account; only older ones are # reported as stale. @@ -42,452 +71,154 @@ CURRENT_ENOUGH_SECONDS = 900 def config_dir() -> Path: - return expand_path(os.environ.get("CLAUDE_CONFIG_DIR") or "~/.claude") + return claude_code_dir() def expand_path(value: str) -> Path: return Path(os.path.expandvars(os.path.expanduser(value))).resolve() -def cache_root() -> Path: - root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "omarchy" / "agent-usage" - root.mkdir(parents=True, exist_ok=True) - return root - - -def 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() - - # Claude JSONL timestamps are usually ISO-8601. Python accepts offsets but - # not a trailing Z until we normalize it to +00:00. - 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 usage_token(usage: dict[str, Any], snake_key: str, camel_key: str) -> int: - value = usage.get(snake_key, usage.get(camel_key, 0)) - try: - return round(float(value or 0)) - except Exception: - return 0 - - -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, - } - - # ---------------------------------------------------------------- local scan -# Transcripts are append-only, and most of the corpus never changes between -# ticks, so the scan keeps a per-file index: the unique usage records already -# parsed out of each file plus the byte offset they end at. A file whose size -# and mtime match the index is not opened; a file that grew is read from the -# stored offset; anything else (truncated, rewritten, replaced, unknown) is -# read from the start. Without this the whole history was re-parsed every -# refresh, and the cost grew for the life of the machine. - -INDEX_VERSION = 4 - -# How much of each end of what was read the index remembers a hash of, so a -# file replaced or rewritten with more than it had is not mistaken for an -# append. -INDEX_HEAD_BYTES = 4096 +# Claude Code's transcripts are read by the shared reader, through the index +# every collector that reads them shares. Claude Code fronts whatever +# ANTHROPIC_BASE_URL names, and Z.ai, Ollama, MiniMax, DeepSeek, Kimi, +# OpenRouter and Alibaba's Model Studio each document pointing it at their +# own Anthropic-compatible API; +# only what Anthropic served is Claude's, and claude_code_owner() hands the +# rest to the collector that bills it. -# Records hold local days, so an index is only true in the timezone that -# wrote it; after a change every transcript is read again. -def index_zone() -> list[Any]: - return [os.environ.get("TZ"), list(time.tzname), time.timezone, time.altzone] - - -def load_index(index_file: Path) -> dict[str, Any]: - try: - data = json.loads(index_file.read_text(encoding="utf-8")) - if ( - isinstance(data, dict) - and data.get("version") == INDEX_VERSION - and data.get("zone") == index_zone() - and isinstance(data.get("files"), dict) - ): - return data["files"] - except Exception: - pass - return {} - - -# What was already read, by its first and last few kilobytes: a transcript -# rewritten in place rather than appended to changes one or the other. -def read_digest(path: Path, length: int) -> str: - digest = hashlib.sha1() - with path.open("rb") as handle: - digest.update(handle.read(min(length, INDEX_HEAD_BYTES))) - handle.seek(max(0, length - INDEX_HEAD_BYTES)) - digest.update(handle.read(min(length, INDEX_HEAD_BYTES))) - return digest.hexdigest() - - -def parse_usage_records(path: Path, entry: dict[str, Any] | None) -> dict[str, Any]: - """Return the index entry for one transcript, reusing `entry` where the file - only grew. Records are [key, model, day, input, output, cacheRead, - cacheWrite, session]; `day` is None when the line carries no timestamp.""" - stat = path.stat() - size = stat.st_size - mtime = stat.st_mtime_ns - - if entry and entry.get("size") == size and entry.get("mtime") == mtime: - return entry - - # Only the same file, grown past what was read with what was read still in - # place, is an append. A transcript replaced by a larger one is a new file, - # and one rewritten in place has a different start or end. - if ( - entry - and size > number(entry.get("size")) - and entry.get("file") == [stat.st_dev, stat.st_ino] - and entry.get("head") == read_digest(path, number(entry.get("offset"))) - ): - records = list(entry.get("records") or []) - offset = number(entry.get("offset")) - line_number = number(entry.get("lines")) - else: - records = [] - offset = 0 - line_number = 0 - - # A streamed response is several lines sharing one message id, and the - # first line's output_tokens is a placeholder while the last carries the - # real count. Keep each message's line with the highest output count (the - # last one read wins a tie), whole: a response that fell back to another - # model mid-stream has lines that are separate snapshots, and mixing fields - # across them over-counts cache tokens and credits the wrong model. - positions = {record[0]: index for index, record in enumerate(records)} - with path.open("rb") as handle: - handle.seek(offset) - for raw in handle: - if not raw.endswith(b"\n"): - # A line still being written; pick it up whole next time. - break - offset += len(raw) - line_number += 1 - - # Cheap pre-filter before JSON parsing keeps files with unrelated - # lines inexpensive. - if b'"usage":' not in raw: - continue - - try: - parsed = json.loads(raw.decode("utf-8", errors="replace")) - except Exception: - continue - - message = parsed.get("message") if isinstance(parsed.get("message"), dict) else {} - if parsed.get("type") != "assistant" and message.get("role") != "assistant": - continue - - usage = message.get("usage") or parsed.get("usage") - if not isinstance(usage, dict): - continue - - message_id = message.get("id") or parsed.get("messageId") or "" - key = str(message_id) if message_id else f"{path}:{parsed.get('uuid') or parsed.get('requestId') or line_number}" - - input_tokens = usage_token(usage, "input_tokens", "inputTokens") - output_tokens = usage_token(usage, "output_tokens", "outputTokens") - cache_read = usage_token(usage, "cache_read_input_tokens", "cacheReadInputTokens") - cache_write = usage_token(usage, "cache_creation_input_tokens", "cacheCreationInputTokens") - if input_tokens + output_tokens + cache_read + cache_write <= 0: - continue - - held = positions.get(key) - if held is not None and output_tokens < records[held][4]: - continue - - timestamp = parsed.get("timestamp") or message.get("timestamp") - record = [ - key, - str(message.get("model") or parsed.get("model") or "claude"), - local_date_from_timestamp(timestamp) if timestamp is not None else None, - input_tokens, - output_tokens, - cache_read, - cache_write, - str(parsed.get("sessionId") or path), - ] - if held is None: - positions[key] = len(records) - records.append(record) - else: - records[held] = record - - return { - "size": size, - "file": [stat.st_dev, stat.st_ino], - "mtime": mtime, - "head": read_digest(path, offset), - "offset": offset, - "lines": line_number, - "records": records, - } - - -def scan_projects(projects_path: Path, index_file: Path | None = None) -> dict[str, Any]: - previous = load_index(index_file) if index_file else {} - current: dict[str, Any] = {} - - # Directory order, as before: when one API message shows up in several - # transcripts with the same output count, the last file visited wins, and - # the numbers must not move because the scan got faster. - order: list[str] = [] - unreadable_count = 0 - unreadable_examples: list[str] = [] - files = projects_path.rglob("*.jsonl") if projects_path.is_dir() else [] - for path in files: - key = str(path) - try: - current[key] = parse_usage_records(path, previous.get(key)) - except Exception as exc: - unreadable_count += 1 - if len(unreadable_examples) < 3: - unreadable_examples.append(f"{path}: {exc}") - # What was read before is still true of the file. - if key in previous: - current[key] = previous[key] - if key in current: - order.append(key) - - if unreadable_count: - examples = "; ".join(unreadable_examples) - print( - f"Ignoring {unreadable_count} unreadable Claude project files" - f" (examples: {examples})", - file=sys.stderr, - ) - - if index_file and current != previous: - try: - write_json(index_file, {"version": INDEX_VERSION, "zone": index_zone(), "files": current}) - except Exception as exc: - print(f"Could not write Claude scan index {index_file}: {exc}", file=sys.stderr) - - return summarize_records(current, order) - - -def summarize_records(indexed: dict[str, Any], order: list[str]) -> dict[str, Any]: - today = local_date_string() - recent_dates = recent_date_strings() - recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} - - # One API message can show up in several transcripts (a resumed session - # carries the earlier ones). As within a file, the record with the highest - # output count wins, the last one visited breaking a tie. - seen: dict[str, list[Any]] = {} - 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]] = {} - prompts = 0 - today_prompt_count = 0 - today_token_total = 0 - - for key in order: - for record in indexed[key].get("records") or []: - held = seen.get(record[0]) - if held is not None and record[4] < held[4]: - continue - seen[record[0]] = record - - for _, model, day, input_tokens, output_tokens, cache_read, cache_write, session_key in seen.values(): - if day is None: - day = today - total = input_tokens + output_tokens + cache_read + cache_write - - 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 - - return { - "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), - # Days with any recorded usage, for the all-time "N days" summary. The - # dates travel too: merging snapshots from several machines needs their - # union, which a count alone cannot give. - "activeDays": len(active_days), - "activeDates": sorted(active_days), - } - - -def scan_cache_paths(projects_path: Path) -> tuple[Path, Path, Path]: - digest = hashlib.sha1(str(projects_path).encode("utf-8")).hexdigest()[:16] - root = cache_root() - return root / f"claude-scan-{digest}.json", root / f"claude-scan-{digest}.lock", root / f"claude-index-{digest}.json" - - -def read_json(path: Path) -> dict[str, Any] | None: - try: - return json.loads(path.read_text(encoding="utf-8")) - except Exception: - return None +def scan_projects(projects_path: Path, index_file: Path | None = None, base_url: str = "", reread: bool = False) -> Stats: + return scan_projects_read(projects_path, index_file, base_url, reread)[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: - # A negative age means the mtime is in the future: the clock moved - # backwards since the write, so the cache's freshness cannot be trusted. - age = time.time() - path.stat().st_mtime - if 0 <= age <= max_age_seconds: - return read_json(path) - except Exception: - return None - return None +def scan_projects_read(projects_path: Path, index_file: Path | None = None, base_url: str = "", reread: bool = False) -> tuple[Stats, bool]: + """The transcripts' stats, and whether every transcript was read: a scan + that missed one isn't cached.""" + stats = Stats() + rows, read_all = claude_code_usage(AGENT_ID, projects_path, index_file, base_url, reread) + for row in rows: + stats.add_row(row, row.model) + return stats, read_all -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 - - -def cached_scan(projects_path: Path, max_age_seconds: float, use_index: bool = True) -> dict[str, Any]: - cache_file, lock_file, index_file = scan_cache_paths(projects_path) - - cached = read_fresh_json(cache_file, max_age_seconds) - if cached is not None: - return cached - - with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - cached = read_fresh_json(cache_file, max_age_seconds) - if cached is not None: - return cached - if not use_index: - index_file.unlink(missing_ok=True) - summary = scan_projects(projects_path, index_file) - write_json(cache_file, summary) - return summary +def anthropic_served(model: Any, base_url: str) -> bool: + return claude_code_owner(model, base_url) == AGENT_ID # ------------------------------------------------------------- local fallback # # A machine without transcripts on disk can still know its history: Claude # Code keeps aggregate counters in stats-cache.json and per-prompt history in -# history.jsonl. Only consulted when the project scan comes back empty. +# history.jsonl. Only consulted when projects/ holds no transcript at all: a +# scan that gave every reply to another collector has counted each prompt +# already, under whoever answered it. Neither file says who answered, so +# while ANTHROPIC_BASE_URL names another provider's host their prompts and +# sessions are that provider's, which reads no fallback, and count nowhere. + +# Claude Code's built-in commands that run on their own and ask the model +# nothing (Claude Code 2.1.308's `local` and `local-jsx` commands and their +# aliases), which history.jsonl keeps as typed. The transcript reader knows +# one by the local output answering it; history.jsonl keeps no output, so +# only the name tells. A custom command, a skill, and a built-in that sends a +# prompt (/init, /statusline) are prompts. +CLAUDE_CODE_LOCAL_COMMANDS = frozenset(( + "add-dir", "advisor", "agents", "allowed-tools", "android", "app", "artifacts", "autocompact", "autofix-pr", + "auto-mode-setup", "background", "bashes", "bg", "branch", "break-reminder", "breaks", "brief", "btw", "bug", "cd", + "checkpoint", "clear", "cloud-plugins", "color", "compact", "config", "context", "copy", "cost", "daemon", + "design-consent", "design-login", "design-revoke", "desktop", "diff", "downtime", "effort", "exit", "export", + "extra-usage", "fast", "feedback", "focus", "fork", "goal", "heapdump", "help", "hooks", "ide", "import", "install", + "install-github-app", "install-slack-app", "ios", "list-agents", "login", "logout", "loops", "marketplace", "mcp", + "memory", "memory-pause", "mobile", "model", "name", "new", "output-style", "passes", "pause-memory", "peers", + "permissions", "plan", "plugin", "plugins", "powerup", "privacy-settings", "pro-trial-expired", "quit", "radio", + "rate-limit-options", "rc", "recap", "reload-plugins", "reload-skills", "remote", "remote-control", "remote-env", + "rename", "reset", "restart", "resume", "rewind", "scroll-speed", "session", "settings", "setup-bedrock", + "setup-vertex", "skill-doctor", "skills", "stats", "status", "stickers", "stop", "subtask", "tasks", "teleport", + "terminal-setup", "theme", "toggle-memory", "tp", "tui", "ultraplan", "ultrareview", "undo", "update", "upgrade", + "usage", "usage-credits", "version", "voice", "web-setup", "wellbeing", "workflow-launch-exec", "workflows", +)) + + +def history_prompt(display: Any) -> bool: + """Whether a history.jsonl entry's text was a prompt: anything but a + built-in local command, with or without arguments.""" + text = str(display or "").strip() + if not text.startswith("/"): + return True + name = text[1:].split(None, 1)[0].lower() if len(text) > 1 else "" + return name not in CLAUDE_CODE_LOCAL_COMMANDS + + +def has_transcripts(projects: Path) -> bool: + try: + return next(projects.rglob("*.jsonl"), None) is not None + except OSError: + return False -def stats_cache_fallback(claude_dir: Path) -> dict[str, Any] | None: +def stats_cache_fallback(claude_dir: Path, base_url: str = "") -> dict[str, Any] | None: try: data = json.loads((claude_dir / "stats-cache.json").read_text(encoding="utf-8")) except Exception: return None + # A file or a field of the wrong shape reads as nothing, never a crash. + if not isinstance(data, dict): + return None - today = local_date_string() - daily_model_tokens = data.get("dailyModelTokens") or [] - today_tokens = {} - for entry in daily_model_tokens: - if isinstance(entry, dict) and entry.get("date") == today: - today_tokens = entry.get("tokensByModel") or {} - break + def entries(key: str) -> list[dict[str, Any]]: + value = data.get(key) + return [entry for entry in value if isinstance(entry, dict)] if isinstance(value, list) else [] + + # Tokens by model per day, the first entry for a day standing, as the + # record's recentDays and today's tokens read them. Claude Code keys these + # by UTC date (toISOString) and stops at yesterday's, so each is read as + # the local day of the same date, the one it mostly overlaps: off by the + # UTC offset at either end, and today's tokens are usually none. + tokens_by_day: dict[str, dict[str, int]] = {} + for entry in entries("dailyModelTokens"): + date = entry.get("date") + if isinstance(date, str) and date not in tokens_by_day: + tokens_by_day[date] = { + str(model): number(tokens) for model, tokens in obj(entry.get("tokensByModel")).items() + if number(tokens) > 0 and anthropic_served(model, base_url) + } + today_tokens = tokens_by_day.get(today_string(), {}) + recent_days = [{"date": date, "messageCount": sum(tokens_by_day.get(date, {}).values())} for date in recent_dates()] + + # The four token kinds of each model, each a count number() trusts; the + # file's other fields (costUSD, webSearchRequests) aren't the record's. + model_usage = { + str(model): {field: number(bucket.get(field)) for field in empty_bucket()} + for model, bucket in obj(data.get("modelUsage")).items() + if isinstance(bucket, dict) and anthropic_served(model, base_url) + } - daily_activity = [day for day in (data.get("dailyActivity") or []) if isinstance(day, dict)] + # Prompts and sessions name no model, so they are Claude's unless Claude + # Code is sent to another provider's host. + prompts_here = anthropic_served(None, base_url) + daily_activity = entries("dailyActivity") if prompts_here else [] active_dates = sorted({str(day.get("date")) for day in daily_activity if number(day.get("messageCount")) > 0 and day.get("date")}) - today_prompts, today_sessions = today_prompts_from_history(claude_dir) + today_prompts, today_sessions = today_prompts_from_history(claude_dir, base_url) return { "todayPrompts": today_prompts, "todaySessions": today_sessions, - "todayTotalTokens": sum(number(v) for v in today_tokens.values()), + "todayTotalTokens": sum(today_tokens.values()), "todayTokensByModel": today_tokens, - "recentDays": daily_activity[-7:], - "modelUsage": data.get("modelUsage") or {}, - "totalPrompts": number(data.get("totalMessages")), - "totalSessions": number(data.get("totalSessions")), + "recentDays": recent_days, + "modelUsage": model_usage, + "totalPrompts": number(data.get("totalMessages")) if prompts_here else 0, + "totalSessions": number(data.get("totalSessions")) if prompts_here else 0, "activeDays": len(active_dates), "activeDates": active_dates, } -def today_prompts_from_history(claude_dir: Path) -> tuple[int, int]: +def today_prompts_from_history(claude_dir: Path, base_url: str = "") -> tuple[int, int]: + if not anthropic_served(None, base_url): + return 0, 0 prompts = 0 sessions: set[str] = set() start_of_day = dt.datetime.combine(dt.datetime.now().date(), dt.time.min).timestamp() * 1000 @@ -505,259 +236,93 @@ def today_prompts_from_history(claude_dir: Path) -> tuple[int, int]: entry = json.loads(line) except Exception: continue + if not isinstance(entry, dict): + continue if number(entry.get("timestamp")) < start_of_day: break + if not history_prompt(entry.get("display")): + continue prompts += 1 if entry.get("sessionId"): sessions.add(str(entry.get("sessionId"))) return prompts, len(sessions) -# --------------------------------------------------------------- pi and omp + +# ------------------------------------------------------- pi, omp, and opencode # # These agents can consume a Claude subscription without writing native -# Claude Code transcripts. Their compatible JSONL session formats carry the -# provider, model, and token usage on every assistant message. - - -# `omp --profile=` (and pi's equivalent) relocates the whole agent tree -# under /profiles//, so a subscription driven entirely through a -# profile leaves the default root empty and its usage uncounted. Sessions are -# keyed by file path, so a profile adds sessions rather than double-counting -# the default root. -def pi_session_roots() -> list[Path]: - roots: list[Path] = [] - for base in (Path.home() / ".pi", Path.home() / ".omp"): - roots.append(base / "agent" / "sessions") - profiles = base / "profiles" - try: - # Sorted so the scan order does not depend on directory order. - roots.extend(sorted(child / "agent" / "sessions" for child in profiles.iterdir() if child.is_dir())) - except OSError: - # No profiles directory, or it is unreadable: the default root stands. - pass - return roots - - -def scan_pi_usage(max_age_seconds: float) -> dict[str, Any] | None: - roots = pi_session_roots() - cache_file = cache_root() / "claude-pi-sessions.json" - cached = read_fresh_json(cache_file, max_age_seconds) - if cached is not None: - return cached.get("stats") - - today = local_date_string() - recent_dates = recent_date_strings() - recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} - sessions: set[str] = set() - active_days: set[str] = set() - today_sessions: set[str] = set() - today_tokens: dict[str, int] = {} - usage_by_model: dict[str, dict[str, int]] = {} - seen: set[str] = set() - prompts = 0 - today_prompt_count = 0 - today_token_total = 0 - - for root in roots: - files = root.rglob("*.jsonl") if root.is_dir() else [] - for path in files: - try: - with path.open("r", encoding="utf-8", errors="replace") as handle: - for line_number, line in enumerate(handle, 1): - if '"usage"' not in line or '"assistant"' 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 - provider = str(message.get("provider") or "") - if provider != "anthropic": - 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 = usage_token(usage, "input", "inputTokens") - output_tokens = usage_token(usage, "output", "outputTokens") - cache_read = usage_token(usage, "cacheRead", "cache_read_input_tokens") - cache_write = usage_token(usage, "cacheWrite", "cache_creation_input_tokens") - total = input_tokens + output_tokens + cache_read + cache_write - if total <= 0: - total = number(usage.get("totalTokens")) - input_tokens = total - if total <= 0: - continue - model = str(message.get("model") or "claude") - 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: - recent[day]["messageCount"] += total - if day == today: - today_prompt_count += 1 - today_sessions.add(session_key) - today_token_total += total - today_tokens[model] = today_tokens.get(model, 0) + total - except OSError: - continue - - stats = None - if prompts > 0: - stats = { - "todayPrompts": today_prompt_count, - "todaySessions": len(today_sessions), - "todayTotalTokens": today_token_total, - "todayTokensByModel": today_tokens, - "recentDays": [recent[day] for day in recent_dates], - "modelUsage": usage_by_model, - "totalPrompts": prompts, - "totalSessions": len(sessions), - "activeDays": len(active_days), - "activeDates": sorted(active_days), - } - write_json(cache_file, {"stats": stats}) - return stats - - -# ---------------------------------------------------------------- opencode -# -# A Claude subscription burned entirely through opencode never writes a -# transcript under ~/.claude, but opencode records per-message provider, -# model, and token usage in its own database. Scan it for Anthropic-provider -# messages and merge the result into whatever the transcript scan found. - - -def scan_opencode_usage(max_age_seconds: float) -> dict[str, Any] | None: - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" - if not db.is_file(): - return None - - # Same freshness contract as the transcript scan: --limits-only promises to - # reuse recent local stats, and a big opencode history walked on every panel - # open would break that promise. - cache_file = cache_root() / f"claude-opencode-{hashlib.sha1(str(db).encode('utf-8')).hexdigest()[:16]}.json" - cached = read_fresh_json(cache_file, max_age_seconds) - if cached is not None: - return cached.get("stats") - - today = local_date_string() - recent_dates = recent_date_strings() - recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} - sessions: set[str] = set() - active_days: set[str] = set() - today_sessions: set[str] = set() - today_tokens: dict[str, int] = {} - usage_by_model: dict[str, dict[str, int]] = {} - prompts = 0 - today_prompt_count = 0 - today_token_total = 0 - - try: - # Read-only: opencode may be writing right now. - conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) - except sqlite3.Error: - return None +# Claude Code transcripts: pi's and omp's JSONL sessions, and opencode's +# message database, carry the provider, model, and token usage of every +# reply. The shared readers find them (pi and omp with `rg --no-ignore`, so a +# $HOME that is a git checkout hides nothing) and this keeps the Anthropic +# ones. + +PI_PATTERN = r'"provider"\s*:\s*"anthropic"' + + +# Exact match: provider ids are free-form, and a custom "anthropic-proxy" +# gateway is not this subscription. +def billed_here(provider: str, model: str, entry: dict) -> bool: + return provider == "anthropic" + + +def opencode_model(model: str) -> str: + return str(model or "claude").rstrip("/").split("/")[-1] or "claude" + + +def add_shared_rows(stats: Stats, pi_rows: list, db_rows: list) -> None: + for row in pi_rows: + stats.add_row(row, row.model or "claude") + for row in db_rows: + stats.add_row(row, opencode_model(row.model)) + + +def scan_local(claude_dir: Path, index_file: Path | None, reread: bool = False) -> tuple[dict[str, Any], bool]: + """Local stats from every source, and whether every source was read in full + (safe to cache). Transcripts come first; only a machine without any falls + back on Claude Code's own aggregate counters.""" + base_url = claude_code_base_url(claude_dir) + projects = claude_dir / "projects" + stats, transcripts_complete = scan_projects_read(projects, index_file, base_url, reread) + pi_rows, pi_complete = scan_pi_sessions(PI_PATTERN, billed_here) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode(['anthropic"'], billed_here) + complete = transcripts_complete and pi_complete and db_complete + + if stats.total_prompts > 0 or stats.model_usage or has_transcripts(projects): + add_shared_rows(stats, pi_rows, db_rows) + return stats.snapshot(), complete + + local = stats_cache_fallback(claude_dir, base_url) + if local is None: + local = stats.snapshot() + # No transcripts and no aggregate cache, but history.jsonl alone can + # still put numbers on today. + today_prompts, today_sessions = today_prompts_from_history(claude_dir, base_url) + if today_prompts or today_sessions: + local.update(todayPrompts=today_prompts, todaySessions=today_sessions) + extra = Stats() + add_shared_rows(extra, pi_rows, db_rows) + if extra.total_prompts > 0: + local = merge_stats(local, extra.snapshot()) + return local, complete + + +def local_stats(claude_dir: Path, max_age: float, use_index: bool = True) -> dict[str, Any]: + # Without the index every transcript is read again, and the index rebuilt + # in place: the other collectors reading Claude Code share it. try: - conn.execute("PRAGMA query_only = ON") - tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} - # opencode v1 keeps messages in `message`, with the role, provider and model - # in each row's data. v2 writes `session_message` instead: the role moves to - # the `type` column and the provider and model nest under data.model. A - # database can hold either or both. - queries = [] - if "message" in tables: - queries.append(("opencode:", "SELECT session_id, data, NULL FROM message")) - if "session_message" in tables: - queries.append(("opencode-v2:", "SELECT session_id, data, type FROM session_message")) - rows = ((prefix, row) for prefix, query in queries for row in conn.execute(query)) - for session_prefix, (session_id, raw, row_type) in rows: - # One malformed row must not abort the scan, so every shape assumption - # lives inside the try. - try: - entry = json.loads(raw) - if not isinstance(entry, dict): - continue - if row_type is None: - role = entry.get("role") - provider = entry.get("providerID") - model_id = entry.get("modelID") - else: - role = row_type - model_info = entry.get("model") or {} - provider = model_info.get("providerID") - model_id = model_info.get("id") - # Exact match: opencode provider ids are free-form, and a custom - # "anthropic-proxy" gateway is not this subscription. - if role != "assistant" or str(provider or "") != "anthropic": - 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")) - total = input_tokens + output_tokens + cache_read + cache_write - if total <= 0: - continue - - created = number((entry.get("time") or {}).get("created")) - day = dt.datetime.fromtimestamp(created / 1000).strftime("%Y-%m-%d") if created > 0 else today - model = str(model_id or "claude").rstrip("/").split("/")[-1] - except Exception: - continue - session_key = session_prefix + str(session_id) - 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: - 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 sqlite3.Error: - return None - finally: - conn.close() - - stats = None - if prompts > 0: - stats = { - "todayPrompts": today_prompt_count, - "todaySessions": len(today_sessions), - "todayTotalTokens": today_token_total, - "todayTokensByModel": today_tokens, - "recentDays": [recent[day] for day in recent_dates], - "modelUsage": usage_by_model, - "totalPrompts": prompts, - "totalSessions": len(sessions), - "activeDays": len(active_days), - "activeDates": sorted(active_days), - } - write_json(cache_file, {"stats": stats}) - return stats + index_file = claude_code_index_path(claude_dir / "projects") + except Exception as exc: + warn(f"scan index unavailable ({exc})") + index_file = None + # The digest covers every data path the scan reads: the transcripts, the + # config home's fallbacks, the pi/omp session roots, and the opencode DB. + sources = [claude_dir / "projects", claude_dir, *pi_session_roots(), opencode_db()] + # Where Claude Code is sent decides whose its turns are, so a scan only + # answers for the endpoint it was made under. + return cached_scan(AGENT_ID, sources, max_age, lambda: scan_local(claude_dir, index_file, not use_index), match=claude_code_scan_match(claude_dir)) def merge_stats(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: @@ -801,11 +366,7 @@ def merge_stats(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: # nowhere but the Authorization header of the limits probe, and only the # plan label may travel into the printed record. def oauth_login(claude_dir: Path) -> tuple[str, int, int, str]: - try: - data = json.loads((claude_dir / ".credentials.json").read_text(encoding="utf-8")) - except Exception: - return "", 0, 0, "" - login = data.get("claudeAiOauth") + login = read_json(claude_dir / ".credentials.json").get("claudeAiOauth") if not isinstance(login, dict): return "", 0, 0, "" # The profile's tier is asked first, but only one that names a multiplier @@ -828,10 +389,7 @@ def oauth_login(claude_dir: Path) -> tuple[str, int, int, str]: # an account moved to Max 20x keeps reading Max 5x. The CLI's own profile is # refreshed against the account, so the tier it states there is asked first. def profile_tier(claude_dir: Path) -> str: - try: - account = json.loads(profile_path(claude_dir).read_text(encoding="utf-8")).get("oauthAccount") - except Exception: - return "" + account = read_json(profile_path(claude_dir)).get("oauthAccount") if not isinstance(account, dict): return "" # A seat with no tier of its own is limited by its organization, which is @@ -876,17 +434,14 @@ def parse_utilization(value: Any) -> float: return float("nan") -def normalize_utilization(value: Any, percent_scale: bool) -> float: +# Every utilization and percent this endpoint answers is a percentage: Claude +# Code's own /usage renders each as `${Math.floor(utilization)}% used`, so 0.6 +# is 0.6%, not 60%. +def normalize_utilization(value: Any) -> float: n = parse_utilization(value) if not (n >= 0): return -1.0 - # Anthropic's OAuth usage endpoint currently reports percentages (for - # example 37.0 or 1.0). Older payloads sometimes used fractions (0.37). - # A payload containing any value >= 1 is percent-scaled, so 1.0 renders - # as 1%, not 100%. - if percent_scale or n > 1: - return min(1.0, n / 100.0) - return min(1.0, n) + return min(1.0, n / 100.0) def normalize_reset_at(value: Any) -> str: @@ -940,7 +495,7 @@ def scoped_window(kind: str) -> str: # actually spending against. A model can hold more than one scoped window, and # only the pair of model and window tells them apart, so both make the title # and both make the key that keeps a repeat out. -def scoped_limits(payload: dict[str, Any], percent_scale: bool) -> list[dict[str, Any]]: +def scoped_limits(payload: dict[str, Any]) -> list[dict[str, Any]]: entries = payload.get("limits") if not isinstance(entries, list): return [] @@ -959,7 +514,7 @@ def scoped_limits(payload: dict[str, Any], percent_scale: bool) -> list[dict[str kind = str(entry.get("kind") or "").strip() if name == "" or (name, kind) in seen: continue - percent = normalize_utilization(entry.get("percent"), percent_scale) + percent = normalize_utilization(entry.get("percent")) if percent < 0: continue seen.add((name, kind)) @@ -974,6 +529,7 @@ def scoped_limits(payload: dict[str, Any], percent_scale: bool) -> list[dict[str return out + def probe_limits(access_token: str) -> dict[str, Any]: request = urllib.request.Request( USAGE_ENDPOINT, @@ -984,96 +540,74 @@ def probe_limits(access_token: str) -> dict[str, Any]: }, ) 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: - retry_after = error.headers.get("retry-after", "") if error.headers else "" - if error.code == 429: - help_text = "Anthropic's usage endpoint is rate limiting checks right now" + ( - f" (retry after {retry_after}s)" if retry_after else "" - ) + ". Local Claude Code stats are still shown." - else: - help_text = f"Anthropic's usage endpoint returned status {error.code}. Local Claude Code stats are still shown." - return {"ok": False, "helpText": help_text} - 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 Anthropic's usage endpoint. Retrying shortly. Local Claude Code stats are still shown.", - } - - weekly = usage_bucket(payload, "seven_day_oauth_apps") or usage_bucket(payload, "seven_day") + payload = open_json(request, timeout=10) + except HttpError as error: + if error.status: + retry_after = error.headers.get("retry-after", "") if error.headers else "" + if error.status == 429: + help_text = "Anthropic's usage endpoint is rate limiting checks right now" + ( + f" (retry after {retry_after}s)" if retry_after else "" + ) + ". Local Claude Code stats are still shown." + return {"ok": False, "rateLimited": True, "helpText": help_text} + if error.status in (401, 403): + return { + "ok": False, + "refused": True, + "helpText": "Anthropic refused Claude Code's sign-in. Start Claude Code, or run `claude auth login`, to sign in again.", + } + help_text = f"Anthropic's usage endpoint returned status {error.status}. Local Claude Code stats are still shown." + return {"ok": False, "helpText": help_text} + if error.transport: + # 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 Anthropic's usage endpoint. Retrying shortly. Local Claude Code stats are still shown.", + } + return {"ok": False, "helpText": "Anthropic's usage endpoint returned an invalid answer. Local Claude Code stats are still shown."} + if not isinstance(payload, dict): + payload = {} + + # The account's week is `seven_day`, the row Claude Code's own /usage titles + # "Current week (all models)". `seven_day_oauth_apps` is a side bucket it + # groups with cowork and omelette and never shows as the weekly limit. + weekly = usage_bucket(payload, "seven_day") session = usage_bucket(payload, "five_hour") - raw = [session.get("utilization") if session else None, weekly.get("utilization") if weekly else None] - # One payload speaks one convention, so the scoped entries settle the scale - # alongside the buckets rather than assuming their own. - entries = payload.get("limits") - if isinstance(entries, list): - raw += [entry.get("percent") for entry in entries if isinstance(entry, dict)] - percent_scale = any(parse_utilization(v) >= 1 for v in raw) limits = [] if session is not None: - percent = normalize_utilization(session.get("utilization"), percent_scale) + percent = normalize_utilization(session.get("utilization")) if percent >= 0: limits.append({"label": "Session (5-hour)", "percent": percent, "resetsAt": normalize_reset_at(session.get("resets_at"))}) if weekly is not None: - percent = normalize_utilization(weekly.get("utilization"), percent_scale) + percent = normalize_utilization(weekly.get("utilization")) if percent >= 0: limits.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": normalize_reset_at(weekly.get("resets_at"))}) - limits.extend(scoped_limits(payload, percent_scale)) + limits.extend(scoped_limits(payload)) if not limits: return {"ok": False, "helpText": "Anthropic's usage endpoint returned no limits. Local Claude Code stats are still shown."} return {"ok": True, "limits": limits} -# Cached windows reset to zero when their measured period ends. Without a -# fresh probe, the next rolling window has no known reset time yet. -def limit_window_open(entry: dict[str, Any], now: dt.datetime) -> bool: - raw = str(entry.get("resetsAt") or "") - if raw == "": - return True - try: - resets_at = dt.datetime.fromisoformat(raw.replace("Z", "+00:00")) - except Exception: - return True - if resets_at.tzinfo is None: - resets_at = resets_at.replace(tzinfo=dt.timezone.utc) - return resets_at > now - - -def usable_cached_limits(cached: dict[str, Any]) -> list[dict[str, Any]]: - entries = cached.get("limits") - if not isinstance(entries, list): - return [] - now = dt.datetime.now(dt.timezone.utc) - return [ - entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") - for entry in entries if isinstance(entry, dict) - ] - - def collect_limits(access_token: str, expires_at_ms: int, refresh_expires_at_ms: int, force: bool, cache_name: str = "claude-limits.json") -> dict[str, Any]: # `live` says the numbers came from Anthropic just now (or within the reuse # window) rather than from a cache kept past a failed or impossible probe. - result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "live": False} - - # When the numbers were measured, so a stale account can say how old - # they are. - result["fetchedAtMs"] = 0 + # `fetchedAtMs` says when they were measured, so a stale account can say + # how old they are. + result = {"limits": [], "usageStatusText": "", "authHelpText": AUTH_HELP, "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. Every account # keeps its own, so one account's allowance never stands in for another's. - probe_cache = cache_root() / cache_name - cached = read_json(probe_cache) or {} - fallback = usable_cached_limits(cached) - result["fetchedAtMs"] = number(cached.get("fetchedAtMs")) + # Elapsed windows come back at zero. A home that says who it's signed in as + # keeps claude-limits-.json, and one that doesn't the + # single-account claude-limits.json. + cache = LimitsCache.at(cache_root() / cache_name) + fallback, result["fetchedAtMs"] = cache.kept() # Probing needs a live token and only the Claude Code CLI can mint one: it # refreshes the credential file when it runs, so a machine left alone long @@ -1110,9 +644,7 @@ def collect_limits(access_token: str, expires_at_ms: int, refresh_expires_at_ms: # --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 - fetched_age = time.time() - fetched_at - if fallback and not force and 0 <= fetched_age < PROBE_MIN_INTERVAL_SECONDS: + if cache.fresh(force): result["limits"] = fallback result["live"] = True result["authHelpText"] = "" @@ -1123,8 +655,7 @@ def collect_limits(access_token: str, expires_at_ms: int, refresh_expires_at_ms: result["limits"] = probe["limits"] result["live"] = True result["authHelpText"] = "" - result["fetchedAtMs"] = round(time.time() * 1000) - write_json(probe_cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": probe["limits"]}) + result["fetchedAtMs"] = cache.save(probe["limits"]) return result # The first probe after login often fires before DHCP has handed out a @@ -1133,47 +664,25 @@ def collect_limits(access_token: str, expires_at_ms: int, refresh_expires_at_ms: result["retryAdvised"] = True if fallback: result["limits"] = fallback - fetched_age = time.time() - result["fetchedAtMs"] / 1000 - result["live"] = 0 <= fetched_age < CURRENT_ENOUGH_SECONDS + result["authHelpText"] = probe["helpText"] + # Anthropic rate-limits this endpoint readily, and a 429 only says "ask + # later": numbers measured minutes ago still describe the account. Any + # other failure leaves them stale, as the collector contract has it. + result["live"] = bool(probe.get("rateLimited")) and cache.younger_than(CURRENT_ENOUGH_SECONDS) else: result["usageStatusText"] = "Claude limits unavailable" result["authHelpText"] = probe["helpText"] + if probe.get("refused"): + result["usageStatusText"] = "Sign-in expired" return result # ---------------------------------------------------------------- accounts -# The accounts `omarchy agent account` registered, read straight from its -# registry. Only a registry holding a second account matters: with one, the -# record is exactly what it always was. -def registered_accounts() -> list[dict[str, Any]]: - state = Path(os.environ.get("XDG_STATE_HOME") or Path.home() / ".local" / "state") - try: - registry = json.loads((state / "omarchy" / "agents" / "accounts" / "claude.json").read_text(encoding="utf-8")) - except Exception: - return [] - accounts = [a for a in registry.get("accounts") or [] if isinstance(a, dict) and a.get("id")] - if len(accounts) < 2: - return [] - active_id = registry.get("active") or accounts[0]["id"] - if not any(a["id"] == active_id for a in accounts): - active_id = accounts[0]["id"] - for account in accounts: - account["active"] = account["id"] == active_id - account["switch"] = { - "mode": "auto" if registry.get("switch") == "auto" else "manual", - "threshold": registry.get("threshold") or 95, - } - return accounts - - def current_account_id(home: Path) -> str: - try: - account = json.loads(profile_path(home).read_text(encoding="utf-8")).get("oauthAccount") or {} - except Exception: - return "" - return str(account.get("accountUuid") or "") + account = read_json(profile_path(home)).get("oauthAccount") + return str(account.get("accountUuid") or "") if isinstance(account, dict) else "" def account_limits(account: dict[str, Any], force: bool) -> dict[str, Any]: @@ -1193,20 +702,7 @@ def account_limits(account: dict[str, Any], force: bool) -> dict[str, Any]: else: cache_name = f"claude-limits-{key}.json" limits = collect_limits(access_token, expires_at_ms, refresh_expires_at_ms, force, cache_name) - return { - "id": account["id"], - "label": str(account.get("label") or account["id"]), - "email": str(account.get("email") or ""), - "plan": plan or str(account.get("plan") or ""), - "active": account["active"], - "primary": bool(account.get("primary")), - "limits": limits["limits"], - "stale": not limits["live"], - "fetchedAt": limits["fetchedAtMs"], - "usageStatusText": limits["usageStatusText"], - "authHelpText": limits["authHelpText"], - "retryAdvised": bool(limits.get("retryAdvised")), - } + return account_entry(account, limits, plan) # -------------------------------------------------------------------- record @@ -1216,33 +712,14 @@ def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--force", action="store_true", help="rescan transcripts from scratch and re-probe limits, ignoring caches") parser.add_argument("--limits-only", action="store_true", help="reuse any recent transcript scan; only the limits probe must be fresh") - parser.add_argument("--cache-seconds", type=float, default=20) + parser.add_argument("--cache-seconds", type=float, default=SCAN_REUSE_SECONDS) args = parser.parse_args() claude_dir = config_dir() - scan_age = 0 if args.force else (900 if args.limits_only else args.cache_seconds) - stats = cached_scan(claude_dir / "projects", scan_age, use_index=not args.force) - - if number(stats.get("totalPrompts")) <= 0: - fallback = stats_cache_fallback(claude_dir) - if fallback is not None: - stats = fallback - else: - # No transcripts and no aggregate cache, but history.jsonl alone can - # still put numbers on today. - today_prompts, today_sessions = today_prompts_from_history(claude_dir) - if today_prompts or today_sessions: - stats = dict(stats, todayPrompts=today_prompts, todaySessions=today_sessions) - - pi_usage = scan_pi_usage(scan_age) - if pi_usage is not None: - stats = merge_stats(stats, pi_usage) - - opencode = scan_opencode_usage(scan_age) - if opencode is not None: - stats = merge_stats(stats, opencode) - - registered = registered_accounts() + scan_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else args.cache_seconds) + stats = local_stats(claude_dir, scan_age, use_index=not args.force) + + registered = registered_accounts(AGENT_ID) accounts = [account_limits(account, args.force) for account in registered] current = next((a for a in accounts if a["active"]), None) if current: @@ -1253,32 +730,49 @@ def main() -> int: stale, fetched_at = current["stale"], current["fetchedAt"] else: access_token, expires_at_ms, refresh_expires_at_ms, plan = oauth_login(claude_dir) - limits = collect_limits(access_token, expires_at_ms, refresh_expires_at_ms, args.force) - stale, fetched_at = not limits["live"], limits["fetchedAtMs"] - - record = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": dt.datetime.now(dt.timezone.utc).isoformat(), - "ready": number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0, - "hasLocalStats": True, - "tierLabel": plan, - "usageStatusText": limits["usageStatusText"], - "authHelpText": limits["authHelpText"], - "limits": limits["limits"], + # Keyed by who the home is signed in as, as a registered primary is, so + # signing in to another subscription never shows the last one's numbers + # and registering the home later keeps its cache. + identity = re.sub(r"[^A-Za-z0-9_-]", "", current_account_id(claude_dir)) + cache_name = f"claude-limits-{identity}.json" if identity else "claude-limits.json" + limits = collect_limits(access_token, expires_at_ms, refresh_expires_at_ms, args.force, cache_name) + stale, fetched_at = not limits["live"] and bool(limits["limits"]), limits["fetchedAtMs"] + # Asking for a sign-in only makes sense where Claude Code is in use: its + # config home or profile exists, or there is local usage. A machine that + # never ran it prints the unused record, which the panel skips. + in_use = ( + claude_dir.exists() + or profile_path(claude_dir).exists() + or number(stats.get("totalPrompts")) > 0 + or bool(stats.get("modelUsage")) + or number(stats.get("todayPrompts")) > 0 + or len(limits["limits"]) > 0 + ) + if not in_use: + limits = {"limits": [], "usageStatusText": "", "authHelpText": ""} + stale, fetched_at = False, 0 + + record = base_record( + AGENT_ID, + AGENT_NAME, + ready=number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")) or len(limits["limits"]) > 0, + hasLocalStats=True, + tierLabel=plan, + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + limits=limits["limits"], # Numbers kept past a failed probe, and when they were measured, so the # panel can tell them apart from ones checked just now. - "limitsStale": stale, - "limitsFetchedAt": fetched_at, - } + limitsStale=stale, + limitsFetchedAt=fetched_at, + ) if limits.get("retryAdvised"): record["retryAdvised"] = True if accounts: record["accountSwitch"] = registered[0]["switch"] record["accounts"] = accounts record.update(stats) - print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + emit(record, sort_keys=True) return 0 diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index 2b958f5fe10..a5c593cd721 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -11,70 +11,71 @@ panel only ever reads the JSON this prints. """ import argparse -import fcntl import hashlib import json +import math import os import select -import shutil -import sqlite3 import subprocess import sys import tempfile import time import tomllib -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone from pathlib import Path +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + LimitsCache, + Stats, + account_entry, + base_record, + cached_scan, + codex_provider_urls, + codex_rollout_cache, + codex_rollouts, + codex_route_owner, + codex_scan_match, + emit, + find_command, + jwt_claims, + now_ms, + number, + obj, + opencode_db, + pi_session_roots, + registered_accounts, + runtime_env, + scan_max_age, + scan_opencode, + scan_pi_sessions, + warn, +) + AGENT_ID = "codex" AGENT_NAME = "Codex" AUTH_HELP = "Run `codex login` to authenticate." +API_KEY_HELP = "Signed in with an API key, which has no plan limits." # session_meta.model_provider for the built-in OpenAI backend, which is the # only one this subscription pays for. Any other id points somewhere else: # `--oss`'s ollama or lmstudio, or a custom provider in config.toml. NATIVE_PROVIDER = "openai" -# A scan this recent is only reused to dedup concurrent collector runs (the -# update command backgrounds one per agent while the panel refreshes on its -# own); every periodic widget refresh lands a real rescan, however low -# refreshIntervalSec is set. --limits-only promises only fresh limits, so it -# may reuse a scan for up to 15 minutes. -SCAN_REUSE_SECONDS = 20 -LIMITS_ONLY_REUSE_SECONDS = 900 - -# Per-file totals of native Codex sessions survive between runs, so a rescan -# only reads the session files that changed. Bumping the schema retires every -# record written by an older layout. -FILE_CACHE_SCHEMA = 1 - - -def local_day(value): - if value is None: - return datetime.now().strftime("%Y-%m-%d") - if isinstance(value, (int, float)): - # pi message timestamps are milliseconds; Codex timestamps are usually seconds. - if value > 10_000_000_000: - value = value / 1000 - return datetime.fromtimestamp(value).strftime("%Y-%m-%d") - text = str(value) - try: - if text.endswith("Z"): - dt = datetime.fromisoformat(text[:-1] + "+00:00") - else: - dt = datetime.fromisoformat(text) - if dt.tzinfo is not None: - dt = dt.astimezone() - return dt.strftime("%Y-%m-%d") - except Exception: - return datetime.now().strftime("%Y-%m-%d") - +# What a dead app-server's stderr says when it never reached OpenAI at all, +# rather than being turned away by it. +TRANSPORT_ERRORS = ( + "error sending request", + "dns error", + "failed to lookup address", + "connection refused", + "connection reset", + "network is unreachable", + "timed out", +) -def number(value): - try: - return int(value or 0) - except Exception: - return 0 def model_name(raw): @@ -82,26 +83,11 @@ def model_name(raw): return value if value else "codex" -def runtime_env(): - home = str(Path.home()) - path_parts = [ - os.environ.get("PATH", ""), - f"{home}/.local/bin", - f"{home}/.npm-global/bin", - f"{home}/.local/share/mise/shims", - ] - env = os.environ.copy() - env["PATH"] = os.pathsep.join(part for part in path_parts if part) - return env ENV = runtime_env() -def find_command(name): - return shutil.which(name, path=ENV.get("PATH")) - - def mise_shims_dir(): data_home = ENV.get("MISE_DATA_DIR") or os.path.join( ENV.get("XDG_DATA_HOME") or os.path.join(str(Path.home()), ".local", "share"), "mise") @@ -148,604 +134,110 @@ def find_codex_binary(): return None -now = datetime.now() -today = now.strftime("%Y-%m-%d") -recent_dates = [(now - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1)] -recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} -today_tokens_by_model = {} -model_usage = {} -today_sessions = set() -active_days = set() - -today_prompts = 0 -today_total_tokens = 0 -total_prompts = 0 -total_sessions = set() -seen_pi_messages = set() - - -# One session file's totals, the unit the file cache stores: day -> model -> -# [input, output, cacheRead, cacheWrite, prompts]. Day-keyed and with no -# "today" in it, so a record stays true however long it is kept. -def new_file_record(stat_result): - return {"mtime": stat_result.st_mtime, "size": stat_result.st_size, "days": {}} - - -def record_usage(record, day, model, input_tokens, output_tokens, cache_read, cache_write): - totals = record["days"].setdefault(day, {}).setdefault(model, [0, 0, 0, 0, 0]) - totals[0] += input_tokens - totals[1] += output_tokens - totals[2] += cache_read - totals[3] += cache_write - totals[4] += 1 - - -def merge_file_record(session_key, record): - """Fold one file's totals, freshly read or replayed from cache, into the run.""" - global today_prompts, today_total_tokens, total_prompts - for day, models in record["days"].items(): - for model, totals in models.items(): - input_tokens, output_tokens, cache_read, cache_write, prompts = totals - total = input_tokens + output_tokens + cache_read + cache_write - total_prompts += prompts - total_sessions.add(session_key) - active_days.add(day) - - bucket = model_usage.setdefault(model, { - "inputTokens": 0, - "outputTokens": 0, - "cacheReadInputTokens": 0, - "cacheCreationInputTokens": 0, - }) - bucket["inputTokens"] += input_tokens - bucket["outputTokens"] += output_tokens - bucket["cacheReadInputTokens"] += cache_read - bucket["cacheCreationInputTokens"] += cache_write - - if day in recent: - recent[day]["messageCount"] += total - - if day == today: - today_prompts += prompts - today_sessions.add(session_key) - today_total_tokens += total - today_tokens_by_model[model] = today_tokens_by_model.get(model, 0) + total - - -def add_usage(day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): - record = {"days": {}} - record_usage(record, day, model, input_tokens, output_tokens, cache_read, cache_write) - merge_file_record(session_key, record) - - -def pi_parent_session(path): - try: - with open(path, encoding="utf-8", errors="replace") as session_file: - header = json.loads(session_file.readline()) - except (OSError, json.JSONDecodeError): - return None - if not isinstance(header, dict): - return None - parent = header.get("parentSession") - if not isinstance(parent, str) or not parent: - return None - return os.path.abspath(os.path.expanduser(parent)) - - -# `omp --profile=` (and pi's equivalent) relocates the whole agent tree -# under /profiles//, so a subscription driven entirely through a -# profile leaves the default root empty and its usage uncounted. Each root is -# scanned separately and sessions are keyed by path, so a profile adds -# sessions rather than double-counting the default one. -def pi_session_roots(): - roots = [] - for base in (Path.home() / ".pi", Path.home() / ".omp"): - roots.append(base / "agent" / "sessions") - profiles = base / "profiles" - try: - # Sorted so the scan order does not depend on directory order. - roots.extend(sorted(child / "agent" / "sessions" for child in profiles.iterdir() if child.is_dir())) - except OSError: - # No profiles directory, or it is unreadable: the default root stands. - pass - return roots +def codex_home(): + return Path(ENV.get("CODEX_HOME") or (Path.home() / ".codex")) -def scan_pi_sessions(): - roots = pi_session_roots() - rg = find_command("rg") or "rg" - session_keys = {} - today_session_keys = {} - for root in roots: - if not root.exists(): - continue - try: - proc = subprocess.Popen( - # --no-ignore: these are data files, not a source tree. When $HOME is - # itself a git checkout (a common dotfiles setup), the parent repo's - # ignore rules would otherwise silently exclude every session file and - # the scan would count zero usage. - [rg, "--json", "--no-ignore", "-e", r'"provider"\s*:\s*"openai-codex"', "-e", r'"api"\s*:\s*"openai-codex', str(root)], - stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, - text=True, - errors="replace", - env=ENV, - ) - except FileNotFoundError: - return - - assert proc.stdout is not None - for raw in proc.stdout: - try: - event = json.loads(raw) - if event.get("type") != "match": - continue - line = event.get("data", {}).get("lines", {}).get("text", "") - raw_path = event.get("data", {}).get("path", {}).get("text", "pi-session") - path = os.path.abspath(os.path.expanduser(raw_path)) - entry = json.loads(line) - except Exception: - continue +# Every source folds into one run's totals. +STATS = Stats() - if entry.get("type") != "message": - continue - message_id = str(entry.get("id") or "") - message_timestamp = str(entry.get("timestamp") or "") - # Forks retain both fields; IDs alone can collide across sessions. - if message_id and message_timestamp: - message_key = ("entry", message_id, message_timestamp) - else: - message_key = ("path", path, message_id) - duplicate = message_key in seen_pi_messages - seen_pi_messages.add(message_key) - message = entry.get("message") or {} - if message.get("role") != "assistant": - continue - provider = str(message.get("provider") or "") - api = str(message.get("api") or "") - if provider != "openai-codex" and not api.startswith("openai-codex"): - continue - usage = message.get("usage") or {} - if not usage: - continue - total = number(usage.get("totalTokens")) - input_tokens = number(usage.get("input")) - output_tokens = number(usage.get("output")) - cache_read = number(usage.get("cacheRead")) - cache_write = number(usage.get("cacheWrite")) - if total and not (input_tokens or output_tokens or cache_read or cache_write): - input_tokens = total - if not (input_tokens or output_tokens or cache_read or cache_write): - continue +# pi and omp sessions that ran through openai-codex, by provider or by the +# Responses API pi drives for it. Their compatible JSONL transcripts carry the +# provider, model, and token usage on every assistant message. +PI_PATTERN = r'"provider"\s*:\s*"openai-codex"|"api"\s*:\s*"openai-codex' - day = local_day(entry.get("timestamp") or message.get("timestamp")) - session_keys.setdefault(path, set()).add(message_key) - if day == today: - today_session_keys.setdefault(path, set()).add(message_key) - if duplicate: - continue - add_usage(day, path, model_name(message.get("model")), input_tokens, output_tokens, cache_read, cache_write) - try: - proc.wait(timeout=1) - except Exception: - proc.kill() - - # Resolve session ownership after scanning so rg's file order cannot affect it. - parents = {path: pi_parent_session(path) for path in session_keys} - total_sessions.difference_update(session_keys) - today_sessions.difference_update(session_keys) - for path, keys in session_keys.items(): - inherited = set() - visited = {path} - parent = parents.get(path) - while parent and parent not in visited: - visited.add(parent) - inherited.update(session_keys.get(parent, ())) - parent = parents.get(parent) - - unique_keys = keys - inherited - if unique_keys: - total_sessions.add(path) - if unique_keys & today_session_keys.get(path, set()): - today_sessions.add(path) +def pi_billed_here(provider, model, entry): + return provider == "openai-codex" or str(entry.get("api") or "").startswith("openai-codex") -def scan_opencode_sessions(): - # A subscription burned entirely through opencode leaves no native session - # files, but opencode records per-message provider, model, and token usage - # in its own database. Read-only: opencode may be writing right now. - # Returns whether the scan ran to completion: a scan cut short by a - # database error still contributes what it read, but must not be cached - # as if it were the whole story. - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" - if not db.is_file(): - return True - try: - conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) - except sqlite3.Error: - return False - try: - conn.execute("PRAGMA query_only = ON") - # OpenCode DBs grow huge (every historical message, JSON included), and - # Python-side json.loads of every row wasted ~600 MB of RSS on machines - # whose subscription never ran on OpenAI. The json_extract conditions are - # the authority for the rows that reach them: role == "assistant" and - # providerID == "openai", the same exact-match values the Python filter - # below checks. The guards in front are pure acceleration, not a perfect - # proxy for the Python filter: - # - The LIKE gates skip rows whose JSON cannot contain the two - # key/value pairs, avoiding the JSON parse of tens-of-MB blobs. They - # can differ from json.loads on duplicate keys (Python keeps the - # last, SQLite json_extract keeps the first) and ASCII-escaped - # values (\"ass\\u0069stant\" decodes for Python but not for LIKE), - # so a row the authority would accept can be gated out. Both cases - # are vanishingly rare in real opencode data. - # - json_valid guards the parse itself: json_extract() RAISES on - # malformed JSON instead of returning NULL, and one such row would - # otherwise abort the whole scan. SQLite does not promise that AND - # terms evaluate left to right, so the guard is a CASE around each - # json_extract rather than a separate AND term. Rows that are not - # well-formed JSON are skipped here; the per-row try/except below - # stays as the final safety net for rows that pass the SQL filter - # but fail json.loads. - for session_id, raw in conn.execute( - "SELECT session_id, data FROM message" - " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" - " AND data LIKE '%\"providerID\"%:%\"openai\"%'" - " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" - " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END = 'openai'" - ): - # One malformed row must not abort the scan, so every shape assumption - # lives inside the try. - try: - entry = json.loads(raw) - # Exact match: opencode provider ids are free-form, and a custom - # "openai-local" gateway is not this subscription. - if not isinstance(entry, dict) or entry.get("role") != "assistant": - continue - if str(entry.get("providerID") or "") != "openai": - continue - tokens = entry.get("tokens") or {} - cache = tokens.get("cache") or {} - input_tokens = number(tokens.get("input")) - # opencode keeps thinking tokens out of output; both are generated. - output_tokens = number(tokens.get("output")) + number(tokens.get("reasoning")) - cache_read = number(cache.get("read")) - cache_write = number(cache.get("write")) - if not (input_tokens or output_tokens or cache_read or cache_write): - continue - day = local_day((entry.get("time") or {}).get("created")) - model = model_name(str(entry.get("modelID") or "").rstrip("/").split("/")[-1]) - except Exception: - continue - add_usage(day, "opencode:" + str(session_id), model, input_tokens, output_tokens, cache_read, cache_write) - except sqlite3.Error: - # Transient lock, schema migration, corruption: the numbers stop here, - # incomplete. - return False - finally: - conn.close() - return True +# A subscription burned entirely through opencode leaves no native session +# files, but opencode records per-message provider, model, and token usage in +# its own database. Exact match: opencode provider ids are free-form, and a +# custom "openai-local" gateway is not this subscription. +def opencode_billed_here(provider, model, entry): + return provider == "openai" -def scan_native_codex_sessions(cached_files, scanned_files): - codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) - roots = [codex_home / "sessions", codex_home / "archived_sessions"] - files = [] - cutoff = time.time() - 30 * 24 * 60 * 60 - for root in roots: - if not root.exists(): - continue - for path in root.rglob("*.jsonl"): - try: - stat_result = path.stat() - except OSError: - continue - if stat_result.st_mtime >= cutoff: - files.append((path, stat_result)) - - for path, stat_result in files: - key = str(path) - record = reusable_record(cached_files, key, stat_result) - if record is None: - record = read_native_codex_session(path, stat_result) - if record["mtime"] is not None: - scanned_files[key] = record - merge_file_record(key, record) - - -def read_native_codex_session(path, stat_result): - record = new_file_record(stat_result) - current_model = "codex" - seen_meta = False - previous_total_usage = None - try: - with path.open(errors="replace") as handle: - for raw in handle: - # Cheap pre-filter before JSON parsing keeps files with unrelated - # lines inexpensive. session_meta still has to get through: it says - # which provider served the rollout. - if '"token_count"' not in raw and '"turn_context"' not in raw and '"session_meta"' not in raw: - continue - try: - entry = json.loads(raw) - except Exception: - continue - if entry.get("type") == "session_meta": - # A forked rollout copies its parent's session_meta after its own, - # so only the first one, as Codex reads it, says who served this. - if seen_meta: - continue - seen_meta = True - # Codex CLI happily fronts any OpenAI-compatible backend (`--oss`, - # or a custom model_provider in config.toml). Those turns bill the - # local box or a third party, never this subscription, so drop the - # whole rollout the way the pi and opencode scans drop foreign - # providers. - meta = entry.get("payload") or {} - provider = str(meta.get("model_provider") or "") - # Rollouts predating the field carry no provider at all; count - # those rather than silently lose the history. - if provider and provider != NATIVE_PROVIDER: - break - continue - if entry.get("type") == "turn_context": - payload = entry.get("payload") or {} - current_model = model_name(payload.get("model") or payload.get("model_slug") or current_model) - continue - payload = entry.get("payload") or entry - if entry.get("type") == "response_item" and isinstance(payload, dict): - payload = payload.get("payload") or payload - if not isinstance(payload, dict): - continue - if payload.get("type") != "token_count": - continue - info = payload.get("info") or {} - # total_token_usage is cumulative for the session. Adding every - # snapshot makes usage grow quadratically, so count the last turn. - usage = info.get("last_token_usage") or {} - cache_read = number(usage.get("cached_input_tokens")) - cache_write = number(usage.get("cache_write_input_tokens")) - # Cached tokens are included in input_tokens, and reasoning tokens - # are included in output_tokens. Keep the cache split without - # counting either category twice. - input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write) - output_tokens = number(usage.get("output_tokens")) - if not (input_tokens or output_tokens or cache_read or cache_write): - continue - # Quota updates can repeat the previous request's last_token_usage. - # Only an unchanged cumulative snapshot proves this is a repeat: - # separate requests may have identical last-token counts, and older - # records may not carry cumulative counters at all. Keep this state - # per rollout and accept changed counters, including counter resets. - total_usage = info.get("total_token_usage") - if isinstance(total_usage, dict) and total_usage: - if total_usage == previous_total_usage: - continue - previous_total_usage = total_usage - day = local_day(entry.get("timestamp") or stat_result.st_mtime) - record_usage(record, day, current_model, input_tokens, output_tokens, cache_read, cache_write) - except Exception: - # Keep what was read before the failure, but caching a partial record would - # hide the rest until the file changes, so hand back one the cache refuses. - record["mtime"] = None - return record - - -def cache_root(): - root = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" - root.mkdir(parents=True, exist_ok=True) - return root - - -def scan_cache_digest(): - codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) - db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" - # The digest covers every data path the scan reads: the codex session - # roots, the opencode DB, and (via Path.home()) the pi/omp session roots. - return hashlib.sha1((str(Path.home()) + "\n" + str(codex_home) + "\n" + str(db)).encode("utf-8")).hexdigest()[:16] - - -def scan_cache_paths(): - digest = scan_cache_digest() - root = cache_root() - return root / f"codex-scan-{digest}.json", root / f"codex-scan-{digest}.lock" - - -# Records hold local days, so they are only true in the timezone that wrote them. -def file_cache_zone(): - return [os.environ.get("TZ"), list(time.tzname), time.timezone, time.altzone] - +def scan_pi_sessions_into_stats(): + # Whether every root was read in full, so a partial scan is never cached + # as the whole story. + rows, complete = scan_pi_sessions(PI_PATTERN, pi_billed_here) + for row in rows: + STATS.add_row(row, model_name(row.model)) + return complete -def file_cache_path(): - return cache_root() / f"codex-files-{scan_cache_digest()}.json" - - -def valid_file_record(value): - try: - counts = [totals for models in value["days"].values() for totals in models.values()] - return ( - type(value["mtime"]) in (int, float) and type(value["size"]) is int - and all(len(totals) == 5 and all(type(n) is int for n in totals) for totals in counts) - ) - except Exception: - return False - - -def read_file_cache(): - """Per-file totals of native sessions from earlier scans, keyed by path. - - Session files are append-only, so a file whose mtime and size both match the - record cannot have grown a turn since: its totals can be replayed instead of - read. That is what keeps a refresh off the gigabytes of history it already - counted. Unlike the aggregate scan cache these records carry no "today" and - never expire on their own -- only the file they describe or a change of - timezone can invalidate one, so midnight and the clock moving leave them true. - - pi and omp sessions stay out of it: a forked pi session repeats its parent's - messages, so they are only deduplicated across the whole tree at once. - """ - try: - payload = json.loads(file_cache_path().read_text(encoding="utf-8")) - except Exception: - return {} - if not isinstance(payload, dict) or payload.get("schemaVersion") != FILE_CACHE_SCHEMA: - return {} - if payload.get("zone") != file_cache_zone(): - return {} - files = payload.get("files") - if not isinstance(files, dict): - return {} - return {path: record for path, record in files.items() if valid_file_record(record)} - - -def write_file_cache(files): - # Only the files this scan saw are written back, so deleted sessions and - # sessions that aged out of the 30-day window leave the cache instead of - # growing it forever. - try: - write_json(file_cache_path(), {"schemaVersion": FILE_CACHE_SCHEMA, "zone": file_cache_zone(), "files": files}) - except Exception as exc: - print(f"omarchy-agent-usage-codex: could not write file cache ({exc})", file=sys.stderr) - - -def reusable_record(cached_files, key, stat_result): - record = cached_files.get(key) - if record is None: - return None - if record["mtime"] != stat_result.st_mtime or record["size"] != stat_result.st_size: - return None - return record - - -def read_fresh_json(path, max_age_seconds): - if max_age_seconds <= 0 or not path.exists(): - return None - try: - # A negative age means the mtime is in the future: the clock moved - # backwards since the write, so the cache's freshness cannot be trusted. - age = time.time() - path.stat().st_mtime - if 0 <= age <= max_age_seconds: - return json.loads(path.read_text(encoding="utf-8")) - except Exception: - return None - return None +def scan_opencode_sessions(): + # The closing quote makes the database's prefilter match the whole id. + rows, complete = scan_opencode(['openai"'], opencode_billed_here) + for row in rows: + STATS.add_row(row, model_name(str(row.model or "").rstrip("/").split("/")[-1])) + return complete + + + +def scan_native_codex_sessions(reuse): + """Fold every native rollout into the run, as the shared reader reads it, + and say whether each was read in full. A rollout that only grew is read on + from where the last read stopped, by whichever collector got to it first. + + Codex CLI happily fronts any OpenAI-compatible backend (`--oss`, or a + custom model_provider in config.toml). Those turns bill the local box or a + third party, never this subscription, so only rollouts on the built-in + provider count, and ones from before Codex recorded a provider at all. Where + openai_base_url sends the built-in provider decides which of their models + are this subscription's; the models codex_route_owner() gives another + collector are left to it, each with the prompts it answered.""" + urls = codex_provider_urls() + records, complete = codex_rollouts(codex_rollout_cache(), reuse=reuse, native=True) + owners = {} + for record in records.values(): + if record["provider"] not in ("", NATIVE_PROVIDER): + continue + for day, models in record["days"].items(): + for model, (input_tokens, output_tokens, cache_read, cache_write, prompts) in models.items(): + if model not in owners: + owners[model] = codex_route_owner(NATIVE_PROVIDER, model, urls) == AGENT_ID + if not owners[model]: + continue + # One prompt can take many requests, one per tool call it makes, so + # the requests count no prompts of their own. + if input_tokens or output_tokens or cache_read or cache_write: + STATS.add(day, record["session"], model, input_tokens, output_tokens, cache_read, cache_write, prompts=0) + if prompts: + STATS.add_prompts(day, record["session"], prompts) + return complete -def write_json(path, payload): - # A temp name unique to this writer, not derived from the target: several - # collectors can run at once (the update command backgrounds one per agent, - # the panel refreshes on its own), and a shared temp path means the second - # replace finds the first one's file already moved away. - handle_fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp") - tmp = Path(tmp_name) - try: - with os.fdopen(handle_fd, "w", encoding="utf-8") as handle: - handle.write(json.dumps(payload, separators=(",", ":")) + "\n") - # mkstemp opens at 0600; nothing in the cache is sensitive, so open it - # up to the usual 0644. - tmp.chmod(0o644) - tmp.replace(path) - except BaseException: - tmp.unlink(missing_ok=True) - raise - - -# The cache payload is a versioned envelope around the local-stats dict, so a -# corrupted or foreign-shaped file is a cache miss (rescan + rewrite) instead -# of a crash or a garbage record. -# Version 2 invalidates totals counted before native notification deduplication. -def read_cached_stats(cache_file, max_age_seconds): - cached = read_fresh_json(cache_file, max_age_seconds) - if not isinstance(cached, dict) or cached.get("schemaVersion") != 2: - return None - # today* fields only mean "today" on the day they were scanned. A cache - # from another local date (midnight passed, or the clock moved) is a miss, - # not merely old, whatever its mtime says. - if cached.get("scanDate") != today: - return None - stats = cached.get("stats") - if not isinstance(stats, dict): - return None - if not all(key in stats for key in ("todayPrompts", "todayTotalTokens", "recentDays", "activeDates", "modelUsage")): - return None - return stats -def write_cached_stats(cache_file, stats): - try: - write_json(cache_file, {"schemaVersion": 2, "scanDate": today, "stats": stats}) - except Exception as exc: - print(f"omarchy-agent-usage-codex: could not write usage cache ({exc})", file=sys.stderr) +# Every data path the scan reads: the codex session roots, the opencode DB, +# and the pi/omp session roots. +def scan_sources(): + return [Path.home(), codex_home(), opencode_db(), *pi_session_roots()] def local_stats(): """Snapshot the aggregated local usage into the record's stats dict.""" - return { - "todayPrompts": today_prompts, - "todaySessions": len(today_sessions), - "todayTotalTokens": today_total_tokens, - "todayTokensByModel": today_tokens_by_model, - "recentDays": [recent[day] for day in recent_dates], - "totalPrompts": total_prompts, - "totalSessions": len(total_sessions), - # Days with any recorded usage, for the all-time "N days" summary. The - # dates travel too: merging snapshots from several machines needs their - # union, which a count alone cannot give. - "activeDays": len(active_days), - "activeDates": sorted(active_days), - "modelUsage": model_usage, - } + return STATS.snapshot() # A forced scan (no reuse at all) reads every session file again. def run_local_scans(max_age): - cached_files = read_file_cache() if max_age > 0 else {} - scanned_files = {} - scan_pi_sessions() - scan_native_codex_sessions(cached_files, scanned_files) - complete = scan_opencode_sessions() - write_file_cache(scanned_files) + global STATS + STATS = Stats() + complete = scan_pi_sessions_into_stats() + complete = scan_native_codex_sessions(max_age > 0) and complete + complete = scan_opencode_sessions() and complete return local_stats(), complete def cached_local_stats(max_age): - """Local stats, with the cache as a pure optimization. - - The cache must never take the collector down: any cache-layer failure - (unwritable cache root, lock errors, disk full) degrades to a direct scan - and a warning on stderr. The JSON record is the contract; the cache is not. - """ - try: - return _cached_local_stats(max_age) - except Exception as exc: - print(f"omarchy-agent-usage-codex: cache unavailable ({exc}); scanning directly", file=sys.stderr) - stats, _ = run_local_scans(max_age) - return stats - + # The route config can change who bills a rollout without it changing. + return cached_scan(AGENT_ID, scan_sources(), max_age, lambda: run_local_scans(max_age), match=codex_scan_match()) -def _cached_local_stats(max_age): - cache_file, lock_file = scan_cache_paths() - - cached = read_cached_stats(cache_file, max_age) - if cached is not None: - return cached - - with lock_file.open("w") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - cached = read_cached_stats(cache_file, max_age) - if cached is not None: - return cached - stats, complete = run_local_scans(max_age) - # An interrupted scan still serves this run, but caching it would - # suppress the missing usage for every reader until the cache expires. - if complete: - write_cached_stats(cache_file, stats) - return stats def rpc_send(proc, payload, operation): @@ -802,8 +294,12 @@ def rpc_request(proc, request_id, method, params=None, timeout=8): def limit_window(window): if not isinstance(window, dict): return None - used = window.get("usedPercent") - if used is None: + try: + used = float(window.get("usedPercent")) + except (TypeError, ValueError): + return None + # A percentage that isn't one (NaN, infinite, a boolean) says nothing. + if isinstance(window.get("usedPercent"), bool) or not math.isfinite(used): return None mins = number(window.get("windowDurationMins")) if mins == 10080: @@ -814,34 +310,145 @@ def limit_window(window): label = f"{mins}m window" else: label = "Limit" - reset = window.get("resetsAt") + reset = number(window.get("resetsAt")) + try: + resets_at = datetime.fromtimestamp(reset, timezone.utc).isoformat() if reset > 0 else "" + except (OverflowError, OSError, ValueError): + resets_at = "" return { "label": label, - "percent": float(used) / 100.0, - "resetsAt": datetime.fromtimestamp(number(reset), timezone.utc).isoformat() if reset else "", + "percent": max(0.0, min(1.0, used / 100.0)), + "resetsAt": resets_at, + "minutes": mins, } +# The panel reads a window from its title, and a model's own allowance titled +# " Session" or " Weekly" sits beside the general window of the +# same span rather than as a row of its own. +def window_title(minutes): + if minutes >= 28 * 24 * 60: + return "Monthly" + if minutes >= 24 * 60: + return "Weekly" + if minutes > 0: + return "Session" + return "" + + +def rate_limit_buckets(result): + # New app-servers return every metered bucket; rateLimits is only the + # backward-compatible single-bucket view. Never count both views. + buckets = result.get("rateLimitsByLimitId") + if isinstance(buckets, dict): + entries = [(str(key), value) for key, value in buckets.items() if isinstance(value, dict)] + if entries: + # Keep the subscription's general limits and plan ahead of model limits. + return sorted(entries, key=lambda entry: entry[0] != "codex") + legacy = result.get("rateLimits") + return [("codex", legacy)] if isinstance(legacy, dict) else [] + + +def bucket_limits(buckets): + limits = [] + for limit_id, bucket in buckets: + # The general bucket's windows are the plain ones. Any other bucket is a + # model's own allowance, named the way the app-server names it. + general = limit_id == "codex" or (len(buckets) == 1 and not bucket.get("limitName")) + name = str(bucket.get("limitName") or limit_id).strip() + for window in (bucket.get("primary"), bucket.get("secondary")): + entry = limit_window(window) + if not entry: + continue + minutes = entry.pop("minutes") + if not general: + span = window_title(minutes) + entry["title"] = f"{name} {span}" if span else name + limits.append(entry) + return limits + + +def finite_number(value): + # A boolean, NaN or infinity is no amount. + if isinstance(value, bool): + return None + try: + amount = float(value) + except (TypeError, ValueError): + return None + return amount if math.isfinite(amount) else None + + # A spend cap can be returned without any rate windows, as observed on a -# Business seat. Its snapshot gives a reset time but no period, so don't -# infer a frequency from the plan or date. Use the amounts for the share; -# remainingPercent only has whole-percent precision. +# Business seat. Its snapshot gives a reset time but no period, so no +# frequency is inferred from the plan or date, and the meter is titled for +# what it is. The amounts give the share; remainingPercent only has +# whole-percent precision, so it stands in only for a cap of zero. def spend_limit(snapshot): if not isinstance(snapshot, dict): return None + cap = finite_number(snapshot.get("limit")) + used = finite_number(snapshot.get("used")) + if cap is not None and cap > 0 and used is not None: + share = used / cap + else: + remaining = finite_number(snapshot.get("remainingPercent")) + if remaining is None: + return None + share = 1 - remaining / 100 + reset = number(snapshot.get("resetsAt")) try: - cap = float(snapshot.get("limit")) - used = float(snapshot.get("used")) / cap if cap > 0 else 1 - float(snapshot.get("remainingPercent")) / 100 - except (TypeError, ValueError): - return None - reset = snapshot.get("resetsAt") + resets_at = datetime.fromtimestamp(reset, timezone.utc).isoformat() if reset > 0 else "" + except (OverflowError, OSError, ValueError): + resets_at = "" return { "label": "Spend limit", - "percent": max(0.0, min(1.0, used)), - "resetsAt": datetime.fromtimestamp(number(reset), timezone.utc).isoformat() if reset else "", + "percent": max(0.0, min(1.0, share)), + "resetsAt": resets_at, + "title": "Spend limit", } +# The cap is the account's, so it is read once: from the general bucket ahead +# of a model's, and from the backward-compatible view when no bucket has it. +def individual_spend_limit(buckets, rpc_result): + legacy = rpc_result.get("rateLimits") + for bucket in [bucket for _, bucket in buckets] + [legacy if isinstance(legacy, dict) else {}]: + if isinstance(bucket.get("individualLimit"), dict): + return spend_limit(bucket["individualLimit"]) + return None + + +# The plans as Codex's own /status names them (SubscriptionDisplay::Status in +# Codex 0.163's TUI), keyed by the wire planType: "pro" is the 200 plan, and +# a "team" workspace is what Codex sells as Business. +PLAN_NAMES = { + "free": "Free", + "go": "Go", + "plus": "Plus", + "pro": "Pro 200", + "prolite": "Pro 100", + "promax": "Pro 500", + "team": "Business", + "self_serve_business_usage_based": "Business", + "self_serve_business_prolite": "Business Premium", + "business": "Enterprise", + "enterprise_cbp_automation": "Enterprise (Automation)", + "enterprise_cbp_usage_based": "Enterprise", + "ent26": "Enterprise", + "enterprise": "Enterprise", + "edu": "Edu", + "edu_plus": "Edu Plus", + "edu_pro": "Edu Pro", +} + + +def plan_name(plan): + # A plan newer than this table keeps its wire name, made readable. + plan = str(plan or "") + return PLAN_NAMES.get(plan, plan.replace("_", " ")) + + def codex_rpc_help(proc, stderr_file, exc): """Explain an app-server RPC failure: a live server's stall, a dead one's own stderr, or the login hint when it died saying nothing.""" @@ -869,16 +476,20 @@ def codex_rpc_help(proc, stderr_file, exc): return AUTH_HELP +# Credits are an extra: a malformed field loses them, never the limits. def reset_credits(result): - granted = (result.get("rateLimitResetCredits") or {}).get("credits") or [] + granted = obj(result.get("rateLimitResetCredits")).get("credits") + if not isinstance(granted, list): + return None available = [c for c in granted if isinstance(c, dict) and c.get("status") == "available"] if not available: return None expiries = [number(c.get("expiresAt")) for c in available if number(c.get("expiresAt")) > 0] - return { - "available": len(available), - "nextExpiresAt": datetime.fromtimestamp(min(expiries), timezone.utc).isoformat() if expiries else "", - } + try: + next_expiry = datetime.fromtimestamp(min(expiries), timezone.utc).isoformat() if expiries else "" + except (OverflowError, OSError, ValueError): + next_expiry = "" + return {"available": len(available), "nextExpiresAt": next_expiry} # Keyring and auto stores keep credentials outside auth.json, so only the file @@ -899,12 +510,49 @@ def has_codex_credentials(codex_home, env): return (codex_home / "auth.json").is_file() -def fetch_codex_rpc(home=None): +# An API key in auth.json signs Codex in to the API, which bills by the token +# and has no plan limits: account/rateLimits/read only says it wants ChatGPT. +# Codex reads auth_mode first, else any API key, ChatGPT tokens beside it or +# not, as an older login wrote them both. +def api_key_login(codex_home, env): + if env.get("CODEX_ACCESS_TOKEN"): + return False + try: + with (codex_home / "config.toml").open("rb") as handle: + config = tomllib.load(handle) + except FileNotFoundError: + config = {} + except Exception: + return False + if str(config.get("cli_auth_credentials_store") or "file") != "file": + return False + try: + auth = json.loads((codex_home / "auth.json").read_text(encoding="utf-8")) + except Exception: + return False + if not isinstance(auth, dict): + return False + if auth.get("auth_mode") is not None: + return auth["auth_mode"] == "apikey" + return bool(auth.get("OPENAI_API_KEY")) + + +# The environment a home's app-server runs in. CODEX_ACCESS_TOKEN signs in +# the primary home alone: handed to every registered home, it would make each +# of them answer for the same account. +def home_env(home=None, primary=True): + env = dict(ENV, CODEX_HOME=str(home)) if home else dict(ENV) + if not primary: + env.pop("CODEX_ACCESS_TOKEN", None) + return env + + +def fetch_codex_rpc(home=None, primary=True): result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": AUTH_HELP} - env = dict(ENV, CODEX_HOME=str(home)) if home else ENV + env = home_env(home, primary) codex = find_codex_binary() if not codex: - result["usageStatusText"] = "Codex unavailable" + result["usageStatusText"] = "Codex limits unavailable" result["authHelpText"] = "codex not found in PATH" return result @@ -918,6 +566,11 @@ def fetch_codex_rpc(home=None): result["usageStatusText"] = "Waiting for auth" result["authHelpText"] = AUTH_HELP return result + # Signed in, with no limits to read: no app-server, and nothing to sign in to. + if api_key_login(Path(env.get("CODEX_HOME") or (Path.home() / ".codex")), env): + result["authHelpText"] = API_KEY_HELP + result["ok"] = True + return result # stderr goes to a regular file, not a PIPE: an unread pipe can fill and # deadlock a chatty app-server. It is only read once the process has exited. @@ -938,8 +591,8 @@ def fetch_codex_rpc(home=None): except Exception as exc: if stderr_file: stderr_file.close() - result["usageStatusText"] = "Codex unavailable" - result["authHelpText"] = str(exc) + result["usageStatusText"] = "Codex limits unavailable" + result["authHelpText"] = str(exc)[:300] return result try: @@ -960,9 +613,18 @@ def fetch_codex_rpc(home=None): result["usageStatusText"] = "Codex limits unavailable" result["authHelpText"] = message return result - limits = (limits_msg.get("result") or {}).get("rateLimits") or {} + rpc_result = limits_msg.get("result") if isinstance(limits_msg.get("result"), dict) else {} + buckets = rate_limit_buckets(rpc_result) + # An answer with neither view of the limits isn't a plan without any; it + # fails like any other probe, so the last good windows stand, stale. + if not buckets: + result["usageStatusText"] = "Codex limits unavailable" + result["authHelpText"] = "Codex's app-server answered without its rate limits." + return result result["authHelpText"] = "" - plan = limits.get("planType") or "" + plan = next((bucket.get("planType") for _, bucket in buckets if bucket.get("planType")), "") + if not plan and isinstance(rpc_result.get("rateLimits"), dict): + plan = rpc_result["rateLimits"].get("planType") or "" # The limits name the plan themselves. account/read is only a fallback # for when they don't, and never a reason to lose them: Codex 0.158's @@ -970,28 +632,35 @@ def fetch_codex_rpc(home=None): if not plan: try: account_msg = rpc_request(proc, 3, "account/read", timeout=2) - account = (account_msg.get("result") or {}).get("account") or {} + account = obj(obj(account_msg.get("result")).get("account")) plan = account.get("planType") or account.get("type") or "" except (TimeoutError, RuntimeError): pass - result["tierLabel"] = str(plan) if plan else "" + result["tierLabel"] = plan_name(plan) - for window in (limits.get("primary"), limits.get("secondary")): - entry = limit_window(window) - if entry: - result["limits"].append(entry) - entry = spend_limit(limits.get("individualLimit")) - if entry: - result["limits"].append(entry) + result["limits"] = bucket_limits(buckets) + spend = individual_spend_limit(buckets, rpc_result) + if spend: + result["limits"].append(spend) # OpenAI hands out free "full reset" credits that wipe the rate limits on # demand and lapse after a month. Worth knowing about before a limit bites. - credits = reset_credits(limits_msg.get("result") or {}) + credits = reset_credits(rpc_result) if credits: result["resetCredits"] = credits + result["ok"] = True except Exception as exc: result["usageStatusText"] = "Codex limits unavailable" result["authHelpText"] = codex_rpc_help(proc, stderr_file, exc) + # A live app-server that never answered, or one that died saying it + # couldn't reach OpenAI, is a network that isn't up yet (the first probe + # after login often beats DHCP), so the panel asks again soon. One that + # was turned away, or died for its own reasons, would only fail again. + help_text = result["authHelpText"].lower() + if isinstance(exc, TimeoutError) or ( + help_text.startswith("codex app-server exited:") and any(error in help_text for error in TRANSPORT_ERRORS) + ): + result["transport"] = True finally: try: proc.terminate() @@ -1006,49 +675,130 @@ def fetch_codex_rpc(home=None): return result -# The accounts `omarchy agent account` registered, read straight from its -# registry. Only a registry holding a second account matters: with one, the -# record is exactly what it always was. -def registered_accounts(): - state = Path(os.environ.get("XDG_STATE_HOME") or (Path.home() / ".local" / "state")) +# Who a home's probe answers for, or "" when only Codex could say: keyring and +# auto stores keep the login out of auth.json. +def signed_in_identity(codex_home, primary=True): + token = ENV.get("CODEX_ACCESS_TOKEN") if primary else "" + if token: + return "token:" + hashlib.sha1(token.encode("utf-8", "surrogatepass")).hexdigest() try: - registry = json.loads((state / "omarchy" / "agents" / "accounts" / "codex.json").read_text(encoding="utf-8")) + with (codex_home / "config.toml").open("rb") as handle: + config = tomllib.load(handle) + except FileNotFoundError: + config = {} + except Exception: + return "" + if str(config.get("cli_auth_credentials_store") or "file") != "file": + return "" + # account_id names the workspace, whose members each have their own limits, + # so the user comes from the ID token's claims, read as Codex reads them. + try: + tokens = json.loads((codex_home / "auth.json").read_text(encoding="utf-8"))["tokens"] + claims = jwt_claims(tokens["id_token"])["https://api.openai.com/auth"] + workspace = str(tokens["account_id"] or "") + user = str(claims.get("chatgpt_user_id") or claims.get("user_id") or "") except Exception: - return [] - accounts = [a for a in registry.get("accounts") or [] if isinstance(a, dict) and a.get("id")] - if len(accounts) < 2: - return [] - active_id = registry.get("active") or accounts[0]["id"] - if not any(a["id"] == active_id for a in accounts): - active_id = accounts[0]["id"] - for account in accounts: - account["active"] = account["id"] == active_id - account["switch"] = { - "mode": "auto" if registry.get("switch") == "auto" else "manual", - "threshold": registry.get("threshold") or 95, + return "" + return f"{workspace}\n{user}" if workspace and user else "" + + +# Keyed by the home and the account it is signed in to, so a home signed in to +# someone else never falls back on the last one's numbers. None when the +# account can't be told, since then no numbers can be vouched for. +def limits_cache(home, primary=True): + home = Path(home or codex_home()) + identity = signed_in_identity(home, primary) + if not identity: + return None + try: + return LimitsCache(AGENT_ID, str(home), identity) + except Exception as exc: + warn(f"limits cache unavailable ({exc})") + return None + + +# The app-server's answer, with the last limits it gave kept as the answer of +# record when a later probe fails: one slow or dropped call is no reason to +# blank limits that were right a refresh ago. `live` says the limits came from +# this probe, or one recent enough to stand in for it. Elapsed windows come +# back at zero. +def collect_limits(home=None, force=False, primary=True): + cache = limits_cache(home, primary) + fallback, fetched_at_ms = cache.kept() if cache else ([], 0) + + # Each probe starts an app-server, which syncs the plugin list on the way + # up, so a panel opened and shut repeatedly reuses a recent answer instead. + # --force is a person asking for fresh numbers, so it skips the reuse + # window. A plan without windows is an answer like any other, so it stands + # in too. + if cache and cache.fresh(force, allow_empty=True): + result = { + "limits": fallback, + "tierLabel": str(cache.data.get("tierLabel") or ""), + "usageStatusText": "", + "authHelpText": "", + "live": True, + "fetchedAtMs": cache.fetched_at, } - return accounts + if isinstance(cache.data.get("resetCredits"), dict): + result["resetCredits"] = cache.data["resetCredits"] + return result + + rpc = fetch_codex_rpc(home, primary) + ok = rpc.pop("ok", False) + transport = rpc.pop("transport", False) + # A login that changed while the probe ran leaves no telling whose these are. + after = limits_cache(home, primary) if cache else None + same_login = cache is not None and after is not None and after.path == cache.path + + if ok: + rpc["live"] = True + if same_login: + rpc["fetchedAtMs"] = cache.save(rpc["limits"], tierLabel=rpc["tierLabel"], resetCredits=rpc.get("resetCredits")) + else: + rpc["fetchedAtMs"] = now_ms() + return rpc + + rpc["live"] = False + rpc["fetchedAtMs"] = 0 + if transport: + rpc["retryAdvised"] = True + if fallback and same_login: + # A sign-in to restore still says so; any other failure leaves the kept + # numbers to speak for themselves, with their age, and the detail of what + # went wrong in the help text. + rpc.update(limits=fallback, fetchedAtMs=fetched_at_ms, tierLabel=rpc["tierLabel"] or str(cache.data.get("tierLabel") or "")) + if rpc["usageStatusText"] == "Codex limits unavailable": + rpc["usageStatusText"] = "" + elif same_login and cache.fetched_at > 0: + # The last answer was a plan without windows, which has no limits to + # lose: a failed recheck is no news, though a sign-in to restore still is. + rpc["tierLabel"] = rpc["tierLabel"] or str(cache.data.get("tierLabel") or "") + if rpc["usageStatusText"] == "Codex limits unavailable": + rpc.update(usageStatusText="", authHelpText="") + return rpc # Each account's app-server runs in that account's home, so it answers for that # login and refreshes that login's tokens there. -def account_limits(account): +def account_limits(account, force): # The primary is ~/.codex by definition, not whatever CODEX_HOME this run # happened to inherit — that would report another account's limits as Main's. - rpc = fetch_codex_rpc(account.get("home") or (Path.home() / ".codex")) - return { - "id": account["id"], - "label": str(account.get("label") or account["id"]), - "email": str(account.get("email") or ""), - "plan": rpc["tierLabel"] or str(account.get("plan") or ""), - "active": account["active"], - "primary": bool(account.get("primary")), - "limits": rpc["limits"], - "stale": not rpc["limits"], - "usageStatusText": rpc["usageStatusText"], - "authHelpText": rpc["authHelpText"], - "resetCredits": rpc.get("resetCredits"), - } + rpc = collect_limits(account.get("home") or (Path.home() / ".codex"), force, primary=bool(account.get("primary"))) + return account_entry(account, rpc, rpc["tierLabel"], resetCredits=rpc.get("resetCredits")) + + +# Codex is in use here once its CLI or a launcher for it is on PATH, its home +# exists, a token is handed to it, or something has local usage to show. +# Only then is a missing sign-in or binary worth saying. +def codex_in_use(stats): + return ( + bool(find_command("codex")) + or bool(find_codex_binary()) + or codex_home().exists() + or bool(ENV.get("CODEX_ACCESS_TOKEN")) + or number(stats.get("totalSessions")) > 0 + ) def main(): @@ -1062,10 +812,9 @@ def main(): parser.add_argument("--limits-only", action="store_true") args = parser.parse_args() - max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) - stats = cached_local_stats(max_age) - registered = registered_accounts() - accounts = [account_limits(account) for account in registered] + stats = cached_local_stats(scan_max_age(args.force, args.limits_only)) + registered = registered_accounts(AGENT_ID) + accounts = [account_limits(account, args.force) for account in registered] current = next((a for a in accounts if a["active"]), None) if current: # The top-level fields keep describing the account new sessions use. @@ -1077,23 +826,37 @@ def main(): } if current.get("resetCredits"): rpc["resetCredits"] = current["resetCredits"] + stale, fetched_at, retry = current["stale"], current["fetchedAt"], current["retryAdvised"] + elif codex_in_use(stats): + rpc = collect_limits(force=args.force) + stale, fetched_at, retry = not rpc.pop("live"), rpc.pop("fetchedAtMs"), rpc.pop("retryAdvised", False) else: - rpc = fetch_codex_rpc() - - record = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": True, - "hasLocalStats": True, - } + # Nothing here has ever run Codex, so there is no sign-in to ask for and + # no binary to miss: the unused record, which the panel skips. + rpc = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": ""} + stale, fetched_at, retry = False, 0, False + + record = base_record( + AGENT_ID, + AGENT_NAME, + # In use here once it has local usage or limits to show; a machine that + # never ran Codex prints an empty record the panel skips. + # Requests a typed prompt can't be found for still count as use. + ready=number(stats.get("totalSessions")) > 0 or bool(rpc["limits"]) or bool(accounts), + hasLocalStats=True, + ) record.update(stats) record.update(rpc) + # Numbers kept past a failed probe, and when they were measured, so the + # panel can tell them apart from ones checked just now. + record["limitsStale"] = stale + record["limitsFetchedAt"] = fetched_at + if retry: + record["retryAdvised"] = True if accounts: record["accountSwitch"] = registered[0]["switch"] record["accounts"] = accounts - print(json.dumps(record, separators=(",", ":"))) + emit(record) if __name__ == "__main__": diff --git a/bin/omarchy-agent-usage-copilot b/bin/omarchy-agent-usage-copilot new file mode 100755 index 00000000000..2ec86c91b8d --- /dev/null +++ b/bin/omarchy-agent-usage-copilot @@ -0,0 +1,945 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the GitHub Copilot usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect GitHub Copilot usage into one display-ready JSON record. + +GitHub exposes no supported personal-account usage API — the Copilot REST +endpoints are organization- or enterprise-scoped — but the Copilot CLI keeps +two local sources this collector reads. The primary one is the session store +(~/.copilot/session-store.db): its assistant_usage_events table records every +billed request with day, model, session, token split, and cost in nano AI +units (1e9 nano units = 1 AI credit at 1 cent). The fallback, for CLI versions +whose store predates that table, is the session transcript directory +(~/.copilot/session-state//events.jsonl), whose shutdown events carry the +same per-model token split and cost. The record prefers whichever source +produced numbers, never both. + +The monthly meter is an allowance of AI credits: GitHub's own account-wide +figure from the internal quota endpoint (the same one the editor plugins +use), asked with a token from COPILOT_QUOTA_TOKEN, GH_TOKEN, GITHUB_TOKEN, or +`gh auth token`. It names the plan, the allowance, and when it resets, and a +failed check keeps the last answer, marked stale. An answer under 15 seconds +old is reused rather than asked again (--force always asks), and `gh` is only +run when GitHub is actually asked. Only when it can't be asked at all does a +monthlyCredits budget from ~/.config/omarchy/agents/copilot.json give a local +estimate, labeled as one, of this machine's CLI spend: + + { + "monthlyCredits": 1500, + "remote": true + } + +remote turns the quota probe off when false, and on even without a Copilot +CLI home when true; by default it runs only where the CLI has been used. + +The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import calendar +import datetime as dt +import hashlib +import json +import math +import os +import sqlite3 +import subprocess +import sys +import urllib.request +from pathlib import Path +from typing import Any + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + NUMBER_CEILING, + HttpError, + LimitsCache, + Stats, + base_record, + cache_root, + cached_scan, + config_home, + emit, + empty_bucket, + empty_stats, + find_command, + local_day, + number, + open_json, + read_json, + scan_max_age, + stamp, + token_count, + write_json, +) + +AGENT_ID = "copilot" +AGENT_NAME = "GitHub Copilot" +AUTH_HELP = ( + "Run `gh auth login`, or set COPILOT_QUOTA_TOKEN, to show Copilot's monthly allowance." +) + +ULTRA_SECONDS = 1_000_000_000 +QUOTA_URL = "https://api.github.com/copilot_internal/user" +REMOTE_TIMEOUT = 6.0 + +# Plan names as the quota endpoint spells them. +API_PLAN_LABELS = { + "individual": "Pro", + "individual_pro": "Pro+", + "business": "Business", + "enterprise": "Enterprise", + "free": "Free", + "max": "Max", +} + +def expand_path(value: str) -> Path: + return Path(os.path.expandvars(os.path.expanduser(value))).resolve() + + +def read_config() -> dict[str, Any]: + return read_json(config_home() / "omarchy" / "agents" / "copilot.json") + + +TOKEN_KEYS = tuple(empty_bucket()) + + +def month_start(now: dt.datetime) -> dt.datetime: + return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + + +def next_month_start(now: dt.datetime) -> dt.datetime: + days = calendar.monthrange(now.year, now.month)[1] + return month_start(now) + dt.timedelta(days=days) + + +def positive_float(value: Any) -> float: + """A cost from disk, 0.0 when it can't be trusted: like number(), an + infinite or huge figure is a damaged field, not spend that fills the + allowance on its own.""" + try: + n = float(value or 0) + except (TypeError, ValueError, OverflowError): + return 0.0 + return n if math.isfinite(n) and 0 < n < NUMBER_CEILING else 0.0 + + +def credits_from_nano(nano: Any) -> float: + return positive_float(nano) / ULTRA_SECONDS + + +def token_split( + prompt_tokens: Any, output_tokens: Any, cache_read: Any, cache_write: Any +) -> dict[str, int]: + """Map Copilot's token counts onto the panel's four disjoint buckets. + + Copilot's input count is the whole prompt, cache reads and writes included, + and its output count already includes reasoning tokens. The buckets are + disjoint, like the other collectors', so the cached part comes off input + and reasoning is never added a second time. + """ + read = token_count(cache_read) + write = token_count(cache_write) + return { + "inputTokens": max(0, token_count(prompt_tokens) - read - write), + "outputTokens": token_count(output_tokens), + "cacheReadInputTokens": read, + "cacheCreationInputTokens": write, + } + + +def clean_name(value: Any) -> str: + return str(value or "").strip() + + +def parse_event_time(value: Any) -> dt.datetime | None: + """Parse a store or transcript timestamp into timezone-aware local time. + + Accepts ISO-8601 (with or without Z) and epoch seconds or milliseconds. + As in local_day(), zero, a huge or negative number, NaN, a boolean, a date + outside 2000-2100, and anything else collapse to None; the caller falls + back to the file's time or the local date. + """ + if isinstance(value, str): + value = value.strip() + if value.isdigit(): + # ASCII digits only, and few enough to be a time: isdigit() also passes + # other scripts' digits and int() refuses strings past 4300 digits. + value = int(value) if value.isascii() and len(value) <= 20 else None + if value is None or value == "" or isinstance(value, bool): + return None + if isinstance(value, (int, float)): + if number(value) <= 0: + return None + seconds = value / 1000 if value > 10_000_000_000 else value + try: + parsed = dt.datetime.fromtimestamp(seconds, dt.timezone.utc) + except (OverflowError, OSError, ValueError): + return None + else: + try: + parsed = dt.datetime.fromisoformat(str(value).replace("Z", "+00:00")) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=dt.timezone.utc) + if not 2000 <= parsed.year <= 2100: + return None + return parsed.astimezone() + + +def utc_month_key(value: Any) -> str: + parsed = parse_event_time(value) + if parsed is None: + return "" + return parsed.astimezone(dt.timezone.utc).strftime("%Y-%m") + + +# ------------------------------------------------------------------ database + + +def open_store(path: Path) -> sqlite3.Connection | None: + """Open the session store read-only, or None when it can't be read. + + The CLI writes to this database continuously in WAL mode. A read-only + connection reads it in place: each query sees one consistent state of the + store, without blocking the writer and without copying a store that grows + without bound. Every scan aggregates as it streams, so the record's size + depends on days and models, never on the number of rows. + """ + if not path.is_file(): + return None + try: + conn = sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True, timeout=5.0) + except sqlite3.Error: + return None + try: + conn.execute("PRAGMA query_only = ON") + # A file that isn't a database only says so when it is first read. + conn.execute("SELECT 1 FROM sqlite_master LIMIT 1").fetchall() + except sqlite3.Error: + conn.close() + return None + conn.row_factory = sqlite3.Row + return conn + + +def has_table(conn: sqlite3.Connection, name: str) -> bool: + try: + row = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=? LIMIT 1", (name,) + ).fetchone() + return row is not None + except sqlite3.Error: + return False + + +def blank_stats() -> dict[str, Any]: + return {**empty_stats(), "monthCredits": 0.0} + + +# --------------------------------------------------------------- aggregate + + +class UsageTally(Stats): + """Running aggregate shared by both scans, on the shared stats fields. + + Unlike a pi or opencode reply, a Copilot prompt and the tokens it spent + arrive apart (a transcript books a segment's tokens at its shutdown), so + each is counted on its own, along with the month's spend. It grows with + days, sessions, and models, never with raw events, so a scan can stream an + entire store or transcript directory through it. + """ + + def __init__(self) -> None: + super().__init__() + self.current_utc_month = dt.datetime.now(dt.timezone.utc).strftime("%Y-%m") + self.month_credits = 0.0 + + def usage(self, model: str, day: str, tokens: dict[str, int]) -> None: + # Tokens no prompt and no one session is known for: a segment's spend. + self.add(day, None, model, *(tokens[key] for key in TOKEN_KEYS), prompts=0) + + def spend(self, month_key: str, credits: float) -> None: + # The allowance window is the current UTC month, matching the reset + # boundary; older spend belongs to windows that are already over. + if month_key == self.current_utc_month: + self.month_credits += credits + + def stats(self) -> dict[str, Any]: + return {**self.snapshot(), "monthCredits": round(self.month_credits, 6)} + + +# ---------------------------------------------------------- store scan +# +# The session store's assistant_usage_events table records one row per billed +# assistant request with day, model, session, the token split, and the cost in +# total_nano_aiu. Rows stream through the tally rather than being fetched, so +# lifetime totals cover the whole store at constant memory. A row with no time +# is today's, and one whose time can't be trusted counts nowhere. +# +# Current CLIs also say who started each request: initiator "user" is a +# prompt typed into the session, and "agent" the follow-up calls the agent +# made on its own (a tool loop's next step, a subagent). And a request GitHub +# served carries GitHub's usage report, its copilot_usage_model and cost; one +# sent to a custom provider (BYOK, COPILOT_PROVIDER_BASE_URL) has neither, so +# that provider billed it, not Copilot, and it stays out of this record. The +# store doesn't keep where such a request went. A store from before these +# columns counts every row as Copilot's, and its prompts can't be told apart. + + +def table_columns(conn: sqlite3.Connection, name: str) -> set[str]: + try: + return {str(row[1]) for row in conn.execute(f"PRAGMA table_info({name})")} + except sqlite3.Error: + return set() + + +# The stats, and how many dated rows the store held: a store with none falls +# back on the transcripts. +def scan_store(conn: sqlite3.Connection) -> tuple[dict[str, Any], int]: + tally = UsageTally() + rows_seen = 0 + columns = table_columns(conn, "assistant_usage_events") + knows_initiator = "initiator" in columns + knows_billing = "copilot_usage_model" in columns + optional = [name for name in ("initiator", "agent_id", "copilot_usage_model") if name in columns] + extra = "".join(", " + name for name in optional) + try: + rows = conn.execute( + "SELECT session_id, model, created_at, total_nano_aiu, input_tokens," + " output_tokens, cache_read_tokens, cache_write_tokens" + extra + + " FROM assistant_usage_events" + ) + for row in rows: + session_id = clean_name(row["session_id"]) + day = local_day(row["created_at"]) + if day is None: + continue + rows_seen += 1 + if knows_billing and row["copilot_usage_model"] is None and row["total_nano_aiu"] is None: + continue + if knows_initiator: + typed = clean_name(row["initiator"]) == "user" and not ("agent_id" in optional and clean_name(row["agent_id"])) + tally.add_prompts(day, session_id or "unknown", 1 if typed else 0) + else: + tally.add_prompts(day, session_id or "unknown") + tally.spend( + utc_month_key(row["created_at"]), + credits_from_nano(row["total_nano_aiu"]), + ) + tally.usage( + clean_name(row["model"]) or "copilot", + day, + token_split( + row["input_tokens"], + row["output_tokens"], + row["cache_read_tokens"], + row["cache_write_tokens"], + ), + ) + except sqlite3.Error: + return blank_stats(), 0 + # Without an initiator a row is a billed model request, and one prompt can + # make many, so its count stays out of the panel's prompt totals. + return dict(tally.stats(), hasPromptStats=knows_initiator), rows_seen + + +# ------------------------------------------------------ transcript fallback +# +# A CLI whose store predates the assistant_usage_events table still writes one +# events.jsonl per session, one {"type", "timestamp", "data"} event per line: +# +# - user.message is a prompt. data.source marks messages a skill, a subagent, +# or the CLI itself injected; those are not the user's prompts. +# - session.shutdown carries data.modelMetrics: per-model token usage and cost +# (totalNanoAiu), cumulative for the session across resumes. The growth since +# the previous shutdown is the usage of the segment that just ended, spread +# over the days that segment's replies landed on. +# - assistant.message carries data.outputTokens per reply. It dates a segment's +# usage, and stands in for a segment still running with no shutdown yet. + + +def is_prompt(event_type: Any, data: dict[str, Any]) -> bool: + return event_type == "user.message" and clean_name(data.get("source")) in ( + "", + "user", + ) + + +def shutdown_metrics(data: dict[str, Any]) -> dict[str, dict[str, float]]: + metrics = data.get("modelMetrics") + if not isinstance(metrics, dict): + return {} + totals: dict[str, dict[str, float]] = {} + for model, entry in metrics.items(): + if not isinstance(entry, dict): + continue + usage = entry.get("usage") + if not isinstance(usage, dict): + usage = {} + totals[clean_name(model) or "copilot"] = { + **token_split( + usage.get("inputTokens"), + usage.get("outputTokens"), + usage.get("cacheReadTokens"), + usage.get("cacheWriteTokens"), + ), + "credits": credits_from_nano(entry.get("totalNanoAiu")), + } + return totals + + +def apportion(total: int, weights: list[int]) -> list[int]: + """Split a whole total in proportion to positive weights, summing exactly.""" + whole = sum(weights) + shares: list[int] = [] + given = 0 + running = 0 + for weight in weights: + running += weight + share = total * running // whole - given + shares.append(share) + given += share + return shares + + +Reply = tuple[str, str, str, int] # model, local day, UTC month, output tokens + + +def spread_segment( + tally: UsageTally, + model: str, + delta: dict[str, int], + credits: float, + replies: list[Reply], + shutdown_slot: tuple[str, str], +) -> None: + """Book one model's usage for a finished segment on the days it happened. + + A segment can span midnight, and its shutdown only reports totals. Output + splits by each day's reply output, which the transcript records exactly. + Input, cache, and cost split by reply count: every reply is one model call + that resends the context, so calls, not output, track what was read. + """ + slots: dict[tuple[str, str], list[int]] = {} + for reply_model, day, month_key, output in replies: + if reply_model == model: + slot = slots.setdefault((day, month_key), [0, 0]) + slot[0] += 1 + slot[1] += output + if not slots: + slots[shutdown_slot] = [1, 0] + + calls = [counts[0] for counts in slots.values()] + outputs = [counts[1] for counts in slots.values()] + if sum(outputs) <= 0: + outputs = calls + split = { + key: apportion(delta[key], outputs if key == "outputTokens" else calls) + for key in TOKEN_KEYS + } + all_calls = sum(calls) + for index, (day, month_key) in enumerate(slots): + tally.usage(model, day, {key: split[key][index] for key in TOKEN_KEYS}) + tally.spend(month_key, credits * calls[index] / all_calls) + + +def scan_transcript(event_file: Path, session_key: str, tally: UsageTally) -> None: + try: + session_mtime = event_file.stat().st_mtime + except OSError: + session_mtime = 0.0 + + # Session totals as of the last shutdown, and the replies since then. + reported: dict[str, dict[str, float]] = {} + replies: list[Reply] = [] + + with event_file.open("r", encoding="utf-8", errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except ValueError: + continue + if not isinstance(entry, dict): + continue + event_type = entry.get("type") + data = entry.get("data") + if not isinstance(data, dict): + continue + + # A line with no time is the file's; one whose time can't be trusted + # counts nowhere, as in every other collector. + stamped = entry.get("timestamp") + parsed = parse_event_time(stamped) + if parsed is None and stamped is not None: + continue + if parsed is None and session_mtime > 0: + parsed = dt.datetime.fromtimestamp( + session_mtime, dt.timezone.utc + ).astimezone() + day = parsed.date().isoformat() if parsed else tally.today + month_key = ( + parsed.astimezone(dt.timezone.utc).strftime("%Y-%m") if parsed else "" + ) + + if is_prompt(event_type, data): + tally.add_prompts(day, session_key) + elif event_type == "assistant.message": + model = clean_name(data.get("model")) or "copilot" + replies.append( + (model, day, month_key, token_count(data.get("outputTokens"))) + ) + elif event_type == "session.shutdown": + for model, totals in shutdown_metrics(data).items(): + before = reported.get(model, {}) + spread_segment( + tally, + model, + { + key: max(0, int(totals[key] - before.get(key, 0))) + for key in TOKEN_KEYS + }, + max(0.0, totals["credits"] - before.get("credits", 0.0)), + replies, + (day, month_key), + ) + reported[model] = totals + # The shutdown totals already count these replies. + replies.clear() + + for model, day, _month_key, output in replies: + tally.usage(model, day, {**empty_bucket(), "outputTokens": output}) + + +# The stats, and whether every transcript was read: one that couldn't be +# leaves the totals short, so they must not be cached. +def scan_transcripts(session_root: Path) -> tuple[dict[str, Any], bool]: + if not session_root.is_dir(): + return blank_stats(), True + tally = UsageTally() + complete = True + try: + session_dirs = sorted(p for p in session_root.iterdir() if p.is_dir()) + except OSError: + return blank_stats(), False + for session_dir in session_dirs: + event_file = session_dir / "events.jsonl" + if not event_file.is_file(): + continue + try: + scan_transcript(event_file, session_dir.name, tally) + except OSError: + complete = False + return tally.stats(), complete + + +# ------------------------------------------------------------------- scan + + +# The stats, and whether they are complete: a store that is there but could +# not be read leaves only the transcripts, which newer CLIs may not write. +def direct_scan(home_dir: Path) -> tuple[dict[str, Any], bool]: + store = home_dir / "session-store.db" + conn = open_store(store) + if conn is None: + stats, complete = scan_transcripts(home_dir / "session-state") + return stats, complete and not store.exists() + try: + if has_table(conn, "assistant_usage_events"): + stats, rows_seen = scan_store(conn) + if rows_seen: + return stats, True + # A store without the usage table, or one that is still empty, falls back + # to the transcripts exactly like a store-less machine. + return scan_transcripts(home_dir / "session-state") + finally: + conn.close() + + +# ------------------------------------------------------------ remote quota + + +def env_token() -> str: + for var in ("COPILOT_QUOTA_TOKEN", "GH_TOKEN", "GITHUB_TOKEN"): + value = os.environ.get(var) + if value and value.strip(): + return value.strip() + return "" + + +def gh_token() -> str: + exe = find_command("gh") + if not exe: + return "" + try: + result = subprocess.run( + # gh follows GH_HOST without --hostname, and an Enterprise token must + # never go to api.github.com. + [exe, "auth", "token", "--hostname", "github.com"], + capture_output=True, + text=True, + timeout=REMOTE_TIMEOUT, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return "" + return result.stdout.strip() if result.returncode == 0 else "" + + +def gh_user() -> str: + """Who `gh` is signed in to github.com as, from its hosts.yml. + + Read as plain lines rather than YAML: only the host's own `user:` key + matters, and reading it costs no process, so a reused answer never has to + run `gh` to know whose it is. + """ + config = os.environ.get("GH_CONFIG_DIR") or str(Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "gh") + try: + lines = (Path(config) / "hosts.yml").read_text(encoding="utf-8").splitlines() + except (OSError, UnicodeDecodeError): + return "" + in_host, child_indent = False, None + for line in lines: + stripped = line.strip() + if not stripped or stripped.startswith("#"): + continue + indent = len(line) - len(line.lstrip(" ")) + if indent == 0: + in_host, child_indent = stripped == "github.com:", None + continue + if not in_host: + continue + child_indent = indent if child_indent is None else child_indent + if indent == child_indent and stripped.startswith("user:"): + return stripped.split(":", 1)[1].strip().strip("'\"") + return "" + + +def github_credential() -> tuple[str, str]: + """Who the quota is asked as, and the token when it is already in hand. + + Environment first, so a user can scope exactly what this sees, then the + `gh` CLI's own credential store, which is where a machine that has signed + in already keeps it. An exported token keys by itself (the limits cache + hashes whatever it is keyed by). gh keys by the account its hosts.yml + names, and its token is only fetched when GitHub is asked; a hosts.yml + without a user falls back to fetching it now. Nothing here is ever stored + or logged. + """ + token = env_token() + if token: + return token, token + user = gh_user() if find_command("gh") else "" + if user: + return f"gh:github.com:{user}", "" + token = gh_token() + return (token, token) if token else ("", "") + + +# A quota amount as GitHub sends it, or None when it isn't a readable one. +def quota_amount(value: Any) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + return None + try: + amount = float(value) + except ValueError: + return None + return amount if math.isfinite(amount) and amount >= 0 else None + + +def ask_quota(token: str) -> dict[str, Any]: + """Ask GitHub for the account-wide credit balance.""" + request = urllib.request.Request( + QUOTA_URL, + headers={ + "Authorization": f"Bearer {token}", + "Accept": "application/json", + "User-Agent": "omarchy-agent-usage", + }, + ) + try: + payload = open_json(request, timeout=REMOTE_TIMEOUT, limit=2_000_000) + except HttpError as error: + if error.status: + return {"available": False, "reason": f"http-{error.status}"} + # No route, no DNS, refused, reset or timed out: nothing answered. + if error.transport: + return {"available": False, "reason": "unreachable"} + # A reply that came back cut short or unreadable is still an answer from + # GitHub. + return {"available": False, "reason": "bad-payload"} + + if not isinstance(payload, dict): + return {"available": False, "reason": "bad-payload"} + + snapshots = payload.get("quota_snapshots") + snapshot = snapshots.get("premium_interactions") if isinstance(snapshots, dict) else None + if not isinstance(snapshot, dict): + return {"available": False, "reason": "no-quota"} + + # Under token-based billing the snapshot counts AI credits directly; only + # the local store and transcripts count in nano units. GitHub's snapshot + # (the CLI's CopilotUserResponseQuotaSnapshotsPremiumInteractions) says + # what is left, as remaining, quota_remaining or percent_remaining, so the + # amount used is the entitlement less that; a credits_used amount is still + # read first where one is sent. A limited allowance needs both amounts: one + # missing or unreadable would pass for no meter or an untouched allowance + # and replace the last good answer, so it is a failed check instead. + unlimited = snapshot.get("unlimited") is True + entitlement = quota_amount(snapshot.get("entitlement")) + used = quota_amount(snapshot.get("credits_used")) + if used is None and entitlement is not None: + remaining = quota_amount(snapshot.get("remaining")) + if remaining is None: + remaining = quota_amount(snapshot.get("quota_remaining")) + percent_remaining = quota_amount(snapshot.get("percent_remaining")) + if remaining is not None: + used = max(0.0, entitlement - remaining) + elif percent_remaining is not None: + used = entitlement * max(0.0, 100.0 - percent_remaining) / 100.0 + if not unlimited and (entitlement is None or used is None): + return {"available": False, "reason": "bad-payload"} + return { + "available": True, + "unlimited": unlimited, + "entitlement": entitlement or 0.0, + "used": used or 0.0, + "resetsAt": str(payload.get("quota_reset_date_utc") or ""), + "planLabel": API_PLAN_LABELS.get(str(payload.get("copilot_plan") or "").strip().lower(), ""), + } + + +# Before the limits caches were named copilot-limits-*, each token kept its +# answer as copilot-quota-. Read once, when this account has no +# answer under the new name, so the last allowance survives the rename. +def adopt_legacy_cache(cache: Path, token: str) -> dict[str, Any]: + legacy = cache_root() / f"copilot-quota-{hashlib.sha256(token.encode('utf-8')).hexdigest()[:16]}.json" + cached = read_json(legacy) + if not cached.get("available"): + return {} + try: + write_json(cache, cached) + legacy.unlink(missing_ok=True) + except OSError: + pass + return cached + + +def fetch_remote_quota(force: bool = False) -> dict[str, Any]: + """The live quota, or the last good one marked stale. + + Every failure is soft: a missing token, an offline machine, or a changed + payload keeps the last answer this account got, marked stale, or leaves + the record on its local stats rather than failing the whole collection. + Each account keeps its own cache, so a machine signed in to another + account never shows the last one's allowance. The cache is read whatever + its file time, so a clock that went back doesn't lose it. + """ + identity, token = github_credential() + if not identity: + return {"available": False, "reason": "no-token"} + + # The cache keeps GitHub's answer itself rather than meter rows: the meter + # is drawn from it at print time, budget and reset date included. + cache = LimitsCache(AGENT_ID, identity) + cached = cache.data + # A panel opened and shut repeatedly must not turn into a request (and a + # gh process) per flick, so an answer this recent is the answer. + if cached.get("available") and cache.fresh(force, allow_empty=True): + return open_quota(dict(cached, stale=False)) + + # hosts.yml still naming the account while `gh auth token` fails is gh or + # its credential store having a bad moment, not a sign-out (signing out + # drops the account from hosts.yml), so that account's last allowance + # stands, stale, like any other failed check. + token = token or gh_token() + if not token: + quota = {"available": False, "reason": "no-token"} + else: + if not cached: + cached = adopt_legacy_cache(cache.path, token) + quota = ask_quota(token) + if quota.get("available"): + quota["stale"] = False + quota["fetchedAtMs"] = cache.save([], **quota) + return quota + + if cached.get("available"): + return open_quota(dict(cached, stale=True, failure=quota.get("reason", ""))) + return quota + + +# A kept answer describes its own allowance window and nothing after it: +# once GitHub's reset date passes, its usage is back at zero and the next +# reset is unknown until GitHub is asked again. One without a reset date is +# kept as it is. +def open_quota(cached: dict[str, Any]) -> dict[str, Any]: + resets_at = parse_event_time(str(cached.get("resetsAt") or "")) + if resets_at is not None and resets_at <= dt.datetime.now(dt.timezone.utc): + cached.update(used=0.0, resetsAt="") + return cached + + +# What went wrong asking GitHub, and what to do about it. +def quota_trouble(reason: str) -> str: + if reason == "unreachable": + return "Couldn't reach GitHub's Copilot quota endpoint. Retrying shortly." + if reason == "no-token": + return "The gh CLI didn't hand over its GitHub token; check `gh auth status`." + if reason in ("http-401", "http-403"): + return "GitHub refused the token for Copilot's quota. " + AUTH_HELP + if reason.startswith("http-"): + return f"GitHub's Copilot quota endpoint returned status {reason[5:]}." + return "GitHub's Copilot quota endpoint sent no allowance." + + +# ---------------------------------------------------------------- allowance + + +def allowance_limits( + month_credits: float, + config: dict[str, Any], + remote: dict[str, Any] | None, + now: dt.datetime, + has_stats: bool, +) -> dict[str, Any]: + """The monthly allowance meter. + + The live figure from the quota endpoint counts every machine, editor, and + agent on the entitlement, plan included. When it can't be asked, a + monthlyCredits budget from the configuration gives a local estimate of this + machine's CLI spend, labeled as such; without one there is no meter. + """ + result = { + "limits": [], + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limitsStale": False, + "limitsFetchedAt": 0, + "retryAdvised": False, + } + remote = remote if isinstance(remote, dict) else {} + reason = str(remote.get("failure") or remote.get("reason") or "") + result["retryAdvised"] = reason == "unreachable" + + if remote.get("available"): + result["tierLabel"] = str(remote.get("planLabel") or "") + if remote.get("stale"): + result["authHelpText"] = quota_trouble(reason) + " Showing the last known allowance." + # A refused token is a sign-in to redo, kept allowance or not. + if reason in ("http-401", "http-403"): + result["usageStatusText"] = "Sign-in expired" + if remote.get("unlimited") or float(remote.get("entitlement") or 0) <= 0: + return result + # GitHub's own reset date, or none: a kept window that has reset has no + # known next one. + parsed = parse_event_time(str(remote.get("resetsAt") or "")) + result["limits"] = [{ + "label": "Monthly allowance", + "percent": min(1.0, float(remote.get("used") or 0.0) / float(remote["entitlement"])), + "resetsAt": parsed.isoformat() if parsed is not None and parsed > now else "", + }] + result["limitsStale"] = bool(remote.get("stale")) + result["limitsFetchedAt"] = stamp(remote.get("fetchedAtMs")) + return result + + if reason == "no-token": + result["usageStatusText"] = "Waiting for auth" + result["authHelpText"] = AUTH_HELP + elif reason in ("http-401", "http-403"): + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = quota_trouble(reason) + elif reason and reason != "disabled": + result["usageStatusText"] = "Copilot limits unavailable" + result["authHelpText"] = quota_trouble(reason) + + # A budget must not advertise a meter on a machine the CLI has never + # touched: the estimate needs local numbers to mean anything. Its window is + # the calendar month in UTC, the boundary GitHub documents for Copilot's + # allowance. + budget = config.get("monthlyCredits") + if isinstance(budget, (int, float)) and budget > 0 and has_stats: + result["limits"] = [{ + "label": "Monthly allowance (est.)", + "percent": min(1.0, month_credits / float(budget)), + "resetsAt": next_month_start(now).isoformat(), + }] + return result + + +# ------------------------------------------------------------------ record + + +def main() -> int: + parser = argparse.ArgumentParser() + # --force rescans everything and rewrites the cache. --limits-only is kept + # for CLI compatibility with the panel's refreshLimits() call: only the + # limits probe must be fresh, so it may reuse a scan for far longer than a + # normal run, whose short window exists purely to dedup concurrent + # collector runs. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.add_argument("--copilot-home", default=os.environ.get("COPILOT_HOME") or "~/.copilot") + args = parser.parse_args() + + home_dir = expand_path(args.copilot_home) + # The digest covers the one data path the scan reads, the CLI's home. + stats = cached_scan(AGENT_ID, [home_dir], scan_max_age(args.force, args.limits_only), lambda: direct_scan(home_dir)) + + config = read_config() + now = dt.datetime.now(dt.timezone.utc) + + # The quota probe hands a GitHub token to GitHub, so it only runs where the + # Copilot CLI has been used, or where the configuration asks for it. A + # machine with gh signed in but no Copilot stays an empty record. + remote = {"available": False, "reason": "disabled"} + if config.get("remote") is True or (config.get("remote", True) is not False and home_dir.is_dir()): + remote = fetch_remote_quota(args.force) + + has_stats = number(stats.get("totalPrompts")) > 0 or float(stats.get("monthCredits") or 0.0) > 0 + allowance = allowance_limits(float(stats.get("monthCredits") or 0.0), config, remote, now, has_stats) + ready = number(stats.get("totalPrompts")) > 0 or len(allowance["limits"]) > 0 + + record = base_record( + AGENT_ID, + AGENT_NAME, + ready=ready, + # A store of custom-provider requests alone leaves nothing of Copilot's. + hasLocalStats=number(stats.get("totalSessions")) > 0 or bool(stats.get("modelUsage")), + hasPromptStats=stats.get("hasPromptStats") is not False, + tierLabel=allowance["tierLabel"], + usageStatusText=allowance["usageStatusText"], + authHelpText=allowance["authHelpText"], + limits=allowance["limits"], + limitsStale=allowance["limitsStale"], + limitsFetchedAt=allowance["limitsFetchedAt"], + ) + # The first probe after login often fires before the network is up. + if allowance["retryAdvised"]: + record["retryAdvised"] = True + for key in ( + "todayPrompts", + "todaySessions", + "todayTotalTokens", + "totalPrompts", + "totalSessions", + "activeDays", + ): + record[key] = number(stats.get(key)) + record["todayTokensByModel"] = stats.get("todayTokensByModel") or {} + record["recentDays"] = stats.get("recentDays") or [] + record["modelUsage"] = stats.get("modelUsage") or {} + record["activeDates"] = stats.get("activeDates") or [] + emit(record) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-cursor b/bin/omarchy-agent-usage-cursor new file mode 100755 index 00000000000..4b9b60f7eb1 --- /dev/null +++ b/bin/omarchy-agent-usage-cursor @@ -0,0 +1,473 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Cursor usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Cursor usage into one display-ready JSON record. + +Cursor keeps no token counts on disk, so everything here comes from the same +dashboard RPCs the Cursor CLI itself calls: the plan and the period meters +from GetCurrentPeriodUsage and GetPlanInfo, tokens per model and per day from +GetAggregatedUsageEvents. They are asked with the sign-in the Cursor CLI +stores, or failing that the one the Cursor editor keeps in its state.vscdb. +The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from datetime import date, datetime, time, timedelta, timezone +from pathlib import Path +from typing import Any + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + account_key, + base_record as shared_base_record, + cached_scan, + config_home, + cursor_auth_path, + cursor_token, + emit, + empty_bucket, + empty_stats, + finite, + number, + open_json, + read_json, + scan_cache_paths, + scan_max_age, + token_lapsed, + trusted_endpoint, +) + +AGENT_ID = "cursor" +AGENT_NAME = "Cursor" +AUTH_HELP = "Sign in to the Cursor editor, or run `cursor-agent login`." +API_BASE_URL = "https://api2.cursor.sh" +SERVICE = "aiserver.v1.DashboardService" +REQUEST_TIMEOUT = 20 +RECENT_DAYS = 7 + + +class CursorError(Exception): + def __init__(self, message: str, auth: bool = False, transport: bool = False) -> None: + super().__init__(message) + self.auth = auth + # Nothing answered at all, so the panel should try again soon. + self.transport = transport + + +def fraction(value: Any) -> float: + # Cursor reports whole percentages; the panel's meters take 0..1. + return min(1.0, max(0.0, finite(value) / 100)) + + +# The Cursor CLI keeps its config beside auth.json and in ~/.cursor, the editor +# under ~/.config/Cursor. Any of them means Cursor is used here, so a missing +# sign-in is worth saying; without them the machine doesn't use Cursor at all. +def cursor_in_use(auth_path: Path) -> bool: + return any(path.is_dir() for path in (auth_path.parent, config_home() / "Cursor", Path.home() / ".cursor")) + + +# ------------------------------------------------------------------ record + + +def iso(epoch_millis: Any) -> str: + millis = number(epoch_millis) + if millis <= 0: + return "" + try: + return datetime.fromtimestamp(millis / 1000, timezone.utc).isoformat() + except (OverflowError, OSError, ValueError): + return "" + + +def base_record(**overrides: Any) -> dict[str, Any]: + # Dashboard numbers are account-global, not machine-local: every synced + # device reports the same truth, so aggregation must not sum them. The + # dashboard reports usage events, never prompt or session counts; the panel + # keeps those numbers out of today's tooltip when hasPromptStats is false. + return shared_base_record(AGENT_ID, AGENT_NAME, scope="account", hasPromptStats=False, **overrides) + + +class CursorClient: + def __init__(self, token: str, api_base_url: str) -> None: + self.token = token + self.api_base_url = api_base_url.rstrip("/") + + def call(self, method: str, payload: dict[str, Any]) -> dict[str, Any]: + request = urllib.request.Request( + f"{self.api_base_url}/{SERVICE}/{method}", + data=json.dumps(payload).encode("utf-8"), + headers={ + "Content-Type": "application/json", + "Connect-Protocol-Version": "1", + "Authorization": f"Bearer {self.token}", + "x-cursor-client-type": "cli", + }, + method="POST", + ) + try: + parsed = open_json(request, timeout=REQUEST_TIMEOUT) + except HttpError as error: + if error.status in (401, 403): + raise CursorError(AUTH_HELP, auth=True) from error + if error.status: + raise CursorError(f"Cursor's dashboard returned status {error.status}.") from error + # A reply cut short came from a server that answered, so it is no + # reason to ask the panel for a sooner retry. + if error.cut_short: + raise CursorError("Cursor's dashboard cut its reply short.") from error + # No route, no DNS, refused, reset or timed out: nothing answered. + if error.transport: + raise CursorError("Couldn't reach Cursor's dashboard. Retrying shortly.", transport=True) from error + raise CursorError("Cursor's dashboard returned malformed JSON.") from error + if not isinstance(parsed, dict): + raise CursorError("Cursor's dashboard returned malformed JSON.") + return parsed + + def period(self) -> dict[str, Any]: + return self.call("GetCurrentPeriodUsage", {}) + + # None when the plan couldn't be asked, so the last known one stands. + def plan_info(self) -> dict[str, Any] | None: + try: + plan = self.call("GetPlanInfo", {}).get("planInfo") + except CursorError: + return None + return plan if isinstance(plan, dict) else {} + + def aggregations(self, start_millis: int, end_millis: int) -> list[Any]: + payload = {"startDate": str(start_millis), "endDate": str(end_millis)} + rows = self.call("GetAggregatedUsageEvents", payload).get("aggregations") + return rows if isinstance(rows, list) else [] + + +# ------------------------------------------------------------------ tokens + + +def day_windows(today: date, count: int) -> list[tuple[str, int, int]]: + """The last `count` local days, oldest first, as (date, startMs, endMs).""" + windows = [] + for offset in range(count - 1, -1, -1): + day = today - timedelta(days=offset) + # Each midnight converts on its own: a day that springs forward or falls + # back is 23 or 25 hours long, never a fixed 24. + start = datetime.combine(day, time.min).astimezone() + end = datetime.combine(day + timedelta(days=1), time.min).astimezone() + windows.append((day.isoformat(), int(start.timestamp() * 1000), int(end.timestamp() * 1000))) + return windows + + +def summarize_day(rows: list[Any]) -> dict[str, dict[str, int]]: + """One day's aggregations under the token names every collector prints.""" + usage: dict[str, dict[str, int]] = {} + for row in rows: + if not isinstance(row, dict): + continue + model = str(row.get("modelIntent") or "cursor") + bucket = usage.setdefault(model, empty_bucket()) + # Connect's JSON writes 64-bit counts as strings, so number() rather than + # the strict token_count(). + bucket["inputTokens"] += number(row.get("inputTokens")) + bucket["outputTokens"] += number(row.get("outputTokens")) + bucket["cacheReadInputTokens"] += number(row.get("cacheReadTokens")) + bucket["cacheCreationInputTokens"] += number(row.get("cacheWriteTokens")) + return usage + + +def summarize_usage(days: list[tuple[str, dict[str, dict[str, int]]]]) -> dict[str, Any]: + recent_days = [] + active_dates = [] + model_usage: dict[str, dict[str, int]] = {} + today_by_model: dict[str, int] = {} + + for index, (day, usage) in enumerate(days): + tokens = 0 + for model, bucket in usage.items(): + carried = model_usage.setdefault(model, empty_bucket()) + for name, value in bucket.items(): + carried[name] += value + tokens += sum(bucket.values()) + recent_days.append({"date": day, "messageCount": tokens}) + if tokens > 0: + active_dates.append(day) + if index == len(days) - 1: + today_by_model = {model: sum(bucket.values()) for model, bucket in usage.items() if sum(bucket.values()) > 0} + + return { + "hasLocalStats": bool(model_usage), + "todayTotalTokens": sum(today_by_model.values()), + "todayTokensByModel": today_by_model, + "recentDays": recent_days, + "activeDays": len(active_dates), + "activeDates": active_dates, + "modelUsage": model_usage, + } + + +# The token scan is cached per account rather than per data path: Cursor's +# numbers live on its dashboard, and the account is who they belong to. +def stats_cache(key: str) -> Path: + return scan_cache_paths(AGENT_ID, [key])[0] + + +def carried_stats(key: str, today: date) -> dict[str, Any]: + """The last scan of this account, for runs that fetch no fresh tokens. + + Another account's scan never stands in for this one's. A scan from an + earlier day is moved onto this week's dates, and its today figures are + cleared, since nothing in it happened today. + """ + stats = {"hasLocalStats": False, **empty_stats()} + cached = read_json(stats_cache(key)) + previous = cached.get("stats") + if not isinstance(previous, dict): + return stats + for name in stats: + if name in previous: + stats[name] = previous[name] + if cached.get("scanDate") != today.isoformat(): + week = [window[0] for window in day_windows(today, RECENT_DAYS)] + kept_days = {str(day): usage for day, usage in previous.get("days") or [] if isinstance(usage, dict)} + if kept_days: + # Summed again over the new week, so the day that fell out leaves the + # model totals and active days along with the chart. + stats.update(summarize_usage([(day, kept_days.get(day, {})) for day in week]), hasLocalStats=stats["hasLocalStats"]) + else: + counts = {str(row.get("date")): row.get("messageCount", 0) for row in stats["recentDays"] if isinstance(row, dict)} + stats["recentDays"] = [{"date": day, "messageCount": counts.get(day, 0)} for day in week] + stats["activeDates"] = [day for day in stats.get("activeDates") or [] if day in week] + stats["activeDays"] = len(stats["activeDates"]) + stats["todayTotalTokens"] = 0 + stats["todayTokensByModel"] = {} + return stats + + +# This week's tokens: a scan of this account recent enough to reuse, else a +# fresh one, else the last one when the fetch fails. +def week_stats(client: CursorClient, today: date, key: str, max_age_seconds: float) -> dict[str, Any]: + def scan() -> tuple[dict[str, Any], bool]: + windows = day_windows(today, RECENT_DAYS) + with ThreadPoolExecutor(max_workers=4) as pool: + rows = list(pool.map(lambda window: client.aggregations(window[1], window[2]), windows)) + days = [(window[0], summarize_day(day_rows)) for window, day_rows in zip(windows, rows)] + # Each day's usage is kept too, so a scan carried past midnight can be + # summed again over the new week. + return {**empty_stats(), **summarize_usage(days), "days": days}, True + + try: + stats = cached_scan(AGENT_ID, [key], max_age_seconds, scan) + except Exception: + return carried_stats(key, today) + return {name: value for name, value in stats.items() if name != "days"} + + +# ------------------------------------------------------------------ limits + + +# One section of the period reply. Connect's JSON leaves an empty message out +# (or sends null), so a missing section is a plan with nothing in it; one of +# another type is a reply this doesn't read, which must not pass for no meters +# and replace the known limits. +def usage_section(period: dict[str, Any], key: str) -> dict[str, Any]: + section = period.get(key) + if section is None: + return {} + if not isinstance(section, dict): + raise CursorError("Cursor's dashboard sent usage in a form this doesn't read.") + return section + + +# PlanInfo.included_usage_period says whether the included allowance runs on +# the billing cycle or resets each week, at included_usage_resets_at. The +# enum may arrive as its number or its name, prefixed or not. +def weekly_included(plan: dict[str, Any] | None) -> bool: + value = (plan or {}).get("includedUsagePeriod") + if isinstance(value, bool): + return False + return value == 2 or (isinstance(value, str) and value.strip().upper() in ("WEEKLY", "INCLUDED_USAGE_PERIOD_WEEKLY")) + + +# The weekly reset as an epoch in milliseconds, like billingCycleEnd; one +# small enough to be seconds is read as seconds rather than as 1970. +def weekly_reset(value: Any) -> str: + stamp = number(value) + return iso(stamp * 1000 if 0 < stamp < 100_000_000_000 else stamp) + + +def summarize_limits(period: dict[str, Any], plan_info: dict[str, Any] | None = None) -> list[dict[str, Any]]: + """Cursor's own percentages, never a percentage derived from spend. + + The dollar figures do not divide into the allowance: bonus usage the + providers hand out is spend without an allowance behind it, so + totalSpend/limit reads well past 100% while Cursor's own meters do not. + + Every meter runs on the billing cycle, unless the plan says its included + allowance is weekly. The included total is the Monthly (or Weekly) row; + the model meters and the on-demand budget are titled after it, so the + panel attaches them to that row rather than showing four rows that each + read "Monthly". The on-demand budget always runs on the billing cycle. + """ + plan = usage_section(period, "planUsage") + spend = usage_section(period, "spendLimitUsage") + resets_at = iso(period.get("billingCycleEnd")) + window, included_resets_at = "Monthly", resets_at + if weekly_included(plan_info): + window, included_resets_at = "Weekly", weekly_reset((plan_info or {}).get("includedUsageResetsAt")) + limits = [] + + for label, title, key in ( + ("Included total", window, "totalPercentUsed"), + ("Auto models", f"Auto models {window}", "autoPercentUsed"), + ("Named models", f"Named models {window}", "apiPercentUsed"), + ): + if key in plan: + limits.append({"label": label, "title": title, "percent": fraction(plan.get(key)), "resetsAt": included_resets_at}) + + limit_cents = number(spend.get("individualLimit")) + if limit_cents > 0: + used_cents = number(spend.get("individualUsed")) + limits.append(dict(on_demand_names(used_cents, limit_cents), percent=min(1.0, used_cents / limit_cents), resetsAt=resets_at, onDemandLimitCents=limit_cents)) + return limits + + +# The on-demand meter names its dollars, so a kept one that resets to zero +# is renamed with the spend it now stands at. +def on_demand_names(used_cents: float, limit_cents: float) -> dict[str, str]: + label = f"On-demand ${used_cents / 100:.2f} / ${limit_cents / 100:.2f}" + return {"label": label, "title": f"{label} Monthly"} + + +def limits_cache(key: str) -> LimitsCache: + return LimitsCache(AGENT_ID, key) + + +# The kept limits as they read now: one whose billing cycle has ended is back +# at zero, with the next cycle's reset unknown until a check succeeds, and an +# on-demand meter among them is renamed for its spend of zero. +def kept_limits(cache: LimitsCache) -> list[dict[str, Any]]: + raw = cache.data.get("limits") + raw = [entry for entry in raw if isinstance(entry, dict)] if isinstance(raw, list) else [] + kept = [] + for entry, opened in zip(raw, cache.limits): + cents = number(entry.get("onDemandLimitCents")) + if cents > 0 and entry.get("resetsAt") and not opened.get("resetsAt"): + opened = dict(opened, **on_demand_names(0, cents)) + kept.append(opened) + return kept + + +# ---------------------------------------------------------------- assembly + + +def scan(api_base_url: str, auth_path: Path, limits_only: bool, state_db: Path | None = None, force: bool = False) -> dict[str, Any]: + token = cursor_token(auth_path, state_db) + if not token: + if cursor_in_use(auth_path): + return base_record(ready=True, usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + return base_record() + + client = CursorClient(token, api_base_url) + key = account_key(token) + cache = limits_cache(key) + kept = kept_limits(cache) + fetched_at = cache.fetched_at + kept_tier = str(cache.data.get("tierLabel") or "") + today = datetime.now().astimezone().date() + # --limits-only runs while a meter is near its limit, which can last hours, + # so the tokens it carries are only so old before they are fetched again. + scan_age = scan_max_age(force, limits_only) + + # A failed check keeps the last good limits, marked stale so the panel + # gives their age on hover, and the week of tokens already on disk. + def failed(status: str, detail: str, retry: bool = False) -> dict[str, Any]: + record = base_record( + ready=True, + tierLabel=kept_tier, + usageStatusText=status, + authHelpText=detail + (" Showing the last known limits." if kept else ""), + limits=kept, + limitsStale=True, + limitsFetchedAt=fetched_at if kept else 0, + ) + if retry: + record["retryAdvised"] = True + record.update(carried_stats(key, today)) + return record + + # CURSOR_API_ENDPOINT moves the endpoint, never the sign-in onto a + # connection anyone on the way can read. + if not trusted_endpoint(api_base_url): + return failed("" if kept else "Cursor limits unavailable", "CURSOR_API_ENDPOINT must use HTTPS (plain HTTP only for loopback).") + + if token_lapsed(token): + return failed( + "Limits paused", + "Cursor's access token has lapsed. Open the Cursor editor, or run `cursor-agent`, to renew it; signing in again is not needed.", + ) + + # A panel opened and shut repeatedly must not turn into a request per + # flick, so an answer this recent is the answer, even one with no meters. + if cache.fresh(force, allow_empty=True): + record = base_record(ready=True, tierLabel=kept_tier, limits=kept, limitsFetchedAt=fetched_at) + record.update(week_stats(client, today, key, scan_age)) + return record + + try: + period = client.period() + plan = client.plan_info() + # Only the plan says the included allowance is weekly, and the week's + # reset with it, so one last known weekly keeps its limits, stale, + # rather than reading them as Monthly on the billing cycle. + if plan is None and cache.data.get("weekly") is True: + raise CursorError("Cursor's plan couldn't be checked.") + limits = summarize_limits(period, plan) + except CursorError as error: + if error.auth: + return failed("Sign-in expired", "Cursor refused the saved sign-in. " + AUTH_HELP) + return failed("" if kept else "Cursor limits unavailable", str(error), retry=error.transport) + except Exception as error: + print(f"omarchy-agent-usage-cursor: {type(error).__name__}", file=sys.stderr) + return failed("" if kept else "Cursor limits unavailable", "Cursor's usage check failed.") + + tier = kept_tier if plan is None else str(plan.get("planName") or "") + weekly = cache.data.get("weekly") is True if plan is None else weekly_included(plan) + # The cache only backs later runs; failing to keep it never costs this + # run's live answer. + record = base_record(ready=True, tierLabel=tier, limits=limits, limitsFetchedAt=cache.save(limits, tierLabel=tier, weekly=weekly)) + record.update(week_stats(client, today, key, scan_age)) + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the Cursor usage record as JSON") + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.add_argument("--auth-path", default=os.environ.get("CURSOR_AUTH_PATH", "")) + parser.add_argument("--api-base-url", default=os.environ.get("CURSOR_API_ENDPOINT", API_BASE_URL)) + args = parser.parse_args() + + auth_path = Path(args.auth_path).expanduser() if args.auth_path else cursor_auth_path() + try: + record = scan(args.api_base_url, auth_path, args.limits_only, force=args.force) + # scan() keeps the last limits through any failed check; only a failure + # before it knows the account ends up here. + except Exception as error: + record = base_record(usageStatusText="Cursor limits unavailable", authHelpText="Cursor's usage check failed.") + print(f"omarchy-agent-usage-cursor: {type(error).__name__}", file=sys.stderr) + emit(record) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-deepseek b/bin/omarchy-agent-usage-deepseek new file mode 100755 index 00000000000..7f30bfc7430 --- /dev/null +++ b/bin/omarchy-agent-usage-deepseek @@ -0,0 +1,450 @@ +#!/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, from opencode sessions that ran on a deepseek provider, and +from DeepSeek Harness (dsh) session logs on its DeepSeek routes. DeepSeek +billing is prepaid, so the record carries the live credit balance from +GET https://api.deepseek.com/user/balance when a key is available, plus the +optional funded/spent detail derived from ~/.config/omarchy/agents/deepseek.json +(the same shape the fireworks collector uses). The agents panel only ever reads +the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +import urllib.request +from decimal import Decimal, InvalidOperation +from pathlib import Path +from typing import Any + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record as shared_base_record, + cached_scan, + config_home, + data_home, + dsh_home, + emit, + limits_cache_path, + now_ms, + number, + obj, + open_json, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_dsh_sessions, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, + trusted_endpoint, + warn, +) + +AGENT_ID = "deepseek" +AGENT_NAME = "DeepSeek" +API_BASE_URL = "https://api.deepseek.com" +OUT_OF_CREDITS_HELP = "DeepSeek says the balance no longer covers API calls. Top up at platform.deepseek.com." + +# DeepSeek Harness bills DeepSeek through these provider routes: an API key +# and a DeepSeek account sign-in (its DeepSeek adapter), and its pi-ai +# adapter's route for pi-ai's own `deepseek` provider, the id pi and opencode +# name too. Its other pi-ai routes belong to whoever bills them: OpenRouter's +# to the openrouter collector, a hand-declared gateway's to nobody here. +DSH_PROVIDERS = ("deepseek-official", "deepseek-account", "deepseek") + + +class DeepSeekError(Exception): + def __init__(self, message: str, transport: bool = False, refused: bool = False) -> None: + super().__init__(message) + self.transport = transport + self.refused = refused + + +def model_name(raw: Any) -> str: + value = str(raw or AGENT_ID) + return value if value else AGENT_ID + + +# ---------------------------------------------------- pi, omp, and opencode + +# pi, omp, and opencode name the provider alike. Exact match: a custom +# gateway named like DeepSeek is not this account. +def billed_here(provider: str, model: str, entry: dict) -> bool: + return provider == "deepseek" + + +def add_harness_rows(stats: Stats) -> bool: + """Add pi/omp and opencode messages on the deepseek provider. Returns + whether every source was read in full.""" + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"deepseek"', billed_here) + for row in pi_rows: + stats.add_row(row, model_name(row.model)) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode(['deepseek"'], billed_here) + for row in db_rows: + stats.add_row(row, model_name(row.model.rstrip("/").split("/")[-1])) + return pi_complete and db_complete + + +# -------------------------------------------------------- DeepSeek Harness + +def dsh_billed_here(provider: str, model: str, request: dict) -> bool: + return provider in DSH_PROVIDERS + + +def scan_dsh(stats: Stats) -> bool: + rows, complete = scan_dsh_sessions(dsh_billed_here) + for row in rows: + stats.add_row(row, model_name(row.model)) + return complete + + +# ------------------------------------------------------------- scan cache + +def run_local_scans(reuse: bool = True) -> tuple[dict[str, Any], bool]: + stats = Stats() + complete = add_harness_rows(stats) + complete = scan_dsh(stats) and complete + # Claude Code pointed at DeepSeek's Anthropic API, or answered by a + # DeepSeek model, and a MiniMax Code custom provider at DeepSeek's host. + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + for row in routed_rows: + stats.add_row(row, model_name(row.model)) + return stats.snapshot(), complete and routed_complete + + +def cached_local_stats(max_age: float) -> dict[str, Any]: + # The digest covers every data path the scan reads. + sources = [opencode_db(), dsh_home(), *pi_session_roots(), *routed_sources()] + # A forced scan reads every Codex rollout again too. + return cached_scan(AGENT_ID, sources, max_age, lambda: run_local_scans(reuse=max_age > 0), match=routed_match()) + + +# -------------------------------------------------------------- credentials + +def read_key_from_dict(parsed: Any, provider: str) -> str: + if not isinstance(parsed, dict): + return "" + entry = parsed.get(provider) + if isinstance(entry, dict): + entry = entry.get("key") or entry.get("apiKey") or entry.get("api_key") + key = entry.strip() if isinstance(entry, str) else "" + # Only a literal key is taken: a pi `$VAR` or `!command` entry is resolved + # by pi alone, and sent as it stands it would read as a refused key. + return "" if key.startswith(("!", "$")) else key + + +def pi_agent_dir() -> Path: + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + +def read_json_key(path: Path, provider: str) -> str: + return read_key_from_dict(read_json(path), provider) + + +def unquote(value: str) -> str: + value = value.strip() + if value[:1] in ("'", '"'): + end = value.find(value[0], 1) + return value[1:end] if end > 0 else "" + # An unquoted scalar ends where a comment begins. + return value.split(" #", 1)[0].strip() + + +def read_dsh_credentials_key(path: Path) -> str: + """DEEPSEEK_API_KEY from dsh's managed `.credentials.yaml`. + + Only the one plain scalar is needed, so this reads the line rather than + pulling in a YAML parser: `refs:` holds it in the current layout, and the + pre-release layout kept it at the top level. + """ + try: + text = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError): + return "" + section = "" + for line in text.splitlines(): + if not line.strip() or line.lstrip().startswith("#"): + continue + if not line[0].isspace(): + section = line.split(":", 1)[0].strip() + if section == "DEEPSEEK_API_KEY": + return unquote(line.split(":", 1)[1]) + continue + if section == "refs": + name, _, value = line.strip().partition(":") + if name.strip() == "DEEPSEEK_API_KEY": + return unquote(value) + return "" + + +def read_dotenv_key(path: Path) -> str: + try: + text = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError): + return "" + for line in text.splitlines(): + line = line.strip() + if line.startswith("export "): + line = line[len("export "):].strip() + name, sep, value = line.partition("=") + if sep and name.strip() == "DEEPSEEK_API_KEY": + return unquote(value) + return "" + + +def api_key() -> str: + # The environment first, then the agents that keep a DeepSeek key, read + # only: pi and omp, opencode, and dsh's managed store with its `.env` + # fallback. + env_key = str(os.environ.get("DEEPSEEK_API_KEY", "")).strip() + if env_key: + return env_key + for read in ( + lambda: read_json_key(pi_agent_dir() / "auth.json", "deepseek"), + lambda: read_json_key(Path.home() / ".omp" / "agent" / "auth.json", "deepseek"), + lambda: read_json_key(data_home() / "opencode" / "auth.json", "deepseek"), + lambda: read_dsh_credentials_key(dsh_home() / ".credentials.yaml"), + lambda: read_dotenv_key(dsh_home() / ".env"), + ): + key = read() + if key: + return key + return "" + + +def read_config() -> dict[str, Any]: + return read_json(config_home() / "omarchy" / "agents" / "deepseek.json") + + +# ----------------------------------------------------------------- balance + +def representable(amount: Decimal) -> bool: + # Finite as a Decimal can still overflow the float the record prints. + return amount.is_finite() and math.isfinite(float(amount)) + + +def money(value: Any) -> Decimal: + if isinstance(value, bool): + return Decimal("0") + try: + amount = Decimal(str(value or "0")) + except InvalidOperation: + return Decimal("0") + return amount if representable(amount) else Decimal("0") + + +def fetch_balance(api_base_url: str, key: str) -> dict[str, Any]: + """GET /user/balance: total_balance, granted_balance, topped_up_balance.""" + url = api_base_url.rstrip("/") + "/user/balance" + # DEEPSEEK_API_BASE_URL moves the endpoint, never the key onto a + # connection anyone on the way can read. + if not trusted_endpoint(url): + raise DeepSeekError("DEEPSEEK_API_BASE_URL must use HTTPS (plain HTTP only for loopback).") + request = urllib.request.Request( + url, + headers={"Authorization": "Bearer " + key, "Accept": "application/json"}, + ) + try: + payload = open_json(request, timeout=15) + except HttpError as error: + if error.status in (401, 403): + raise DeepSeekError("DeepSeek rejected the saved API key. Set a current one in DEEPSEEK_API_KEY or the agent that holds it.", refused=True) from error + if error.status: + raise DeepSeekError(f"DeepSeek's balance endpoint returned HTTP {error.status}.") from error + if error.transport: + # No answer at all: no route, no DNS, a timeout. Worth a retry soon. + raise DeepSeekError("Couldn't reach the DeepSeek API. Retrying shortly.", transport=True) from error + 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: + raise DeepSeekError("DeepSeek's balance response carried no balance_infos.") + + # One entry per currency (CNY or USD). The one holding credit is the + # balance, so an empty CNY entry never hides a funded USD one; with none + # funded, the first. An entry without a readable total says nothing and is + # skipped, and an answer with none is a failed check that keeps the last + # good balance, never a balance of zero. + readable = [] + for info in infos: + if not isinstance(info, dict): + continue + total = info.get("total_balance") + try: + remaining = None if isinstance(total, bool) or total is None else Decimal(str(total).strip()) + except InvalidOperation: + remaining = None + if remaining is not None and representable(remaining): + readable.append((remaining, info)) + if not readable: + raise DeepSeekError("DeepSeek's balance response carried no total_balance.") + remaining, info = next((entry for entry in readable if entry[0] > 0), readable[0]) + + return { + "remaining": remaining, + "granted": money(info.get("granted_balance")), + "toppedUp": money(info.get("topped_up_balance")), + "currency": str(info.get("currency") or "CNY"), + # DeepSeek's own word on whether the balance still covers API calls. + "available": payload.get("is_available") is not False and remaining > 0, + } + + +def balance_record(remaining: Decimal, currency: str) -> dict[str, Any]: + config = read_config() + funded = money(config.get("fundedAmount")) + if funded < 0: + funded = Decimal("0") + 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(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, + } + + +# One cache per key and endpoint, so a machine signed in to another account +# never shows the last one's balance. +def balance_cache_path(key: str, api_base_url: str) -> Path: + return limits_cache_path(AGENT_ID, api_base_url.rstrip("/"), key) + + +def out_of_credits(result: dict[str, Any], available: Any) -> None: + # With no funded figure the balance gauge can't raise this alarm itself, + # so the status does. + if available is False: + result.update(usageStatusText="Out of credits", authHelpText=OUT_OF_CREDITS_HELP) + + +def collect_balance(key: str, api_base_url: str, force: bool) -> dict[str, Any]: + result: dict[str, Any] = {"balance": None, "stale": False, "fetchedAtMs": 0, "usageStatusText": "", "authHelpText": "", "retryAdvised": False} + cache = None + try: + cache = LimitsCache.at(balance_cache_path(key, api_base_url)) + except OSError as exc: + warn(f"could not read the balance cache ({exc})") + kept = obj(cache.data.get("balance")) if cache else {} + fetched_at = cache.fetched_at if cache else 0 + + # A panel opened and closed repeatedly must not become one request per + # flick, so a recent answer is reused, though a balance has no limit + # windows; --force skips the wait. + if kept and cache.fresh(force, allow_empty=True): + result.update(balance=balance_record(money(kept.get("remaining")), str(kept.get("currency") or "")), fetchedAtMs=fetched_at) + out_of_credits(result, kept.get("available")) + return result + + try: + live = fetch_balance(api_base_url, key) + except DeepSeekError as error: + result["retryAdvised"] = error.transport + result["authHelpText"] = str(error) + if kept: + # The last good balance, marked stale with when it was read, beats an + # empty section while DeepSeek can't be asked. + result.update(balance=balance_record(money(kept.get("remaining")), str(kept.get("currency") or "")), stale=True, fetchedAtMs=fetched_at) + out_of_credits(result, kept.get("available")) + result["authHelpText"] = str(error) + if error.refused: + result["usageStatusText"] = "Sign-in expired" + elif not kept: + result["usageStatusText"] = "Balance unavailable" + return result + + balance = {"remaining": str(live["remaining"]), "currency": live["currency"], "available": live["available"]} + fetched_at = cache.save([], only_if_newer=True, balance=balance) if cache is not None else now_ms() + result.update(balance=balance_record(live["remaining"], live["currency"]), fetchedAtMs=fetched_at) + out_of_credits(result, live["available"]) + return result + + +# ------------------------------------------------------------------ record + +# Whether a scan found any DeepSeek usage on this machine: a session, or +# tokens no session is known for. +def has_local_stats(stats: dict[str, Any]) -> bool: + return number(stats.get("totalSessions")) > 0 or bool(stats.get("modelUsage")) + + +def base_record(stats: dict[str, Any] | None = None) -> dict[str, Any]: + # Prompt and session counts are real: every pi session message carries + # them, so today's tooltip may show "N prompts · M sessions". + record = shared_base_record(AGENT_ID, AGENT_NAME) + if stats is not None: + record.update(stats, hasLocalStats=has_local_stats(stats)) + return record + + +def collect(api_base_url: str, max_age: float, force: bool) -> dict[str, Any]: + record = base_record(cached_local_stats(max_age)) + record["ready"] = record["hasLocalStats"] + + # 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 not key: + return record + record.update(ready=True, tierLabel="Prepaid") + balance = collect_balance(key, api_base_url, force) + if balance["balance"] is not None: + record["balance"] = balance["balance"] + record.update( + usageStatusText=balance["usageStatusText"], + authHelpText=balance["authHelpText"], + limitsStale=balance["stale"], + limitsFetchedAt=balance["fetchedAtMs"], + ) + if balance["retryAdvised"]: + record["retryAdvised"] = True + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the DeepSeek usage record as JSON") + # --force rescans everything and asks DeepSeek past the reuse window. + # --limits-only (the panel's quick refresh) may reuse a scan for 15 + # minutes, but asks for the balance like a normal run. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.add_argument("--api-base-url", default=os.environ.get("DEEPSEEK_API_BASE_URL", API_BASE_URL)) + args = parser.parse_args() + + try: + record = collect(args.api_base_url, scan_max_age(args.force, args.limits_only), args.force) + except Exception as error: + # A bug, not a state the panel should explain: an empty record keeps + # the panel steady, and stderr says what broke. + record = base_record() + warn(f"{type(error).__name__}: {error}") + emit(record) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-devin b/bin/omarchy-agent-usage-devin new file mode 100755 index 00000000000..701bdbe62d2 --- /dev/null +++ b/bin/omarchy-agent-usage-devin @@ -0,0 +1,497 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Devin usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Devin usage into one display-ready JSON record. + +Token stats come from the Devin CLI session database at +$XDG_DATA_HOME/devin/cli/sessions.db, where each assistant message carries +metadata.metrics with its input/output/cache token counts and the model that +generated it, and from pi/omp sessions and opencode's message database (both +of its stores) that ran on the devin provider. + +The daily and weekly quota comes from the user-status cache the Devin CLI +writes at $XDG_CACHE_HOME/devin/cli/user_status.*.bin every time it starts: +the GetUserStatus answer it fetched, base64 protobuf. The collector reads that +cache, the same data the CLI's /usage command renders, and makes no network +call of its own. The agents panel only ever reads the JSON this prints. +""" + +import argparse +import base64 +import json +import math +import os +import sqlite3 +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + Stats, + base_record, + cache_root, + cached_scan, + data_home, + emit, + local_day, + number, + obj, + open_limits, + opencode_db, + pi_session_roots, + read_json, + scan_max_age, + scan_opencode, + scan_pi_sessions, + token_count, + warn, + write_json, +) + +AGENT_ID = "devin" +AGENT_NAME = "Devin" +PROVIDER = "devin" +AUTH_HELP = "Run `devin auth login` so Devin can record its quota." +REFRESH_HELP = "Devin hasn't recorded its quota since it was signed in. Start Devin to refresh it." + +# The quota is a file read, so it is read on every run; only the local scan +# is cached. + +# The quota comes from a cache the CLI writes when it starts, never from a +# check of ours, and cloud sessions keep spending it meanwhile. Numbers this +# recent still describe the account; older ones are reported stale, with +# their age. +LIMITS_CURRENT_ENOUGH_SECONDS = 900 + +# The plan tiers GetUserStatus numbers, named as Devin names them. Only asked +# when the plan carries no name of its own. +TEAMS_TIER = { + 1: "Teams", 2: "Pro", 3: "Enterprise Saas", 4: "Hybrid", 5: "Enterprise Self Hosted", + 6: "Waitlist Pro", 7: "Teams Ultimate", 8: "Pro Ultimate", 9: "Trial", 10: "Enterprise Self Serve", + 11: "Enterprise Saas Pooled", 12: "Devin Enterprise", 14: "Devin Teams", 15: "Devin Teams V2", 16: "Devin Pro", + 17: "Devin Max", 18: "Max", 19: "Devin Free", 20: "Devin Trial", +} + +# The CLI's status cache dates its fetch in epoch seconds; past these bounds +# the stamp is garbage, not a time. +EARLIEST_EPOCH = 946_684_800 # 2000-01-01 +LATEST_EPOCH = 4_102_444_800 # 2100-01-01 + + +def epoch_seconds(value): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + # Stamps are seconds or milliseconds. + seconds = value / 1000 if value > 100_000_000_000 else value + if not EARLIEST_EPOCH <= seconds < LATEST_EPOCH: + return None + return seconds + + +def iso_from_unix(value): + seconds = epoch_seconds(value) + return datetime.fromtimestamp(seconds, timezone.utc).isoformat() if seconds else "" + + +def stored_day(value): + # The Devin CLI dates every row it writes, so a row without a time is + # unreadable rather than today's. + return None if value is None else local_day(value) + + +def model_name(raw): + return str(raw or "").strip() or AGENT_ID + + +# ------------------------------------------------------------- Devin CLI + + +def data_dir(): + return Path(os.environ.get("DEVIN_DATA_DIR") or (data_home() / "devin" / "cli")) + + +def credentials_path(): + return data_dir().parent / "credentials.toml" + + +def session_databases(): + """The CLI's own database, or failing that the older layout some builds + wrote straight under $XDG_DATA_HOME/devin. Never both: a machine that moved + to the newer layout may still have the old files, holding the same work.""" + primary = data_dir() / "sessions.db" + if primary.is_file(): + return [primary] + root = data_home() / "devin" + return [path for path in (root / "cli_sessions.db", root / "sessions.db") if path.is_file()] + + +def table_columns(conn, table): + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + + +def scan_message_nodes(conn, stats, prefix): + # Every assistant message records what it spent under metadata.metrics. + # The LIKE gate skips rows that cannot hold metrics before any parse. + for session_id, created_at, raw in conn.execute( + "SELECT session_id, created_at, chat_message FROM message_nodes WHERE chat_message LIKE '%\"metrics\"%'" + ): + try: + metadata = obj(obj(json.loads(raw)).get("metadata")) + except (TypeError, ValueError): + continue + metrics = obj(metadata.get("metrics")) + day = stored_day(created_at) + if not metrics or day is None: + continue + # The Devin CLI records what a reply spent and what a person asked in + # separate rows: prompts are counted from prompt_history, where each one + # is a row. + stats.add( + day, prefix + str(session_id), model_name(metadata.get("generation_model")), + token_count(metrics.get("input_tokens")), token_count(metrics.get("output_tokens")), + token_count(metrics.get("cache_read_tokens")), token_count(metrics.get("cache_creation_tokens")), + prompts=0, + ) + + +def scan_prompt_history(conn, stats, prefix): + for session_id, timestamp in conn.execute("SELECT session_id, timestamp FROM prompt_history"): + day = stored_day(timestamp) + if day is not None: + stats.add_prompts(day, prefix + str(session_id)) + + +def scan_session_totals(conn, stats, prefix, tables): + # The older layout keeps one lifetime total per session and no split by + # kind, so the tokens land on the day the session was last active and in + # inputTokens, the one bucket that cannot overstate a split it doesn't have. + for session_id, model, created_at, last_activity_at, total_tokens in conn.execute( + "SELECT id, model, created_at, last_activity_at, total_tokens FROM sessions" + ): + day = stored_day(last_activity_at) or stored_day(created_at) + if day is not None: + stats.add(day, prefix + str(session_id), model_name(model), token_count(total_tokens), 0, 0, 0, prompts=0) + if "messages" in tables and {"role", "timestamp", "session_id"} <= table_columns(conn, "messages"): + for session_id, role, timestamp in conn.execute("SELECT session_id, role, timestamp FROM messages WHERE role = 'assistant'"): + day = stored_day(timestamp) + if day is not None: + stats.add_prompts(day, prefix + str(session_id)) + + +def scan_devin_database(path, stats): + """Add one session database. Returns whether it was read to the end.""" + try: + # Read-only: Devin may be writing right now. + conn = sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return False + prefix = "devin:" + path.name + ":" + try: + conn.execute("PRAGMA query_only = ON") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + known = False + if "message_nodes" in tables and {"session_id", "created_at", "chat_message"} <= table_columns(conn, "message_nodes"): + known = True + scan_message_nodes(conn, stats, prefix) + if "prompt_history" in tables and {"session_id", "timestamp"} <= table_columns(conn, "prompt_history"): + known = True + scan_prompt_history(conn, stats, prefix) + if not known and "sessions" in tables and {"id", "model", "created_at", "last_activity_at", "total_tokens"} <= table_columns(conn, "sessions"): + known = True + scan_session_totals(conn, stats, prefix, tables) + if not known: + # A layout this collector doesn't know: reading it again won't help. + warn(f"{path} has no session tables this collector reads") + except sqlite3.Error as error: + # Transient lock, schema migration, corruption: the numbers stop here, + # incomplete. + warn(f"{path}: {error}") + return False + finally: + conn.close() + return True + + +# ------------------------------------------------------- pi, omp, opencode + + +# pi, omp, and opencode name the provider alike, by its exact id. +def billed_here(provider, model, entry): + return provider == PROVIDER + + +# -------------------------------------------------------------- scan cache + + +def scan_sources(): + return [Path.home(), data_dir(), data_home(), *pi_session_roots(), opencode_db()] + + +# The local stats and whether every source was read; a session database that +# couldn't be read is added to `unreadable`. +def run_local_scans(unreadable): + stats = Stats() + pi_rows, complete = scan_pi_sessions(r'"provider"\s*:\s*"devin"', billed_here) + for path in session_databases(): + if not scan_devin_database(path, stats): + complete = False + unreadable.append(path) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode([PROVIDER + '"'], billed_here) + for row in pi_rows + db_rows: + stats.add_row(row, model_name(row.model)) + return stats.snapshot(), complete and db_complete + + +# ------------------------------------------------------------------ quota +# +# GetUserStatus is protobuf, and the record needs a handful of its fields, so +# they are walked by hand rather than through generated code. In the cached +# UserStatus, field 10 is the plan tier and field 13 the PlanStatus; inside +# that, field 1 is the PlanInfo, 3 the plan's end, 5-7 the credits used, +# 14/15 the daily/weekly percent remaining and 17/18 their reset times. + + +def decode_varint(data, pos): + result = 0 + shift = 0 + while True: + byte = data[pos] + result |= (byte & 0x7F) << shift + pos += 1 + if not byte & 0x80: + return result, pos + shift += 7 + if shift > 63: + raise ValueError("varint too long") + + +def proto_fields(data): + """field number -> first value: an int for a varint, bytes for a message.""" + found = {} + pos = 0 + try: + while pos < len(data): + tag, pos = decode_varint(data, pos) + number_, wire = tag >> 3, tag & 0x07 + if wire == 0: + value, pos = decode_varint(data, pos) + elif wire == 1: + value, pos = None, pos + 8 + elif wire == 2: + length, pos = decode_varint(data, pos) + value, pos = bytes(data[pos:pos + length]), pos + length + elif wire == 5: + value, pos = None, pos + 4 + else: + break + if value is not None: + found.setdefault(number_, value) + except (IndexError, ValueError): + pass + return found + + +def message(fields, number_): + value = fields.get(number_) + return proto_fields(value) if isinstance(value, bytes) else {} + + +def varint(fields, number_): + value = fields.get(number_) + if not isinstance(value, int): + return None + return value - (1 << 64) if value >= (1 << 63) else value + + +def text(fields, number_): + value = fields.get(number_) + return value.decode("utf-8", errors="replace").strip() if isinstance(value, bytes) else "" + + +def quota_window(label, remaining, reset): + # The percent is what's left; the panel wants what's spent. Proto3 leaves a + # zero off the wire, so a window with nothing left carries no percent. + remaining = max(0, min(100, remaining or 0)) + return {"label": label, "percent": (100 - remaining) / 100.0, "resetsAt": iso_from_unix(reset)} + + +def parse_user_status(payload): + status = proto_fields(payload) + plan_status = message(status, 13) + plan = message(plan_status, 1) + tier = text(plan, 2) or TEAMS_TIER.get(varint(plan, 1)) or TEAMS_TIER.get(varint(status, 10)) or "" + + limits = [] + if plan_status: + # A credit plan leaves the quota percents at their defaults, so a window + # only shows when the server dates its reset or the plan bills by quota, + # and never when the plan hides it. + quota_plan = varint(plan, 35) == 2 + windows = (("Daily", 36, 14, 17), ("Weekly", 37, 15, 18)) + for label, hidden, remaining, reset in windows: + if varint(plan, hidden): + continue + if (varint(plan_status, reset) or 0) > 0 or quota_plan: + limits.append(quota_window(label, varint(plan_status, remaining), varint(plan_status, reset))) + + # Monthly credit buckets exist only on credit plans, and last until the + # plan's period ends; a quota plan reports no grant and is skipped. + plan_end = iso_from_unix(varint(message(plan_status, 3), 1)) + for label, limit_field, used_field in (("Prompt credits", 12, 6), ("Flow credits", 13, 5), ("Flex credits", 14, 7)): + limit = varint(plan, limit_field) or 0 + used = max(0, varint(plan_status, used_field) or 0) + if limit > 0: + limits.append({"label": label, "percent": min(1.0, used / limit), "resetsAt": plan_end}) + return {"tierLabel": tier, "limits": limits} + + +def user_status_dir(): + xdg_cache_home = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) + return Path(os.environ.get("DEVIN_CACHE_DIR") or (xdg_cache_home / "devin" / "cli")) + + +# The CLI rewrites credentials.toml on sign-in, so a status fetched before +# that write describes the account signed in before. The new account's first +# fetch can land a beat before the file is flushed, hence the allowance. +SIGN_IN_ALLOWANCE_SECONDS = 30 + + +def sign_in_time(): + try: + return credentials_path().stat().st_mtime + except OSError: + return None + + +def sign_in_state_path(): + return cache_root() / "devin-sign-in.json" + + +def sign_in_state(signed_in_at): + """Which status caches belonged to the sign-in before this one. + + The allowance alone can't tell the old account's last fetch from the new + one's first: both can land in the seconds before the rewrite. The cache + this collector served under the previous sign-in is known to be the old + account's, so it is remembered and kept out of that gap. + """ + state = read_json(sign_in_state_path()) + current = [name for name in state.get("current") or [] if isinstance(name, str)] + previous = [name for name in state.get("previous") or [] if isinstance(name, str)] + if state.get("signedInAt") != signed_in_at: + return {"signedInAt": signed_in_at, "current": [], "previous": current if "signedInAt" in state else []} + return {"signedInAt": signed_in_at, "current": current, "previous": previous} + + +def remember_served(state, name): + if name in state["current"]: + return + state["current"] = state["current"] + [name] + try: + write_json(sign_in_state_path(), state) + except OSError as exc: + warn(f"could not write the sign-in state ({exc})") + + +def read_quota(): + """The newest user-status cache fetched under the current sign-in, or None. + + The CLI keeps one cache per sign-in and no marker of which is active, so + the newest fetch wins, and fetches from before the last sign-in are dropped: + after a switch with no fetch since, nothing qualifies rather than the + previous account's numbers. Without a credentials file to date the sign-in, + the newest fetch stands. + """ + signed_in_at = sign_in_time() + cutoff = None if signed_in_at is None else signed_in_at - SIGN_IN_ALLOWANCE_SECONDS + state = None if signed_in_at is None else sign_in_state(signed_in_at) + best = None + try: + paths = sorted(user_status_dir().glob("user_status.*.bin")) + except OSError: + paths = [] + for path in paths: + cache = read_json(path) + if not isinstance(cache.get("payload"), str): + continue + fetched_at = epoch_seconds(cache.get("fetched_at_secs")) + if fetched_at is None or (cutoff is not None and fetched_at < cutoff): + continue + if state and fetched_at < signed_in_at and path.name in state["previous"]: + continue + try: + quota = parse_user_status(base64.b64decode(cache["payload"], validate=True)) + except ValueError: + continue + if not quota["limits"] and not quota["tierLabel"]: + continue + if best is None or fetched_at > best[0]: + best = (fetched_at, quota, path.name) + if best is None: + return None + fetched_at, quota, name = best + if state: + remember_served(state, name) + age = time.time() - fetched_at + return { + "tierLabel": quota["tierLabel"], + # A window whose reset has passed is untouched, not unknown: it comes back + # at zero with no reset time until the CLI fetches the next one. + "limits": open_limits(quota["limits"]), + # Floored, never rounded: a stamp half a millisecond ahead reads as a + # clock that went back. + "limitsFetchedAt": int(fetched_at * 1000), + "limitsStale": not 0 <= age < LIMITS_CURRENT_ENOUGH_SECONDS, + } + + +# ------------------------------------------------------------------ record + + +def build_record(stats, unreadable): + quota = read_quota() + signed_in = credentials_path().is_file() + in_use = number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")) + + # A sign-in with no sessions here has no local stats to show. + local = in_use or number(stats.get("totalSessions")) > 0 + record = base_record(AGENT_ID, AGENT_NAME, ready=in_use or signed_in or bool(quota and quota["limits"]), hasLocalStats=local) + if quota: + record.update(quota) + elif signed_in: + record["usageStatusText"] = "Devin limits unavailable" + record["authHelpText"] = REFRESH_HELP + elif in_use: + record["usageStatusText"] = "Waiting for auth" + record["authHelpText"] = AUTH_HELP + elif unreadable: + # Nothing else to say, and a database that can't be read is trouble the + # panel shows rather than a Devin it hides as unused. + record["ready"] = True + record["usageStatusText"] = "Devin unavailable" + record["authHelpText"] = "Couldn't read Devin's session database (" + "; ".join(map(str, unreadable)) + ")." + record.update(stats) + return record + + +def main(): + parser = argparse.ArgumentParser(description="Print the Devin usage record as JSON") + parser.add_argument("--force", action="store_true", help="rescan local sessions, ignoring the scan cache") + parser.add_argument("--limits-only", action="store_true", help="reuse a recent local scan") + args = parser.parse_args() + + # A scan served from the cache read every database. + unreadable = [] + stats = cached_scan(AGENT_ID, scan_sources(), scan_max_age(args.force, args.limits_only), lambda: run_local_scans(unreadable)) + emit(build_record(stats, unreadable), sort_keys=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-fireworks b/bin/omarchy-agent-usage-fireworks index fe04bc1df38..fa16c3d24f5 100755 --- a/bin/omarchy-agent-usage-fireworks +++ b/bin/omarchy-agent-usage-fireworks @@ -9,6 +9,9 @@ the last 30 days. Fireworks does not expose its prepaid ledger, so the record carries an estimated balance instead of rate limits: the credits configured in ~/.config/omarchy/agents/fireworks.json minus rated account costs since the funding date. The agents panel only ever reads the JSON this prints. + +Each API key keeps its account's last good usage rows and balance under the +agent-usage cache, so a failed check shows them stale rather than blank. """ from __future__ import annotations @@ -16,10 +19,10 @@ from __future__ import annotations import argparse import configparser import json +import math import os import re import sys -import urllib.error import urllib.parse import urllib.request from datetime import date, datetime, time, timedelta, timezone @@ -27,6 +30,25 @@ from decimal import Decimal, InvalidOperation from pathlib import Path from typing import Any +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record as shared_base_record, + cache_root, # noqa: F401 - the tests find the cached answers through it + config_home, + data_home, + emit, + number, + open_json, + read_json, + scan_max_age, + trusted_endpoint, +) + AGENT_ID = "fireworks" AGENT_NAME = "Fireworks" AUTH_HELP = "Set FIREWORKS_API_KEY, run `firectl set-api-key`, or sign in to Fireworks in pi or opencode." @@ -37,22 +59,38 @@ class FireworksError(Exception): pass -def number(value: Any) -> int: - try: - return max(0, round(float(value or 0))) - except (TypeError, ValueError): - return 0 +# No server answered at all (DNS, connection, timeout): the one failure worth +# a sooner retry. +class FireworksTransportError(FireworksError): + pass + + +# The key itself was refused, so only a new one helps. +class FireworksAuthError(FireworksError): + pass + + +# The key may not read this: a billing role it lacks. +class FireworksDeniedError(FireworksError): + pass -def money_value(value: Any) -> Decimal: +def representable(amount: Decimal) -> bool: + # Finite as a Decimal can still overflow the float the record prints. + return amount.is_finite() and math.isfinite(float(amount)) + + +def money_value(value: Any, unreadable: Decimal | None = Decimal("0")) -> Decimal | None: if not isinstance(value, dict): - return Decimal("0") + return unreadable try: units = Decimal(str(value.get("units", 0) or 0)) nanos = Decimal(str(value.get("nanos", 0) or 0)) / Decimal("1000000000") - return units + nanos + total = units + nanos except (InvalidOperation, TypeError, ValueError): - return Decimal("0") + return unreadable + # An amount the record can't print as a number is no amount. + return total if representable(total) else unreadable def model_id(row: dict[str, Any]) -> str: @@ -73,44 +111,21 @@ def row_date(row: dict[str, Any]) -> str: try: parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) except ValueError: - return raw[:10] if len(raw) >= 10 else "" + try: + return date.fromisoformat(raw[:10]).isoformat() + except ValueError: + return "" if parsed.tzinfo is None: parsed = parsed.replace(tzinfo=timezone.utc) - return parsed.astimezone().date().isoformat() - - -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": {}, - } + try: + return parsed.astimezone().date().isoformat() + except (OverflowError, ValueError): + # A time at the edge of the calendar has no local day. + return "" 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, + return shared_base_record(AGENT_ID, AGENT_NAME, **{ # Billing-API numbers are account-global, not machine-local: every synced # device reports the same truth, so aggregation must not sum them. "scope": "account", @@ -118,22 +133,16 @@ def base_record(**overrides: Any) -> dict[str, Any]: # panel keeps those numbers out of today's tooltip when this is false. "hasPromptStats": False, "tierLabel": "Prepaid", - "usageStatusText": "", - "authHelpText": "", - "limits": [], - } - record.update(empty_stats()) - record.update(overrides) - return record + **overrides, + }) +# The billing API's rows are account-wide day-and-model buckets, not replies: +# no prompt or session to count, a bucket without a day still counts toward +# its model, and `today` is the caller's. So the record carries only their +# token fields. def summarize_usage(payload: dict[str, Any], today: date | None = None) -> dict[str, Any]: - today = today or datetime.now().astimezone().date() - recent_dates = [(today - timedelta(days=offset)).isoformat() for offset in range(6, -1, -1)] - recent = {day: 0 for day in recent_dates} - today_by_model: dict[str, int] = {} - model_usage: dict[str, dict[str, int]] = {} - active_dates: set[str] = set() + stats = Stats(today) rows = payload.get("serverlessCosts") if not isinstance(rows, list): @@ -144,36 +153,18 @@ def summarize_usage(payload: dict[str, Any], today: date | None = None) -> dict[ continue day = row_date(raw_row) model = model_id(raw_row) + # Counts from the API are untrusted, and come as strings: number() reads + # a boolean, a negative, a non-finite or an absurdly large one as nothing. prompt = number(raw_row.get("promptTokens")) cached = min(prompt, number(raw_row.get("cachedPromptTokens"))) uncached = number(raw_row.get("uncachedPromptTokens")) if "uncachedPromptTokens" not in raw_row: uncached = max(0, prompt - cached) output = number(raw_row.get("completionTokens")) - total = uncached + cached + output - if total <= 0: - continue - - bucket = model_usage.setdefault(model, empty_bucket()) - bucket["inputTokens"] += uncached - bucket["outputTokens"] += output - bucket["cacheReadInputTokens"] += cached - - if day: - active_dates.add(day) - if day in recent: - recent[day] += total - if day == today.isoformat(): - today_by_model[model] = today_by_model.get(model, 0) + total + stats.add(day or None, None, model, uncached, output, cached, 0, prompts=0) - return { - "todayTotalTokens": sum(today_by_model.values()), - "todayTokensByModel": today_by_model, - "recentDays": [{"date": day, "messageCount": recent[day]} for day in recent_dates], - "activeDays": len(active_dates), - "activeDates": sorted(active_dates), - "modelUsage": model_usage, - } + snapshot = stats.snapshot() + return {key: snapshot[key] for key in ("todayTotalTokens", "todayTokensByModel", "recentDays", "activeDays", "activeDates", "modelUsage")} def read_auth_file(path: Path) -> tuple[str, str]: @@ -183,7 +174,7 @@ def read_auth_file(path: Path) -> tuple[str, str]: parser = configparser.ConfigParser(interpolation=None) try: parser.read(path) - except configparser.Error: + except (configparser.Error, UnicodeDecodeError): return "", "" api_key = "" @@ -202,8 +193,7 @@ def pi_auth_path() -> Path: 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" + return data_home() / "opencode" / "auth.json" PI_ENV_REFERENCE = re.compile(r"\$(?:\{([A-Za-z_][A-Za-z0-9_]*)\}|([A-Za-z_][A-Za-z0-9_]*))") @@ -235,11 +225,7 @@ def resolve_pi_key(raw: Any) -> str: def read_pi_key(path: Path) -> str: - try: - parsed = json.loads(path.read_text()) - except (OSError, json.JSONDecodeError): - return "" - entry = parsed.get("fireworks") if isinstance(parsed, dict) else None + entry = read_json(path).get("fireworks") if not isinstance(entry, dict): return "" if entry.get("type") not in (None, "api_key"): @@ -248,27 +234,18 @@ def read_pi_key(path: Path) -> str: def read_opencode_key(path: Path) -> str: - try: - parsed = json.loads(path.read_text()) - except (OSError, json.JSONDecodeError): - return "" - entry = parsed.get("fireworks-ai") if isinstance(parsed, dict) else None + entry = read_json(path).get("fireworks-ai") if not isinstance(entry, dict): return "" return str(entry.get("key") or "").strip() def config_path() -> Path: - config_home = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) - return config_home / "omarchy" / "agents" / "fireworks.json" + return config_home() / "omarchy" / "agents" / "fireworks.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 {} + return read_json(config_path()) def credentials(auth_path: Path, config: dict[str, Any]) -> tuple[str, str]: @@ -343,6 +320,10 @@ class FireworksClient: body: dict[str, Any] | None = None, ) -> dict[str, Any]: url = self.base_url + path + # FIREWORKS_API_BASE_URL moves the endpoint, never the key onto a + # connection anyone on the way can read. + if not trusted_endpoint(url): + raise FireworksError("FIREWORKS_API_BASE_URL must use HTTPS (plain HTTP only for loopback)") if query: url += "?" + urllib.parse.urlencode(query, doseq=True) data = None if body is None else json.dumps(body).encode("utf-8") @@ -357,25 +338,31 @@ class FireworksClient: }, ) try: - with urllib.request.urlopen(request, timeout=15) as response: - decoded = json.load(response) - return decoded if isinstance(decoded, dict) else {} - except urllib.error.HTTPError as error: - if error.code == 401: - raise FireworksError("Fireworks rejected the API key") - if error.code == 403: - raise FireworksError("The Fireworks API key cannot read billing data") - if error.code == 404: - raise FireworksError("Fireworks account not found") - raise FireworksError(f"Fireworks API returned HTTP {error.code}") - except urllib.error.URLError as error: - raise FireworksError("Could not reach the Fireworks API") from error - except (json.JSONDecodeError, TimeoutError) as error: + decoded = open_json(request, timeout=15) + except HttpError as error: + if error.status == 401: + raise FireworksAuthError("Fireworks rejected the API key") from error + if error.status == 403: + raise FireworksDeniedError("The Fireworks API key cannot read billing data") from error + if error.status == 404: + raise FireworksError("Fireworks account not found") from error + if error.status: + raise FireworksError(f"Fireworks API returned HTTP {error.status}") from error + if error.transport: + raise FireworksTransportError("Could not reach the Fireworks API") from error raise FireworksError("Fireworks returned an invalid billing response") from error + # Every answer is an object; anything else would read as an account with + # nothing in it and replace the last good numbers. + if not isinstance(decoded, dict): + raise FireworksError("Fireworks returned an invalid billing response") + return decoded def discover_account(self) -> tuple[str, dict[str, Any]]: payload = self.request("/v1/accounts", query={"pageSize": 100}) - accounts = [item for item in payload.get("accounts", []) if isinstance(item, dict)] + accounts = payload.get("accounts", []) + if not isinstance(accounts, list): + raise FireworksError("Fireworks returned an invalid billing response") + accounts = [item for item in accounts if isinstance(item, dict)] if len(accounts) == 1: account = accounts[0] return normalize_account_id(str(account.get("name") or "")), account @@ -401,9 +388,12 @@ class FireworksClient: rows: list[Any] = [] for _ in range(20): payload = self.request(f"/v1/accounts/{quoted}/billingUsage", query=query) - page = payload.get("serverlessCosts") - if isinstance(page, list): - rows.extend(page) + # The API leaves out a list with nothing in it, so a missing one is an + # idle month; one that is there but isn't a list is a broken answer. + page = payload.get("serverlessCosts", []) + if not isinstance(page, list): + raise FireworksError("Fireworks returned an invalid billing response") + rows.extend(page) token = str(payload.get("nextPageToken") or "") if not token: break @@ -419,9 +409,10 @@ class FireworksClient: } try: payload = self.request(f"/v1/accounts/{quoted}/usageCosts:query", body=body) - if not isinstance(payload.get("subtotal"), dict): + subtotal = money_value(payload.get("subtotal"), None) + if subtotal is None: raise FireworksError("Fireworks cost response did not include a subtotal") - return money_value(payload.get("subtotal")) + return subtotal except FireworksError: parsed_end = datetime.fromisoformat(end_at.replace("Z", "+00:00")) summary_end = (parsed_end.date() + timedelta(days=1)).isoformat() + "T00:00:00Z" @@ -429,10 +420,10 @@ class FireworksClient: f"/v1/accounts/{quoted}/billing/summary", query={"startTime": start_at, "endTime": summary_end}, ) - return sum( - (money_value(item.get("totalCost")) for item in payload.get("lineItems", []) if isinstance(item, dict)), - Decimal("0"), - ) + items = payload.get("lineItems", []) + if not isinstance(items, list): + raise FireworksError("Fireworks returned an invalid billing response") + return sum((money_value(item.get("totalCost")) for item in items if isinstance(item, dict)), Decimal("0")) def live_balance(client: FireworksClient, account_id: str) -> Decimal | None: @@ -441,14 +432,20 @@ def live_balance(client: FireworksClient, account_id: str) -> Decimal | None: # is the best we can do. The response shape is undocumented, so accept a # Money object at the top level or under any plausible field name. quoted = urllib.parse.quote(normalize_account_id(account_id), safe="") + # Any other failure means the balance couldn't be checked at all, which a + # missing estimate must not paper over as "no balance": it propagates, and + # the last good balance stands. try: payload = client.request(f"/v1/accounts/{quoted}:getBalance") - except FireworksError: + except FireworksDeniedError: return None candidates = [payload] + [payload.get(field) for field in ("balance", "creditBalance", "prepaidBalance", "amount")] for value in candidates: if isinstance(value, dict) and ("units" in value or "nanos" in value): - return money_value(value) + amount = money_value(value, None) + if amount is None: + raise FireworksError("Fireworks returned an unreadable balance") + return amount return None @@ -462,7 +459,7 @@ def estimated_balance( funded = Decimal(str(config.get("fundedAmount") or "0")) except InvalidOperation: raise FireworksError("Fireworks fundedAmount must be a number") - if not funded.is_finite(): + if not representable(funded): raise FireworksError("Fireworks fundedAmount must be a finite number") if funded <= 0: return None @@ -486,71 +483,162 @@ def estimated_balance( } -def scan(api_base_url: str, auth_path: Path) -> dict[str, Any]: +def current_balance( + client: FireworksClient, + account_id: str, + account: dict[str, Any], + config: dict[str, Any], +) -> dict[str, Any] | None: + live = live_balance(client, account_id) + if live is None: + return estimated_balance(client, account_id, account, config) + try: + funded = Decimal(str(config.get("fundedAmount") or "0")) + if not representable(funded) or funded < 0: + funded = Decimal("0") + except InvalidOperation: + funded = Decimal("0") + # An overdrawn account has nothing left, and spent no more than it funded. + live = max(Decimal("0"), live) + return { + "remaining": float(live), + "funded": float(funded), + "spent": float(max(Decimal("0"), funded - live)), + "currency": "USD", + "estimated": False, + } + + +# ---------------------------------------------------------------- caches + + +# Both answers are keyed by the API key: the balance under the limits cache +# name every collector shares, the usage rows beside it. Neither holds limit +# windows, so either stands in for a probe on its stamp alone. +def caches(api_key: str) -> tuple[LimitsCache, LimitsCache]: + return LimitsCache(AGENT_ID, api_key, kind="usage"), LimitsCache(AGENT_ID, api_key) + + +# The last good answer of one kind (the account's usage rows, or its balance) +# for this API key and account, whatever its file time: a clock that went +# back doesn't lose it. An answer for another account never stands in. +def cached_part(cache: LimitsCache, account_id: str) -> dict[str, Any] | None: + if not cache.fetched_at: + return None + if account_id and normalize_account_id(str(cache.data.get("accountId") or "")) != account_id: + return None + return dict(cache.data, fetchedAtMs=cache.fetched_at) + + +# Refreshes can overlap (the panel's own and a sooner retry), so an answer +# only replaces one fetched before it. +def remember(cache: LimitsCache, **part: Any) -> dict[str, Any]: + cache.save([], only_if_newer=True, **part) + return dict(cache.data) + + +# ---------------------------------------------------------------- record + + +def scan(api_base_url: str, auth_path: Path, force: bool = False, limits_only: bool = False) -> dict[str, Any]: config = read_config() api_key, account_id = credentials(auth_path, config) + # No key anywhere: Fireworks isn't used here, and the panel skips it. if not api_key: - return base_record(usageStatusText="Fireworks unavailable", authHelpText=AUTH_HELP) - - client = FireworksClient(api_key, api_base_url) - account: dict[str, Any] = {} - if account_id: - account_id = normalize_account_id(account_id) - else: - account_id, account = client.discover_account() - + return base_record(tierLabel="") + + account_id = normalize_account_id(account_id) if account_id else "" + usage_cache, balance_cache = caches(api_key) + usage_part = cached_part(usage_cache, account_id) + balance_part = cached_part(balance_cache, account_id) + + # The usage rows are this record's stats, so they are reused like a local + # scan: briefly on a normal run, for 15 minutes on the panel's near-limit + # refresh. That refresh is about the balance, which comes from separate + # calls, so it is probed like a normal run. --force reuses nothing. + usage_fresh = usage_part is not None and usage_cache.fresh(force, scan_max_age(force, limits_only), allow_empty=True) + balance_fresh = balance_part is not None and balance_cache.fresh(force, allow_empty=True) + + errors: list[FireworksError] = [] + failed: set[str] = set() today = datetime.now().astimezone().date() - usage = client.usage(account_id, today - timedelta(days=29), today + timedelta(days=1)) - record = base_record(ready=True, hasLocalStats=True) - record.update(summarize_usage(usage, today)) - - live = live_balance(client, account_id) - if live is not None: + if not (usage_fresh and balance_fresh): + client = FireworksClient(api_key, api_base_url) + account: dict[str, Any] = {} try: - funded = Decimal(str(config.get("fundedAmount") or "0")) - if not funded.is_finite() or funded < 0: - funded = Decimal("0") - except InvalidOperation: - funded = Decimal("0") - record["balance"] = { - "remaining": float(live), - "funded": float(funded), - "spent": float(max(Decimal("0"), funded - live)), - "currency": "USD", - "estimated": False, - } - return record - - try: - balance = estimated_balance(client, account_id, account, config) - if balance: - record["balance"] = balance - except FireworksError as error: + if not account_id: + account_id, account = client.discover_account() + # A key that now reaches another account starts over. + if usage_part and normalize_account_id(str(usage_part.get("accountId") or "")) != account_id: + usage_part, usage_fresh = None, False + if balance_part and normalize_account_id(str(balance_part.get("accountId") or "")) != account_id: + balance_part, balance_fresh = None, False + except FireworksError as error: + errors.append(error) + failed.update(name for name, fresh in (("usage", usage_fresh), ("balance", balance_fresh)) if not fresh) + else: + if not usage_fresh: + try: + usage = client.usage(account_id, today - timedelta(days=29), today + timedelta(days=1)) + usage_part = remember(usage_cache, accountId=account_id, serverlessCosts=usage.get("serverlessCosts") or []) + except FireworksError as error: + errors.append(error) + failed.add("usage") + if not balance_fresh: + try: + balance_part = remember(balance_cache, accountId=account_id, balance=current_balance(client, account_id, account, config)) + except FireworksError as error: + errors.append(error) + failed.add("balance") + + record = base_record(ready=True) + # Kept rows are summarized against today, so a day-old answer still files + # its tokens under the right days. + if usage_part is not None: + rows = usage_part.get("serverlessCosts") + record.update(summarize_usage({"serverlessCosts": rows if isinstance(rows, list) else []}, today)) + record["hasLocalStats"] = True + if balance_part is not None and isinstance(balance_part.get("balance"), dict): + record["balance"] = balance_part["balance"] + kept = [part for part in (usage_part, balance_part) if part is not None] + if kept: + record["limitsFetchedAt"] = min(part["fetchedAtMs"] for part in kept) + # Anything that didn't come back fresh marks the record stale, kept or not; + # the panel shows the kept numbers' age on hover. + record["limitsStale"] = bool(failed) + if any(isinstance(error, FireworksTransportError) for error in errors): + # Usually the seconds after login before the network is up; the panel + # tries once more soon instead of waiting out its refresh interval. + record["retryAdvised"] = True + + if any(isinstance(error, FireworksAuthError) for error in errors): + record["usageStatusText"] = "Sign-in expired" + record["authHelpText"] = "Fireworks rejected the API key. " + AUTH_HELP + elif "usage" in failed and usage_part is None: + record["usageStatusText"] = "Fireworks unavailable" + record["authHelpText"] = str(errors[0]) + elif "balance" in failed and balance_part is None: record["usageStatusText"] = "Balance unavailable" - record["authHelpText"] = str(error) - + record["authHelpText"] = str(errors[0]) + elif errors: + record["authHelpText"] = str(errors[0]) return record def main() -> int: parser = argparse.ArgumentParser(description="Print the Fireworks usage record as JSON") - # Stats and balance come from the same few API calls, so there is no cache - # to force past and no faster limits-only path. The flags exist so every - # collector accepts the same invocation. - parser.add_argument("--force", action="store_true") - parser.add_argument("--limits-only", action="store_true") + parser.add_argument("--force", action="store_true", help="re-fetch usage and balance, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse any recent usage answer; only the balance must be fresh") parser.add_argument("--auth-path", default=os.environ.get("FIREWORKS_AUTH_PATH", "~/.fireworks/auth.ini")) parser.add_argument("--api-base-url", default=os.environ.get("FIREWORKS_API_BASE_URL", API_BASE_URL)) args = parser.parse_args() try: - record = scan(args.api_base_url, Path(args.auth_path).expanduser()) - except FireworksError as error: - record = base_record(usageStatusText="Fireworks unavailable", authHelpText=str(error)) + record = scan(args.api_base_url, Path(args.auth_path).expanduser(), args.force, args.limits_only) except Exception as error: record = base_record(usageStatusText="Fireworks unavailable", authHelpText="Fireworks usage scan failed") print(f"omarchy-agent-usage-fireworks: {type(error).__name__}", file=sys.stderr) - print(json.dumps(record, separators=(",", ":"))) + emit(record) return 0 diff --git a/bin/omarchy-agent-usage-grok b/bin/omarchy-agent-usage-grok index f2d4d7fe340..f001e21fcda 100755 --- a/bin/omarchy-agent-usage-grok +++ b/bin/omarchy-agent-usage-grok @@ -14,53 +14,45 @@ by model. The agents panel only ever reads the JSON this prints. import argparse import json +import math import os import re -import tempfile +import sys import time -import urllib.error import urllib.request from datetime import datetime, timedelta, timezone from pathlib import Path +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + account_entry, + base_record, + cached_scan, + emit, + empty_bucket, + local_day, + obj, + open_json, + read_json, + registered_accounts, + scan_max_age, + token_count, +) + AGENT_ID = "grok" AGENT_NAME = "Grok" AUTH_HELP = "Start Grok, or run `grok login`, to sign in." CREDITS_URL = "https://cli-chat-proxy.grok.com/v1/billing?format=credits" -# Grok's own CLI watches its subscription once a minute; a panel opened and -# shut repeatedly reuses an answer for about that long. -PROBE_REUSE_SECONDS = 60 +# A unified-billing account counts its usage against a monthly quota, which +# the same endpoint answers without the credits format. +MONTHLY_URL = "https://cli-chat-proxy.grok.com/v1/billing" # Grok's docs give a saved sign-in 30 days; a refresh token older than that # can't renew the access token, so it no longer counts as signed in. REFRESH_LIFETIME_SECONDS = 30 * 24 * 3600 -# A limits-only refresh needs fresh limits, not a fresh history scan. -LIMITS_ONLY_SCAN_REUSE_SECONDS = 900 - - -def read_json(path): - try: - return json.loads(Path(path).read_text(encoding="utf-8")) - except Exception: - return None - - -# A temporary file of its own per write, so overlapping runs can't rename -# each other's away. -def write_json(path, payload): - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - fd, tmp = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent)) - try: - with os.fdopen(fd, "w", encoding="utf-8") as handle: - json.dump(payload, handle) - os.replace(tmp, path) - except Exception: - Path(tmp).unlink(missing_ok=True) - raise - - -def cache_root(): - return Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" def parse_time(value): @@ -77,12 +69,30 @@ def default_home(): return Path(os.environ.get("GROK_HOME") or (Path.home() / ".grok")) +# The local day a session or turn was active on, or None to skip it: a +# missing time is no activity here, and one local_day() can't trust (out of +# range, before 2000 or after 2100) counts nowhere. A time without a zone is +# UTC, as parse_time() reads it. +def active_day(value): + if value is None: + return None + parsed = parse_time(value) if isinstance(value, str) else None + return local_day(parsed.isoformat() if parsed else value) + + # Every session under /sessions///: its summary says -# when it was last active and how many prompts it had, and its usage.json -# lists each finished turn with its tokens by model. Added accounts link -# their sessions to the primary's, so this covers all of them. Grok counts -# cached input inside inputTokens; the record keeps the two apart, as the -# other collectors do. +# when it was last active, and its usage.json lists each finished turn with +# its tokens by model, one turn per typed prompt. Added accounts link their +# sessions to the primary's, so this covers all of them. Grok counts cached +# input inside inputTokens; the record keeps the two apart, as the other +# collectors do. Counts are read as every collector reads them: one that +# isn't a plain number is no count. +# +# A fork's usage.json repeats its source's turns verbatim before its own, so +# a turn already seen (same end time, same tokens) counts once. A subagent's +# turns are tokens it spent, but its prompts came from the parent agent, not +# the person typing. The summary's num_messages counts system and assistant +# messages, inherited ones too, so it is no prompt count. def local_stats(home): now = datetime.now().astimezone() today = now.date().isoformat() @@ -92,14 +102,15 @@ def local_stats(home): by_day = {day: 0 for day in week} today_by_model = {} models = {} + seen = set() for session in (Path(home) / "sessions").glob("*/*/"): summary = read_json(session / "summary.json") + subagent = False if isinstance(summary, dict): - active = parse_time(summary.get("last_active_at")) - if active: - day = active.astimezone().date().isoformat() + subagent = str(summary.get("session_kind") or "").startswith("subagent") + day = active_day(summary.get("last_active_at")) + if day: sessions += 1 - prompts += int(summary.get("num_messages") or 0) days.add(day) if day == today: today_sessions += 1 @@ -107,24 +118,29 @@ def local_stats(home): usage = read_json(session / "usage.json") turns = usage.get("turns") if isinstance(usage, dict) else None for turn in turns if isinstance(turns, list) else []: - ended = parse_time(turn.get("endedAt")) if isinstance(turn, dict) else None - if not ended: + day = active_day(turn.get("endedAt")) if isinstance(turn, dict) else None + if not day: continue - day = ended.astimezone().date().isoformat() + copy = json.dumps([turn.get("endedAt"), turn.get("modelUsage")], sort_keys=True, default=str) + if copy in seen: + continue + seen.add(copy) days.add(day) - if day == today: - today_prompts += 1 - for model, bucket in (turn.get("modelUsage") or {}).items(): + if not subagent: + prompts += 1 + if day == today: + today_prompts += 1 + for model, bucket in obj(turn.get("modelUsage")).items(): if not isinstance(bucket, dict): continue # Grok's totalTokens leaves cache writes out; counting them here keeps # the day's total equal to the sum of its model's buckets. - cached = int(bucket.get("cachedReadTokens") or 0) - uncached = max(0, int(bucket.get("inputTokens") or 0) - cached) - output = int(bucket.get("outputTokens") or 0) - created = int(bucket.get("cacheCreationTokens") or 0) + cached = token_count(bucket.get("cachedReadTokens")) + uncached = max(0, token_count(bucket.get("inputTokens")) - cached) + output = token_count(bucket.get("outputTokens")) + created = token_count(bucket.get("cacheCreationTokens")) tokens = uncached + cached + output + created - totals = models.setdefault(model, {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}) + totals = models.setdefault(model, empty_bucket()) totals["inputTokens"] += uncached totals["outputTokens"] += output totals["cacheReadInputTokens"] += cached @@ -150,16 +166,11 @@ def local_stats(home): } -# A scan from before midnight would count yesterday's sessions as today's. +# Cached per home and per day, and rebuilt under a lock, so overlapping runs +# (the update command backgrounds one per agent while the panel refreshes on +# its own) scan once between them. def cached_local_stats(home, max_age): - cache = cache_root() / "grok-stats.json" - cached = read_json(cache) or {} - today = datetime.now().astimezone().date().isoformat() - if cached.get("home") == str(home) and cached.get("day") == today and time.time() - float(cached.get("at") or 0) < max_age: - return cached["stats"] - stats = local_stats(home) - write_json(cache, {"home": str(home), "day": today, "at": time.time(), "stats": stats}) - return stats + return cached_scan(AGENT_ID, [home], max_age, lambda: (local_stats(home), True)) # One login per issuer in auth.json; the first is the CLI's own. @@ -171,42 +182,64 @@ def login(home): def plan(home): - cached = read_json(Path(home) / "settings_cache.json") or {} + cached = read_json(Path(home) / "settings_cache.json") try: - settings = json.loads(cached.get("payload") or "{}").get("settings") or {} + settings = obj(obj(json.loads(cached.get("payload") or "{}")).get("settings")) except Exception: return "" return str(settings.get("subscription_tier_display") or "") +# A percentage or amount from xAI's answer: a number, or a money field's +# {"val": n}, whose val protobuf JSON leaves out when it's zero. +def amount(value): + if isinstance(value, dict): + value = value.get("val") or 0 + if isinstance(value, bool): + raise ValueError("not a number") + value = float(value) + # NaN would pass any clamp as 100%; a number that isn't one is no answer. + if not math.isfinite(value): + raise ValueError("not a number") + return value + + # The credits answer as one limit window: how much of this period's included # usage is gone, and when the period ends. The answer is protobuf JSON, which # leaves out a field holding zero, so a period with nothing used yet has no -# percentage at all. +# percentage at all. A unified-billing account leaves it out too, its usage +# counted against the monthly quota instead, so it has no window here (None); +# and a period already over without one says nothing about this one. def credits_limit(data): config = data["config"] period = config["currentPeriod"] + ends = parse_time(period["end"]) + used = config.get("creditUsagePercent") + if used is None: + if config.get("isUnifiedBillingUser") is True: + return None + if ends.timestamp() <= time.time(): + raise ValueError("an ended period without a percentage") + used = 0 return { "label": "Weekly" if "WEEK" in str(period.get("type")) else "Monthly", - "percent": max(0.0, min(1.0, float(config.get("creditUsagePercent") or 0) / 100)), - "resetsAt": parse_time(period["end"]).isoformat(), + "percent": max(0.0, min(1.0, amount(used) / 100)), + "resetsAt": ends.isoformat(), } -# A window whose reset has passed starts over empty. A weekly one's next -# reset is a whole number of weeks later; a monthly one's isn't known. -def rolled_forward(limits): - now = datetime.now(timezone.utc) - rolled = [] - for limit in limits: - resets = parse_time(limit.get("resetsAt")) - if resets and resets <= now: - limit = dict(limit, percent=0.0, resetsAt="") - if limit.get("label") == "Weekly": - weeks = (now - resets) // timedelta(weeks=1) + 1 - limit["resetsAt"] = (resets + weeks * timedelta(weeks=1)).isoformat() - rolled.append(limit) - return rolled +# The monthly answer as one limit window: how much of the month's included +# quota is used, and when the billing period ends. +def monthly_limit(data): + config = data["config"] + quota = amount(config["monthlyLimit"]) + if quota <= 0: + raise ValueError("no monthly quota") + return { + "label": "Monthly", + "percent": max(0.0, min(1.0, amount(config["used"]) / quota)), + "resetsAt": parse_time(config["billingPeriodEnd"]).isoformat(), + } def collect_limits(home, key, force): @@ -217,74 +250,81 @@ def collect_limits(home, key, force): result["usageStatusText"] = "Waiting for auth" return result - cache = cache_root() / f"grok-limits-{key}.json" - cached = read_json(cache) or {} - fetched_at = float(cached.get("fetchedAtMs") or 0) - result["fetchedAtMs"] = fetched_at + # A kept window whose reset has passed starts over empty. Its next reset is + # xAI's to say, so it stays unknown until a probe answers. + cache = LimitsCache(AGENT_ID, key) + fallback, result["fetchedAtMs"] = cache.kept() expires = parse_time(entry.get("expires_at")) if expires and expires.timestamp() <= time.time(): - result["limits"] = cached.get("limits") or [] + result["limits"] = fallback # The access token lives six hours and Grok mints a new one from its # refresh token whenever it starts, so a lapsed one is routine, not a - # sign-out. And while Grok hasn't run, nothing here has spent its - # allowance, so the last numbers still stand, rolled past any reset. + # sign-out: the limits are paused, the last numbers kept as stale, until + # Grok runs again. Only a refresh token past its 30 days is a sign-in to + # redo. if entry.get("refresh_token") and time.time() - expires.timestamp() < REFRESH_LIFETIME_SECONDS: - result.update(limits=rolled_forward(result["limits"]), live=True) - if not result["limits"]: - result["usageStatusText"] = "Limits paused" - result["authHelpText"] = "Start Grok to see its limits; signing in again is not needed." + result["usageStatusText"] = "Limits paused" + result["authHelpText"] = ( + "Grok's access token has lapsed" + + (" — showing the last known limits." if fallback else ".") + + " Start Grok to refresh it; signing in again is not needed." + ) else: result["usageStatusText"] = "Sign-in expired" result["authHelpText"] = "Grok's saved sign-in expired. Start Grok to refresh it." return result - if cached.get("limits") and not force and time.time() - fetched_at / 1000 < PROBE_REUSE_SECONDS: - result.update(limits=cached["limits"], live=True, authHelpText="") + # A panel opened and shut repeatedly reuses an answer this recent rather + # than asking xAI again. + if cache.fresh(force): + result.update(limits=fallback, live=True, authHelpText="") return result - request = urllib.request.Request(CREDITS_URL, headers={"Authorization": f"Bearer {token}", "Accept": "application/json"}) + headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"} + limits = [] try: - with urllib.request.urlopen(request, timeout=8) as response: - limit = credits_limit(json.loads(response.read())) - except urllib.error.HTTPError as error: - limit = None - if error.code in (401, 403): - result["usageStatusText"] = "Waiting for auth" + credits = open_json(urllib.request.Request(CREDITS_URL, headers=headers), timeout=8) + weekly = credits_limit(credits) + if weekly is None: + limits = [monthly_limit(open_json(urllib.request.Request(MONTHLY_URL, headers=headers), timeout=8))] + else: + limits = [weekly] + # A unified-billing account can carry both: the weekly credits and the + # monthly included quota, which Grok's own /usage shows side by side. + # The weekly window stands on its own when the monthly one won't come. + if credits["config"].get("isUnifiedBillingUser") is True and weekly["label"] == "Weekly": + try: + limits.append(monthly_limit(open_json(urllib.request.Request(MONTHLY_URL, headers=headers), timeout=8))) + except Exception: + pass + except HttpError as error: + if error.status in (401, 403): + # The token was turned away before it lapsed: a sign-in to redo. + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = "xAI refused Grok's saved sign-in. Start Grok, or run `grok login`, to sign in again." + elif error.status: + result["authHelpText"] = f"Grok's usage endpoint returned status {error.status}. Local Grok stats are still shown." + elif error.transport: + # No server answered at all (no route, no DNS, a timeout), often a + # network not up yet after login, so the panel asks again soon. Any real + # answer, an error status included, is a server to stop pestering. + result["retryAdvised"] = True + result["authHelpText"] = "Couldn't reach Grok's usage endpoint. Retrying shortly. Local Grok stats are still shown." + else: + result["authHelpText"] = "Grok's usage endpoint answered without a usable limit. Local Grok stats are still shown." except Exception: - limit = None + result["authHelpText"] = "Grok's usage endpoint answered without a usable limit. Local Grok stats are still shown." - if limit: - result.update(limits=[limit], live=True, fetchedAtMs=round(time.time() * 1000), authHelpText="") - write_json(cache, {"fetchedAtMs": result["fetchedAtMs"], "limits": result["limits"]}) + if limits: + result.update(limits=limits, live=True, fetchedAtMs=cache.save(limits), authHelpText="") else: - result["limits"] = cached.get("limits") or [] + result["limits"] = fallback if not result["usageStatusText"] and not result["limits"]: result["usageStatusText"] = "Grok limits unavailable" return result -# The accounts `omarchy agent account` registered, read straight from its -# registry. Only a registry holding a second account matters: with one, the -# record describes ~/.grok alone. -def registered_accounts(): - state = Path(os.environ.get("XDG_STATE_HOME") or (Path.home() / ".local" / "state")) - registry = read_json(state / "omarchy" / "agents" / "accounts" / "grok.json") or {} - accounts = [a for a in registry.get("accounts") or [] if isinstance(a, dict) and a.get("id")] - if len(accounts) < 2: - return [] - active_id = registry.get("active") or accounts[0]["id"] - if not any(a["id"] == active_id for a in accounts): - active_id = accounts[0]["id"] - for account in accounts: - account["active"] = account["id"] == active_id - account["switch"] = { - "mode": "auto" if registry.get("switch") == "auto" else "manual", - "threshold": registry.get("threshold") or 95, - } - return accounts - - def home_of(account): return Path(account.get("home") or (Path.home() / ".grok")) @@ -299,19 +339,7 @@ def cache_key(home, fallback): def account_limits(account, force): home = home_of(account) limits = collect_limits(home, cache_key(home, account["id"]), force) - return { - "id": account["id"], - "label": str(account.get("label") or account["id"]), - "email": str(account.get("email") or ""), - "plan": plan(home) or str(account.get("plan") or ""), - "active": account["active"], - "primary": bool(account.get("primary")), - "limits": limits["limits"], - "stale": not limits["live"], - "fetchedAt": limits["fetchedAtMs"], - "usageStatusText": limits["usageStatusText"], - "authHelpText": limits["authHelpText"], - } + return account_entry(account, limits, plan(home)) def main(): @@ -320,7 +348,7 @@ def main(): parser.add_argument("--limits-only", action="store_true") args = parser.parse_args() - registered = registered_accounts() + registered = registered_accounts(AGENT_ID) accounts = [account_limits(account, args.force) for account in registered] current = next((a for a in accounts if a["active"]), None) if current: @@ -331,30 +359,37 @@ def main(): home = default_home() tier = plan(home) limits = collect_limits(home, cache_key(home, "main"), args.force) - stale, fetched_at = not limits["live"], limits["fetchedAtMs"] + stale, fetched_at = not limits["live"] and bool(limits["limits"]), limits["fetchedAtMs"] - # Nobody signed in to Grok here: an empty record, which the panel skips. # With accounts registered, the primary is ~/.grok by definition. stats_home = Path.home() / ".grok" if accounts else default_home() - signed_in = bool(accounts) or bool(login(stats_home)) - record = { - "schemaVersion": 1, - "id": AGENT_ID, - "name": AGENT_NAME, - "updatedAt": datetime.now(timezone.utc).isoformat(), - "ready": signed_in, - "tierLabel": tier, - "usageStatusText": limits["usageStatusText"], - "authHelpText": limits["authHelpText"], - "limits": limits["limits"], - "limitsStale": stale, - "limitsFetchedAt": fetched_at, - } - record.update(cached_local_stats(stats_home, LIMITS_ONLY_SCAN_REUSE_SECONDS if args.limits_only and not args.force else 0)) + stats = cached_local_stats(stats_home, scan_max_age(args.force, args.limits_only)) + # In use here once someone signed in or ran a session; otherwise an empty + # record, which the panel skips. + in_use = bool(accounts) or bool(login(stats_home)) or stats["totalSessions"] > 0 + # A missing sign-in is only news where Grok has been set up: a machine + # with no Grok home at all says nothing. + if not in_use and not stats_home.exists() and not limits["limits"]: + limits = {"limits": [], "usageStatusText": "", "authHelpText": ""} + stale, fetched_at = False, 0 + record = base_record( + AGENT_ID, + AGENT_NAME, + ready=in_use, + tierLabel=tier, + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + limits=limits["limits"], + limitsStale=stale, + limitsFetchedAt=fetched_at, + ) + if limits.get("retryAdvised"): + record["retryAdvised"] = True + record.update(stats) if accounts: record["accountSwitch"] = registered[0]["switch"] record["accounts"] = accounts - print(json.dumps(record, separators=(",", ":"))) + emit(record) if __name__ == "__main__": diff --git a/bin/omarchy-agent-usage-grokbot b/bin/omarchy-agent-usage-grokbot new file mode 100755 index 00000000000..5c71a105b5d --- /dev/null +++ b/bin/omarchy-agent-usage-grokbot @@ -0,0 +1,173 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Grok Bot usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Grok Bot usage into one display-ready JSON record. + +Grok Bot runs on a Cursor account, but its weekly allowance is metered apart +from Cursor's own plan: the dashboard's GetSandUsageStatus RPC (Sand is Grok +Bot's internal name) says how much of this week's included usage is gone and +when it resets. It is asked as the same account the Cursor collector reads, +found the same way, so the two never describe different people. Grok Bot +keeps nothing countable on disk, so the record carries the limit alone, and +an account without a Grok Bot allowance gives an unused record, which the +panel skips. The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import os +import sys +import urllib.request +from datetime import datetime, timezone +from typing import Any + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + account_key, + base_record, + cursor_token, + emit, + finite, + open_json, + token_lapsed, +) + +AGENT_ID = "grokbot" +AGENT_NAME = "Grok Bot" +AUTH_HELP = "Sign in to the Cursor editor, or run `cursor-agent login`, with your Grok Bot account." +USAGE_URL = "https://api2.cursor.sh/aiserver.v1.DashboardService/GetSandUsageStatus" + + +def parse_time(value: Any) -> datetime | None: + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except ValueError: + return None + return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc) + + +# ------------------------------------------------------------------ limit + + +def fetch_status(token: str) -> dict[str, Any]: + request = urllib.request.Request( + USAGE_URL, + data=b"{}", + headers={ + "Content-Type": "application/json", + "Connect-Protocol-Version": "1", + "Authorization": f"Bearer {token}", + "x-cursor-client-type": "cli", + }, + method="POST", + ) + status = open_json(request, timeout=15) + if not isinstance(status, dict): + raise ValueError("not an object") + return status + + +# The answer is protobuf JSON, which leaves out a field holding zero or +# false: no included limit means no Grok Bot allowance on this account, and +# a week with nothing used yet has no percentage at all. +def weekly_limits(status: dict[str, Any]) -> list[dict[str, Any]]: + if status.get("hasNonZeroIncludedLimit") is not True: + return [] + resets = parse_time(status.get("nextResetTimestampUtc")) + return [{ + "label": "Weekly", + "percent": max(0.0, min(1.0, finite(status.get("usagePercent")) / 100)), + "resetsAt": resets.isoformat() if resets else "", + }] + + +def collect(force: bool = False) -> dict[str, Any]: + record = base_record(AGENT_ID, AGENT_NAME, hasPromptStats=False) + # Without a sign-in there is no telling whether this account has Grok Bot + # at all, so the machine reads as not using it; the Cursor collector is the + # one that asks for a sign-in. + token = cursor_token() + if not token: + return record + + cache = LimitsCache(AGENT_ID, account_key(token)) + # A kept week that has ended is a fresh allowance, so it comes back at zero + # with no reset time until a check says when the next one ends. One with no + # known end is kept as it is. + kept = cache.limits + fetched_at = cache.fetched_at + kept_tier = str(cache.data.get("tierLabel") or "") + + # A panel opened and shut repeatedly must not turn into a request per + # flick, so an answer this recent is the answer, even one saying the + # account has no Grok Bot. + if cache.fresh(force, allow_empty=True): + record.update(ready=bool(kept), tierLabel=kept_tier, limits=kept, limitsFetchedAt=fetched_at) + return record + + retry = False + if token_lapsed(token): + status = "Limits paused" + detail = "Cursor's access token has lapsed. Open the Cursor editor, or run `cursor-agent`, to renew it; signing in again is not needed." + else: + try: + answer = fetch_status(token) + limits = weekly_limits(answer) + except HttpError as error: + if error.status in (401, 403): + status, detail = "Sign-in expired", "Cursor refused the saved sign-in. " + AUTH_HELP + elif error.status: + status, detail = "", f"Cursor's dashboard returned status {error.status}." + # No route, no DNS, refused, reset or timed out: nothing answered, so + # the panel tries again soon. + elif error.transport: + status, detail, retry = "", "Couldn't reach Cursor's dashboard. Retrying shortly.", True + # A reply cut short or one that isn't the expected JSON came from a + # server that answered. + else: + status, detail = "", "Cursor's dashboard sent a reply that couldn't be read." + except (ValueError, TypeError, AttributeError): + status, detail = "", "Cursor's dashboard sent a reply that couldn't be read." + else: + tier = str(answer.get("grokPlanLabel") or "") + # The cache only backs later runs; failing to keep it never costs this + # run's live answer. + record.update(ready=bool(limits), tierLabel=tier, limits=limits, limitsFetchedAt=cache.save(limits, tierLabel=tier)) + return record + + # A failed check keeps the last good limit, marked stale so the panel + # gives its age on hover. With nothing kept the record stays unused, so a + # Cursor account that never had Grok Bot doesn't appear on a blip. + record.update( + ready=bool(kept), + tierLabel=kept_tier if kept else "", + usageStatusText=status or ("" if kept else "Grok Bot limits unavailable"), + authHelpText=detail + (" Showing the last known limit." if kept else ""), + limits=kept, + limitsStale=True, + limitsFetchedAt=fetched_at if kept else 0, + ) + if retry: + record["retryAdvised"] = True + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the Grok Bot usage record as JSON") + parser.add_argument("--force", action="store_true") + # Grok Bot has no local scan to reuse, so a near-limit refresh is a normal + # run; the flag keeps the invocation every collector takes. + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + emit(collect(force=args.force)) + 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..85872e97a4b --- /dev/null +++ b/bin/omarchy-agent-usage-hermes @@ -0,0 +1,278 @@ +#!/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 its root's profiles/, and in ~/.hermes whole. Hermes +has no subscription of its own, so the record has no limits. + +Tokens count under the provider that bills them, so this record holds what +Hermes ran on Nous Portal, and what it ran before it recorded a route. A +session sent to a server you run is the local collector's, which reads it +through harness_rows(); one sent to any other provider is not counted here. + +Chat sessions that arrived over messaging gateways count too: Hermes treats +them as first-class sessions and so does this record. +""" + +from __future__ import annotations + +import argparse +import os +import sqlite3 +import sys +from datetime import datetime +from pathlib import Path + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + NUMBER_CEILING, + Stats, + UsageRow, + base_record, + cached_scan, + claude_code_owner, + date_string, + emit, + local_day, + number, + scan_max_age, + url_host, +) + +AGENT_ID = "hermes" +AGENT_NAME = "Hermes" + +# Nous Portal's hosts (portal., inference-api., welcome-api.), the service +# Nous Research, who make Hermes, bill. +NOUS_HOST = "nousresearch.com" + + +# Hermes' home is HERMES_HOME, its ~ and $VARs expanded, else +# ~/.hermes. Named profiles sit under the root's +# profiles/: the default home, unless HERMES_HOME lies outside it, where a +# /profiles/ home makes the root and any other home is its +# own (hermes_constants.get_default_hermes_root). The root is the default +# profile, so a profile home still counts the root's own state.db +# (hermes_cli.profiles.list_profiles). The default home ~/.hermes is read +# whatever the root, as before, and whole: its own state.db is its default +# profile, beside its profiles/. A profile linked to the home counts once. +def state_paths(): + native = Path.home() / (".hermes" + os.environ.get("HERMES_DATA_DIR_SUFFIX", "")) + configured = os.environ.get("HERMES_HOME", "").strip() + home = Path(os.path.expanduser(os.path.expandvars(configured))) if configured else native + root = native + if not home.resolve().is_relative_to(native.resolve()): + root = home.parent.parent if home.parent.name == "profiles" else home + roots = dict.fromkeys((root, native, Path.home() / ".hermes")) + candidates = [home / "state.db", *(base / "state.db" for base in roots)] + for profiles in (base / "profiles" for base in roots): + 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 + + +# Hermes records where it sent each session and each request: the provider +# it billed and the base URL it reached. Nous Portal, Hermes' own maker's +# service, is Hermes'; any other URL is read the way every collector reads +# Claude Code's (claude_code_owner): a provider's host is that provider's, +# and a server you run is local's unless the model is a Claude or an Ollama +# cloud one it only forwards. A gateway nothing here can place, or a +# provider's own API no collector reads Hermes for (Anthropic's, OpenAI's), +# is nobody's. A row from before Hermes kept the route ran where nothing +# says, and stays Hermes', as every row was read before. A session Hermes +# keeps routes for that recorded none never reached a provider, since Hermes +# records the route with the first request it accounts, so no one bills it. +def route_owner(provider, base_url, model): + url = str(base_url or "").strip() + if url: + host = url_host(url) + if host == NOUS_HOST or host.endswith("." + NOUS_HOST): + return AGENT_ID + return claude_code_owner(str(model or ""), url) + provider = str(provider or "").strip().lower() + return AGENT_ID if provider in ("", "nous") else None + + +def columns(conn, table): + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + + +def route_columns(conn, table, alias): + """The route columns of a table, or NULLs where Hermes didn't keep them yet.""" + if {"billing_provider", "billing_base_url"} <= columns(conn, table): + return f"{alias}.billing_provider, {alias}.billing_base_url" + return "NULL, NULL" + + +# A column's sum over the rows whose count can be trusted: a negative row +# would net out real usage, and text or a count past the ceiling is none. +def trusted_sum(column): + return f"SUM(CASE WHEN typeof({column}) IN ('integer', 'real') AND {column} >= 0 AND {column} < {NUMBER_CEILING} THEN {column} ELSE 0 END)" + + +# A time names a day only when it is a number that local_day trusts: SQLite +# ranks text above every number, and a stray 1 is a damaged stamp, not a day +# in 1970. +def stamp_day(value): + return local_day(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None + + +def scan(path, owner): + """One state.db's rows that `owner` bills: a row per session, dated by its + start, with the tokens of each model it used and the prompts typed on each + day. A session counts where its route says, else where its first request + went, else nowhere when Hermes keeps routes and it made none.""" + conn = sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + try: + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + usage = [] + # A Hermes that keeps no usage table yet still has its sessions and + # prompts to count. + if "session_model_usage" in tables: + usage = conn.execute( + """SELECT smu.session_id, smu.model, {}, {}, {}, {}, {} + FROM session_model_usage smu + GROUP BY 1, 2, 3, 4""".format(route_columns(conn, "session_model_usage", "smu"), + *(trusted_sum("smu." + column) for column in ("input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens"))) + ).fetchall() + first_owner = {} + for session_id, model, provider, url, *_ in usage: + first_owner.setdefault(session_id, route_owner(provider, url, model)) + + model_column = "s.model" if "model" in columns(conn, "sessions") else "NULL" + session_routes = route_columns(conn, "sessions", "s") + unrouted = None if "session_model_usage" in tables and session_routes != "NULL, NULL" else AGENT_ID + sessions = {} + for session_id, started, model, provider, url in conn.execute(f"SELECT s.id, s.started_at, {model_column}, {session_routes} FROM sessions s"): + recorded = str(provider or "").strip() or str(url or "").strip() + sessions[session_id] = (stamp_day(started), route_owner(provider, url, model) if recorded else first_owner.get(session_id, unrouted)) + + rows = [] + database = path.resolve() + key = lambda session_id: f"{database}:{session_id}" + for session_id, (day, billed) in sessions.items(): + if billed == owner: + rows.append(UsageRow(day, key(session_id), "", "", 0, 0, 0, 0, prompts=0)) + + # Usage is kept per session and model, so a session's tokens land on the + # day it started. + for session_id, model, provider, url, *counts in usage: + if session_id not in sessions or route_owner(provider, url, model) != owner: + continue + # Counts from disk are untrusted: a negative, non-finite, absurdly + # large or non-numeric value counts as nothing. + counts = [number(count) for count in counts] + if sum(counts): + rows.append(UsageRow(sessions[session_id][0], key(session_id), str(provider or ""), str(model or "hermes"), *counts, prompts=0)) + + # A prompt from the future is no day's yet. + today = date_string(datetime.now()) + for session_id, stamp, count in conn.execute( + "SELECT session_id, CASE WHEN typeof(timestamp) IN ('integer', 'real') THEN timestamp END, COUNT(*) FROM messages WHERE role = 'user' GROUP BY 1, 2" + ): + if sessions.get(session_id, (None, AGENT_ID))[1] != owner: + continue + day = stamp_day(stamp) + rows.append(UsageRow(day if day is not None and day <= today else None, key(session_id), "", "", 0, 0, 0, 0, prompts=count)) + return rows + finally: + conn.close() + + +def harness_sources(): + return [path.resolve() for path in state_paths()] + + +# What harness_rows() decided owners under, for cached_scan(match=...): +# nothing beyond the databases, which record every route themselves. +def harness_match(): + return {} + + +def scan_all(paths, owner): + """Every state.db's rows that `owner` bills, and the paths that couldn't + be read.""" + rows, unreadable = [], [] + for path in paths: + try: + rows.extend(scan(path, owner)) + except sqlite3.Error as error: + unreadable.append(f"{path}: {error}") + return rows, unreadable + + +def harness_rows(owner): + """Every state.db's rows that `owner` bills, and whether all were read.""" + rows, unreadable = scan_all(state_paths(), owner) + return rows, not unreadable + + +# The stats from every state.db, with the paths that couldn't be read. +def collect_stats(paths): + stats = Stats() + rows, unreadable = scan_all(paths, AGENT_ID) + for row in rows: + stats.add_row(row, row.model) + return stats.snapshot(), unreadable + + +# The scan is cached per day and per set of databases. A scan that couldn't +# read every database isn't cached, so a scan served from the cache read them +# all. +def cached_stats(paths, max_age): + unreadable = [] + + def scan_all(): + stats, missed = collect_stats(paths) + unreadable.extend(missed) + return stats, not missed + + stats = cached_scan(AGENT_ID, [path.resolve() for path in paths], max_age, scan_all) + return stats, unreadable + + +def collect(force=False, limits_only=False): + record = base_record(AGENT_ID, AGENT_NAME) + + # No state.db: Hermes isn't used on this machine, which is not trouble -- + # with sync on, another machine's Hermes still shows through its stats. + paths = state_paths() + if not paths: + return record + + stats, unreadable = cached_stats(paths, scan_max_age(force, limits_only)) + record.update(stats) + # Sessions all billed elsewhere leave Hermes nothing of its own to show. + record["hasLocalStats"] = bool(stats.get("modelUsage") or stats.get("totalSessions")) + record["ready"] = record["hasLocalStats"] or bool(unreadable) + if unreadable: + record["usageStatusText"] = "Hermes unavailable" + record["authHelpText"] = "Couldn't read Hermes' state.db (" + "; ".join(unreadable) + ")." + return record + + +def main(): + parser = argparse.ArgumentParser(description="Print the Hermes Agent usage record as JSON") + parser.add_argument("--force", action="store_true", help="rescan every state.db, ignoring the cache") + parser.add_argument("--limits-only", action="store_true", help="reuse any recent scan; Hermes has no limits to probe") + args = parser.parse_args() + emit(collect(args.force, args.limits_only)) + + +if __name__ == "__main__": + main() diff --git a/bin/omarchy-agent-usage-kimchi b/bin/omarchy-agent-usage-kimchi new file mode 100755 index 00000000000..9118408e00a --- /dev/null +++ b/bin/omarchy-agent-usage-kimchi @@ -0,0 +1,401 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Kimchi usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Kimchi usage into one display-ready JSON record. + +The plan, prepaid credits, and spend budgets come from the gateway endpoints +Kimchi's own status line reads (/v1/credits and /v1/budget), asked +with the API key Kimchi signs in with. Tokens and prompts come from its pi +session files: one line per message, each assistant reply carrying its model +and token usage. pi and omp sessions on Kimchi's provider ids count too, and +so do the opencode, Claude Code and Codex setups Kimchi writes, which send +to its gateway. The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +import urllib.request +from datetime import datetime, timezone +from pathlib import Path + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record, + cached_scan, + emit, + kimchi_sessions_root, + open_json, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, + stamp, + trusted_endpoint, +) + +AGENT_ID = "kimchi" +AGENT_NAME = "Kimchi" +AUTH_HELP = "Start Kimchi and log in, or set KIMCHI_API_KEY." +REJECTED_HELP = "Kimchi rejected its API key. Log in again in Kimchi, or update KIMCHI_API_KEY." +OUT_OF_CREDITS_HELP = "Kimchi's gateway is rate limiting requests until the account has credits again." +REGION_LLM_BASE = {"us": "https://llm.kimchi.dev", "eu": "https://llm.eu.kimchi.dev"} +# The region a self-hosted login saves; its gateway is the host it names. +SELF_HOSTED = "self-hosted" +PLAN_LABELS = {"community": "Community", "coder": "Coder", "teams": "Teams", "enterprise": "Enterprise"} +BUDGET_LABELS = { + "API_KEY": "API key budget", + "USER": "Budget", + "TEAM_PER_USER": "Team budget", + "TEAM_POOLED": "Team pool", + "ORGANIZATION_SOFT": "Org budget", + "ORGANIZATION_HARD": "Org budget", +} +CREDITS_FIELDS = ("serverless", "remaining", "has_credits", "hasCredits", "tier", "tier_name", "tierName") + + +# Kimchi sends dollar amounts as strings, sometimes with a sign or unit. +def usd(value): + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + parsed = float(value) + else: + text = str(value or "").lower().replace("$", "").replace("usd", "").replace(",", ".").strip() + try: + parsed = float(text) if text else None + except ValueError: + return None + return parsed if parsed is not None and math.isfinite(parsed) else None + + +def parse_time(value): + if not isinstance(value, str) or not value.strip(): + return None + try: + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return None + return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc) + + +# Kimchi's own gateway routes, as its isKimchiProvider names them: kimchi-dev, +# kimchi-dev/ and kimchi-experimental, and `kimchi`, the provider id +# its setup commands write into opencode and other agents and a pi or omp +# models.json names it by. `--provider` runs a session on any other pi +# provider (OpenRouter, a custom one in harness/models.json), which that +# provider bills: the openrouter collector reads these sessions for its own, +# and a custom provider's count nowhere. +def kimchi_route(provider, model, entry): + name = provider.lower() + return name in ("kimchi", "kimchi-dev", "kimchi-experimental") or name.startswith("kimchi-dev/") + + +# Kimchi's harness also wrote replies naming no provider, which only it can +# have sent: those are its own there, and no one's anywhere else. +def unnamed(provider, model, entry): + return provider == "" + + +# Every pi-format session on a Kimchi route, in Kimchi's harness and in pi's +# and omp's trees, read in one pass so a root inside another counts once and +# a message a fork or resumed session repeats counts once; the harness's +# replies naming no provider; opencode replies on Kimchi's provider; and the +# Claude Code, MiniMax Code and Codex turns sent to Kimchi's gateway. Returns +# the stats and whether every source was read. +def local_stats(sessions_root, reuse=True): + stats = Stats() + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"kimchi', kimchi_route, [*pi_session_roots(), sessions_root]) + harness_rows, harness_complete = scan_pi_sessions(r'"usage"', unnamed, [sessions_root]) + # Every id starting kimchi gets past the database's prefilter; the accept + # keeps Kimchi's own. + db_rows, db_complete = scan_opencode(["kimchi"], kimchi_route) + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + for row in pi_rows + harness_rows + db_rows + routed_rows: + stats.add_row(row, row.model or AGENT_ID) + complete = pi_complete and harness_complete and db_complete and routed_complete + return {"hasLocalStats": bool(stats.sessions), **stats.snapshot()}, complete + + +# The key and gateway Kimchi itself would use: KIMCHI_API_KEY, KIMCHI_BASE_URL, +# and KIMCHI_REGION over ~/.config/kimchi/config.json, the region's gateway +# when nothing names one. A self-hosted login's gateway is /llm +# (KIMCHI_SELF_HOSTED_URL over the saved one), and without one there is none: +# its key is never sent to Kimchi's public gateway. +def gateway(): + config = read_json(Path.home() / ".config" / "kimchi" / "config.json") + key = os.environ.get("KIMCHI_API_KEY") or config.get("apiKey") or config.get("api_key") or "" + base = os.environ.get("KIMCHI_BASE_URL") or "" + region = next((value.strip().lower() for value in (os.environ.get("KIMCHI_REGION"), config.get("region")) + if isinstance(value, str) and value.strip().lower() in (*REGION_LLM_BASE, SELF_HOSTED)), "us") + if not base and region == SELF_HOSTED: + self_hosted = os.environ.get("KIMCHI_SELF_HOSTED_URL") or config.get("selfHostedUrl") or "" + if not isinstance(self_hosted, str) or not self_hosted.strip(): + return str(key).strip(), "" + base = self_hosted.strip().rstrip("/") + "/llm" + if not base: + base = config.get("llmEndpoint") or REGION_LLM_BASE[region] + base = str(base).strip().rstrip("/") + if base.lower().endswith("/openai/v1"): + base = base[: -len("/openai/v1")] + return str(key).strip(), base + + +def get_json(url, key): + request = urllib.request.Request(url, headers={"Authorization": f"Bearer {key}", "Accept": "application/json", "User-Agent": "omarchy-agent-usage-kimchi"}) + return open_json(request, timeout=8) + + +def plan_label(raw): + text = str(raw or "").strip().lower().replace("_", "-").replace(" ", "-") + if text in ("community", "free", "free-tier", "free-slow"): + return PLAN_LABELS["community"] + if text in ("starter", "individual"): + return PLAN_LABELS["coder"] + if text == "team": + return PLAN_LABELS["teams"] + for plan in ("coder", "teams", "enterprise"): + if text == plan or text.startswith(plan + "-"): + return PLAN_LABELS[plan] + return str(raw or "") + + +# The credits answer names the plan and what's left of the prepaid balance. +# It says how much remains but not how much was bought, so the gauge shows +# the remaining figure alone. +def credits_part(data): + # A reply that names none of the credits fields isn't an answer, so it + # fails the check rather than replacing the kept plan and balance. + if not isinstance(data, dict) or not any(k in data for k in CREDITS_FIELDS): + raise ValueError("not a credits answer") + if data.get("serverless") is False: + return {} + tier = next((data[k] for k in ("tier", "tier_name", "tierName") if k in data), "") + part = {"tierLabel": plan_label(tier)} + remaining = usd(data.get("remaining")) + if remaining is not None and remaining >= 0: + part["balance"] = {"remaining": remaining, "funded": 0, "spent": 0, "currency": "USD", "estimated": False} + has_credits = next((data[k] for k in ("has_credits", "hasCredits") if k in data), None) + # The gateway throttles an account with no credits left, and with no funded + # figure the balance gauge has no way to raise that alarm itself. + if has_credits is False or (remaining is not None and remaining <= 0): + part["outOfCredits"] = True + return part + + +# The panel titles a window by its clock, so a budget period that spans a +# month or a week says so rather than leaving the panel to guess from a label. +def period_title(period): + start, end = parse_time(period.get("startTime")), parse_time(period.get("endTime")) + if not (start and end): + return "" + days = (end - start).days + if 27 <= days <= 31: + return "Monthly" + if days == 7: + return "Weekly" + return "" + + +# Each spend budget is a limit window: spend against its cap until the +# period ends. A provider capped inside it is its own allowance on the same +# clock, which the panel draws as a tick on that budget's meter; in a period +# the panel has no clock name for, it gets a line of its own. A cap is what +# may be spent, not what was bought, so it never funds the balance gauge. +# +# The panel hangs a provider's tick on the meter its title names, so with +# several budgets on one clock each is titled by its scope ("Team pool +# Monthly") and its providers get lines of their own that name it, rather +# than every cap landing on whichever budget came last. +def budget_part(data): + # The budgets list is the answer, and an empty or null one still is; a + # reply without it fails the check rather than erasing the kept budgets. + if not isinstance(data, dict) or "budgets" not in data or not isinstance(data["budgets"] or [], list): + raise ValueError("not a budget answer") + period = data.get("period") if isinstance(data.get("period"), dict) else {} + # The gateway's own period end, or none: a reset time is never made up. + resets = str(period.get("endTime") or "") if parse_time(period.get("endTime")) else "" + title = period_title(period) + budgets = [] + for entry in data["budgets"] or []: + if not isinstance(entry, dict): + continue + cap = usd(entry.get("budgetLimitUsd")) + spent = usd(entry.get("totalSpendUsd")) + if cap and cap > 0 and spent is not None: + budgets.append((entry, spent / cap)) + shared = len(budgets) > 1 + limits = [] + for entry, share in budgets: + label = BUDGET_LABELS.get(entry.get("scope"), "Budget") + limit = {"label": label, "percent": max(0.0, min(1.0, share)), "resetsAt": resets} + if title: + limit["title"] = f"{label} {title}" if shared else title + limits.append(limit) + for provider in entry.get("providerBudgets") or []: + if not isinstance(provider, dict) or not str(provider.get("limitType") or "").endswith("CAPPED"): + continue + provider_cap = usd(provider.get("budgetLimitUsd")) + used = usd(provider.get("usageUsd")) + name = str(provider.get("provider") or "").strip() + if not (provider_cap and provider_cap > 0 and used is not None and name): + continue + name = name[:1].upper() + name[1:] + limits.append({ + "label": f"{name} budget", + "title": f"{name} budget ({label})" if shared else (f"{name} {title}" if title else f"{name} budget"), + "percent": max(0.0, min(1.0, used / provider_cap)), + "resetsAt": resets, + }) + return {"limits": limits} + + +# Keyed by the API key and the gateway it was asked of, so a machine signed +# in to someone else, or pointed at another region, never shows the last +# one's numbers. One file per endpoint, so a run only ever writes the answer +# it actually got: the budget's limits under the limits cache name every +# collector shares, and the credits' plan and balance beside them. An answer +# without limits (no budget, or the credits) still stands in for a probe. +# Returns the cache to save to and the one holding the last answer. +def billing_cache(key, base, resource): + if resource == "credits": + cache = LimitsCache(AGENT_ID, base, key, kind="credits") + return cache, cache + cache = LimitsCache(AGENT_ID, base, key) + if not cache.data: + # Budgets were cached as kimchi-budget-.json before they joined the + # limits cache name every collector shares. + return cache, LimitsCache(AGENT_ID, base, key, kind="budget") + return cache, cache + + +# The two endpoints answer on their own, and so does each fall back: a part +# that didn't come back is the last one that did, so one endpoint failing +# never blanks or overwrites what the other said. Any part that didn't come +# back fresh marks the record stale, cached or not. +def collect_billing(force): + result = {"tierLabel": "", "limits": [], "usageStatusText": "", "authHelpText": "", "signedIn": False, "limitsStale": False, "limitsFetchedAt": 0} + key, base = gateway() + if not key: + return result + result["signedIn"] = True + if not base: + result["usageStatusText"] = "Balance unavailable" + result["authHelpText"] = "Kimchi is set to its self-hosted region but names no self-hosted URL. Local Kimchi stats are still shown." + return result + # KIMCHI_BASE_URL, llmEndpoint and a self-hosted URL move the gateway, + # never the key onto a connection anyone on the way can read. + if not trusted_endpoint(base): + result["usageStatusText"] = "Balance unavailable" + result["authHelpText"] = "Kimchi's gateway must use HTTPS (plain HTTP only for loopback). Local Kimchi stats are still shown." + return result + + parts = {} + failed = set() + rejected = unreachable = False + for resource, parse in (("credits", credits_part), ("budget", budget_part)): + cache, kept = billing_cache(key, base, resource) + cached = None + if kept.fetched_at: + # A kept budget whose period ended is an untouched allowance, not an + # unknown one: it comes back at zero, with no reset time until the + # gateway names the next period. + cached = dict(kept.data, limits=kept.limits, fetchedAtMs=kept.fetched_at) + if kept.fresh(force, allow_empty=True): + parts[resource] = cached + continue + try: + data = get_json(f"{base}/v1/{resource}", key) + part = parse(data) + # Refreshes can overlap (the panel's own and a limits retry), so an + # answer only replaces one fetched before it. + cache.save(part.pop("limits", []), only_if_newer=True, **part) + parts[resource] = dict(cache.data) + continue + except HttpError as error: + rejected = rejected or error.status in (401, 403) + # No answer at all: no route, no DNS, a timeout, or a dropped + # connection. A reply cut short is an answer that failed to read. + unreachable = unreachable or error.transport + except Exception: + pass + failed.add(resource) + if cached is not None: + parts[resource] = cached + + credits = parts.get("credits") or {} + result["tierLabel"] = credits.get("tierLabel") or "" + if credits.get("balance"): + result["balance"] = credits["balance"] + result["limits"] = (parts.get("budget") or {}).get("limits") or [] + result["limitsStale"] = bool(failed) + if parts: + result["limitsFetchedAt"] = min(stamp(part.get("fetchedAtMs")) for part in parts.values()) + + if rejected: + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = REJECTED_HELP + elif credits.get("outOfCredits") or credits.get("usageStatusText"): + # With no funded figure the balance gauge can't raise this alarm itself, + # so the status does. A credits answer kept before outOfCredits marked it + # with a status of its own. + result["usageStatusText"] = "Out of credits" + result["authHelpText"] = OUT_OF_CREDITS_HELP + elif "credits" not in parts: + result["usageStatusText"] = "Balance unavailable" + result["authHelpText"] = ( + "Couldn't reach Kimchi's gateway. Local Kimchi stats are still shown." if unreachable + else "Kimchi's gateway didn't return the credits. Local Kimchi stats are still shown." + ) + if unreachable and not rejected: + # Usually the seconds after login before the network is up; the panel + # tries once more soon instead of waiting out its refresh interval. + result["retryAdvised"] = True + return result + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + # A scan is reused on the day it was made only, since one from before + # midnight would count yesterday's sessions as today's. A limits-only + # refresh needs fresh credits, not a fresh history scan. The digest covers + # every source the scan reads, and the match where Claude Code, MiniMax + # Code and Codex are sent, which decides whose their turns are. A forced + # scan reads every Codex rollout again too. + sessions_root = kimchi_sessions_root() + sources = [sessions_root, *pi_session_roots(), opencode_db(), *routed_sources()] + max_age = scan_max_age(args.force, args.limits_only) + stats = cached_scan(AGENT_ID, sources, max_age, lambda: local_stats(sessions_root, reuse=max_age > 0), match=routed_match()) + billing = collect_billing(args.force) + signed_in = billing.pop("signedIn") + # Sessions on disk with no key to ask the gateway with: Kimchi is in use + # here, its plan and credits just can't be read yet. + if not signed_in and stats.get("hasLocalStats"): + billing["usageStatusText"] = "Waiting for auth" + billing["authHelpText"] = AUTH_HELP + record = base_record(AGENT_ID, AGENT_NAME, ready=signed_in or bool(stats.get("hasLocalStats"))) + record.update(billing) + record.update(stats) + emit(record) + + +if __name__ == "__main__": + main() diff --git a/bin/omarchy-agent-usage-kimi b/bin/omarchy-agent-usage-kimi index f7d251743d1..61bdefe0fd7 100755 --- a/bin/omarchy-agent-usage-kimi +++ b/bin/omarchy-agent-usage-kimi @@ -2,7 +2,8 @@ # omarchy:summary=Print the Kimi usage record as JSON # omarchy:args=[--force] [--limits-only] # omarchy:hidden=true -"""Collect Kimi account quota into one display-ready JSON record.""" +"""Collect Kimi account quota and Kimi Code's local usage into one +display-ready JSON record.""" from __future__ import annotations @@ -12,11 +13,12 @@ import hashlib import json import math import os +import stat import sys import tempfile import threading import time -import urllib.error +import tomllib import urllib.parse import urllib.request from datetime import datetime, timezone @@ -24,6 +26,35 @@ from contextlib import contextmanager from pathlib import Path from typing import Any +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + UsageRow, + base_record as shared_base_record, + cached_scan, + claude_code_owner, + emit, + local_day, + obj, + open_json, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, + token_count, + warn, + write_json, +) + AGENT_ID = "kimi" AGENT_NAME = "Kimi" AUTH_HELP = "Set KIMI_API_KEY or sign in to Kimi (run `kimi login`)." @@ -39,13 +70,23 @@ OAUTH_CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098" ALLOWED_DOMAINS = ("kimi.com", "kimi.ai") REFRESH_SKEW_S = 60 LOCK_WAIT_SECONDS = 15 -PROBE_REUSE_SECONDS = 60 +# proper-lockfile's own stale age: a lease left this long unrenewed is dead. +LOCK_STALE_SECONDS = 10 + + +# The legacy kimi-cli's share directory: KIMI_SHARE_DIR, else ~/.kimi +# (kimi_cli/share.py), a relative one read from here as Kimi Code's own +# migration reads it. +def legacy_home() -> Path: + if share := os.environ.get("KIMI_SHARE_DIR", "").strip(): + return Path(os.path.abspath(share)) + return Path.home() / ".kimi" def credentials_dir() -> Path: if home := os.environ.get("KIMI_CODE_HOME"): return Path(home).expanduser() / "credentials" - candidates = [Path.home() / ".kimi-code" / "credentials", Path.home() / ".kimi" / "credentials"] + candidates = [Path.home() / ".kimi-code" / "credentials", legacy_home() / "credentials"] for candidate in candidates: if (candidate / "kimi-code.json").is_file(): return candidate @@ -57,16 +98,15 @@ CREDENTIALS_PATH = CREDENTIALS_DIR / "kimi-code.json" class KimiError(Exception): - def __init__(self, message: str, retry: bool = False): + # `retry` marks a service that couldn't be reached at all, so the panel + # tries again soon; `refused` marks a credential Kimi turned down, which + # only signing in again (or a new key) repairs; `missing` marks no + # credential at all. + def __init__(self, message: str, retry: bool = False, refused: bool = False, missing: bool = False): super().__init__(message) self.retry = retry - - -class NoRedirectHandler(urllib.request.HTTPRedirectHandler): - # urllib would copy every header (including the bearer token) onto the - # redirect target, potentially leaking it cross-host. Never follow. - def redirect_request(self, req, fp, code, msg, headers, newurl): - return None + self.refused = refused + self.missing = missing def trusted_url(url: str) -> tuple[str, bool]: @@ -84,57 +124,21 @@ def trusted_url(url: str) -> tuple[str, bool]: return url, any(host == domain or host.endswith("." + domain) for domain in ALLOWED_DOMAINS) -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, - # Quota numbers are account-global, not machine-local: every synced device - # reports the same truth, so aggregation must not sum them. - "scope": "account", - "hasPromptStats": False, - "tierLabel": "", - "usageStatusText": "", - "authHelpText": "", - "limits": [], - "limitsStale": False, - "limitsFetchedAt": 0, - } - record.update(empty_stats()) - record.update(overrides) - return record + return shared_base_record(AGENT_ID, AGENT_NAME, **overrides) # --------------------------------------------------------------- kimi auth def read_credentials() -> dict[str, Any]: - try: - credentials = json.loads(CREDENTIALS_PATH.read_text()) - except (OSError, json.JSONDecodeError): - return {} - return credentials if isinstance(credentials, dict) else {} + return read_json(CREDENTIALS_PATH) -def write_json(path: Path, payload: dict[str, Any]) -> None: +# The Kimi CLI's own credential file, rewritten only under its lock: private +# from the first byte and flushed before it replaces the old one, unlike the +# shared cache writer. +def write_private_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) fd, temporary = tempfile.mkstemp(prefix=path.name + ".", dir=path.parent) try: @@ -150,7 +154,7 @@ def write_json(path: Path, payload: dict[str, Any]) -> None: def write_credentials(credentials: dict[str, Any]) -> None: CREDENTIALS_DIR.mkdir(parents=True, exist_ok=True, mode=0o700) CREDENTIALS_DIR.chmod(0o700) - write_json(CREDENTIALS_PATH, credentials) + write_private_json(CREDENTIALS_PATH, credentials) def token_fresh(credentials: dict[str, Any]) -> bool: @@ -177,17 +181,14 @@ def request_refresh(refresh_token: str) -> dict[str, Any]: headers = {"Content-Type": "application/x-www-form-urlencoded"} request = urllib.request.Request(url, data=form, headers=headers) try: - with urllib.request.build_opener(NoRedirectHandler()).open(request, timeout=15) as response: - decoded = json.load(response) - except urllib.error.HTTPError as error: - if error.code in (400, 401, 403): - raise KimiError("Kimi rejected the sign-in; sign in again with `kimi login`") from error - raise KimiError(f"Kimi token refresh failed (HTTP {error.code})") from error - except urllib.error.URLError as error: - raise KimiError("Could not reach the Kimi login service", retry=True) from error - except TimeoutError as error: - raise KimiError("Kimi token refresh timed out", retry=True) from error - except json.JSONDecodeError as error: + decoded = open_json(request, timeout=15) + except HttpError as error: + if error.status in (400, 401, 403): + raise KimiError("Kimi rejected the sign-in; sign in again with `kimi login`", refused=True) from error + if error.status: + raise KimiError(f"Kimi token refresh failed (HTTP {error.status})") from error + if error.transport: + raise KimiError("Could not reach the Kimi login service", retry=True) from error raise KimiError("Kimi returned an invalid token response") from error if not isinstance(decoded, dict) or any( not isinstance(decoded.get(key), str) or not decoded[key] for key in ("access_token", "refresh_token") @@ -205,7 +206,7 @@ def request_refresh(refresh_token: str) -> dict[str, Any]: @contextmanager def refresh_lock(): deadline = time.monotonic() + LOCK_WAIT_SECONDS - legacy = CREDENTIALS_DIR == Path.home() / ".kimi" / "credentials" and not os.environ.get("KIMI_CODE_HOME") + legacy = CREDENTIALS_DIR == legacy_home() / "credentials" and not os.environ.get("KIMI_CODE_HOME") if legacy: with (CREDENTIALS_DIR / "kimi-code.lock").open("a") as handle: while True: @@ -227,7 +228,23 @@ def refresh_lock(): lock.mkdir(mode=0o700) break except FileExistsError: - if not lock.is_dir() or time.monotonic() >= deadline: + # The CLI may let go between the mkdir and this look: a lock that is + # gone is free, so try again. Only something that isn't a directory + # blocks the lock for good. + try: + info = lock.lstat() + except FileNotFoundError: + continue + blocked = not stat.S_ISDIR(info.st_mode) + # A lease nobody renewed is left by a CLI that died holding it, and is + # taken over the way proper-lockfile itself does. + if not blocked and time.time() - info.st_mtime > LOCK_STALE_SECONDS: + try: + lock.rmdir() + continue + except OSError: + pass + if blocked or time.monotonic() >= deadline: raise KimiError("Kimi sign-in is busy; start Kimi and try again") time.sleep(0.1) @@ -269,7 +286,7 @@ def refresh_credentials(auth: dict[str, str] | None = None) -> dict[str, Any]: return credentials refresh_token = str(credentials.get("refresh_token") or "") if not refresh_token: - raise KimiError("Kimi sign-in is missing; sign in with `kimi login`") + raise KimiError("Kimi sign-in expired; sign in again with `kimi login`", refused=True) granted = request_refresh(refresh_token) check() latest = read_credentials() @@ -305,7 +322,7 @@ def api_key(auth: dict[str, str] | None = None) -> str: if auth is not None: auth["key"] = str(credentials.get("access_token") or "") if not credentials.get("access_token"): - raise KimiError(AUTH_HELP) + raise KimiError(AUTH_HELP, missing=True) if not token_fresh(credentials): credentials = refresh_credentials(auth) return str(credentials.get("access_token") or "") @@ -321,23 +338,24 @@ def quota(key: str) -> dict[str, Any]: headers["Authorization"] = "Bearer " + key request = urllib.request.Request(url, headers=headers) try: - with urllib.request.build_opener(NoRedirectHandler()).open(request, timeout=15) as response: - decoded = json.load(response) - except urllib.error.HTTPError as error: - if error.code in (401, 403): + decoded = open_json(request, timeout=15) + except HttpError as error: + if error.status in (401, 403): remedy = "Update or unset KIMI_API_KEY." if os.environ.get("KIMI_API_KEY", "").strip() else "Sign in again with `kimi login`." - raise KimiError("Kimi rejected the credentials. " + remedy) from error - if error.code == 404: + raise KimiError("Kimi rejected the credentials. " + remedy, refused=True) from error + if error.status == 404: raise KimiError("Kimi usage endpoint not available. Try Kimi for Coding.") from error - raise KimiError(f"Kimi API returned HTTP {error.code}") from error - except urllib.error.URLError as error: - raise KimiError("Could not reach the Kimi API", retry=True) from error - except TimeoutError as error: - raise KimiError("Kimi usage request timed out", retry=True) from error - except json.JSONDecodeError as error: + if error.status: + raise KimiError(f"Kimi API returned HTTP {error.status}") from error + if error.transport: + raise KimiError("Could not reach the Kimi API", retry=True) from error raise KimiError("Kimi returned an invalid usage response") from error if not isinstance(decoded, dict): raise KimiError("Kimi returned an empty usage response") + # A reply naming none of the quota shapes isn't an answer, so it fails the + # check rather than replacing the kept limits with none. + if not any(key in decoded for key in ("usages", "usage", "limits")): + raise KimiError("Kimi returned an invalid usage response") return decoded @@ -353,14 +371,16 @@ def reset_at(detail: dict[str, Any]) -> str: raw = str(detail.get(key) or "").strip() if not raw: continue + # A time that won't parse names no reset, rather than passing through as + # one the panel can't read. if raw.isdigit(): - stamp = int(raw) - if stamp < 1_000_000_000_000: - stamp *= 1000 try: + stamp = int(raw) + if stamp < 1_000_000_000_000: + stamp *= 1000 return datetime.fromtimestamp(stamp / 1000, tz=timezone.utc).isoformat().replace("+00:00", "Z") except (OverflowError, OSError, ValueError): - return raw + return "" # Kimi stamps carry nanosecond fractions Python cannot parse; trim them. normalized = raw.replace("Z", "+00:00") if "." in normalized: @@ -375,7 +395,7 @@ def reset_at(detail: dict[str, Any]) -> str: try: return datetime.fromisoformat(normalized).isoformat().replace("+00:00", "Z") except ValueError: - return raw + return "" return "" @@ -448,75 +468,315 @@ def to_limits(data: dict[str, Any]) -> list[dict[str, Any]]: return limits -def probe_cache(key: str) -> Path: - identity = str(CREDENTIALS_PATH) + "\0" + USAGE_URL + "\0" + key - digest = hashlib.sha256(identity.encode()).hexdigest() - root = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") / "omarchy" / "agent-usage" - return root / f"kimi-limits-{digest}.json" +# -------------------------------------------------------------- local stats -def read_probe(path: Path) -> dict[str, Any]: +# Kimi Code's home, as resolveKimiHome finds it. +def code_home() -> Path: + if home := os.environ.get("KIMI_CODE_HOME"): + return Path(home).expanduser() + return Path.home() / ".kimi-code" + + +# Kimi as pi, omp and opencode name it: the Kimi Code plan (pi's own +# `kimi-coding` provider and omp's `kimi-code`, both api.kimi.com/coding/v1, +# and models.dev's `kimi-code-plan-global` and `kimi-code-plan-cn`), and +# Moonshot's pay-as-you-go API (pi's and opencode's `moonshotai` and +# `moonshotai-cn`, omp's `moonshot`), which these limits don't meter but +# which counts here, as Claude Code's turns to Moonshot do. Exact ids. +PROVIDERS = ("kimi-coding", "kimi-code", "kimi-code-plan-global", "kimi-code-plan-cn", "moonshotai", "moonshotai-cn", "moonshot") + + +def billed_here(provider: str, model: str, entry: dict) -> bool: + return provider in PROVIDERS + + +# Kimi Code logs each agent of a session, subagents included, to its own +# append-only sessions///agents//wire.jsonl: a +# usage.record for every model request it made, compactions included, the +# input its user gave (turn.prompt, or turn.steer into a running turn), the +# model alias an agent is bound to (profile.bind) and each request is sent on +# (llm.request), answered or not, and, at the end of a fork's copy of its +# source's records, a forked marker. +WIRE_TYPES = (b'"usage.record"', b'"turn.prompt"', b'"turn.steer"', b'"forked"', b'"profile.bind"', b'"llm.request"') +USAGE_FIELDS = ("inputOther", "output", "inputCacheRead", "inputCacheCreation") + + +# What someone typed, as Kimi Code's isUserVisibleTurnInputRecord reads it, +# less a shell command, which asks the model nothing. +def typed(origin: dict[str, Any]) -> bool: + kind = origin.get("kind") + return kind == "user" or kind in ("skill_activation", "plugin_command") and origin.get("trigger") == "user-slash" + + +def read_wire(path: Path) -> tuple[bool, list[tuple]]: + """Whether a wire.jsonl is a fork's, and its rows: ("u", key, day, model, + counts) for a request's usage, ("p", key, day) for a typed input, and + ("b", alias) and ("r", alias) for an alias bound and a request sent on + it. A usage or input row whose time can't be trusted is skipped.""" + forked, rows = False, [] + with path.open("rb") as handle: + for raw in handle: + if not any(marker in raw for marker in WIRE_TYPES): + continue + try: + record = json.loads(raw) + except (ValueError, RecursionError): + continue + if not isinstance(record, dict): + continue + kind, time_ms = record.get("type"), record.get("time") + if kind == "forked": + forked = True + continue + alias = record.get("modelAlias") + if kind in ("profile.bind", "llm.request") and isinstance(alias, str) and alias.strip(): + rows.append(("b" if kind == "profile.bind" else "r", alias.strip())) + continue + day = local_day(time_ms) if isinstance(time_ms, (int, float)) else None + if day is None: + continue + if kind == "usage.record": + model = record.get("model") + usage = obj(record.get("usage")) + counts = tuple(token_count(usage.get(field)) for field in USAGE_FIELDS) + if isinstance(model, str) and model.strip() and sum(counts) > 0: + # A fork copies its source's records as they were, so a record's + # agent, time, model and counts together name it. + rows.append(("u", (str(record.get("agentId")), time_ms, model, counts), day, model.strip(), counts)) + elif kind in ("turn.prompt", "turn.steer") and typed(obj(record.get("origin"))): + key = record.get("promptId" if kind == "turn.prompt" else "messageId") + if not isinstance(key, str) or not key: + key = (str(record.get("agentId")), time_ms) + rows.append(("p", (kind, key), day)) + return forked, rows + + +# A usage.record names the model alias it ran on, and config.toml says where +# that alias goes: its [models.] provider, whose type and base_url +# (else the *_BASE_URL in the provider's env table) name the endpoint, unless +# the model carries a base_url of its own. A provider with no base_url is its +# type's official endpoint: Kimi's for a `kimi` one, no collector's for +# Anthropic's, OpenAI's or Google's. +def code_routes() -> dict[str, tuple[str, str, str]]: + """model alias -> (provider type, base URL, model id), from Kimi Code's + config.toml.""" try: - data = json.loads(path.read_text()) - if isinstance(data, dict) and isinstance(data.get("limits"), list): - stamp = float(data.get("fetchedAt")) - if math.isfinite(stamp) and stamp > 0: - return data - except (OSError, ValueError, TypeError): + with (code_home() / "config.toml").open("rb") as handle: + config = tomllib.load(handle) + except (OSError, UnicodeDecodeError, tomllib.TOMLDecodeError): return {} - return {} + providers, models = obj(config.get("providers")), obj(config.get("models")) + routes = {} + for alias, model in models.items(): + model = obj(model) + provider = obj(providers.get(str(model.get("provider") or ""))) + env = obj(provider.get("env")) + url = model.get("base_url") if model.get("protocol") else None + url = url or provider.get("base_url") or next((value for key, value in env.items() if str(key).endswith("_BASE_URL")), "") + routes[str(alias)] = (str(provider.get("type") or ""), str(url or "").strip(), str(model.get("model") or alias)) + return routes + + +# The prefix of the aliases /login provisions for the Kimi Code plan +# (managedModelKey). +MANAGED_PREFIX = "kimi-code/" + + +def route_owner(alias: str, routes: dict[str, tuple[str, str, str]]) -> str | None: + """The collector that bills a request on `alias`, or None for one nothing + here can place. A base URL is read the way every collector reads Claude + Code's (claude_code_owner): a provider's host is that provider's, and a + server you run is local's, unless the model is a Claude or an Ollama cloud + one it only forwards. Kimi Code sends nothing on an alias config.toml + doesn't name, so one missing from it is either Kimi's own, the alias + KIMI_MODEL_NAME makes up for a run (no provider prefix) or one of the + plan's `kimi-code/` models /logout or a catalog refresh dropped, or a + provider's since removed, which nothing here can place.""" + if alias not in routes: + return AGENT_ID if "/" not in alias or alias.startswith(MANAGED_PREFIX) else None + kind, url, model = routes[alias] + if url: + return claude_code_owner(model, url) + return AGENT_ID if kind == "kimi" else None + + +def harness_sources() -> list[Path]: + return [code_home()] + + +def harness_match() -> dict[str, str]: + """What harness_rows() decided owners under, for cached_scan(match=...): + config.toml's routes, so a scan made before a route changed never stands in + for one after it, in this collector or in local.""" + routes = json.dumps(sorted(code_routes().items())) + return {"kimiCodeRoutes": hashlib.sha256(routes.encode("utf-8", "surrogatepass")).hexdigest()[:16]} + + +def harness_rows(owner: str) -> tuple[list[UsageRow], bool]: + """Kimi Code's usage that `owner` bills, each request counted once though + forks copy them, with each typed prompt as a row of its own billed with the + next request sent for it, answered or not, else with the alias its agent + was bound to, and whether every session was read. A prompt in a log that + names neither, from before Kimi Code recorded them, is Kimi Code's own.""" + routes = code_routes() + complete = True + wires = [] + for path in sorted((code_home() / "sessions").glob("*/*/agents/*/wire.jsonl")): + try: + forked, rows = read_wire(path) + except OSError as exc: + warn(f"skipping {path}: {type(exc).__name__}") + complete = False + continue + wires.append((forked, path, rows)) + # A source's rows go to the source's session, not to a fork read first. + wires.sort(key=lambda wire: wire[0]) + seen = set() + found = [] + for _, path, rows in wires: + session = path.parents[2].name + bound, unanswered = AGENT_ID, [] + for row in rows: + if row[0] == "b": + bound = route_owner(row[1], routes) + unanswered.append(bound) + requested, answered = False, None + owners = [] + for row, fallback in zip(reversed(rows), reversed(unanswered)): + if row[0] in ("u", "r"): + requested, answered = True, route_owner(row[3] if row[0] == "u" else row[1], routes) + owners.append(answered if requested else fallback) + for row, billed in zip(rows, reversed(owners)): + if row[0] not in ("u", "p") or (row[0], row[1]) in seen: + continue + seen.add((row[0], row[1])) + if billed != owner: + continue + if row[0] == "u": + found.append(UsageRow(row[2], session, "", row[3], *row[4], prompts=0)) + else: + found.append(UsageRow(row[2], session, "", "", 0, 0, 0, 0)) + return found, complete + + +def scan(reuse: bool = True) -> tuple[dict[str, Any], bool]: + """Kimi Code's usage and typed prompts on Kimi's routes, pi, omp and + opencode sessions on the Kimi Code plan, Claude Code's turns on Kimi + (pointed at Kimi's or Moonshot's API, or answered by a Kimi model) and + MiniMax Code's custom providers at those hosts, and whether every session + was read.""" + stats = Stats() + code_rows, complete = harness_rows(AGENT_ID) + for row in code_rows: + stats.add_row(row, row.model) + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"(' + "|".join(PROVIDERS) + r')"', billed_here) + # The database's prefilter matches every plan id. + db_rows, db_complete = scan_opencode(["kimi-cod", "moonshot"], billed_here) + for row in routed_rows + pi_rows + db_rows: + stats.add_row(row, row.model.rstrip("/").split("/")[-1] or AGENT_ID) + return stats.snapshot(), complete and routed_complete and pi_complete and db_complete + + +# Keyed by the credential store, the endpoint, and the token, so another +# sign-in or another Kimi service never shows this one's quota. +def limits_cache(key: str) -> LimitsCache: + return LimitsCache(AGENT_ID, str(CREDENTIALS_PATH), USAGE_URL, key) + + +def has_credential() -> bool: + return bool(os.environ.get("KIMI_API_KEY", "").strip() or read_credentials().get("access_token")) + + +# Kimi counts as in use on this machine when it has a credential or a Kimi +# home of its own, where Kimi Code keeps its sessions: a home without a +# sign-in is waiting for one, while a machine with neither never ran Kimi. +def in_use() -> bool: + if has_credential() or CREDENTIALS_PATH.exists(): + return True + if home := os.environ.get("KIMI_CODE_HOME"): + return Path(home).expanduser().is_dir() + return (Path.home() / ".kimi-code").is_dir() or legacy_home().is_dir() + + +def failure_record(error: Exception, cache: LimitsCache | None) -> dict[str, Any]: + # A kept window whose reset has passed comes back at zero, with no reset + # time until Kimi names the next one. + fallback, fetched_at = cache.kept() if cache else ([], 0) + help_text = str(error) if isinstance(error, KimiError) else "Kimi usage scan failed" + if isinstance(error, KimiError) and error.missing or not has_credential(): + status, help_text = "Waiting for auth", AUTH_HELP + elif isinstance(error, KimiError) and error.refused: + status = "Sign-in expired" + elif fallback: + # A failed check that kept the last good limits isn't trouble to flag; + # limitsStale and the age on hover say the rest. + status = "" + else: + status = "Kimi limits unavailable" + record = base_record( + ready=True, + usageStatusText=status, + authHelpText=help_text, + limits=fallback, + limitsStale=bool(fallback), + limitsFetchedAt=fetched_at, + ) + if isinstance(error, KimiError) and error.retry: + record["retryAdvised"] = True + return record def main() -> int: parser = argparse.ArgumentParser(description="Print the Kimi usage record as JSON") - # Flags accepted so the collector answers the same invocation as the others. parser.add_argument("--force", action="store_true") parser.add_argument("--limits-only", action="store_true") args = parser.parse_args() - cached = {} + # The digest covers Kimi Code's home, so another home never reuses this + # one's numbers, pi's and omp's session roots and opencode's database, and + # Claude Code's and MiniMax Code's sessions, whose owner depends on where + # they are sent, as Kimi Code's does on its routes. + sources = [code_home(), *pi_session_roots(), opencode_db(), *routed_sources()] + # A forced scan reads every Codex rollout again too. + max_age = scan_max_age(args.force, args.limits_only) + stats = cached_scan(AGENT_ID, sources, max_age, lambda: scan(reuse=max_age > 0), match={**routed_match(), **harness_match()}) + if not in_use() and not stats.get("modelUsage") and not stats.get("totalPrompts"): + emit(base_record()) + return 0 + + cache = None auth = {} try: key = api_key(auth) if not key: - record = base_record(usageStatusText="Kimi unavailable", authHelpText=AUTH_HELP) - else: - path = probe_cache(key) - cached = read_probe(path) - if not cached and auth.get("previous_key"): - cached = read_probe(probe_cache(auth["previous_key"])) - if cached: - try: - write_json(path, cached) - except OSError: - pass - age = time.time() * 1000 - float(cached.get("fetchedAt") or 0) - if not args.force and 0 <= age < PROBE_REUSE_SECONDS * 1000: - limits, fetched_at = cached["limits"], cached["fetchedAt"] - else: - limits, fetched_at = to_limits(quota(key)), int(time.time() * 1000) + raise KimiError(AUTH_HELP, missing=True) + cache = limits_cache(key) + if not cache.fetched_at and auth.get("previous_key"): + # A refresh just rotated the token: the last answer under the old one + # is still this account's, so it moves to the new key. + previous = limits_cache(auth["previous_key"]) + if previous.fetched_at: try: - write_json(path, {"limits": limits, "fetchedAt": fetched_at}) + write_json(cache.path, previous.data) except OSError: pass - record = base_record( - ready=True, - limits=limits, - limitsFetchedAt=fetched_at, - ) - except KimiError as error: - if not cached: - cached = read_probe(probe_cache(auth.get("key", ""))) - record = base_record(usageStatusText="Kimi unavailable", authHelpText=str(error), - limits=cached.get("limits", []), limitsStale=bool(cached.get("limits")), - limitsFetchedAt=cached.get("fetchedAt", 0), retryAdvised=error.retry) - except Exception: - if not cached: - cached = read_probe(probe_cache(auth.get("key", ""))) - record = base_record(usageStatusText="Kimi unavailable", authHelpText="Kimi usage scan failed", - limits=cached.get("limits", []), limitsStale=bool(cached.get("limits")), - limitsFetchedAt=cached.get("fetchedAt", 0)) - print(json.dumps(record, separators=(",", ":"))) + cache.reload() + # A plan without windows is an answer too, and stands in like one. + if cache.fresh(args.force, allow_empty=True): + limits, fetched_at = cache.limits, cache.fetched_at + else: + limits = to_limits(quota(key)) + fetched_at = cache.save(limits) + record = base_record(ready=True, limits=limits, limitsFetchedAt=fetched_at) + except Exception as error: + if (cache is None or not cache.fetched_at) and auth.get("key"): + cache = limits_cache(auth["key"]) + record = failure_record(error, cache) + record.update(stats) + record["hasLocalStats"] = bool(stats.get("modelUsage") or stats.get("totalSessions")) + emit(record) return 0 diff --git a/bin/omarchy-agent-usage-local b/bin/omarchy-agent-usage-local new file mode 100755 index 00000000000..28b36f43b6f --- /dev/null +++ b/bin/omarchy-agent-usage-local @@ -0,0 +1,319 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the self-hosted model usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect self-hosted model usage into one display-ready JSON record. + +Every other collector counts what a subscription or a provider bills. This +one counts the turns that ran on hardware you own: Codex CLI pointed at +Ollama, opencode on a local gateway, pi on llama.cpp, Kimi Code, Qwen Code, +Hermes or Mistral Vibe on a server of your own, whether the endpoint is this +machine or a box down the hall. It is the exact complement of the +provider collectors: a provider id another collector counts is never counted +here, and the Codex collector's own `openai` rollouts never are either. + +A provider is self-hosted when it is a well-known local runtime, an id named +in ~/.config/omarchy/agents/local.json, or a Codex model_providers entry whose +base_url is on a loopback, private, link-local or LAN address: + + {"providers": ["homelab", "my-gateway"]} + +Ollama's cloud models (a `:cloud` or `-cloud` tag) run on Ollama's servers +even when the local daemon fronts them, so they are never counted here, and +nor are opencode's google rows on antigravity- models, which the agy +collector counts. There +are no limits and no balance: the meter on self-hosted work is the +electricity bill. The agents panel only ever reads the JSON this prints. +""" + +import argparse +import importlib.machinery +import importlib.util +import json +import os +import re +import sys +from pathlib import Path + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + Stats, + base_record, + cached_scan, + cloud_model, + codex_home, + codex_native_key, + codex_native_routed, + codex_rollout_cache, + codex_provider_urls, + codex_route_owner, + codex_rollouts, + config_home, + emit, + number, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, + self_hosted, + warn, +) + +AGENT_ID = "local" +AGENT_NAME = "Local models" + +# Runtimes that serve models from your own hardware, for provider ids with no +# URL to classify: opencode and pi record the id but not the endpoint, and +# Codex's built-in `--oss` providers have no config.toml entry. +KNOWN_LOCAL_PROVIDERS = { + "gpt4all", "jan", "koboldcpp", "llama-cpp", "llamacpp", "llamafile", + "lm-studio", "lmstudio", "local", "localai", "mlx", "ollama", + "openai-local", "oss", "sglang", "tabby", "text-generation-webui", "vllm", +} + +# Provider ids the provider collectors count whole, from pi, omp and opencode, +# each the exact id that collector's filter matches. A local.json naming one +# of them would count its tokens twice, so they are never local, whatever the +# configuration says: the claiming collector wins. Each collector's own +# provider list must be a subset of this one; agent-usage-partition-test.sh +# fails when an id is added there and not here. Two more are claimed only for some models, and only +# those rows are left out: Ollama's cloud-tagged models (cloud_model) and +# opencode's Antigravity plugin (antigravity_row). +CLAIMED_PROVIDERS = { + "anthropic", # claude + "deepseek", # deepseek + "devin", # devin + "fireworks", "fireworks-ai", # fireworks, which counts the whole account + "kimi-coding", "kimi-code", "kimi-code-plan-global", "kimi-code-plan-cn", # kimi + "moonshotai", "moonshotai-cn", "moonshot", # kimi + "google-antigravity", # agy, pi's Antigravity sign-in + "kimchi", "kimchi-dev", "kimchi-experimental", # kimchi + "meta", "muse-code", # muse + "minimax", "minimax-cn", "minimax-coding-plan", "minimax-cn-coding-plan", # minimax + "mistral", # vibe + "ollama-cloud", # ollama + "openai", "openai-codex", # codex + "opencode", "opencode-go", # opencode + "openrouter", # openrouter + "alibaba-coding-plan", "alibaba-coding-plan-cn", "alibaba-token-plan", "alibaba-token-plan-cn", # qwen + "qwen-token-plan", "qwen-token-plan-cn", "qwen-token-plan-individual", "qwen-portal", # qwen + "synthetic", # synthetic + "zai", "zhipu", "zai-coding-cn", "zai-coding-plan", # zai + "zhipu-coding-plan", "zhipuai", "zhipuai-coding-plan", # zai +} + +# Ids a collector claims by prefix rather than whole: Kimchi 1.18 names each +# gateway route kimchi-dev/, and kimchi counts every one. +CLAIMED_PREFIXES = ("kimchi-dev/",) # kimchi + +def model_name(raw): + return str(raw or "").strip() or AGENT_ID + + +# --------------------------------------------------------- classification + + +def config_path(): + return config_home() / "omarchy" / "agents" / "local.json" + + +def configured_providers(): + names = read_json(config_path()).get("providers") + if not isinstance(names, list): + return set() + return {str(name).strip().lower() for name in names if isinstance(name, str) and name.strip()} + + +def antigravity_row(provider, model): + # opencode's Antigravity plugin runs under a "google" provider with model + # ids it prefixes with antigravity-; the agy collector counts exactly + # those opencode rows, and every other google model stays local's to count + # when local.json names it. + return str(provider or "").strip().lower().startswith("google") and str(model or "").startswith("antigravity-") + + +class Classifier: + def __init__(self): + self.configured = configured_providers() + # Every Codex provider's base_url, remembered past config.toml + # dropping it. + self.endpoints = codex_provider_urls() + ids = set(KNOWN_LOCAL_PROVIDERS) | self.configured + ids |= {pid for pid, url in self.endpoints.items() if url and self_hosted(url) and not codex_native_key(pid)} + self.local_ids = {pid for pid in ids - CLAIMED_PROVIDERS if not pid.startswith(CLAIMED_PREFIXES)} + + def provider_is_local(self, provider): + return str(provider or "").strip().lower() in self.local_ids + + def accepts(self, provider, model): + return self.provider_is_local(provider) and not cloud_model(model) + + def fingerprint(self): + return sorted(self.local_ids) + + +CLASSIFIER = None + + +def classifier(): + global CLASSIFIER + if CLASSIFIER is None: + CLASSIFIER = Classifier() + return CLASSIFIER + + +def pi_pattern(): + ids = "|".join(re.escape(pid) for pid in sorted(classifier().local_ids)) + return r'(?i)"provider"\s*:\s*"(?:' + ids + ')"' + + +def accept_pi(provider, model, message): + # The Codex collector also counts any message on its openai-codex API, + # whatever the provider id says. + if str(message.get("api") or "").startswith("openai-codex"): + return False + return classifier().accepts(provider, model) + + +def accept_opencode(provider, model, entry): + return classifier().accepts(provider, model) and not antigravity_row(provider, model) + + +# --------------------------------------------------- harnesses of their own + +# Kimi Code, Qwen Code, Hermes and Mistral Vibe keep their sessions in stores +# of their own. Each one's collector reads its store and says who bills each +# row (harness_rows), from the route Hermes and older Vibe sessions record +# or, for Kimi Code, Qwen Code and Vibe's Unified Harness, from the current +# config, so this one takes the rows a server you run answered from the same +# reader. +HARNESSES = ("hermes", "kimi", "qwen", "vibe") + + +def harness(name): + path = os.path.join(os.environ["OMARCHY_PATH"], "bin", f"omarchy-agent-usage-{name}") + loader = importlib.machinery.SourceFileLoader(f"omarchy_agent_usage_{name}", path) + module = importlib.util.module_from_spec(importlib.util.spec_from_loader(loader.name, loader)) + loader.exec_module(module) + return module + + +HARNESS_MODULES = {} + + +def harness_modules(): + if not HARNESS_MODULES: + HARNESS_MODULES.update((name, harness(name)) for name in HARNESSES) + return HARNESS_MODULES + + +def scan_harnesses(stats): + complete = True + for name, module in harness_modules().items(): + try: + rows, read = module.harness_rows(AGENT_ID) + except Exception as exc: + # One harness's reader failing costs only its rows. + warn(f"could not read {name} sessions ({type(exc).__name__}: {exc})") + complete = False + continue + complete = complete and read + for row in rows: + stats.add_row(row, model_name(row.model)) + return complete + + +def harness_sources(): + return [source for module in harness_modules().values() for source in module.harness_sources()] + + +# Each harness's route config decides which of its rows are self-hosted, and +# can change without its store changing, so a scan is kept only while every +# harness still classifies the way it did. +def harness_match(): + match = {} + for name, module in harness_modules().items(): + try: + match.update(module.harness_match()) + except Exception as exc: + warn(f"could not read {name}'s routes ({exc})") + match[f"{name}Routes"] = "unreadable" + return match + + +# ------------------------------------------------------------------- Codex + + + +def scan_codex(stats, use_file_cache): + endpoints = classifier().endpoints + records, complete = codex_rollouts(codex_rollout_cache(), reuse=use_file_cache, native=codex_native_routed(endpoints)) + for record in records.values(): + for day, models in record["days"].items(): + for model, (input_tokens, output_tokens, cache_read, cache_write, prompts) in models.items(): + # A provider's host, or Ollama Cloud, bills the rollouts + # codex_route_owner() places, whatever the provider's id, and it + # gives the built-in provider's to this one while openai_base_url + # is self-hosted. + owner = codex_route_owner(record["provider"], model, endpoints) + if owner == "local" or (owner is None and classifier().provider_is_local(record["provider"]) and not cloud_model(model)): + # A Codex rollout's totals arrive per day and model, prompts + # counted, under the session that started it. + stats.add(day, record["session"], model, input_tokens, output_tokens, cache_read, cache_write, prompts=prompts) + return complete + + +# -------------------------------------------------------------- scan cache + + +def scan_sources(): + # Every path the scan reads, and which providers it counts as local, so a + # changed local.json or Codex config is a new scan rather than a stale one. + return [Path.home(), codex_home(), opencode_db(), *pi_session_roots(), *routed_sources(), *harness_sources(), json.dumps(classifier().fingerprint())] + + +def run_local_scans(max_age): + stats = Stats() + complete = scan_codex(stats, use_file_cache=max_age > 0) + pi_rows, pi_complete = scan_pi_sessions(pi_pattern(), accept_pi) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode([pid + '"' for pid in sorted(classifier().local_ids)], accept_opencode) + # Claude Code pointed at a self-hosted endpoint, for the models that ran + # there, and MiniMax Code's custom providers on one. + routed_rows, routed_complete = scan_routed(AGENT_ID) + for row in pi_rows + db_rows + routed_rows: + stats.add_row(row, model_name(row.model)) + harnesses_complete = scan_harnesses(stats) + return stats.snapshot(), complete and pi_complete and db_complete and routed_complete and harnesses_complete + + +# ------------------------------------------------------------------ record + + +def main(): + parser = argparse.ArgumentParser(description="Print the self-hosted model usage record as JSON") + # There are no limits to fetch, so --limits-only only means a redraw may + # reuse an older scan. + parser.add_argument("--force", action="store_true", help="rescan every session, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse a recent scan") + args = parser.parse_args() + + max_age = scan_max_age(args.force, args.limits_only) + stats = cached_scan(AGENT_ID, scan_sources(), max_age, lambda: run_local_scans(max_age), match={**routed_match(), **harness_match()}) + # Self-hosted work has no plan, no limits and no balance. + in_use = number(stats.get("totalSessions")) > 0 + record = base_record(AGENT_ID, AGENT_NAME, ready=in_use, hasLocalStats=in_use) + record.update(stats) + emit(record, sort_keys=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-minimax b/bin/omarchy-agent-usage-minimax new file mode 100755 index 00000000000..f58c0e78ab0 --- /dev/null +++ b/bin/omarchy-agent-usage-minimax @@ -0,0 +1,473 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the MiniMax Token Plan usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect MiniMax usage into a display-ready JSON record. + +The 5-hour session and weekly windows come from MiniMax's Token Plan API with +the key MiniMax's own `mmx` CLI, opencode, or pi saved. Each credential's last +good limits are cached, so a failed check keeps showing them, marked stale, +rather than dropping the section. + +Local stats come from MiniMax Code's own session files, and from pi/omp and +opencode sessions that ran on a MiniMax provider. The agents panel only ever +reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record as shared_base_record, + cached_scan, + data_home, + emit, + open_json, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, + trusted_endpoint, + warn, +) + +AGENT_ID = "minimax" +AGENT_NAME = "MiniMax" +DEFAULT_BASE_URL = "https://api.minimax.io" +CN_BASE_URL = "https://api.minimaxi.com" +AUTH_HELP = "Set MINIMAX_API_KEY, or sign in to MiniMax in opencode, to show Token Plan limits." +AUTH_HELP_BAD_KEY = "MiniMax rejected the saved credentials. Set MINIMAX_API_KEY, or sign in to MiniMax again with `mmx` or opencode." +# The plan the quota endpoint describes: only a Token Plan has these windows. +TIER_LABEL = "Token Plan" +# The pi/omp and opencode provider names MiniMax bills, by region. Exact +# names: a custom gateway in front of MiniMax is not this account. +PROVIDERS = ("minimax", "minimax-cn", "minimax-coding-plan", "minimax-cn-coding-plan") +# base_resp codes MiniMax answers a bad key with: 1004 is an authentication +# failure, 2049 an invalid API key. +AUTH_STATUS_CODES = {1004, 2049} +# The quota rows for MiniMax's media models ride on allowances of their own +# (video and image generation, speech, vision) and are not the coding plan. +MEDIA_MODEL_WORDS = ("video", "image", "audio", "vision") + + +class MiniMaxError(Exception): + def __init__(self, message: str, transport: bool = False, rejected: bool = False) -> None: + super().__init__(message) + self.transport = transport + self.rejected = rejected + + +def base_record(**overrides: Any) -> dict[str, Any]: + return shared_base_record(AGENT_ID, AGENT_NAME, **overrides) + + +def limits_cache_paths(api_key: str, base_url: str) -> tuple[Path, Path]: + # Keyed by the credential and the host it signs in to: another key, or the + # same key in the other region, never shows this one's allowance. + cache = LimitsCache(AGENT_ID, base_url.rstrip("/"), api_key) + return cache.path, cache.lock_path + + +# -------------------------------------------------------------------------- +# Credentials +# -------------------------------------------------------------------------- + +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 parsed: + 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: + """The host mmx itself asks: its base_url, then the OAuth session's + resource_url, then the region (MINIMAX_REGION over the saved one).""" + oauth = config.get("oauth") + oauth = oauth if isinstance(oauth, dict) else {} + for value in (config.get("base_url"), oauth.get("resource_url")): + if isinstance(value, str) and value.strip(): + return value.strip() + # mmx refuses a region it doesn't know; here one just doesn't decide. + for region in (os.environ.get("MINIMAX_REGION"), config.get("region"), oauth.get("region")): + if region in ("cn", "global"): + return CN_BASE_URL if region == "cn" else DEFAULT_BASE_URL + return DEFAULT_BASE_URL + + +def pi_agent_dir() -> Path: + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + +def saved_key(path: Path, providers: tuple[str, ...]) -> tuple[str, str]: + """(key, provider) for the first provider `path` holds a literal key for.""" + parsed = read_json(path) + for provider in providers: + entry = parsed.get(provider) + if isinstance(entry, dict): + value = entry.get("key") or entry.get("apiKey") or entry.get("token") + # A pi `$VAR` or `!command` entry is resolved by pi alone. + if isinstance(value, str) and value.strip() and not value.strip().startswith(("!", "$")): + return value.strip(), provider + return "", "" + + +def credentials() -> tuple[str, str]: + """(key, host). A key only answers on its own region's host, so the host + follows wherever the key came from: mmx's region applies to mmx's key + alone. MINIMAX_BASE_URL overrides any of them.""" + override = os.environ.get("MINIMAX_BASE_URL", "").strip() + env_key = os.environ.get("MINIMAX_API_KEY", "").strip() + if env_key: + return env_key, override or (CN_BASE_URL if os.environ.get("MINIMAX_REGION") == "cn" else DEFAULT_BASE_URL) + + config = mmx_config() + if mmx_key := config_key(config): + return mmx_key, override or config_base_url(config) + + # The key opencode or pi signed in to MiniMax with; a China-region + # provider's key only answers on the China host. + for path in (data_home() / "opencode" / "auth.json", pi_agent_dir() / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"): + # Only the provider ids MiniMax bills: a key saved under any other + # name belongs to a gateway, and is never sent to MiniMax. + api_key, provider = saved_key(path, PROVIDERS) + if api_key: + return api_key, override or (CN_BASE_URL if provider.startswith("minimax-cn") else DEFAULT_BASE_URL) + return "", override or DEFAULT_BASE_URL + + +# -------------------------------------------------------------------------- +# Token Plan limits +# -------------------------------------------------------------------------- + +def number(value: Any) -> float | None: + # Booleans are ints to Python; neither they nor NaN or infinity are numbers + # the quota API meant. + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) + + +def reset_at(entry: dict[str, Any], end_key: str, remains_key: str) -> str: + try: + end_ms = number(entry.get(end_key)) + if end_ms and end_ms > 0: + return datetime.fromtimestamp(end_ms / 1000, timezone.utc).isoformat() + # Without an end time MiniMax still says how long the window has left, + # which places its own reset rather than guessing one. + remains_ms = number(entry.get(remains_key)) + if remains_ms and remains_ms > 0: + return (datetime.now(timezone.utc) + timedelta(milliseconds=remains_ms)).isoformat() + except (OverflowError, OSError, ValueError): + pass + return "" + + +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 is_media_row(row: dict[str, Any]) -> bool: + name = str(row.get("model_name") or "").strip().lower() + return any(word in name for word in MEDIA_MODEL_WORDS) + + +# current_*_status per window: 1 normal, 2 exhausted, 3 unlimited. +STATUS_EXHAUSTED = 2 +STATUS_UNLIMITED = 3 + + +def not_in_plan(row: dict[str, Any]) -> bool: + """The shape mmx shows as "not in plan": both windows unlimited with + nothing allotted, a model the plan doesn't cover rather than a full one.""" + return ( + number(row.get("current_interval_total_count")) == 0 and number(row.get("current_weekly_total_count")) == 0 + and number(row.get("current_interval_status")) == STATUS_UNLIMITED and number(row.get("current_weekly_status")) == STATUS_UNLIMITED + ) + + +def window_used(row: dict[str, Any], percent_key: str, status_key: str) -> float | None: + status = number(row.get(status_key)) + # MiniMax's own status outranks the percentage: an exhausted window may + # omit it or keep a stale one, and an unlimited one has nothing to fill. + if status == STATUS_EXHAUSTED: + return 1.0 + if status == STATUS_UNLIMITED: + return None + return used_fraction(row.get(percent_key)) + + +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. Without it the first row that + # isn't a media model's stands for the plan. A model outside the plan is + # never it. + rows = [row for row in rows if isinstance(row, dict) and not not_in_plan(row)] + general = next((row for row in rows if row.get("model_name") == "general"), None) + if general is None: + general = next((row for row in rows if not is_media_row(row)), None) + if general is None: + raise MiniMaxError("MiniMax returned no usable Token Plan quota") + + limits = [] + interval = window_used(general, "current_interval_remaining_percent", "current_interval_status") + weekly = window_used(general, "current_weekly_remaining_percent", "current_weekly_status") + 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" + # MINIMAX_BASE_URL, --base-url and mmx's own base_url or resource_url move + # the endpoint, never the key onto a connection anyone on the way can read. + if not trusted_endpoint(endpoint): + raise MiniMaxError("The MiniMax endpoint must use HTTPS (plain HTTP only for loopback); the key was not sent.") + request = urllib.request.Request( + endpoint, + headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, + ) + try: + payload = open_json(request, timeout=12) + except HttpError as error: + if error.status in (401, 403): + raise MiniMaxError(AUTH_HELP_BAD_KEY, rejected=True) from error + if 300 <= error.status < 400: + # A redirect would carry the bearer to whatever host it names. + raise MiniMaxError("MiniMax's quota endpoint redirected elsewhere; the key was not sent on.") from error + if error.status: + raise MiniMaxError(f"MiniMax's quota endpoint returned HTTP {error.status}.") from error + if error.transport: + # No answer from any server: no route, no DNS, a timeout, a reset. + raise MiniMaxError("Couldn't reach the MiniMax API. Retrying shortly.", transport=True) from 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): + code = number(base_response.get("status_code")) + if code not in (None, 0): + if int(code) in AUTH_STATUS_CODES: + raise MiniMaxError(AUTH_HELP_BAD_KEY, rejected=True) + message = str(base_response.get("status_msg") or "MiniMax rejected the quota request")[:200] + raise MiniMaxError(message) + return payload + + +def collect_limits(api_key: str, base_url: str, force: bool) -> dict[str, Any]: + """Fresh limits when a probe is due, else the last known ones. + + `live` says whether the limits are this check's (or a probe recent enough + to stand in for one), and `fetchedAtMs` when they were measured. + """ + result = {"limits": [], "usageStatusText": "", "authHelpText": "", "retryAdvised": False, "live": False, "fetchedAtMs": 0} + # Keyed by the credential and the host it signs in to: another key, or the + # same key in the other region, never shows this one's allowance. + cache = LimitsCache(AGENT_ID, base_url.rstrip("/"), api_key) + + # --force is a person asking for fresh numbers, so it skips the reuse + # window; the interval absorbs repeated panel opens, not a refresh press. + if cache.fresh(force): + result.update(limits=cache.limits, live=True, fetchedAtMs=cache.fetched_at) + return result + + with cache.lock(): + if not force: + # Another run may have probed while this one waited for the lock. + cache.reload() + if cache.fresh(force): + result.update(limits=cache.limits, live=True, fetchedAtMs=cache.fetched_at) + return result + try: + limits = quota_limits(fetch_quota(api_key, base_url)) + except MiniMaxError as error: + failure = error + except Exception as error: + # An answer shaped like nothing expected is still a failed check: the + # kept limits stand. + failure = MiniMaxError(f"MiniMax's quota check failed ({type(error).__name__}).") + else: + result.update(limits=limits, live=True, fetchedAtMs=cache.save(limits)) + return result + + # A transport failure reached no server — typically the seconds after login + # before the network is up — so ask the shell to retry sooner. A kept + # window whose reset has passed comes back at zero. + fallback, fetched_at = cache.kept() + result["retryAdvised"] = failure.transport + result["authHelpText"] = str(failure) + if fallback: + result.update(limits=fallback, fetchedAtMs=fetched_at) + if failure.rejected: + result["usageStatusText"] = "Sign-in expired" + elif not fallback: + result["usageStatusText"] = "MiniMax limits unavailable" + return result + + +# -------------------------------------------------------------------------- +# Local session stats +# -------------------------------------------------------------------------- + +# pi, omp, and opencode name the provider alike. Exact names: a custom +# gateway in front of MiniMax is not this account. +def billed_here(provider: str, model: str, entry: dict) -> bool: + return provider in PROVIDERS + + +def scan_sessions(reuse: bool = True) -> tuple[dict[str, Any], bool]: + stats = Stats() + # Every source runs, so one unreadable source still leaves the others. + # MiniMax Code on MiniMax's own routes or a custom provider at MiniMax's + # host, and Claude Code pointed at MiniMax or answered by a MiniMax model. + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + for row in routed_rows: + stats.add_row(row, row.model or "minimax") + pattern = r'"provider"\s*:\s*"(' + "|".join(PROVIDERS) + r')"' + pi_rows, pi_complete = scan_pi_sessions(pattern, billed_here) + for row in pi_rows: + # pi counts each assistant turn as the prompt that asked for it. + stats.add_row(row, row.model or "minimax") + # The database's prefilter matches every MiniMax provider id. + db_rows, db_complete = scan_opencode(["minimax"], billed_here) + for row in db_rows: + stats.add_row(row, row.model.rstrip("/").split("/")[-1] or "minimax") + return stats.snapshot(), routed_complete and pi_complete and db_complete + + +def cached_local_stats(max_age_seconds: float) -> dict[str, Any]: + # The digest covers every data path the scan reads. + sources = [opencode_db(), *pi_session_roots(), *routed_sources()] + # A forced scan reads every Codex rollout again too. + return cached_scan(AGENT_ID, sources, max_age_seconds, lambda: scan_sessions(reuse=max_age_seconds > 0), match=routed_match()) + + +# -------------------------------------------------------------------------- +# Record +# -------------------------------------------------------------------------- + +# Nobody using MiniMax here: an empty record, which the panel skips. With a +# key, a failed check keeps the last good limits, marked stale; without one, +# local usage still shows and the limits wait for a key. +def collect(api_key: str, base_url: str, force: bool = False, stats: dict[str, Any] | None = None) -> dict[str, Any]: + record = base_record() + if stats: + record.update(stats) + has_stats = bool(record["modelUsage"]) or record["totalPrompts"] > 0 + record["hasLocalStats"] = has_stats + + if not api_key: + if has_stats: + record.update(ready=True, usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + return record + + limits = collect_limits(api_key, base_url, force) + record.update( + ready=True, + tierLabel=TIER_LABEL if limits["limits"] else "", + limits=limits["limits"], + limitsStale=bool(limits["limits"]) and not limits["live"], + limitsFetchedAt=limits["fetchedAtMs"], + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + ) + if limits["retryAdvised"]: + record["retryAdvised"] = True + return record + + +def main() -> int: + parser = argparse.ArgumentParser(description="Print the MiniMax usage record as JSON") + # --force rescans and re-probes, ignoring both caches. --limits-only reuses + # a recent local scan (only the limits probe must be fresh); a bare run + # reuses a scan only long enough to dedup concurrent collector runs. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + parser.add_argument("--base-url") + args = parser.parse_args() + + try: + stats = cached_local_stats(scan_max_age(args.force, args.limits_only)) + except Exception as error: + # Local stats are an addition; their failure must not cost the limits. + warn(f"local scan failed ({type(error).__name__})") + stats = None + + try: + api_key, saved_base_url = credentials() + record = collect(api_key, args.base_url or saved_base_url, args.force, stats) + except Exception as error: + record = base_record(usageStatusText="MiniMax limits unavailable", authHelpText="MiniMax usage scan failed.") + warn(type(error).__name__) + emit(record, sort_keys=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-muse b/bin/omarchy-agent-usage-muse new file mode 100755 index 00000000000..8b415d83bcc --- /dev/null +++ b/bin/omarchy-agent-usage-muse @@ -0,0 +1,755 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Muse (Meta) usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Muse (Meta Muse Spark) local usage plus subscription windows. + +Local metering (always available, no network): token counts, prompt and +session counts, per-day and per-model breakdowns. Sources (all local): + - native `muse` CLI sessions in $XDG_DATA_HOME/muse/sessions, defaulting to + ~/.local/share/muse/sessions (`model_completed` events carry usage), + with the real prompt and session counts from the CLI's own + session-index.db beside them when it is there, + - opencode sessions with providerID == "meta" (any Muse Spark model). + Muse Spark served through OpenCode Zen (providerID == "opencode") is + Zen usage, and the opencode record counts it. + - pi/omp sessions on the `meta` provider, or omp's `muse-code` (its Muse + Code sign-in), in every session root (profiles and PI_CODING_AGENT_DIR + included). + +Subscription windows (additive, authoritative only): when the CLI is signed +in via `muse login`, the collector probes the client's own key endpoint +(`POST https://api.meta.ai/muse-code/key` with the OAuth token from +`MUSE_AUTH_PATH`, else `~/.config/muse/auth.json`) for `subs_usage` +session (5-hour) + weekly percents with reset times plus `subs_tier_name` — +the same figures the TUI's `/usage` overlay shows. Minting is idempotent so +polling is safe; the token travels only in the Authorization header and +neither it nor the returned API key is persisted, and the login file is only +ever read, never refreshed or rewritten. A failed check keeps the last good +windows, marked stale. Pay-as-you-go accounts +and active plans that return no `subs_usage` (e.g. Muse Code Power Usage) +have no subscription windows, so limits stay empty — the plan name (none for +pay-as-you-go) plus local metering still show. No estimated meters from manual caps: without an +authoritative probe the panel shows measured usage only. The agents panel +only ever reads the JSON this prints. +""" + +import argparse +import hashlib +import json +import math +import os +import re +import sqlite3 +import sys +import time +import urllib.request +from datetime import date, datetime, timedelta, timezone +from pathlib import Path + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + account_key, + base_record, + cache_root, + cached_scan, + config_home, + data_home, + emit, + now_ms, + number, + open_json, + opencode_db, + pi_session_roots, + read_json, + scan_max_age, + scan_opencode, + scan_pi_sessions, + token_count, + trusted_endpoint, + warn, + write_json, +) + +AGENT_ID = "muse" +AGENT_NAME = "Muse" + +# Model ids and session files are read from local data but flow into the +# panel JSON, so bound them: overlong model names fail the row closed, and +# symlinked or gigantic session files are skipped instead of followed. +MAX_MODEL_NAME_LEN = 120 +MAX_SESSION_BYTES = 512 * 1024 * 1024 + +# What the scan counts, stored beside the shared scan cache's own version and +# required to match on reuse. Bump when the counting changes so old caches +# are rescanned instead of trusted. 3: prompts and sessions from the session +# index. 4: pi/omp sessions counted. 5: native sessions the index hasn't seen +# counted. 6: today's prompts from what each session's count grew by today. +# 7: a child's calls known only from the session's usage attribution. +CACHE_SCHEMA_VERSION = 7 + + +def key_endpoint() -> str: + # Overridable for tests, which point it at a local stub instead of Meta's + # key service. + return os.environ.get("MUSE_KEY_ENDPOINT") or "https://api.meta.ai/muse-code/key" + + +def percent_number(value): + if isinstance(value, bool): + return -1.0 + try: + parsed = float(value) + except (TypeError, ValueError): + return -1.0 + return min(1.0, max(0.0, parsed / 100.0)) if math.isfinite(parsed) else -1.0 + + +def model_name(raw): + value = str(raw or "muse-spark").rstrip("/").split("/")[-1] + if not value: + value = "muse-spark" + # Model ids flow into the panel JSON: cap the length so a crafted session + # file or DB row cannot bloat the record or the UI. Callers skip the row + # when this returns None (fail closed). + if len(value) > MAX_MODEL_NAME_LEN: + return None + return value + + +# Muse's own stores date rows in microseconds, which the shared reader of pi +# and opencode times doesn't expect, and Muse can't have written before 2000 +# or from the future. +def local_day(value): + """Normalize assorted timestamps to a local YYYY-MM-DD date string. + + A stamp that isn't a plausible moment (missing, zero, negative, NaN, a + bool, far future, not a time at all) gives None, so the caller can fall + back or leave the row undated rather than file it under today. + """ + if value is None or isinstance(value, bool): + return None + if isinstance(value, (int, float)): + seconds = float(value) + if not math.isfinite(seconds) or seconds <= 0: + return None + # Native muse recorded_at is microseconds; opencode created is millis; + # codex-style seconds also possible. + if seconds > 10_000_000_000_000: + seconds = seconds / 1_000_000 + elif seconds > 10_000_000_000: + seconds = seconds / 1000 + else: + text = str(value).strip() + try: + # Some writers emit the epoch as a string: route micro/milli/second + # magnitudes through the numeric branch instead of misfiling them. + return local_day(float(text)) + except ValueError: + pass + try: + dt = datetime.fromisoformat(text[:-1] + "+00:00" if text.endswith("Z") else text) + # A zoneless time at the edge of the calendar (year 1 or 9999) can't + # be placed in a zone off UTC. + if dt.tzinfo is None: + dt = dt.astimezone() + seconds = dt.timestamp() + except (OverflowError, OSError, ValueError): + return None + # Muse didn't exist before 2000 and can't write from the future. + if not 946_684_800 <= seconds <= time.time() + 2 * 86_400: + return None + try: + return datetime.fromtimestamp(seconds).astimezone().date().isoformat() + except (OverflowError, OSError, ValueError): + return None + + +def muse_data_dir(): + override = (os.environ.get("MUSE_DATA_DIR") or "").strip() + if override: + return Path(os.path.expandvars(os.path.expanduser(override))) + return data_home() / "muse" + + +# The Muse CLI keeps one row per session in session-index.db with the number +# of prompts typed into it, the real counts a completion tally can only +# approximate (one prompt can run many model calls). Returns the ids of the +# sessions the index has a row for, empty without an index, and whether it +# could be read: an index that's there but unreadable falls back on +# completions without caching that scan. A session nobody prompted isn't +# counted, and a row whose creation time isn't a plausible moment counts in +# the totals only. +def scan_session_index(stats): + database = muse_data_dir() / "session-index.db" + if not database.is_file(): + return set(), True + try: + conn = sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return set(), False + try: + conn.execute("PRAGMA query_only = ON") + rows = conn.execute("SELECT session_id, created_at_us, prompt_count FROM sessions").fetchall() + except sqlite3.Error: + return set(), False + finally: + try: + conn.close() + except Exception: + pass + indexed = {str(session_id) for session_id, _, _ in rows} + counts, created = {}, {} + for session_id, created_at_us, prompt_count in rows: + prompts = token_count(prompt_count) + if not prompts: + continue + session_id = str(session_id) + counts[session_id] = prompts + # created_at_us is microseconds, which local_day recognises by size. + created[session_id] = local_day(created_at_us) + grown = prompts_today(database, counts, {key for key, day in created.items() if day == stats.today}, stats.today) + for session_id, prompts in counts.items(): + # A session's prompts land on the day it was created, which may be + # unknown, apart from today's, which are what its count grew by today. + today_prompts = grown.get(session_id, 0) + stats.add_prompts(created[session_id], "muse-index:" + session_id, prompts - today_prompts) + if today_prompts: + stats.add_prompts(stats.today, "muse-index:" + session_id, today_prompts) + return indexed, True + + +# The index keeps a session's running prompt count and the day it was +# created, not when each prompt came, so a session carried on from an +# earlier day would put every prompt on the day it began. Today's prompts are +# instead what each count grew by since the last scan before today, which +# this collector keeps in a file of its own; a session created today counts +# whole. A session from before today that no earlier scan saw counts from +# when it is first seen, since nothing says when its earlier prompts came; +# so does every session after a gap of days since the last scan. +def prompts_today(database, counts, created_today, today): + digest = hashlib.sha1(str(database).encode("utf-8")).hexdigest()[:16] + path = cache_root() / f"muse-prompts-{digest}.json" + kept = read_json(path) + day = kept.get("day") if isinstance(kept.get("day"), str) else "" + if day == today: + start = kept.get("start") + elif day == (date.fromisoformat(today) - timedelta(days=1)).isoformat(): + start = kept.get("last") + else: + # Never scanned, kept from a clock that went back, or from before a gap + # of days: what a count grew by since then wasn't all typed today. + start = None + start = {str(key): token_count(value) for key, value in start.items()} if isinstance(start, dict) else {} + for session_id, prompts in counts.items(): + if session_id not in start and session_id not in created_today: + start[session_id] = prompts + try: + write_json(path, {"day": today, "start": start, "last": counts}) + except Exception as exc: + warn(f"could not keep the session prompt counts ({exc})") + grown = {} + for session_id, prompts in counts.items(): + added = prompts - start.get(session_id, 0) + if added > 0: + grown[session_id] = added + return grown + + +def muse_sessions_root(): + return muse_data_dir() / "sessions" + + +def iter_native_events(path): + """Yield (event_dict, recorded_at, event_id) tuples from a native session.jsonl file. + + Handles both plain event lines and retained_frame wrappers whose children + carry record_json payloads. A file that can't be read raises OSError, so + the scan knows it is incomplete. + """ + with path.open(errors="replace") as handle: + for raw in handle: + # Session logs run to hundreds of megabytes, mostly lines without + # usage. Only completions ("usage", escaped or not inside a + # retained frame) and the model records ("model_id") matter, so + # everything else skips the JSON parse. + if "usage" not in raw and "model_id" not in raw: + continue + raw = raw.strip() + if not raw: + continue + try: + entry = json.loads(raw) + except Exception: + continue + if not isinstance(entry, dict): + continue + candidates = [entry] + children = entry.get("children") + if isinstance(children, list): + for child in children: + if not isinstance(child, dict): + continue + nested = child.get("record_json") + if not nested: + continue + try: + parsed = json.loads(nested) + except Exception: + continue + if isinstance(parsed, dict): + candidates.append(parsed) + for event in candidates: + payload = event.get("payload") or {} + if not isinstance(payload, dict): + continue + yield payload, event.get("recorded_at"), event.get("id") + + +# A session's log credits every billed run to its owner in a +# goal_usage_attribution record: the main agent, and each child it ran, such +# as the reminder agents that watch a prompt. Most children keep a log of +# their own, whose completions are counted from it, but some (a skill +# reminder among them) leave nothing but this record, the only trace of the +# calls they made. Returns (input, output, cache read, cache write) for such a +# child's record, once per usage id, and None for every other. +def attributed_usage(record, logged, seen): + if not isinstance(record, dict): + return None + owner = record.get("owner") if isinstance(record.get("owner"), dict) else {} + quantity = record.get("quantity") if isinstance(record.get("quantity"), dict) else {} + owner_id = owner.get("owner_id") + usage_id = record.get("usage_id") + if not isinstance(owner_id, str) or owner.get("owner_type") == "main_root" or owner_id in logged: + return None + if quantity.get("unit") != "tokens" or quantity.get("reported") is False: + return None + if not isinstance(usage_id, str) or not usage_id or usage_id in seen: + return None + input_tokens = token_count(quantity.get("input_tokens")) + cache_read = min(input_tokens, token_count(quantity.get("cached_tokens"))) + output_tokens = token_count(quantity.get("output_tokens")) + if not (input_tokens or output_tokens): + return None + seen.add(usage_id) + return input_tokens - cache_read, output_tokens, cache_read, 0 + + +def scan_native_muse_sessions(stats, indexed=frozenset()): + """Add native Muse session completions. Returns whether every session file + was read (safe to cache); a deliberately skipped one doesn't count.""" + root = muse_sessions_root() + if not root.exists(): + return True + # Full history, no mtime cutoff: all-time totals and active days need every + # session still on disk (the same tradeoff the claude collector makes). + # Session files are small JSONL and the scan cache dedups concurrent runs, + # so each refresh pays one linear pass at most. + try: + resolved_root = root.resolve() + except OSError: + return False + files = [] + try: + # Session views contain copies of the same history. Check only paths + # below the sessions root, so a custom data parent is not excluded. + files = [p for p in root.rglob("session.jsonl") + if not any(part.startswith(".msp-view") for part in p.relative_to(root).parts[:-1]) + and p.is_file()] + except OSError: + return False + complete = True + seen_event_ids = set() + # Every agent with a log of its own: a session and each subagent are named + # by their ids, so a child attributed usage without one is found here. + logged = {p.parent.name for p in files} + seen_usage_ids = set() + for path in files: + try: + # Fail closed on entries that could escape the sessions tree or stall + # the scan: symlinked files are never followed (even when their target + # sits inside the tree, so a link cannot double-count a session), files + # reached via a symlinked directory are rejected by the resolve check, + # and gigantic files are skipped. + if path.is_symlink(): + continue + if not path.resolve().is_relative_to(resolved_root): + continue + st = path.stat() + if st.st_size > MAX_SESSION_BYTES: + continue + mtime_day = local_day(st.st_mtime) + except OSError: + complete = False + continue + # Session directories are named by the index's session id, and subagent + # trees roll up into their parent session: normalize both the main + # session.jsonl and any subagent/.../session.jsonl to the session + # directory, so one user session counts once however many subagents it + # spawned. + session_dir = re.sub(r"(/subagent/.*)?/session\.jsonl$", "", str(path)) + counted = Path(session_dir).name not in indexed + current_model = "" + try: + for payload, recorded_at, event_id in iter_native_events(path): + kind = payload.get("kind") + if kind == "metadata": + record = payload.get("record") or {} + if isinstance(record, dict) and record.get("model_id"): + current_model = model_name(record.get("model_id")) + continue + if kind == "run_model": + record = payload.get("record") or {} + if isinstance(record, dict) and record.get("model_id"): + current_model = model_name(record.get("model_id")) + continue + # Older logs wrap a completion as payload kind "run"; current ones + # (record_type "event", payload_type "runtime.session") carry the event + # with no payload kind at all. + if kind not in ("run", None): + continue + event = payload.get("event") or {} + if not isinstance(event, dict): + continue + if event.get("kind") == "goal_usage_attribution": + attributed = attributed_usage(event.get("record"), logged, seen_usage_ids) + model = model_name(current_model) if current_model is not None else None + if attributed and model: + # A child's calls, on the session's model, answer no prompt. + stats.add(local_day(recorded_at) or mtime_day, "muse:" + session_dir if counted else None, model, + *attributed, prompts=0) + continue + if event.get("kind") != "model_completed": + continue + usage = event.get("usage") or {} + if not isinstance(usage, dict): + continue + raw_model = event.get("model") + # An oversized metadata model is an explicit rejection, not a reason to + # silently put the following completion in the generic fallback bucket. + if not raw_model and current_model is None: + continue + model = model_name(raw_model or current_model) + if model is None: + continue + input_tokens = token_count(usage.get("input_tokens")) + # Native Muse reports reasoning as a subset of output_tokens. + output_tokens = token_count(usage.get("output_tokens")) + # Cached tokens are a subset of input: keep the split without double + # counting (same convention as the codex collector), and never let a + # cache figure claim more than the input it is part of. + cache_read = min(input_tokens, token_count(usage.get("cache_read_tokens", usage.get("cached_tokens")))) + cache_write = min(input_tokens - cache_read, token_count(usage.get("cache_write_tokens"))) + input_tokens = input_tokens - cache_read - cache_write + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + # Native record IDs survive copies and retained-frame wrapping. Claim + # only valid completions: metadata or malformed usage must not suppress + # a later usable record. Missing/malformed IDs cannot identify duplicates. + if isinstance(event_id, str) and event_id: + if event_id in seen_event_ids: + continue + seen_event_ids.add(event_id) + # An unusable stamp falls back on the file's own time. + day = local_day(recorded_at) or mtime_day + # A session the index has no row for yet counts its completions as + # prompts, as it would with no index at all. One it has a row for + # supplied its real prompt and session counts already: its tokens + # still count, but a completion is a model call, not a prompt, and its + # session is the index's. An undated completion counts in the totals + # only. + stats.add(day, "muse:" + session_dir if counted else None, model, + input_tokens, output_tokens, cache_read, cache_write, prompts=1 if counted else 0) + except OSError: + # Whatever it yielded before failing still counts here, but a scan + # missing the rest of a session isn't kept. + complete = False + return complete + + +# pi, omp, and opencode alike: the meta provider is Meta's own Muse Spark API +# and only serves Muse Spark models, so any model on it counts (muse-spark +# 1.1/1.2/1.3, ...), and so does omp's `muse-code`, the same API reached +# through a Muse Code sign-in. Muse Spark through the Zen gateway +# (providerID "opencode") is OpenCode usage, so the opencode record keeps it. +PROVIDERS = ("meta", "muse-code") + + +def billed_here(provider, model, entry): + return provider in PROVIDERS + + +def scan_shared_stores(stats): + """pi/omp and opencode replies on Meta's providers. Returns whether every + source was read in full (safe to cache).""" + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"(' + "|".join(PROVIDERS) + r')"', billed_here) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode([provider + '"' for provider in PROVIDERS], billed_here) + for row in db_rows + pi_rows: + model = model_name(row.model) + if model is not None: + stats.add_row(row, model) + return pi_complete and db_complete + + +# ------------------------------------------------------------------ limits +# +# Additive subscription windows only. Local metering above always runs; this +# section only fills `limits`/`tierLabel`/status when the authoritative key +# probe answers. Pay-as-you-go and active plans without `subs_usage` have no +# subscription windows: no token, a PAYG probe, or an active probe with no +# subs_usage means limits stay empty (never estimated, never zeroed) and +# local metering still shows. The OAuth token goes nowhere but +# the Authorization header; only plan label + usage figures enter the record. + + +def muse_config_dir() -> Path: + return config_home() / "muse" + + +def read_access_token() -> str: + try: + raw = (os.environ.get("MUSE_AUTH_PATH") or "").strip() + path = Path(os.path.expandvars(os.path.expanduser(raw))) if raw else (muse_config_dir() / "auth.json") + data = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return "" + # Every level checked: a login file that parses to a list, or whose + # `providers` is not an object, is a missing credential, not a crash. + if not isinstance(data, dict): + return "" + providers = data.get("providers") + if not isinstance(providers, dict): + return "" + login = providers.get("meta") + if not isinstance(login, dict): + return "" + return str(login.get("access_token") or "") + + +def normalize_reset_at(value) -> str: + if value is None or isinstance(value, bool): + return "" + try: + ts = float(value) + except Exception: + # Some answers give the reset as an ISO time, which omp reads too. + try: + parsed = datetime.fromisoformat(str(value).strip().replace("Z", "+00:00")) + except ValueError: + return "" + return (parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc)).isoformat() + if ts <= 0: + return "" + # Epoch seconds; accept milliseconds too. + if ts >= 1e12: + ts = ts / 1000 + try: + return datetime.fromtimestamp(ts, timezone.utc).isoformat() + except Exception: + return "" + + +def probe_key(access_token: str) -> dict: + # MUSE_KEY_ENDPOINT moves the endpoint, never the token onto a connection + # anyone on the way can read. + if not trusted_endpoint(key_endpoint()): + return {"ok": False, "helpText": "MUSE_KEY_ENDPOINT must use HTTPS (plain HTTP only for loopback). Local Muse stats are still shown."} + request = urllib.request.Request( + key_endpoint(), + data=json.dumps({}).encode(), + method="POST", + headers={ + "Authorization": "Bearer " + access_token, + "Accept": "application/json", + "Content-Type": "application/json", + }, + ) + try: + payload = open_json(request, timeout=10) + except HttpError as error: + if error.status in (401, 403): + return {"ok": False, "expired": True, + "helpText": "Muse's saved sign-in was rejected. Run `muse login` to refresh it."} + if error.status: + return {"ok": False, + "helpText": f"Muse's usage endpoint returned status {error.status}. Local Muse stats are still shown."} + if error.transport: + return {"ok": False, "transport": True, + "helpText": "Couldn't reach Muse's usage endpoint. Retrying shortly. Local Muse stats are still shown."} + # The server answered but not with JSON: a shape change, not a network + # failure. No retryAdvised — re-asking won't fix the parser. + return {"ok": False, "helpText": "Muse's usage endpoint returned an unreadable response. Local Muse stats are still shown."} + + usage = payload.get("subs_usage") if isinstance(payload, dict) else None + if not isinstance(usage, dict): + # PAYG account: key mints fine but there is no subscription window to + # meter (live shape: is_subs_active=false, subs_tier_name=null, no + # subs_usage). Local metering is the whole story — stay silent, not an + # error. + if isinstance(payload, dict) and payload.get("is_subs_active") is False: + return {"ok": False, "payg": True} + # Active plan without windows (live shape on Muse Code Power Usage: + # is_subs_active=true, subs_tier_name set, no subs_usage key at all). + # That is the steady answer, not a snapshot about to appear: keep the + # tier, show local metering only, stay silent like PAYG. + if isinstance(payload, dict) and payload.get("is_subs_active") is True: + tier_hint = str(payload.get("subs_tier_name") or "") + return {"ok": False, "silent": True, "tierLabel": tier_hint} + tier_hint = str(payload.get("subs_tier_name") or "") if isinstance(payload, dict) else "" + return {"ok": False, "tierLabel": tier_hint, + "helpText": "Muse's usage endpoint returned no usage. Local Muse stats are still shown."} + + limits = [] + window = usage.get("window") or {} + if isinstance(window, dict) and window.get("used_percent") is not None: + percent = percent_number(window.get("used_percent")) + if percent >= 0: + mins = number(window.get("window_duration_mins")) + if mins >= 60 and mins % 60 == 0: + label = f"Session ({mins // 60}-hour)" + else: + label = f"Session ({mins}m)" if mins > 0 else "Session" + limits.append({"label": label, "percent": percent, "resetsAt": normalize_reset_at(window.get("resets_at"))}) + weekly = usage.get("weekly") or {} + if isinstance(weekly, dict) and weekly.get("used_percent") is not None: + percent = percent_number(weekly.get("used_percent")) + if percent >= 0: + limits.append({"label": "Weekly (7-day)", "percent": percent, "resetsAt": normalize_reset_at(weekly.get("resets_at"))}) + + if not limits: + tier_hint = str(payload.get("subs_tier_name") or "") if isinstance(payload, dict) else "" + return {"ok": False, "tierLabel": tier_hint, + "helpText": "Muse's usage endpoint returned no limits. Local Muse stats are still shown."} + tier = "" + if isinstance(payload, dict): + tier = str(payload.get("subs_tier_name") or "") + return {"ok": True, "limits": limits, "tierLabel": tier} + + +def collect_limits(force: bool) -> dict: + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "stale": False, "fetchedAtMs": 0} + + access_token = read_access_token() + if access_token == "": + # No login. Muse Spark used only through opencode's meta provider signs + # in with a Model API key, not `muse login`, so that alone is no reason + # to ask for one; the Muse CLI's own home on this machine is. + if muse_config_dir().is_dir() or muse_sessions_root().is_dir(): + result["usageStatusText"] = "Waiting for auth" + result["authHelpText"] = "Run `muse login` to show Muse's session and weekly limits." + return result + result["signedIn"] = True + + # Keyed by who the token signs in as (the JWT's subject when it is one), so + # a machine signed in to someone else never shows the last one's windows, + # and a refreshed token keeps the same cache. A window kept from an earlier + # check whose reset has passed comes back at zero: the period it measured + # is over. + try: + probe_cache = LimitsCache(AGENT_ID, account_key(access_token)) + cached, fallback, fetched_at = probe_cache.data, probe_cache.limits, probe_cache.fetched_at + except Exception as exc: + warn(f"could not read limits cache ({exc})") + probe_cache, cached, fallback, fetched_at = None, {}, [], 0 + result["tierLabel"] = str(cached.get("tierLabel") or "") + + # The last good windows, marked stale, with when they were measured. + def keep_cached(): + result["limits"] = fallback + result["stale"] = True + result["fetchedAtMs"] = fetched_at if fallback else 0 + + # Minting is idempotent and cheap (the same key comes back every call), so + # the reuse only absorbs repeated panel opens. A plan without windows is an + # answer too, and stands in like any other. + if probe_cache is not None and probe_cache.fresh(force, allow_empty=True): + result["limits"] = fallback + result["fetchedAtMs"] = fetched_at if fallback else 0 + return result + + probe = probe_key(access_token) + if probe["ok"]: + result["limits"] = probe["limits"] + result["tierLabel"] = probe.get("tierLabel") or result["tierLabel"] + if probe_cache is not None: + result["fetchedAtMs"] = probe_cache.save(probe["limits"], tierLabel=result["tierLabel"]) + else: + result["fetchedAtMs"] = now_ms() + return result + + if probe.get("transport"): + result["retryAdvised"] = True + if probe.get("payg") or probe.get("silent"): + # Authenticated PAYG names no plan; an active plan without windows keeps + # its name. Either way local metering stands alone, silent — not an + # error, not expired — and the cached windows go with the old plan, so + # a later failed check can't bring them back. + result["tierLabel"] = "" if probe.get("payg") else probe.get("tierLabel") or result["tierLabel"] + if probe_cache is not None: + probe_cache.save([], tierLabel=result["tierLabel"], noWindows=True) + return result + keep_cached() + if probe.get("expired"): + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = probe["helpText"] + elif not fallback: + result["usageStatusText"] = "Muse limits unavailable" + result["authHelpText"] = probe["helpText"] + result["tierLabel"] = probe.get("tierLabel") or result["tierLabel"] + return result + + +def scan_sources(): + # The digest covers every data path the scan reads. + return [Path.home(), muse_sessions_root(), opencode_db(), *pi_session_roots()] + + +def run_local_scans(): + stats = Stats() + indexed, index_read = scan_session_index(stats) + native_read = scan_native_muse_sessions(stats, indexed) + complete = scan_shared_stores(stats) + return stats.snapshot(), complete and index_read and native_read + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true", help="rescan sessions, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse any recent scan; kept for CLI compatibility with the panel refresh") + args = parser.parse_args() + + max_age = scan_max_age(args.force, args.limits_only) + stats = cached_scan(AGENT_ID, scan_sources(), max_age, run_local_scans, match={"museSchema": CACHE_SCHEMA_VERSION}) + limits = collect_limits(args.force) + + record = base_record( + AGENT_ID, AGENT_NAME, + # A sign-in shows even before its first session or window, and so does + # a Muse CLI waiting for one. + ready=number(stats.get("totalPrompts")) > 0 or len(limits["limits"]) > 0 or bool(limits.get("signedIn")) or limits["usageStatusText"] != "", + # A sign-in with no sessions here has no local stats to show. + hasLocalStats=number(stats.get("totalSessions")) > 0 or number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")), + # Prompts and sessions are counted, so the panel shows them in tooltips. + hasPromptStats=True, + tierLabel=limits["tierLabel"], + limits=limits["limits"], + limitsStale=limits["stale"], + limitsFetchedAt=limits["fetchedAtMs"], + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + ) + record.update(stats) + if limits.get("retryAdvised"): + record["retryAdvised"] = True + emit(record, 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..230d07c22e7 --- /dev/null +++ b/bin/omarchy-agent-usage-ollama @@ -0,0 +1,367 @@ +#!/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, opencode's message +database, Codex's `--oss` rollouts and Claude Code's transcripts that ran on +the `ollama-cloud` provider (or on a cloud model through the local daemon), +and the account's 5-hour session and 7-day weekly +limits (or the monthly included credits) from Ollama's balance 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 math +import os +import sys +import urllib.request +from datetime import datetime +from pathlib import Path +from typing import Any + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record, + cached_scan as shared_cached_scan, + cloud_model, + config_home, + emit, + limits_cache_path as shared_limits_cache_path, + open_json, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, +) + +AGENT_ID = "ollama" +AGENT_NAME = "Ollama Cloud" +PROVIDER = "ollama-cloud" +BALANCE_ENDPOINT = "https://ollama.com/api/balance" +ME_ENDPOINT = "https://ollama.com/api/me" +AUTH_HELP_NO_KEY = "Set OLLAMA_API_KEY, or sign in to pi with an Ollama Cloud key, to show the account limits." +AUTH_HELP_BAD_KEY = ( + "Ollama Cloud rejected the API key. Update OLLAMA_API_KEY or ~/.config/omarchy/agents/ollama.json," + " or sign in to pi again. Local stats are still shown." +) +# GET /api/balance (docs.ollama.com/api/balance) answers a legacy plan with +# `included.session` / `included.weekly` ({remaining_percent 0..100, +# resets_at}) and a current plan with `included.balance_usd`, +# `allowance_usd` and `period.until`. /api/usage only reports token and cost +# history, never the limits. + + +def config_dir() -> Path: + return config_home() / "omarchy" / "agents" + + +def read_config() -> dict[str, Any]: + return read_json(config_dir() / "ollama.json") + + +def pi_agent_dir() -> Path: + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + +def api_key() -> str: + env = os.environ.get("OLLAMA_API_KEY", "").strip() + if env: + 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 balance endpoint accepts the same key. Only a + # literal key is taken: a `$VAR` or `!command` entry is resolved by pi alone. + for auth_path in (pi_agent_dir() / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"): + entry = read_json(auth_path).get(PROVIDER) + key = str(entry.get("key") or "").strip() if isinstance(entry, dict) else "" + if key and not key.startswith(("!", "$")): + return key + return "" + + +# ---------------------------------------------------------------- local scan +# +# pi and omp record every assistant message with its provider, model, and +# token usage in JSONL session files, and opencode in its message database, +# under the same provider ids. The `ollama-cloud` provider is the +# Ollama Cloud API, which serves Claude models alongside the open ones. A +# cloud-tagged model ("glm-4.6:cloud", "gpt-oss:120b-cloud") reached through +# the local daemon's `ollama` provider runs on Ollama's servers too, so it +# counts here; the local collector leaves exactly those out. Codex's `--oss` +# rollouts name the same `ollama` provider, and count here on the same +# models. Claude Code pointed at ollama.com, or answering with a cloud-tagged +# model, is Ollama's too, as is a MiniMax Code custom provider at ollama.com. + + +def billed_by_ollama_cloud(provider: str, model: str, entry: Any = None) -> bool: + return provider == PROVIDER or (provider == "ollama" and cloud_model(model)) + + +# Ollama model ids carry a version tag ("deepseek-v4-pro:0813"); merge tags +# so the chart groups by model, not by release. +def model_name(raw: str) -> str: + return (raw or "ollama").split(":", 1)[0] + + +def scan_sessions(model_prefix: str, reuse: bool = True) -> tuple[dict[str, Any], bool]: + stats = Stats() + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"ollama(-cloud)?"', billed_by_ollama_cloud) + # The database's prefilter matches both provider ids. + db_rows, db_complete = scan_opencode(["ollama"], billed_by_ollama_cloud) + # Codex rollouts come with them, by codex_route_owner(): the `--oss` + # daemon's cloud-tagged models, a model_providers entry named + # `ollama-cloud`, or one whose base_url is ollama.com. + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + for row in pi_rows + db_rows + routed_rows: + model = model_name(row.model) + if model_prefix and not model.lower().startswith(model_prefix): + continue + stats.add_row(row, model) + return stats.snapshot(), pi_complete and db_complete and routed_complete + + +def cached_scan(max_age_seconds: float, model_prefix: str) -> dict[str, Any]: + # The digest covers every data path the scan reads: the pi/omp session + # roots, the opencode DB, Codex's home, and Claude Code's and MiniMax + # Code's sessions. A scan only answers for the model filter it applied, and + # the endpoints those were sent to. + return shared_cached_scan( + AGENT_ID, + [*pi_session_roots(), opencode_db(), *routed_sources()], + max_age_seconds, + lambda: scan_sessions(model_prefix, reuse=max_age_seconds > 0), + match={"modelPrefix": model_prefix, **routed_match()}, + ) + + +# ------------------------------------------------------------------- limits + + +def clamp_fraction(value: Any) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0: + return -1.0 + return min(1.0, float(value)) + + +def reset_time(value: Any) -> str: + raw = value.strip() if isinstance(value, str) else "" + try: + datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return "" + return raw + + +def amount(value: Any) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) + + +def balance_limits(payload: Any) -> list[dict[str, Any]]: + included = payload.get("included") if isinstance(payload, dict) else None + if not isinstance(included, dict): + return [] + out = [] + for key, label in (("session", "Session (5-hour)"), ("weekly", "Weekly (7-day)")): + window = included.get(key) + if not isinstance(window, dict): + continue + # remaining_percent runs 0..100 and says what is left, not what is used. + remaining = amount(window.get("remaining_percent")) + remaining = clamp_fraction(remaining / 100) if remaining is not None else -1.0 + if remaining >= 0: + out.append({"label": label, "percent": 1.0 - remaining, "resetsAt": reset_time(window.get("resets_at"))}) + balance, allowance = amount(included.get("balance_usd")), amount(included.get("allowance_usd")) + if balance is not None and allowance is not None and allowance > 0: + period = included.get("period") + out.append({ + "label": "Included credits", + "title": "Monthly", + "percent": max(0.0, min(1.0, (allowance - balance) / allowance)), + "resetsAt": reset_time(period.get("until")) if isinstance(period, dict) else "", + }) + return out + + +# The credits bought apart from the plan: purchased.balance_usd, what's left +# of them unexpired. The answer says how much remains but not how much was +# bought, so the balance shows the remaining figure alone. A plan that never +# bought any, or used them up, has none to show, and that is no alarm while +# the plan's included credits still serve. +def purchased_balance(payload: Any) -> dict[str, Any] | None: + purchased = payload.get("purchased") if isinstance(payload, dict) else None + remaining = amount(purchased.get("balance_usd")) if isinstance(purchased, dict) else None + if remaining is None or remaining <= 0: + return None + return {"remaining": remaining, "funded": 0, "spent": 0, "currency": "USD", "estimated": False} + + +def kept_balance(cache: LimitsCache) -> dict[str, Any] | None: + balance = cache.data.get("balance") + return balance if isinstance(balance, dict) else None + + +def probe_limits(access_key: str) -> dict[str, Any]: + request = urllib.request.Request( + BALANCE_ENDPOINT, + headers={"Authorization": "Bearer " + access_key, "Accept": "application/json"}, + ) + try: + payload = open_json(request, timeout=10) + except HttpError as error: + if error.status in (401, 403): + return {"ok": False, "auth": True, "helpText": AUTH_HELP_BAD_KEY} + if error.status: + return {"ok": False, "helpText": f"Ollama's balance endpoint returned status {error.status}. Local stats are still shown."} + if error.transport: + # 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 balance endpoint. Retrying shortly. Local stats are still shown.", + } + return {"ok": False, "helpText": "Ollama's balance endpoint returned an invalid answer. Local stats are still shown."} + + out = balance_limits(payload) + balance = purchased_balance(payload) + if not out and balance is None: + return {"ok": False, "helpText": "Ollama's balance endpoint returned no limits. Local stats are still shown."} + return {"ok": True, "limits": out, "balance": balance} + + +def probe_plan(access_key: str) -> str | None: + """The plan Ollama names for the key, "" when it names none, None when it couldn't be asked.""" + request = urllib.request.Request( + ME_ENDPOINT, + headers={"Authorization": "Bearer " + access_key, "Accept": "application/json"}, + method="POST", + ) + try: + payload = open_json(request, timeout=10) + except Exception: + return None + # Ollama's UserResponse tags the field `plan`; Go's decoder takes either + # case, so `Plan` is read too. + plan = str(payload.get("plan") or payload.get("Plan") or "").strip()[:64] if isinstance(payload, dict) else "" + return plan[0].upper() + plan[1:] if plan else "" + + +def limits_cache_path(access_key: str) -> Path: + # Keyed by the credential: another key never shows this one's allowance or + # plan. + return shared_limits_cache_path(AGENT_ID, access_key) + + +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 tell kept limits from fresh ones. + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "live": False, "fetchedAtMs": 0} + + if access_key == "": + result["usageStatusText"] = "Waiting for auth" + result["authHelpText"] = AUTH_HELP_NO_KEY + return result + + # A panel that is opened and shut repeatedly must not turn into a request + # per flick, so recent probe results are reused for a short window — and + # kept as the answer of record when a later probe fails. + cache = LimitsCache(AGENT_ID, access_key) + cached_plan = str(cache.data.get("plan") or "") + + # --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. + if cache.fresh(force, allow_empty=kept_balance(cache) is not None): + result.update(limits=cache.limits, tierLabel=cached_plan, live=True, fetchedAtMs=cache.fetched_at, balance=kept_balance(cache)) + return result + + probe = probe_limits(access_key) + if probe["ok"]: + plan = probe_plan(access_key) + # A plan lookup that failed keeps the plan last named, rather than + # blanking a label the account still has. + plan = cached_plan if plan is None else plan + fetched_at = cache.save(probe["limits"], only_if_newer=True, plan=plan, balance=probe["balance"]) + result.update(limits=probe["limits"], tierLabel=plan, live=True, fetchedAtMs=fetched_at, balance=probe["balance"]) + 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 + result["authHelpText"] = probe["helpText"] + fallback = cache.kept()[0] + balance = kept_balance(cache) + if fallback or balance: + result.update(limits=fallback, tierLabel=cached_plan, fetchedAtMs=cache.fetched_at, balance=balance) + if probe.get("auth"): + result["usageStatusText"] = "Sign-in expired" + elif not fallback and not balance: + result["usageStatusText"] = "Ollama limits unavailable" + 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") + args = parser.parse_args() + + config = read_config() + model_prefix = str(config.get("modelPrefix") or "").strip().lower() + stats = cached_scan(scan_max_age(args.force, args.limits_only), model_prefix) + has_stats = bool(stats.get("modelUsage")) + + key = api_key() + record = base_record(AGENT_ID, AGENT_NAME, hasLocalStats=has_stats) + record.update(stats) + + # A machine with neither a key nor any Ollama Cloud sessions doesn't use it, + # and the panel skips the record; that is no trouble to report. + if key or has_stats: + limits = collect_limits(key, args.force) + record.update( + ready=True, + tierLabel=limits["tierLabel"], + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + limits=limits["limits"], + limitsStale=bool(limits["limits"] or limits.get("balance")) and not limits["live"], + limitsFetchedAt=limits["fetchedAtMs"], + ) + if limits.get("balance"): + record["balance"] = limits["balance"] + if limits.get("retryAdvised"): + record["retryAdvised"] = True + emit(record, sort_keys=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-opencode b/bin/omarchy-agent-usage-opencode new file mode 100755 index 00000000000..315118933d6 --- /dev/null +++ b/bin/omarchy-agent-usage-opencode @@ -0,0 +1,430 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the OpenCode usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect OpenCode usage into one display-ready JSON record. + +Local stats come from pi/omp sessions that ran through an opencode provider +and from opencode's own message database, both of its stores. Only the Zen +and Go providers count: opencode runs on Anthropic or OpenAI belong to the +claude and codex records. The limits are the Go plan's rolling, weekly, and +monthly meters from OpenCode's usage endpoint. The agents panel only ever +reads the JSON this prints. +""" + +import argparse +import json +import math +import os +import sqlite3 +import sys +import time +import urllib.request +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import NamedTuple + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record, + cache_root, # noqa: F401 - the tests reach the cache through it + cached_scan, + data_home, + emit, + limits_cache_path as shared_limits_cache_path, + number, + open_json, + opencode_db, + pi_session_roots, + scan_max_age, + scan_opencode, + scan_pi_sessions, +) + +AGENT_ID = "opencode" +AGENT_NAME = "OpenCode" + +# The subscription meters live on the Zen "go" gateway; the plain gateway has +# no usage endpoint. Limits are per-account, so they come from the same key +# regardless of which provider id a given run used. A Console session asks +# Console's own endpoint for the same numbers. +API_BASE_URL = "https://opencode.ai" +USAGE_PATH = "/zen/go/v1/usage" +CONSOLE_USAGE_PATH = "/inference/go/v1/usage" +CONSOLE_ORG_HEADER = "x-opencode-org-id" +CONSOLE_EXPIRED_HELP = "OpenCode's Console session has expired. Run `opencode auth login` to refresh it." + +# A key with no Go plan behind it (a Zen pay-as-you-go account) gets the same +# answer every time, so it is asked again hourly rather than on every run; +# --force asks at once, for someone who just subscribed. +NO_PLAN_RECHECK_SECONDS = 3600 + +# The subscription has two provider ids: "opencode" for the plain Zen +# gateway and "opencode-go" for the Go plan. Both spend the same account and +# both appear in pi sessions and in opencode's own database. +PROVIDERS = ("opencode", "opencode-go") + + +# Usage percentages arrive as 0-100; the record contract wants 0-1. An +# unreadable value returns -1 so the caller can drop the window. +def fraction(value): + if isinstance(value, bool): + return -1.0 + try: + percent = float(value) / 100.0 + except (TypeError, ValueError): + return -1.0 + return percent if math.isfinite(percent) else -1.0 + + +def model_name(raw): + # opencode stores the model as a {providerID, modelID} object in its own + # database and as a plain string in pi sessions; the shared readers hand + # back the id either way. A variant suffix such as ":high" names the + # thinking level, not the model. + value = str(raw or "").rstrip("/").split("/")[-1].split(":")[0] or "opencode" + # The panel title-cases words but cannot produce all-caps acronyms. + if value[:3].lower() == "glm": + value = "GLM" + value[3:] + return value + + +# Exact match: opencode provider ids are free-form, and a custom +# "opencode-local" gateway is not this subscription. Anthropic and OpenAI runs +# never match, so the claude and codex collectors keep those. +def billed_here(provider, model, entry): + return provider in PROVIDERS + + +def run_local_scans(): + # pi and omp can spend the subscription without opencode's CLI ever + # running, so its own database cannot be the only source. Their JSONL + # transcripts carry the provider, model, and token usage per message. A + # subscription burned entirely through opencode leaves no pi session files, + # but opencode records the same per message in its own database. + stats = Stats() + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"opencode(-go)?"', billed_here) + db_rows, db_complete = scan_opencode(PROVIDERS, billed_here) + for row in pi_rows + db_rows: + stats.add_row(row, model_name(row.model)) + return stats.snapshot(), pi_complete and db_complete + + +def cached_local_stats(max_age): + # The digest covers every data path the scan reads: the opencode DB and the + # pi/omp session roots. + return cached_scan(AGENT_ID, [opencode_db(), *pi_session_roots()], max_age, run_local_scans) + + +class Credential(NamedTuple): + """One way of asking OpenCode for the Go meters.""" + + token: str + url: str + headers: dict + expired: bool = False + + +def read_auth_key(path): + try: + parsed = json.loads(path.read_text(encoding="utf-8")) + except Exception: + # Unreadable, not JSON, or nested past the parser's depth (a + # RecursionError): no key here. + return "" + if not isinstance(parsed, dict): + return "" + # Only the subscription's own provider ids: a custom gateway named like + # "opencode-local" holds a key for somewhere else, never to be sent here. + # An entry is an {type, key} object, or the key as a plain string. + # A pi `$VAR` or `!command` entry is resolved by pi alone, never sent. + entries = [parsed[provider] for provider in PROVIDERS if provider in parsed] + for entry in entries: + if isinstance(entry, dict): + entry = entry.get("key") + if isinstance(entry, str) and entry.strip() and not entry.strip().startswith(("!", "$")): + return entry.strip() + return "" + + +def pi_agent_dir(): + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + +def read_store_credentials(db): + """The active credentials opencode V2 keeps in its database: Go keys, and + the Console session it signs the opencode integration in with.""" + keys, console = [], None + if not db.is_file(): + return keys, console + try: + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return keys, console + try: + conn.execute("PRAGMA query_only = ON") + rows = conn.execute( + "SELECT integration_id, value FROM credential WHERE active = 1 AND integration_id IN ('opencode-go', 'opencode')" + " ORDER BY integration_id = 'opencode-go' DESC, time_updated DESC" + ).fetchall() + except sqlite3.Error: + return keys, console + finally: + conn.close() + for _, raw in rows: + try: + value = json.loads(raw) + except Exception: + # Not JSON, or nested past the parser's depth: skip this row alone. + continue + if not isinstance(value, dict): + continue + if value.get("type") == "key" and str(value.get("key") or "").strip(): + keys.append(str(value["key"]).strip()) + elif value.get("type") == "oauth" and console is None and str(value.get("access") or "").strip(): + # Console traffic is scoped to a workspace, and a wrong one is a hard + # 403, so the header only goes out when the id is there. + metadata = value.get("metadata") + org = str(metadata.get("orgID") or "").strip() if isinstance(metadata, dict) else "" + expired = 0 < number(value.get("expires")) <= time.time() * 1000 + console = Credential(str(value["access"]).strip(), API_BASE_URL + CONSOLE_USAGE_PATH, {CONSOLE_ORG_HEADER: org} if org else {}, expired) + return keys, console + + +def credential_candidates(): + # An explicit key wins, then opencode's own credential stores, then pi's: + # pi drives the same Zen subscription and records it under the same + # provider ids, so a machine that codes through pi never signs the opencode + # CLI in. The Console session is the implicit one, so it comes last. Which + # one the endpoint accepts isn't knowable from here, so a refused one only + # hands over to the next. + store_keys, console = read_store_credentials(opencode_db()) + tokens = [str(os.environ.get(name, "")).strip() for name in ("OPENCODE_API_KEY", "OPENCODE_GO_API_KEY")] + tokens += store_keys + tokens.append(read_auth_key(data_home() / "opencode" / "auth.json")) + tokens.append(read_auth_key(pi_agent_dir() / "auth.json")) + tokens.append(read_auth_key(Path.home() / ".omp" / "agent" / "auth.json")) + candidates = [] + for token in tokens: + if token and token not in [c.token for c in candidates]: + candidates.append(Credential(token, API_BASE_URL + USAGE_PATH, {})) + if console is not None: + candidates.append(console) + return candidates + + +def reset_from_seconds(value): + seconds = number(value) + if seconds <= 0: + return "" + try: + return (datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat() + except OverflowError: + # Further out than a date goes: no reset time to show. + return "" + + +def iso_reset(value): + # Only an ISO time is a reset time; an epoch, a list or garbage is none. + if not isinstance(value, str) or not value.strip(): + return "" + try: + datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return "" + return value.strip() + + +# The live endpoint nests the windows under "usage" with "percent" and +# "resetsAt"; another deployment spells them "Usage" with +# "usagePercent" and "resetInSec". Both are read, so the collector doesn't +# depend on which one answers. +def parse_usage(payload): + usage = payload.get("usage") if isinstance(payload, dict) else None + usage = usage if isinstance(usage, dict) else (payload if isinstance(payload, dict) else {}) + limits = [] + for name, label in (("rolling", "Rolling (5h)"), ("weekly", "Weekly (7-day)"), ("monthly", "Monthly")): + window = usage.get(name) + if not isinstance(window, dict): + window = payload.get(name + "Usage") if isinstance(payload, dict) else None + if not isinstance(window, dict): + continue + percent = fraction(window.get("percent", window.get("usagePercent"))) + if percent < 0: + continue + # A rate-limited window is spent until it resets, whatever it reads. + if str(window.get("status") or "") == "rate-limited": + percent = 1.0 + limits.append({ + "label": label, + "percent": max(0.0, min(1.0, percent)), + "resetsAt": iso_reset(window.get("resetsAt")) or reset_from_seconds(window.get("resetInSec")), + }) + return limits + + +def probe_limits(credential): + request = urllib.request.Request(credential.url, headers={ + "Authorization": "Bearer " + credential.token, + "Accept": "application/json", + # Cloudflare answers the default urllib signature with a 1010 before the + # request reaches the API, so identify the collector instead. + "User-Agent": "omarchy-agent-usage-opencode/1.0", + **credential.headers, + }) + unreachable = {"ok": False, "transport": True, "helpText": "Couldn't reach OpenCode's usage endpoint. Retrying shortly. Local OpenCode stats are still shown."} + try: + payload = open_json(request, timeout=10) + except HttpError as error: + if error.transport: + return unreachable + # A refused credential may just be the wrong one for this account, so the + # next way in still gets its turn. + if error.status == 401: + return {"ok": False, "refused": True, "helpText": "OpenCode rejected the saved key. Reconnect OpenCode Go with /connect in the TUI, or set OPENCODE_API_KEY."} + if error.status == 403: + # A key without a Go plan behind it: Zen pay-as-you-go has no meters, + # which is no trouble at all. + return {"ok": False, "noPlan": True} + if error.status: + return {"ok": False, "helpText": f"OpenCode's usage endpoint returned status {error.status}. Local OpenCode stats are still shown."} + payload = None + except Exception: + # open_json already names every transport failure, so anything else is + # a broken answer, and retrying soon won't mend it. + payload = None + + if not isinstance(payload, dict) or not payload: + return {"ok": False, "helpText": "OpenCode's usage endpoint returned no data. Local OpenCode stats are still shown."} + limits = parse_usage(payload) + if not limits: + return {"ok": False, "helpText": "OpenCode's usage endpoint returned no quota windows. Local OpenCode stats are still shown."} + return {"ok": True, "limits": limits} + + +# One cache per credential, so a machine signed in to another account never +# shows the last one's meters. +def limits_cache_path(token): + return shared_limits_cache_path(AGENT_ID, token) + + +def collect_limits(force): + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False, "stale": False, "fetchedAtMs": 0, "signedIn": False} + # No key at all is a Zen user paying as they go, or someone who only runs + # opencode locally: local stats stand alone, with nothing to fix. + candidates = credential_candidates() + if not candidates: + return result + result["signedIn"] = True + + # Limits only ever show for a Go plan, so a tier follows them. + def answer(limits, fetched_at, stale): + # Stale says limits were kept past a failed check, so none kept is none stale. + result.update(limits=limits, tierLabel="Go" if limits else "", stale=stale and bool(limits), fetchedAtMs=fetched_at if limits else 0) + return result + + refused = None + no_plan = False + for candidate in candidates: + cache = LimitsCache(AGENT_ID, candidate.token) + fallback, fetched_at = cache.limits, cache.fetched_at + + # A key already known to have no Go plan isn't asked again for a while. + if cache.data.get("noPlan") is True and not force and cache.younger_than(NO_PLAN_RECHECK_SECONDS): + no_plan = True + continue + + # A recent answer stands in for a probe; --force skips the wait. + if cache.fresh(force): + return answer(fallback, fetched_at, False) + + probe = probe_limits(candidate) + if probe.get("ok"): + return answer(probe["limits"], cache.save(probe["limits"], only_if_newer=True), False) + + if probe.get("noPlan"): + # The windows of a plan that's gone mean nothing now, so they go too. + cache.save([], only_if_newer=True, noPlan=True) + no_plan = True + continue + + if probe.get("refused"): + # Keep the first refusal that still has open windows, so a later + # credential without a cache can't blank them. A lapsed Console + # session is what to fix, whichever refusal is kept. + expired = candidate.expired or bool(refused and refused[3]) + if refused is None or (not refused[1] and fallback): + refused = (probe, fallback, fetched_at, expired) + else: + refused = refused[:3] + (expired,) + continue + + # Last good numbers, marked stale, beat an empty limits section while + # the probe is down. A key that already said it has no plan makes an + # outage on another one no news. + result["retryAdvised"] = bool(probe.get("transport")) + answer(fallback, fetched_at, True) + if fallback: + result["authHelpText"] = probe["helpText"] + elif not no_plan: + result["usageStatusText"] = "OpenCode limits unavailable" + result["authHelpText"] = probe["helpText"] + return result + + # A working key without a Go plan is the answer: the other ways in only + # failed to say the same. + if no_plan: + return result + + # Every way in was refused: an account problem, not an outage. Keep any + # still-open numbers, but say a new sign-in is needed. + probe, fallback, fetched_at, expired = refused + answer(fallback, fetched_at, True) + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = CONSOLE_EXPIRED_HELP if expired else probe["helpText"] + return result + + +def main(): + parser = argparse.ArgumentParser() + # --force rescans everything and rewrites the cache. --limits-only is kept + # for CLI compatibility with the panel's refreshLimits() call: only the + # limits probe must be fresh, so it may reuse a scan for far longer than a + # normal run, whose short window exists purely to dedup concurrent + # collector runs. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + stats = cached_local_stats(scan_max_age(args.force, args.limits_only)) + limits = collect_limits(args.force) + + # Nobody uses OpenCode here: no usage and no credential leave a record the + # panel skips. + record = base_record( + AGENT_ID, + AGENT_NAME, + ready=number(stats.get("totalPrompts")) > 0 or bool(limits["limits"]) or limits["signedIn"], + # A sign-in with no sessions here has no local stats to show. + hasLocalStats=number(stats.get("totalSessions")) > 0 or number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")), + tierLabel=limits["tierLabel"], + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + limits=limits["limits"], + limitsStale=limits["stale"], + limitsFetchedAt=limits["fetchedAtMs"], + ) + record.update(stats) + if limits["retryAdvised"]: + record["retryAdvised"] = True + emit(record) + + +if __name__ == "__main__": + main() diff --git a/bin/omarchy-agent-usage-openrouter b/bin/omarchy-agent-usage-openrouter new file mode 100755 index 00000000000..dc9f5c776a6 --- /dev/null +++ b/bin/omarchy-agent-usage-openrouter @@ -0,0 +1,627 @@ +#!/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): key usage comes from GET /api/v1/key, and the balance +from GET /api/v1/credits, asked with the management key when there is one +(OpenRouter documents it as management-key only, and an inference key it +refuses isn't asked again for a day); token stats come from the local +opencode SQLite DB +(providerID == "openrouter", role == "assistant") and from pi/omp session +files whose assistant messages ran on the openrouter provider, read-only. +Each key keeps its own last good meter and balance, shown stale when a +check fails. The key is OPENROUTER_API_KEY, opencode's or pi/omp's saved +OpenRouter login, or the config file's. Without any, the local stats still +show and the record waits for auth. + +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 math +import os +import sys +import urllib.request +from datetime import datetime, timedelta +from pathlib import Path +from typing import Any + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + age_below, + base_record as shared_base_record, + cache_root, # noqa: F401 - the tests find the cached answers through it + cached_scan, + config_home, + data_home, + dsh_home, + emit, + empty_bucket, + empty_stats, + kimchi_sessions_root, + limits_cache_path, + now_ms, + number, + open_json, + open_limits, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_cache_paths as shared_scan_cache_paths, + scan_dsh_sessions, + scan_max_age, + scan_opencode as shared_scan_opencode, + scan_pi_sessions, + scan_routed, + stamp, + warn, + write_json, +) + +AGENT_ID = "openrouter" +AGENT_NAME = "OpenRouter" +PROVIDER = "openrouter" +AUTH_HELP = "Set OPENROUTER_API_KEY, or sign in to OpenRouter in opencode, pi or omp. For the 30-day history, add a management key as OPENROUTER_MANAGEMENT_KEY (openrouter.ai/settings/management-keys)." +MANAGEMENT_AUTH_HELP = "Replace OPENROUTER_MANAGEMENT_KEY (or managementKey in ~/.config/omarchy/agents/openrouter.json) with a working one from openrouter.ai/settings/management-keys." +API_BASE = "https://openrouter.ai/api/v1" + + +class OpenRouterError(Exception): + pass + + +def model_name(raw: Any) -> str: + value = str(raw or "openrouter").rstrip("/").split("/")[-1] + return value or "openrouter" + + +# ---------------------------------------------------------------- auth + +def read_config() -> dict[str, Any]: + return read_json(config_home() / "omarchy" / "agents" / "openrouter.json") + + +def read_opencode_key(path: Path) -> str: + entry = read_json(path).get("openrouter") + if not isinstance(entry, dict): + return "" + return str(entry.get("key") or "").strip() + + +def pi_agent_dir() -> Path: + # PI_CODING_AGENT_DIR moves pi's agent dir, as pi itself reads it. + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + +def read_pi_key(path: Path) -> str: + # pi/omp keep a key they were given with /login in auth.json. Only a literal + # key is taken: a `$VAR` or `!command` entry is resolved by pi alone. + key = read_opencode_key(path) + return "" if key.startswith(("!", "$")) else key + + +def api_key(config: dict[str, Any]) -> str: + # Explicit key or an opencode, pi or omp login wins; the config file is + # the last resort for headless/refresh-timer runs without any. + return ( + str(os.environ.get("OPENROUTER_API_KEY", "")).strip() + or read_opencode_key(data_home() / "opencode" / "auth.json") + or read_pi_key(pi_agent_dir() / "auth.json") + or read_pi_key(Path.home() / ".omp" / "agent" / "auth.json") + 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 + +# No server answered at all (DNS, connection, timeout): the one failure worth +# a sooner retry. +class OpenRouterTransportError(OpenRouterError): + pass + + +# The key was refused. 403 from /activity with an inference key is the +# expected tier-1 signal ("Only management keys can perform this operation"), +# not a broken setup -- callers decide what a refusal means. +class OpenRouterAuthError(OpenRouterError): + def __init__(self, code: int): + super().__init__(f"OpenRouter refused the key (status {code})") + self.code = code + + +def api_get(path: str, key: str) -> dict: + req = urllib.request.Request( + API_BASE + path, headers={"Authorization": "Bearer " + key, "Accept": "application/json"} + ) + try: + payload = open_json(req, timeout=15) + except HttpError as error: + if error.status in (401, 403): + raise OpenRouterAuthError(error.status) from error + if error.status: + raise OpenRouterError(f"OpenRouter's API returned status {error.status}.") from error + if error.transport: + # A server that stalls on its headers or body times out the read + # itself; that is a network failure like any other. + raise OpenRouterTransportError("Couldn't reach the OpenRouter API. Retrying shortly.") from error + raise OpenRouterError("OpenRouter returned an invalid response.") from error + return payload if isinstance(payload, dict) else {} + + +# /activity names its day as a date, sometimes with a time after it. Anything +# that isn't one files under no day rather than a made-up one. +def activity_day(value: Any) -> str: + text = str(value or "").strip()[:10] + try: + return datetime.strptime(text, "%Y-%m-%d").strftime("%Y-%m-%d") + except ValueError: + return "" + + +def summarize_activity(items: list[Any], today: datetime) -> dict[str, Any]: + """Fold /activity rows into the panel's stats dict. + + One row per (date, endpoint): prompt tokens count as input, and completion + tokens, reasoning among them, as output (thinking is generated, so it is + output, as the collector contract has it). 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 = activity_day(raw.get("date")) + model = model_name(str(raw.get("model_permaslug") or raw.get("model") or "")) + prompt = number(raw.get("prompt_tokens")) + # Reasoning is part of the completion count, as in OpenRouter's usage + # (its own AI SDK provider reads text = completion - reasoning), so it + # is never added again; a row naming reasoning alone still counts it. + output = max(number(raw.get("completion_tokens")), number(raw.get("reasoning_tokens"))) + total = prompt + output + if total <= 0: + continue + bucket = model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += prompt + bucket["outputTokens"] += output + 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, max_age: float) -> tuple[dict[str, Any] | None, bool]: + """Tier 2: 30-day account history via management key, or None. + + Returns (stats, transport_failed). Stats are None when there is no + management key (plain tier-1 run) or the key is refused -- the caller then + falls back to the local DB scan. The answer is a stats source like the + local scan, so it is reused on the same schedule; when /activity can't be + reached, today's last answer stands in for it. + """ + if not key: + return None, False + scan_date = today.strftime("%Y-%m-%d") + try: + cache = limits_cache_path(AGENT_ID, key, kind="activity") + cached = read_json(cache) + except OSError: + cache, cached = None, {} + kept = cached.get("stats") if cached.get("scanDate") == scan_date and isinstance(cached.get("stats"), dict) else None + if kept is not None and age_below(stamp(cached.get("fetchedAtMs")) / 1000, max_age): + return kept, False + + try: + payload = api_get("/activity", key) + except OpenRouterAuthError: + return None, False + except OpenRouterError as error: + warn(f"activity lookup failed ({error})") + return kept, isinstance(error, OpenRouterTransportError) + items = payload.get("data") + if not isinstance(items, list): + return kept, False + stats = empty_stats() + stats.update(summarize_activity(items, today)) + if cache is not None: + try: + write_json(cache, {"schemaVersion": 1, "scanDate": scan_date, "fetchedAtMs": now_ms(), "stats": stats}) + except OSError as exc: + warn(f"could not write activity cache ({exc})") + return stats, False + + +# ---------------------------------------------------------------- local stats + +# pi, omp, and opencode name the provider alike. Exact match: a custom +# gateway named like OpenRouter is not this account. +def billed_here(provider: str, model: str, entry: dict) -> bool: + return provider == PROVIDER + + +def scan_opencode(stats: Stats) -> bool: + """Add opencode messages on the openrouter provider, from both of its + stores. 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.""" + # The closing quote makes the database's prefilter match the whole id. + rows, complete = shared_scan_opencode([PROVIDER + '"'], billed_here) + for row in rows: + stats.add_row(row, model_name(row.model)) + return complete + + +def pi_roots() -> list[Path]: + # Kimchi's harness keeps pi-format sessions of its own, and one run with + # `--provider openrouter` is OpenRouter's. + return [*pi_session_roots(), kimchi_sessions_root()] + + +def scan_pi(stats: Stats) -> bool: + """Add pi/omp and Kimchi messages on the openrouter provider. Returns + whether every session file could be read, like scan_opencode.""" + rows, complete = scan_pi_sessions(r'"provider"\s*:\s*"openrouter"', billed_here, pi_roots()) + for row in rows: + stats.add_row(row, model_name(row.model)) + return complete + + +def scan_dsh(stats: Stats) -> bool: + """Add DeepSeek Harness attempts on its pi-ai adapter's openrouter route.""" + rows, complete = scan_dsh_sessions(lambda provider, model, request: provider == PROVIDER) + for row in rows: + stats.add_row(row, model_name(row.model)) + return complete + + +def scan_local(stats: Stats, reuse: bool = True) -> bool: + # Every source always runs, so one failing still leaves the others' numbers. + opencode_complete = scan_opencode(stats) + pi_complete = scan_pi(stats) + dsh_complete = scan_dsh(stats) + # Claude Code pointed at OpenRouter's Anthropic-compatible API, and MiniMax + # Code's custom providers (its OpenRouter preset among them) at OpenRouter. + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + for row in routed_rows: + stats.add_row(row, model_name(row.model)) + return opencode_complete and pi_complete and dsh_complete and routed_complete + + +def scan_sources() -> list[Path]: + return [opencode_db(), *pi_roots(), dsh_home(), *routed_sources()] + + +def scan_cache_paths() -> tuple[Path, Path]: + return shared_scan_cache_paths(AGENT_ID, scan_sources()) + + +def cached_local_stats(max_age: float) -> tuple[dict, bool]: + """Local stats, and whether they cover every source. + + The cache is a pure optimization, and a cached scan was a complete one. + Concurrent runs (the update command and the panel's own refresh) wait for + one scan instead of each reading the whole database. + """ + + def scan() -> tuple[dict[str, Any], bool]: + stats = Stats() + # A forced scan reads every Codex rollout again too. + complete = scan_local(stats, reuse=max_age > 0) + return stats.snapshot(), complete + + return cached_scan(AGENT_ID, scan_sources(), max_age, scan, match=routed_match(), with_complete=True) + + +# ---------------------------------------------------------------- record + +def base_record(**overrides: Any) -> dict: + return shared_base_record(AGENT_ID, AGENT_NAME, **overrides) + + +# 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. +# How long an inference key /credits refused goes unasked: OpenRouter +# documents the endpoint as management-key only, so it answers the same +# until the key's permissions change. +CREDITS_REFUSAL_SECONDS = 24 * 3600 + +KEY_LIMIT_PERIODS = {"daily": "Daily", "weekly": "Weekly", "monthly": "Monthly"} + + +# OpenRouter answers `limit: null` for a key without a credit limit, which has +# no meter. An answer missing the limit, or one naming a limit without a +# readable remainder, says nothing about the allowance, so it is a failed check +# that keeps the last good meter rather than an empty one that would replace it. +def key_limit_meter(key_info: dict) -> list[dict]: + if "limit" not in key_info: + raise OpenRouterError("OpenRouter answered /key without a credit limit.") + limit = key_info.get("limit") + if limit is None: + return [] + remaining = key_info.get("limit_remaining") + try: + if isinstance(limit, bool) or isinstance(remaining, bool): + raise TypeError + lim = float(limit) + rem = max(0.0, float(remaining)) + except (TypeError, ValueError): + raise OpenRouterError("OpenRouter answered /key with an unreadable credit limit.") from None + if not (math.isfinite(lim) and math.isfinite(rem)): + raise OpenRouterError("OpenRouter answered /key with an unreadable credit limit.") + if lim <= 0: + return [] + period = KEY_LIMIT_PERIODS.get(str(key_info.get("limit_reset") or "").strip().lower()) + label = f"{period} key limit" if period else "Key limit" + return [{"label": label, "title": label, "percent": min(1.0, (lim - min(rem, lim)) / lim), "resetsAt": ""}] + + +def dollars(value: Any) -> float | None: + if isinstance(value, bool): + return None + try: + parsed = float(value or 0) + except (TypeError, ValueError): + return None + return parsed if math.isfinite(parsed) else None + + +# Balance: total purchased minus total used, or None for an account that never +# bought credits. An answer missing either total says nothing about the +# balance, so it is a failed check that keeps the last good one, not an empty +# balance that would replace it. +def credits_balance(credits: dict) -> dict | None: + total_credits = dollars(credits.get("total_credits")) if credits.get("total_credits") is not None else None + total_usage = dollars(credits.get("total_usage")) if credits.get("total_usage") is not None else None + if total_credits is None or total_usage is None: + raise OpenRouterError("OpenRouter answered /credits without a balance.") + if total_credits <= 0: + return None + return { + "remaining": round(max(0.0, total_credits - total_usage), 6), + "funded": round(total_credits, 6), + "spent": round(max(0.0, total_usage), 6), + "currency": "USD", + "estimated": False, + } + + +# Each credential keeps its own last good answer, so another key never shows +# this one's numbers: the key's meter as openrouter-limits-, the +# balance beside it as openrouter-credits-. +def probe(endpoint: str, key: str, kind: str, parse, force: bool) -> tuple[dict | None, OpenRouterError | None]: + """One endpoint's answer as (part, error). + + Opening and shutting the panel must not turn into three API calls per + flick, so a probe this recent is reused; --force, a person pressing + refresh, skips it. A fresh answer replaces the cached one, and a failed one + falls back on it. part is None when nothing was ever fetched. + A key with no credit limit has an empty meter, and its answer is reused + all the same. + """ + cache = LimitsCache(AGENT_ID, key, kind=kind) + kept = dict(cache.data, fetchedAtMs=cache.fetched_at) if cache.fetched_at else None + if kept is not None and cache.fresh(force, allow_empty=True): + return kept, None + try: + payload = api_get(endpoint, key) + data = payload.get("data") + part = parse(data if isinstance(data, dict) else {}) + except OpenRouterError as error: + return kept, error + # Refreshes can overlap (the panel's own and a sooner retry), so an answer + # only replaces one fetched before it. + cache.save(part.pop("limits", []), only_if_newer=True, **part) + return dict(cache.data), None + + +def collect_billing(key: str, manager: str, force: bool) -> dict[str, Any]: + """The key-limit meter and the balance, each kept past a failed check. + + The meter describes an inference key's own credit limit, so with only a + management key there is none to show; the balance is the account's, and + an inference key reads it only where OpenRouter still lets it. + """ + result: dict[str, Any] = { + "limits": [], "limitsStale": False, "limitsFetchedAt": 0, + "usageStatusText": "", "authHelpText": "", "balance": None, "retryAdvised": False, + } + parts: dict[str, dict | None] = {} + errors: dict[str, OpenRouterError] = {} + if key: + parts["key"], error = probe("/key", key, "limits", lambda data: {"limits": key_limit_meter(data)}, force) + if error: + errors["key"] = error + # OpenRouter documents /credits as management-key only, answering another + # key 403 "Only management keys can perform this operation". So the + # management key asks when there is one, and an inference key alone asks + # until it is refused that way, which is no balance for this key, not a + # refused key: /key says whether the key itself still works. The refusal + # is kept, so the next refreshes don't ask again. A management key that is + # refused itself (revoked, or no management key at all) leaves the + # inference key to ask; when that key has no balance either, the + # management key's refusal stands, with its last good balance kept. + error = manager_error = None + if manager: + parts["credits"], error = probe("/credits", manager, "credits", lambda data: {"balance": credits_balance(data)}, force) + manager_error = error + if key and (not manager or isinstance(error, OpenRouterAuthError)): + managed = parts.get("credits"), error + refusal = LimitsCache(AGENT_ID, key, kind="credits") + if refusal.data.get("refused") is True and refusal.younger_than(CREDITS_REFUSAL_SECONDS): + parts["credits"], error = managed + else: + parts["credits"], error = probe("/credits", key, "credits", lambda data: {"balance": credits_balance(data)}, force) + if isinstance(error, OpenRouterAuthError) and error.code == 403: + refusal.save([], refused=True) + parts["credits"], error = managed + if error: + errors["credits"] = error + + meter = parts.get("key") + if meter is not None: + # A kept window whose reset has passed is an untouched allowance, not an + # unknown one. The key meter only ever names a period, so this guards a + # cache written by a meter that carries a reset time. + result["limits"] = open_limits(meter.get("limits")) + credits = parts.get("credits") + if credits is not None and isinstance(credits.get("balance"), dict): + result["balance"] = credits["balance"] + kept = [part for part in parts.values() if part is not None] + if kept: + result["limitsFetchedAt"] = min(part["fetchedAtMs"] for part in kept) + # Anything that didn't come back fresh marks the record stale, kept or not; + # the panel shows the kept numbers' age on hover. + result["limitsStale"] = bool(errors) + result["retryAdvised"] = any(isinstance(error, OpenRouterTransportError) for error in errors.values()) + + refused = [error for error in errors.values() if isinstance(error, OpenRouterAuthError)] + if refused and refused == [manager_error]: + # Only the management key was refused; the inference key, if any, still + # works, so the help names the key to replace. + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = "OpenRouter refused the management key. " + MANAGEMENT_AUTH_HELP + elif refused: + result["usageStatusText"] = "Sign-in expired" + result["authHelpText"] = "OpenRouter refused the API key. " + AUTH_HELP + elif "key" in errors and meter is None: + result["usageStatusText"] = "OpenRouter limits unavailable" + result["authHelpText"] = str(errors["key"]) + elif "credits" in errors and credits is None: + result["usageStatusText"] = "Balance unavailable" + result["authHelpText"] = str(errors["credits"]) + elif errors: + result["authHelpText"] = str(next(iter(errors.values()))) + return result + + +def scan(args: Any) -> dict[str, Any]: + config = read_config() + key = api_key(config) + manager = management_key(config) + today = datetime.now().astimezone() + max_age = scan_max_age(args.force, args.limits_only) + + if not key and not manager: + # No key anywhere: there is no balance or meter to read, but what pi, + # omp and opencode spent through OpenRouter here is still OpenRouter's + # usage, and no other collector counts it. With none, OpenRouter isn't + # used here and the panel skips it. + stats, complete = cached_local_stats(max_age) + if number(stats.get("totalPrompts")) <= 0 and not stats.get("modelUsage"): + return base_record() + record = base_record(ready=True, hasLocalStats=complete, usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + record.update(stats) + return record + + # Tier 2 first: with a management key the 30-day account history covers + # every machine and client. Without one (or a refused one) fall back to + # the local opencode and pi scan below. + managed, activity_unreachable = activity_stats(manager, today, max_age) + if managed is not None: + stats = dict(managed) + complete = True + has_prompt_stats = False + # Tier-2 numbers are account-global (billing API), like Fireworks: every + # synced device reports the same truth, so aggregation must merge them by + # widest value, not sum them. Tier 1 is machine-local instead. + stats["scope"] = "account" + else: + stats, complete = cached_local_stats(max_age) + has_prompt_stats = True + has_data = number(stats.get("totalPrompts")) > 0 or bool(stats.get("modelUsage")) + + billing = collect_billing(key, manager, args.force) + + record = base_record( + ready=True, + tierLabel="Prepaid", + hasLocalStats=complete, + hasPromptStats=has_prompt_stats, + limits=billing["limits"], + limitsStale=billing["limitsStale"], + limitsFetchedAt=billing["limitsFetchedAt"], + usageStatusText=billing["usageStatusText"], + authHelpText=billing["authHelpText"], + ) + # Usually the seconds after login before the network is up; the panel + # tries once more soon instead of waiting out its refresh interval. + if billing["retryAdvised"] or activity_unreachable: + record["retryAdvised"] = True + record.update(stats) + if billing["balance"]: + record["balance"] = billing["balance"] + # A key nothing has been spent through, with no credits bought and no limit + # set, has nothing to show: the machine doesn't use OpenRouter yet. + if not has_data and not billing["balance"] and not record["limits"] and not record["usageStatusText"] and not record["limitsStale"]: + record["ready"] = False + record["tierLabel"] = "" + return record + + +def main() -> int: + ap = argparse.ArgumentParser(description="Print the OpenRouter usage record as JSON") + # --force rescans and re-probes everything; --limits-only reuses a recent + # local scan or activity history for up to 15 minutes, while the key meter + # and balance are probed like a normal run. + ap.add_argument("--force", action="store_true", help="rescan sessions and re-probe the API, ignoring caches") + ap.add_argument("--limits-only", action="store_true", help="reuse any recent session scan; only the limits probe must be fresh") + args = ap.parse_args() + try: + record = scan(args) + except Exception as error: + record = base_record(usageStatusText="OpenRouter unavailable", authHelpText="OpenRouter usage scan failed") + warn(type(error).__name__) + emit(record) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-qwen b/bin/omarchy-agent-usage-qwen new file mode 100755 index 00000000000..ff23a011b8d --- /dev/null +++ b/bin/omarchy-agent-usage-qwen @@ -0,0 +1,596 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Qwen Code usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Qwen Code and Alibaba Qwen plan usage into one display-ready JSON +record. + +Qwen Code records each session as JSONL under +~/.qwen/projects//chats/.jsonl (QWEN_HOME moves ~/.qwen; +QWEN_RUNTIME_DIR or settings.json's advanced.runtimeOutputDir move the +projects). Every API response it bills lands there as a ui_telemetry record +carrying the model and the provider's token counts, the same records Qwen +Code's own usage statistics are rebuilt from, and every typed prompt as a +user record. + +pi, omp and opencode turns billed by Alibaba's Qwen plans count here too: +Model Studio's Coding Plan and Token Plan, in their international and +mainland-China forms, and omp's Qwen OAuth (portal.qwen.ai) sign-in. So do +Claude Code's and MiniMax Code's turns sent to Model Studio's +Anthropic-compatible API. + +Neither Qwen Code nor Alibaba publishes a way to read a plan's quota: the +Coding Plan's usage shows on its console page alone, and Qwen Code reads none +of its own. So the record carries no meters, and nothing waits on a +credential. The agents panel only ever reads the JSON this prints. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +from pathlib import Path + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + Stats, + UsageRow, + base_record, + cached_scan, + claude_code_owner, + emit, + local_day, + obj, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_cache_paths, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, + token_count, + warn, + write_json, +) + +AGENT_ID = "qwen" +AGENT_NAME = "Qwen Code" + +API_RESPONSE_EVENT = "qwen-code.api_response" +# A request no reply answered (a refused connection, an error status) still +# names the model and auth type it was sent with. +API_ERROR_EVENT = "qwen-code.api_error" + +# Qwen Code's own rule (getEventTotalTokenCount) takes a response's total +# when it has one, and otherwise counts thinking on top of the output only +# where the output leaves it out: through Gemini, but not through an +# OpenAI-style API. Qwen Code's rule leaves out openai-responses, yet it maps +# the Responses API's output_tokens, which include reasoning_tokens, to the +# output count. +THOUGHTS_IN_OUTPUT_AUTH_TYPES = ("openai", "openai-responses", "qwen-oauth") + +# User records Qwen Code writes for text someone typed: a prompt, or a message +# sent while a turn was still running. Goal, mention, and notification records +# carry a subtype of their own. +PROMPT_SUBTYPES = (None, "mid_turn_user_message") + +# Rows keep each record's time as written, so the index outlives the day and +# zone it was built in; version 1 kept parsed seconds, and version 2 kept an +# undated record undated, so it counted as today's on every day after. +# Version 3 counted a reply's thinking twice where its output held it. +# Version 4 kept no reply's auth type and model id, which say where it was +# sent. Version 5 kept no failed request, which says where a prompt went that +# no reply answered. Version 6 keyed a reply by its response id, which a +# server that recycles ids gives different replies. Version 7 kept no mark on +# a fork's copied records, so a fork whose file sorted first took its +# source's records and session. +INDEX_VERSION = 8 +INDEX_HEAD_BYTES = 4096 + +MAX_MODEL_NAME_LEN = 120 + +# The pi, omp and opencode provider ids whose turns Alibaba's Qwen plans bill. +# Exact names: Alibaba's pay-as-you-go `alibaba` and `alibaba-cn` keys carry +# no plan and are left out. +# - opencode (models.dev) and omp: the Coding Plan, `alibaba-coding-plan` and +# opencode's mainland `alibaba-coding-plan-cn`. +# - opencode and omp: the Token Plan, `alibaba-token-plan` and opencode's +# `alibaba-token-plan-cn`; pi names it `qwen-token-plan`, `-cn` and +# `-individual`. +# - omp: `qwen-portal`, the Qwen OAuth sign-in Qwen Code also offered. +PROVIDERS = ( + "alibaba-coding-plan", + "alibaba-coding-plan-cn", + "alibaba-token-plan", + "alibaba-token-plan-cn", + "qwen-token-plan", + "qwen-token-plan-cn", + "qwen-token-plan-individual", + "qwen-portal", +) + + +def model_label(raw): + """A model id fit for the record: its last path segment, or None when it + is too long to be a real id.""" + value = str(raw or "").rstrip("/").split("/")[-1] + if len(value) > MAX_MODEL_NAME_LEN: + return None + return value or "qwen" + + +# ------------------------------------------------------------- locations + + +def resolve_dir(value): + """A directory setting as Qwen Code resolves it, `~` included. A relative + one resolves against wherever Qwen Code was started, which this collector + can't know, so it is left out.""" + if not isinstance(value, str) or not value.strip(): + return None + path = Path(value.strip()).expanduser() + return path if path.is_absolute() else None + + +def qwen_home(): + return resolve_dir(os.environ.get("QWEN_HOME")) or (Path.home() / ".qwen") + + +def runtime_output_dir(home): + return resolve_dir(obj(read_json(home / "settings.json").get("advanced")).get("runtimeOutputDir")) + + +def runtime_dirs(): + """Every directory whose projects/ may hold chats: QWEN_RUNTIME_DIR or the + configured runtimeOutputDir, where Qwen Code writes, and the Qwen home, + where it writes by default and where its own usage rebuild looks.""" + home = qwen_home() + dirs = [] + for candidate in (resolve_dir(os.environ.get("QWEN_RUNTIME_DIR")), runtime_output_dir(home), home): + if candidate is not None and candidate not in dirs: + dirs.append(candidate) + return dirs + + +def chat_files(): + files = [] + for base in runtime_dirs(): + try: + found = sorted((base / "projects").glob("*/chats/*.jsonl")) + except (OSError, ValueError): + # ValueError: a configured directory no path can name (a NUL in it). + continue + # A ledger sits beside its transcript and repeats nothing billable. + files.extend(path for path in found if not path.name.endswith(".ledger.jsonl") and path not in files) + return files + + +# ------------------------------------------------------------- transcripts + + +def trusted_time(*values): + """(time, ok): the first of `values` that is present and names a day. None + of them present is a record of today, and only present but untrusted + times are not ok, skipping the record.""" + present = [value for value in values if value is not None] + for value in present: + if local_day(value) is not None: + return value, True + return None, not present + + +def parse_line(raw, session): + """One transcript line as an index row, or None. + + A prompt is ["p", key, time, session, copy]; a billed response is + ["r", key, time, session, model, input, output, cacheRead, authType, + modelId, copy]; a failed request is ["e", key, time, session, authType, + modelId, copy]. The key, the record's uuid, dedups a record a forked + session carries again, and copy (1 or 0) says the record is such a copy: + Qwen Code's forkSession marks every record it copies with a forkedFrom. + """ + if b'"user"' not in raw and API_RESPONSE_EVENT.encode() not in raw and API_ERROR_EVENT.encode() not in raw: + return None + try: + record = json.loads(raw) + except (ValueError, RecursionError): + # RecursionError: a line nested deeper than the parser will go. + return None + if not isinstance(record, dict): + return None + session = str(record.get("sessionId") or session) + uuid = str(record.get("uuid") or "") + copy = 1 if record.get("forkedFrom") is not None else 0 + + if record.get("type") == "user": + if record.get("subtype") not in PROMPT_SUBTYPES: + return None + stamp, ok = trusted_time(record.get("timestamp")) + if not ok or not uuid: + return None + return ["p", uuid, stamp, session, copy] + + if record.get("type") != "system" or record.get("subtype") != "ui_telemetry": + return None + event = obj(obj(record.get("systemPayload")).get("uiEvent")) + if event.get("event.name") == API_ERROR_EVENT and uuid: + stamp, _ = trusted_time(event.get("event.timestamp"), record.get("timestamp")) + return ["e", uuid, stamp, session, str(event.get("auth_type") or ""), str(event.get("model") or ""), copy] + if event.get("event.name") != API_RESPONSE_EVENT: + return None + stamp, ok = trusted_time(event.get("event.timestamp"), record.get("timestamp")) + model = model_label(event.get("model")) + if not ok or model is None: + return None + # The prompt count includes its cached part; a provider that reports only + # the cached part reports a fully cached prompt, as Qwen Code reads it. + cached = token_count(event.get("cached_content_token_count")) + prompt = token_count(event.get("input_token_count")) or cached + cached = min(cached, prompt) + output = token_count(event.get("output_token_count")) + if str(event.get("auth_type") or "") not in THOUGHTS_IN_OUTPUT_AUTH_TYPES: + thoughts = token_count(event.get("thoughts_token_count")) + total = token_count(event.get("total_token_count")) + if total > 0: + thoughts = min(thoughts, max(0, total - prompt - output)) + output += thoughts + if prompt + output <= 0: + return None + # A fork copies each record verbatim, its uuid included, and a resumed + # session appends to its own transcript, so the uuid is what a copy shares. + # A server's response id is not: one that recycles short ids gives + # different replies the same one. + key = uuid or str(event.get("response_id") or "") + if not key: + return None + return ["r", key, stamp, session, model, prompt - cached, output, cached, str(event.get("auth_type") or ""), str(event.get("model")), copy] + + +def read_digest(path, length): + digest = hashlib.sha1() + with path.open("rb") as handle: + digest.update(handle.read(min(length, INDEX_HEAD_BYTES))) + return digest.hexdigest() + + +def parse_file(path, entry): + """The index entry for one transcript, reusing `entry` where the file only + grew: Qwen Code appends to a transcript and never rewrites what it wrote.""" + stat = path.stat() + if entry and entry.get("size") == stat.st_size and entry.get("mtime") == stat.st_mtime_ns: + return entry + offset = 0 + rows = [] + if ( + entry + and isinstance(entry.get("offset"), int) + and stat.st_size > entry["offset"] + and entry.get("file") == [stat.st_dev, stat.st_ino] + and entry.get("head") == read_digest(path, entry["offset"]) + ): + offset = entry["offset"] + rows = list(entry.get("rows") or []) + with path.open("rb") as handle: + handle.seek(offset) + for raw in handle: + if not raw.endswith(b"\n"): + # A line still being written; pick it up whole next time. + break + offset += len(raw) + row = parse_line(raw, path.stem) + if row is not None: + # A record with no time is dated once, by when the file was last + # written as it is read: the index keeps it past today, and it must + # not count as today's on every day after. In whole milliseconds, + # floored: local_day rounds a float of seconds, which dates a write + # in the last half second of a day as the next. + if row[2] is None: + row[2] = stat.st_mtime_ns // 1_000_000 + rows.append(row) + return { + "size": stat.st_size, + "mtime": stat.st_mtime_ns, + "file": [stat.st_dev, stat.st_ino], + "offset": offset, + "head": read_digest(path, offset), + "rows": rows, + } + + +def load_index(path): + data = read_json(path) + if data.get("version") == INDEX_VERSION and isinstance(data.get("files"), dict): + return data["files"] + return {} + + +# ------------------------------------------------------------------ routes + +# The variable Qwen Code takes an auth type's base URL from when no +# modelProviders route names one (its AUTH_ENV_MAPPINGS). Gemini and Vertex +# AI take none. +BASE_URL_ENV = {"openai": "OPENAI_BASE_URL", "openai-responses": "OPENAI_BASE_URL", "anthropic": "ANTHROPIC_BASE_URL"} + + +def dotenv(path): + """A .env file's variables, as dotenv reads them: KEY=value lines, an + `export ` prefix, quotes and a trailing # comment dropped.""" + values = {} + try: + text = path.read_text(encoding="utf-8", errors="replace") + except (OSError, ValueError): + return values + for line in text.splitlines(): + key, sep, value = line.strip().removeprefix("export ").partition("=") + key, value = key.strip(), value.strip() + if not sep or not key or key.startswith("#"): + continue + if value[:1] in ("'", '"', "`") and value[1:].find(value[0]) >= 0: + value = value[1:1 + value[1:].find(value[0])] + else: + value = value.split(" #", 1)[0].strip() + values[key] = value + return values + + +def env_files(home): + """The .env files Qwen Code loads wherever it starts, first one first: its + home's, ~/.qwen's when QWEN_HOME moves the home, and ~/.env. A project's + own .env applies to the runs started there alone, which nothing here can + tell apart.""" + files = [] + for path in (home / ".env", Path.home() / ".qwen" / ".env", Path.home() / ".env"): + if path not in files: + files.append(path) + return files + + +def base_urls(home): + """auth type -> the base URL Qwen Code sends it to when no modelProviders + route names one, resolved as Qwen Code resolves it: the variable from the + environment it starts in, else from the first .env file setting it, else + from settings.json's env, and without one the security.auth.baseUrl + /auth saved for the auth type it selected.""" + settings = read_json(home / "settings.json") + layers = [dict(os.environ), *(dotenv(path) for path in env_files(home)), obj(settings.get("env"))] + auth = obj(obj(settings.get("security")).get("auth")) + saved = auth.get("baseUrl") if isinstance(auth.get("baseUrl"), str) else "" + urls = {} + for auth_type, key in BASE_URL_ENV.items(): + url = next((str(layer[key]).strip() for layer in layers if isinstance(layer.get(key), str) and layer[key].strip()), "") + if not url and auth.get("selectedType") in (None, auth_type): + url = saved.strip() + if url: + urls[auth_type] = url + return urls + + +def model_routes(home): + """(auth type, model id) -> base URL, from settings.json's modelProviders: + a provider id's routes, each an id and its baseUrl. A model the settings + list twice is sent where its first entry says, as Qwen Code picks it.""" + routes = {} + providers = obj(read_json(home / "settings.json").get("modelProviders")) + for provider, entries in providers.items(): + for entry in entries if isinstance(entries, list) else []: + entry = obj(entry) + url = entry.get("baseUrl") + if isinstance(entry.get("id"), str) and isinstance(url, str) and url.strip(): + # A Responses route is an openai one with its wireApi set. + auth = "openai-responses" if provider == "openai" and entry.get("wireApi") == "responses" else str(provider) + routes.setdefault((auth, entry["id"]), url.strip()) + return routes + + +def selected_route(home): + """(auth type, model id) settings.json selects, where Qwen Code sends a + prompt when nothing else says: /auth's security.auth.selectedType and + /model's model.name.""" + settings = read_json(home / "settings.json") + auth = obj(obj(settings.get("security")).get("auth")).get("selectedType") + model = obj(settings.get("model")).get("name") + return str(auth) if isinstance(auth, str) else "", str(model) if isinstance(model, str) else "" + + +def route_config(home): + """What route_owner() decides with: the modelProviders routes, the base + URLs outside them, and the selected route.""" + return model_routes(home), base_urls(home), selected_route(home) + + +def harness_match(): + """For a scan's cached_scan(match=...): a digest of where Qwen Code's + replies are sent, which its settings and .env files can change without a + transcript changing.""" + routes, urls, selected = route_config(qwen_home()) + config = json.dumps([sorted([*key, url] for key, url in routes.items()), sorted(urls.items()), selected]) + return {"qwenRoutes": hashlib.sha1(config.encode("utf-8", "surrogatepass")).hexdigest()} + + +def route_owner(auth_type, model, routes, urls=None): + """The collector that bills a reply, by the auth type Qwen Code recorded + for it and the base URL it was sent to: its model's modelProviders route, + else the base URL the auth type takes from Qwen Code's environment (urls, + from base_urls()). A URL is read the way every collector reads Claude + Code's (claude_code_owner): Alibaba's hosts are Qwen's, a provider's host + is that provider's, a server you run local's unless the model is a Claude + or an Ollama cloud one it only forwards, and a gateway nobody's. Without a + URL, a reply from before Qwen Code kept the type ran on Qwen Code's + default service, as every reply was read before, while an OpenAI-compatible + one went to OpenAI's API, Qwen Code's default for it (getDefaultBaseUrl), + and Anthropic, Gemini and Vertex AI are their own services, none of which + a collector here bills.""" + if auth_type == "qwen-oauth": + return AGENT_ID + url = routes.get((auth_type, model)) or next((value for (_, routed), value in routes.items() if routed == model), "") + url = url or (urls or {}).get(auth_type, "") + if url: + return claude_code_owner(model, url) + return AGENT_ID if auth_type == "" else None + + +# Where parse_line() keeps each kind of row's copy mark. +COPY_FIELD = {"p": 4, "e": 6, "r": 10} + + +def summarize(rows, owner): + """The usage rows `owner` bills from a transcript's index rows, in file + order. Qwen Code's prompts are its user records, not its replies: a reply + adds its tokens, and a prompt only its session and day, billed with the + next request sent for it, answered or not, else with the route settings.json + selects. Beside them, each row's (kind, key) and whether it is a fork's + copy.""" + routes, urls, selected = route_config(qwen_home()) + parsed = [] + for row in rows: + try: + kind, key, stamp, session = row[:4] + day = local_day(stamp) + if kind == "r": + model = str(row[4]) + counts = [int(value) for value in row[5:8]] + billed = route_owner(str(row[8]), str(row[9]), routes, urls) + elif kind == "e": + model, counts, billed = "", [0, 0, 0], route_owner(str(row[4]), str(row[5]), routes, urls) + elif kind == "p": + model, counts, billed = "", [0, 0, 0], None + else: + continue + copy = bool(row[COPY_FIELD[kind]]) + except (TypeError, ValueError, IndexError, OverflowError, OSError): + continue + parsed.append([kind, key, day, str(session), model, counts, billed, copy]) + answered = route_owner(*selected, routes, urls) + for entry in reversed(parsed): + if entry[0] in ("r", "e"): + answered = entry[6] + else: + entry[6] = answered + parsed = [entry for entry in parsed if entry[0] != "e"] + kept = [entry for entry in parsed if entry[2] is not None and entry[6] == owner] + return [ + UsageRow(day, session, "", model, *counts, 0, prompts=0 if kind == "r" else 1) + for kind, key, day, session, model, counts, billed, copy in kept + ], [((kind, key), copy) for kind, key, day, session, model, counts, billed, copy in kept] + + +def index_path(): + # Beside the scan cache, under its digest of every directory the scan + # reads, so another home or runtime directory never reuses this index. + cache_file, _ = scan_cache_paths(AGENT_ID, runtime_dirs()) + return cache_file.with_name(cache_file.name.replace("-scan-", "-index-", 1)) + + +def harness_sources(): + return runtime_dirs() + + +def harness_rows(owner): + """Qwen Code's usage that `owner` bills, each record counted once though a + resumed or branched session carries it again, and whether every + transcript was read.""" + try: + index_file = index_path() + except Exception: + index_file = None + previous = load_index(index_file) if index_file else {} + current = {} + complete = True + originals, copies = [], [] + for path in chat_files(): + key = str(path) + try: + current[key] = parse_file(path, previous.get(key)) + except OSError as exc: + warn(f"skipping {path}: {type(exc).__name__}") + complete = False + continue + rows, keys = summarize(current[key].get("rows") or [], owner) + for row, (row_key, copy) in zip(rows, keys): + (copies if copy else originals).append((row_key, row)) + # A fork's copy counts only where its source's record is gone, whichever + # file name sorts first, so the record stays the session's that wrote it. + found, seen = [], set() + for row_key, row in originals + copies: + if row_key not in seen: + seen.add(row_key) + found.append(row) + if index_file and current != previous: + try: + write_json(index_file, {"version": INDEX_VERSION, "files": current}) + except Exception as exc: + warn(f"could not write transcript index ({exc})") + return found, complete + + +# pi, omp, and opencode name the provider alike. +def billed_here(provider, model, entry): + return provider.lower() in PROVIDERS + + +def scan(reuse=True): + """Qwen Code's transcripts, the Qwen plan turns in pi/omp sessions and + both of opencode's message stores, and Claude Code's and MiniMax Code's + turns sent to Model Studio. Returns (stats, complete); a scan that + couldn't read every source is not cached.""" + stats = Stats() + # Every source runs, so one unreadable source still leaves the others. + transcript_rows, transcripts_complete = harness_rows(AGENT_ID) + for row in transcript_rows: + stats.add_row(row, row.model) + pattern = r'"provider"\s*:\s*"(' + "|".join(PROVIDERS) + r')"' + pi_rows, pi_complete = scan_pi_sessions(pattern, billed_here) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode([provider + '"' for provider in PROVIDERS], billed_here) + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + for row in pi_rows + db_rows + routed_rows: + # pi and opencode count each reply as the prompt that asked for it; + # Claude Code's rows carry their typed prompts apart. + stats.add_row(row, model_label(row.model) or "qwen") + return stats.snapshot(), transcripts_complete and pi_complete and db_complete and routed_complete + + +def scan_sources(): + # Every data path the scan reads: Qwen Code's directories, the pi/omp + # session roots, the opencode database, and Claude Code's and MiniMax + # Code's sessions. + return [*runtime_dirs(), *pi_session_roots(), opencode_db(), *routed_sources()] + + +def main(): + parser = argparse.ArgumentParser(description="Print the Qwen Code usage record as JSON") + # --force rescans everything. --limits-only only promises fresh limits, and + # Qwen Code reads none, so it reuses a scan far longer than a normal run. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + # The digest covers every path the scan reads, so another home, runtime + # directory, session root or database never reuses this one's numbers. + # Where each harness sends its replies decides whose they are, so a scan + # only answers for the routes it was made under. + match = {**harness_match(), **routed_match()} + # A forced scan reads every Codex rollout again too. + max_age = scan_max_age(args.force, args.limits_only) + stats = cached_scan(AGENT_ID, scan_sources(), max_age, lambda: scan(reuse=max_age > 0), match=match) + + # There is no quota to read, so nothing waits on a sign-in: the status stays + # empty and local stats show alone. A machine that never ran Qwen Code or a + # Qwen plan prints a record the panel skips. + in_use = stats["totalSessions"] > 0 + record = base_record(AGENT_ID, AGENT_NAME, ready=in_use, hasLocalStats=in_use) + record.update(stats) + emit(record) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-synthetic b/bin/omarchy-agent-usage-synthetic new file mode 100755 index 00000000000..35a780a75c9 --- /dev/null +++ b/bin/omarchy-agent-usage-synthetic @@ -0,0 +1,318 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Synthetic usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Synthetic usage into one display-ready JSON record. + +Local stats come from pi/omp sessions and opencode messages that ran on the +synthetic provider, and from Claude Code, MiniMax Code and Codex sessions sent +to Synthetic's API. The account's allowances (the rolling five-hour request +limit, the weekly credits, search and free tool calls) come from Synthetic's +documented GET https://api.synthetic.new/v2/quotas, which doesn't count +against them. The agents panel only ever reads the JSON this prints. +""" + +import argparse +import math +import os +import re +import sys +import urllib.request +from datetime import datetime +from pathlib import Path + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record, + cached_scan, + data_home, + emit, + number, + obj, + open_json, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, + trusted_endpoint, +) + +AGENT_ID = "synthetic" +AGENT_NAME = "Synthetic" +PROVIDER = "synthetic" +API_BASE_URL = "https://api.synthetic.new" +AUTH_HELP = "Set SYNTHETIC_API_KEY, or sign in to Synthetic in pi, omp or opencode, to read its quotas." +REFUSED_HELP = "Synthetic rejected the API key. Set a current one in SYNTHETIC_API_KEY or sign in to Synthetic again." + + +def finite(value): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) + + +def model_name(raw): + # Synthetic serves Hugging Face models as "hf:org/model"; the model is the + # last part. An alias (syn:large:text) is kept whole. + value = str(raw or "").strip().removeprefix("hf:").rstrip("/").split("/")[-1] + return value or AGENT_ID + + +# ------------------------------------------------------------- local stats + + +def pi_agent_dirs(): + pi_agent = os.environ.get("PI_CODING_AGENT_DIR") + dirs = [Path(pi_agent).expanduser() if pi_agent else Path.home() / ".pi" / "agent", Path.home() / ".omp" / "agent"] + # `omp --profile=` (and pi's equivalent) moves the agent tree under + # /profiles//. + for base in (Path.home() / ".pi", Path.home() / ".omp"): + try: + dirs.extend(sorted(child / "agent" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + pass + return dirs + + +# pi, omp, and opencode name the provider alike. Exact match: opencode +# provider ids are free-form. +def billed_here(provider, model, entry): + return provider == PROVIDER + + +def run_local_scans(reuse=True): + stats = Stats() + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"synthetic"', billed_here) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode([PROVIDER + '"'], billed_here) + # Claude Code, MiniMax Code and Codex sent to Synthetic's host, or Claude + # Code answered by an hf:/syn: model. + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + for row in pi_rows + db_rows + routed_rows: + stats.add_row(row, model_name(row.model)) + return stats.snapshot(), pi_complete and db_complete and routed_complete + + +def cached_local_stats(max_age): + # The digest covers every data path the scan reads: the pi/omp session + # roots, the opencode DB, and Claude Code's, MiniMax Code's and Codex's + # sessions, whose owner depends on where they are sent. A forced scan reads + # every Codex rollout again too. + sources = [Path.home(), opencode_db(), *pi_session_roots(), *routed_sources()] + return cached_scan(AGENT_ID, sources, max_age, lambda: run_local_scans(reuse=max_age > 0), match=routed_match()) + + +# ------------------------------------------------------------- credentials + + +PI_ENV_REFERENCE = re.compile(r"\$(?:\{([A-Za-z_][A-Za-z0-9_]*)\}|([A-Za-z_][A-Za-z0-9_]*))") + + +def stored_key(path): + # pi, omp and opencode keep {"synthetic": {"key": ...}}. pi's key may also + # be a $ENV_VAR reference, resolved here, or a !command, which stays pi's. + entry = obj(read_json(path).get(PROVIDER)) + key = str(entry.get("key") or "").strip() + if not key or key.startswith("!"): + return "" + unresolved = [] + + def substitute(match): + value = os.environ.get(match.group(1) or match.group(2), "") + if not value: + unresolved.append(match.group(0)) + return value + + key = PI_ENV_REFERENCE.sub(substitute, key).strip() + return "" if unresolved else key + + +def api_key(): + explicit = str(os.environ.get("SYNTHETIC_API_KEY", "")).strip() + if explicit: + return explicit + # The default pi and omp sign-ins, then opencode's, then each named profile's: + # a profile's sessions are counted, so a sign-in made only there must read + # its quotas too, without outranking the defaults. + agent_dirs = pi_agent_dirs() + paths = [agent_dir / "auth.json" for agent_dir in agent_dirs[:2]] + paths.append(data_home() / "opencode" / "auth.json") + paths.extend(agent_dir / "auth.json" for agent_dir in agent_dirs[2:]) + for path in paths: + key = stored_key(path) + if key: + return key + return "" + + +# ------------------------------------------------------------------ quotas + + +def used_fraction(used, maximum): + used, maximum = finite(used), finite(maximum) + if used is None or maximum is None or maximum <= 0: + return None + return max(0.0, min(1.0, used / maximum)) + + +def reset_time(value): + raw = str(value or "").strip() + if not raw: + return "" + try: + datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return "" + return raw + + +def request_window(label, payload, title=None): + payload = obj(payload) + percent = used_fraction(payload.get("requests"), payload.get("limit")) + if percent is None: + return None + window = {"label": label, "percent": percent, "resetsAt": reset_time(payload.get("renewsAt"))} + if title: + window["title"] = title + return window + + +def parse_quotas(payload): + limits = [] + # The rolling limit and the weekly credits regenerate a little at a time. + # nextTickAt and nextRegenAt say when the next share comes back, not when + # the window starts over, so they are no reset time: reading them as one + # would zero a window that has only regained a sliver. + rolling = obj(payload.get("rollingFiveHourLimit")) + remaining, maximum = finite(rolling.get("remaining")), finite(rolling.get("max")) + if remaining is not None and maximum is not None and maximum > 0: + percent = 1.0 if rolling.get("limited") is True else max(0.0, min(1.0, 1.0 - remaining / maximum)) + limits.append({"label": "Rolling (5h)", "percent": percent, "resetsAt": ""}) + weekly = finite(obj(payload.get("weeklyTokenLimit")).get("percentRemaining")) + if weekly is not None: + limits.append({"label": "Weekly credits", "percent": max(0.0, min(1.0, 1.0 - weekly / 100.0)), "resetsAt": ""}) + + # The documented bucket. Accounts on the credit system still receive it, at + # zero, so it stands in only when neither of those parsed. + if not limits: + window = request_window("Subscription requests", payload.get("subscription")) + if window: + limits.append(window) + for window in ( + request_window("Search (hourly)", obj(payload.get("search")).get("hourly"), "Search"), + request_window("Free tool calls", payload.get("freeToolCalls")), + ): + if window: + limits.append(window) + return limits + + +def probe_quotas(key, base_url): + # SYNTHETIC_API_BASE_URL moves the endpoint, never the key onto a + # connection anyone on the way can read. + if not trusted_endpoint(base_url): + return {"ok": False, "helpText": "SYNTHETIC_API_BASE_URL must use HTTPS (plain HTTP only for loopback). Local Synthetic stats are still shown."} + request = urllib.request.Request( + base_url.rstrip("/") + "/v2/quotas", + headers={ + "Authorization": "Bearer " + key, + "Accept": "application/json", + "User-Agent": "omarchy-agent-usage-synthetic/1.0", + }, + ) + try: + payload = open_json(request, timeout=10) + except HttpError as error: + if error.status in (401, 403): + return {"ok": False, "refused": True, "helpText": REFUSED_HELP} + if error.status: + return {"ok": False, "helpText": f"Synthetic's quota endpoint returned status {error.status}. Local Synthetic stats are still shown."} + if error.transport: + # A transport failure reached no server at all: no route, no DNS. + return {"ok": False, "transport": True, "helpText": "Couldn't reach Synthetic's quota endpoint. Retrying shortly. Local Synthetic stats are still shown."} + return {"ok": False, "helpText": "Synthetic returned an unreadable quota answer. Local Synthetic stats are still shown."} + limits = parse_quotas(payload) if isinstance(payload, dict) else [] + if not limits: + return {"ok": False, "helpText": "Synthetic's quota endpoint returned no quotas. Local Synthetic stats are still shown."} + return {"ok": True, "limits": limits} + + +def collect_limits(key, base_url, force): + result = {"limits": [], "usageStatusText": "", "authHelpText": "", "stale": False, "fetchedAtMs": 0, "retryAdvised": False} + if not key: + return result + + # One cache per key, so another key never shows this one's quotas. + cache = LimitsCache(AGENT_ID, key) + + # --force is a person asking for fresh numbers, so it skips the reuse window. + if cache.fresh(force): + result.update(limits=cache.limits, fetchedAtMs=cache.fetched_at) + return result + + probe = probe_quotas(key, base_url) + if probe["ok"]: + result.update(limits=probe["limits"], fetchedAtMs=cache.save(probe["limits"], only_if_newer=True)) + return result + + # Last good numbers, marked stale, beat an empty section while the check + # fails. A kept window whose reset has passed is untouched, not unknown: it + # comes back at zero with no reset time until the next check. + fallback, fetched_at = cache.kept() + result.update(limits=fallback, stale=bool(fallback), fetchedAtMs=fetched_at, retryAdvised=bool(probe.get("transport"))) + if probe.get("refused"): + result.update(usageStatusText="Sign-in expired", authHelpText=probe["helpText"]) + elif not fallback: + result.update(usageStatusText="Synthetic limits unavailable", authHelpText=probe["helpText"]) + return result + + +# ------------------------------------------------------------------ record + + +def main(): + parser = argparse.ArgumentParser(description="Print the Synthetic usage record as JSON") + parser.add_argument("--force", action="store_true", help="rescan local sessions and re-check quotas, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse a recent local scan; only the quotas must be fresh") + args = parser.parse_args() + + stats = cached_local_stats(scan_max_age(args.force, args.limits_only)) + key = api_key() + limits = collect_limits(key, os.environ.get("SYNTHETIC_API_BASE_URL") or API_BASE_URL, args.force) + in_use = number(stats.get("totalPrompts")) > 0 + + record = base_record( + AGENT_ID, + AGENT_NAME, + ready=in_use or bool(key), + # A key with no sessions here has no local stats to show. + hasLocalStats=number(stats.get("totalSessions")) > 0 or in_use or bool(stats.get("modelUsage")), + # The quotas answer names no plan. + tierLabel="", + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + limits=limits["limits"], + limitsStale=limits["stale"], + limitsFetchedAt=limits["fetchedAtMs"], + ) + if in_use and not key: + record.update(usageStatusText="Waiting for auth", authHelpText=AUTH_HELP) + if limits["retryAdvised"]: + record["retryAdvised"] = True + record.update(stats) + emit(record, sort_keys=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-agent-usage-update b/bin/omarchy-agent-usage-update index 7a79159c3ab..c5c26e9fd2b 100755 --- a/bin/omarchy-agent-usage-update +++ b/bin/omarchy-agent-usage-update @@ -58,6 +58,11 @@ if [[ -d $accounts_dir ]]; then "$OMARCHY_PATH/bin/omarchy-agent-account-state" refresh fi +# Collectors forced together read the Codex rollouts they share once between +# them, whenever each gets to them, by knowing when this refresh began. +# EPOCHREALTIME uses the locale's decimal separator, a comma in many. +export OMARCHY_AGENT_USAGE_REFRESH_AT=${EPOCHREALTIME/,/.} + pids=() for collector in "$OMARCHY_PATH"/bin/omarchy-agent-usage-*; do [[ -x $collector ]] || continue diff --git a/bin/omarchy-agent-usage-vibe b/bin/omarchy-agent-usage-vibe new file mode 100755 index 00000000000..c1e23418c8b --- /dev/null +++ b/bin/omarchy-agent-usage-vibe @@ -0,0 +1,696 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the Mistral Vibe usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect Mistral Vibe usage into one display-ready JSON record. + +Local stats come from the sessions Vibe saves under ~/.vibe/logs/session/ +(VIBE_HOME moves it, and config.toml's session_logging.save_dir moves the +sessions): the Unified Harness store current Vibe keeps under unified/, and +the per-session meta.json files older releases wrote, plus pi/omp and opencode +sessions that ran on the `mistral` provider. Vibe also drives other providers +(a local llama.cpp server out of the box), so a Vibe session only counts when +the model it ran on is served by Mistral; one on a server you run is the +local collector's, which reads it through harness_rows(). The plan comes from the whoami +answer Vibe caches in whoami_cache.json. Mistral publishes no usage limits +to read, so the record carries no meters. The agents panel only ever reads +the JSON this prints. +""" + +import argparse +import hashlib +import json +import os +import re +import sys +import time +import tomllib +from datetime import datetime +from pathlib import Path +from urllib.parse import urlparse + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + Stats, + UsageRow, + base_record, + cached_scan, + claude_code_owner, + emit, + obj, + opencode_db, + pi_session_roots, + read_json, + scan_cache_paths, + scan_max_age, + scan_opencode, + scan_pi_sessions, + token_count, + warn, + write_json, +) + +AGENT_ID = "vibe" +AGENT_NAME = "Mistral Vibe" + +# The provider id pi, omp, and opencode give Mistral's own API. A model reached +# through a gateway (OpenRouter, a proxy) is billed there and counted there. +PROVIDER = "mistral" + +# The alias an unpinned Vibe config resolves to when it is configured, as +# Vibe's own resolve_default_model_alias does. +DEFAULT_MODEL_ALIAS = "mistral-medium-3.5" + +# Version 1 kept only the sessions Mistral served, so it named no owner, +# version 2 gave a mistral-backend provider to Mistral wherever its api_base +# pointed, and version 3 gave a Unified Harness session whose alias config.toml +# dropped to the default model. +INDEX_VERSION = 4 + +# Times read from Vibe's meta.json are untrusted. Anything past these bounds +# is a corrupt or crafted field, not usage. +MIN_EPOCH_SECONDS = 946_684_800 # 2000-01-01 +MAX_MODEL_NAME_LEN = 120 + + +def epoch_seconds(value): + """Seconds since the epoch for an ISO string or a seconds/milliseconds + number, or None when the value can't be a real moment of use.""" + if value is None or isinstance(value, bool): + return None + if isinstance(value, (int, float)): + seconds = float(value) + elif isinstance(value, str) and value.strip(): + text = value.strip() + try: + seconds = float(text) + except ValueError: + try: + parsed = datetime.fromisoformat(text[:-1] + "+00:00" if text.endswith("Z") else text) + # A zoneless time at the edge of the calendar (year 1 or 9999) can't + # be placed in a zone off UTC. + if parsed.tzinfo is None: + parsed = parsed.astimezone() + seconds = parsed.timestamp() + except (OverflowError, OSError, ValueError): + return None + else: + return None + if seconds != seconds: + return None + if seconds > 10_000_000_000: + seconds /= 1000 + # A day of slack past now covers a clock a little behind the writer's. + if seconds < MIN_EPOCH_SECONDS or seconds > time.time() + 86_400: + return None + return seconds + + +def start_day(seconds): + return datetime.fromtimestamp(seconds).strftime("%Y-%m-%d") + + +def model_label(raw): + """A model id fit for the record: its last path segment, without a + thinking-level suffix, or None when it is too long to be a real id.""" + value = str(raw or "").rstrip("/").split("/")[-1].split(":")[0] + if len(value) > MAX_MODEL_NAME_LEN: + return None + return value or "mistral" + + +# ------------------------------------------------------------------ Vibe + + +def vibe_home(): + configured = os.environ.get("VIBE_HOME", "").strip() + if configured: + return Path(configured).expanduser().resolve() + return Path.home() / ".vibe" + + +def vibe_config(): + """Vibe's own config.toml, or {} when there is none to read.""" + try: + with (vibe_home() / "config.toml").open("rb") as handle: + return tomllib.load(handle) + except (OSError, tomllib.TOMLDecodeError, UnicodeDecodeError): + return {} + + +def session_logging(): + """Where Vibe saves sessions and the prefix it names them with, honoring + [session_logging] in Vibe's own config.toml as Vibe does.""" + home = vibe_home() + save_dir = home / "logs" / "session" + prefix = "session" + section = obj(vibe_config().get("session_logging")) + configured_dir = section.get("save_dir") + if isinstance(configured_dir, str) and configured_dir.strip(): + path = Path(configured_dir.strip()).expanduser() + # Vibe resolves a relative save_dir against wherever it was started, which + # this collector can't know; such a config keeps the default location. + if path.is_absolute(): + save_dir = path + configured_prefix = section.get("session_prefix") + if isinstance(configured_prefix, str) and configured_prefix.strip() and "/" not in configured_prefix: + prefix = configured_prefix.strip() + return save_dir, prefix + + +def resolve_model(config): + """The model entry a session ran on, resolved the way Vibe resolves it: the + pinned alias, else the experiment-routed one, else the default alias, else + the first configured model.""" + models = config.get("models") + # Current Vibe dumps the models keyed by alias; older releases as a list. + if isinstance(models, dict): + entries = list(models.values()) + elif isinstance(models, list): + entries = models + else: + return None + by_alias = {} + for entry in entries: + if isinstance(entry, dict): + alias = str(entry.get("alias") or entry.get("name") or "") + if alias and alias not in by_alias: + by_alias[alias] = entry + for alias in (config.get("active_model"), config.get("routed_default_model"), DEFAULT_MODEL_ALIAS): + if isinstance(alias, str) and alias in by_alias: + return by_alias[alias] + return next(iter(by_alias.values()), None) + + +# Vibe's built-in providers, which a config that lists none still has. +DEFAULT_API_BASES = {PROVIDER: "https://api.mistral.ai/v1", "llamacpp": "http://127.0.0.1:8080/v1"} + + +def provider_owner(name, providers, model): + """The collector that bills `model` on the named provider: this one for + Mistral's own API, else whoever runs its api_base, read the way every + collector reads Claude Code's (claude_code_owner): local for a server you + run, unless the model is a Claude or an Ollama cloud one it only forwards, + and None for one nothing here can place. + + Vibe's `mistral` backend sends to the provider's api_base too, so the + api_base decides for it as for a generic one: a mistral.ai host is Mistral's + API, and so is a mistral-backend provider that names none. A provider the + config doesn't describe is one of Vibe's built-in ones, or no one's. + """ + api_base = DEFAULT_API_BASES.get(name, "") + for provider in providers if isinstance(providers, list) else []: + if not isinstance(provider, dict) or str(provider.get("name") or "") != name: + continue + api_base = str(provider.get("api_base") or "").strip() + if not api_base and str(provider.get("backend") or "").lower() == "mistral": + return AGENT_ID + break + host = (urlparse(api_base).hostname or "").lower() + if host == "mistral.ai" or host.endswith(".mistral.ai"): + return AGENT_ID + return claude_code_owner(model, api_base) if api_base else None + + +def session_model(meta): + """The model label a session ran on and the collector that bills it, or + None when it names no model.""" + config = obj(meta.get("config")) + model = resolve_model(config) + if not isinstance(model, dict): + return None + label = model_label(model.get("alias") or model.get("name")) + if label is None: + return None + return label, provider_owner(str(model.get("provider") or ""), config.get("providers"), str(model.get("name") or "")) + + +def count_prompts(messages_path): + """The prompts typed into a session: user messages Vibe did not inject + itself, such as reminders and compaction summaries.""" + prompts = 0 + with messages_path.open("rb") as handle: + for raw in handle: + if b'"user"' not in raw: + continue + try: + message = json.loads(raw) + except ValueError: + continue + if not isinstance(message, dict) or message.get("role") != "user": + continue + if message.get("injected") is True or message.get("context_boundary"): + continue + prompts += 1 + return prompts + + +def file_signature(path): + try: + stat = path.stat() + except FileNotFoundError: + return None + return [stat.st_size, stat.st_mtime_ns] + + +def parse_meta(meta_path, count_session_prompts): + """One session's index row: [startSeconds, model, input, output, cacheRead, + prompts, owner], or None when it names no model or used nothing. + + Vibe keeps running totals per session, so the whole session is booked on + the day it started. + """ + meta = json.loads(meta_path.read_text(encoding="utf-8")) + if not isinstance(meta, dict): + return None + routed = session_model(meta) + started = epoch_seconds(meta.get("start_time")) + if routed is None or started is None: + return None + model, owner = routed + stats = obj(meta.get("stats")) + prompt_tokens = token_count(stats.get("session_prompt_tokens")) + # Mistral reports cached prompt tokens inside prompt_tokens, as OpenAI does. + cached = min(token_count(stats.get("session_cached_tokens")), prompt_tokens) + completion = token_count(stats.get("session_completion_tokens")) + prompts = 0 + messages = meta_path.parent / "messages.jsonl" + if count_session_prompts and messages.is_file(): + prompts = count_prompts(messages) + if prompt_tokens + completion <= 0 and prompts <= 0: + return None + return [started, model, prompt_tokens - cached, completion, cached, prompts, owner] + + +# ------------------------------------------------------- Unified Harness + +# Current Vibe runs every session on the Unified Harness, which keeps it in +# /unified//: CURRENT names the generation in force, +# whose projection-state.json holds the session's public state as of that +# generation, and the journal it names records every change since. The +# session's running token totals live in that state; its meta.json only +# names the model the session ran on. +UNIFIED_DIR = "unified" +DEFAULT_MODEL = {"alias": DEFAULT_MODEL_ALIAS, "name": "mistral-vibe-cli-latest", "provider": "mistral"} +STORE_NAME = re.compile(r"^[A-Za-z0-9._-]+$") +JOURNAL_PATH = re.compile(r"^journal/[0-9]{16}\.jsonl$") +CHUNK_DIGEST = re.compile(r"^[0-9a-f]{64}$") + + +def store_json(path): + return json.loads(path.read_text(encoding="utf-8")) + + +def store_name(value): + """A generation, file or session name from the store, or None when it + would leave the directory it names something in.""" + if isinstance(value, str) and STORE_NAME.match(value) and value not in (".", ".."): + return value + return None + + +def dict_entries(value): + return [entry for entry in value if isinstance(entry, dict)] if isinstance(value, list) else [] + + +def generation_document(generation_dir, described, default): + name = store_name(obj(described).get("path") or default) + if name is None: + raise ValueError(f"manifest names no {default}") + return obj(store_json(generation_dir / name)) + + +def unified_state(session_dir): + """The session's public state and the Harness's own metadata for it, with + its journal folded in, as (envelope, history entries, metadata).""" + current = obj(store_json(session_dir / "CURRENT")) + generation = store_name(current.get("generation")) + if generation is None: + raise ValueError("CURRENT names no generation") + generation_dir = session_dir / "generations" / generation + manifest = obj(store_json(generation_dir / "manifest.json")) + metadata = obj(generation_document(generation_dir, manifest.get("runtime_state"), "runtime-state.json").get("session_metadata")) + described = obj(manifest.get("projection_state")) + envelope = obj(generation_document(generation_dir, described, "projection-state.json").get("snapshot")) + entries = dict_entries(obj(envelope.get("history")).get("entries")) + # A long history is kept apart, in the store's chunk pool. + chunks = described.get("chunks") + if isinstance(chunks, list): + entries = [] + for digest in chunks: + if not isinstance(digest, str) or not CHUNK_DIGEST.match(digest): + raise ValueError("manifest names a bad chunk") + entries.extend(dict_entries(store_json(session_dir / "chunks" / f"{digest}.json"))) + + journal = obj(manifest.get("recovery_journal_segment")).get("path") + if isinstance(journal, str) and JOURNAL_PATH.match(journal) and (session_dir / journal).is_file(): + for raw in (session_dir / journal).read_bytes().splitlines(keepends=True): + # A record still being appended has no newline yet. + if not raw.endswith(b"\n"): + break + try: + record = obj(json.loads(raw)) + except ValueError: + continue + payload = obj(record.get("payload")) + kind = record.get("type") + if kind == "projection_advanced" and isinstance(payload.get("snapshot"), dict): + envelope = payload["snapshot"] + entries = dict_entries(obj(envelope.get("history")).get("entries")) + elif kind in ("projection_delta", "session_relocated") and isinstance(payload.get("delta"), list): + for op in payload["delta"]: + entries, envelope = apply_projection_op(obj(op), entries, envelope) + if kind == "session_relocated" and isinstance(payload.get("metadata"), dict): + metadata = payload["metadata"] + return envelope, entries, metadata + + +def apply_projection_op(op, entries, envelope): + """One step of a projection delta, applied as the Harness replays it.""" + kind = op.get("op") + if kind == "append_entry" and isinstance(op.get("entry"), dict): + entries = [*entries, op["entry"]] + elif kind == "replace_entry" and isinstance(op.get("entry"), dict): + index = next((i for i, entry in enumerate(entries) if entry.get("id") == op.get("id")), None) + if index is not None: + entries = [*entries[:index], op["entry"], *entries[index + 1:]] + elif kind == "remove_entry": + entries = [entry for entry in entries if entry.get("id") != op.get("id")] + elif kind == "set_history_entries": + entries = dict_entries(op.get("entries")) + elif kind == "set_envelope" and isinstance(op.get("state"), dict): + envelope = op["state"] + return entries, envelope + + +def unified_model(alias, config): + """The model label a Unified Harness session ran on and who bills it, as + session_model() reads them. The session names only the model's alias, so + the model and its provider come from config.toml, over Vibe's built-in + default. An alias the config no longer names keeps its label and goes + to no one, rather than to whichever model Vibe would fall back to.""" + models = config.get("models") + if isinstance(models, dict): + models = list(models.values()) + models = [*(models if isinstance(models, list) else []), DEFAULT_MODEL] + if isinstance(alias, str) and alias: + known = {str(entry.get("alias") or entry.get("name") or "") for entry in models if isinstance(entry, dict)} + if alias not in known: + label = model_label(alias) + return None if label is None else (label, None) + return session_model({"config": {"active_model": alias, "models": models, "providers": config.get("providers")}}) + + +def parse_unified(session_dir, config): + """One Unified Harness session's index row: [startSeconds, model, input, + output, cacheRead, prompts, session, owner], or None when it names no + model or used nothing. + + A subagent is a session of its own in the store. It counts toward the + session that spawned it, and its prompts are that session's instructions + rather than anything typed. A fork is a session of its own, starting + from no usage, and the history it copied is not its prompts. + """ + meta_path = session_dir / "meta.json" + meta = obj(store_json(meta_path)) if meta_path.is_file() else {} + envelope, entries, metadata = unified_state(session_dir) + # meta.json names the model Vibe resolved for the session; without one, the + # session's own pin, else the configured model, is the one it ran on. + alias = obj(meta.get("config")).get("active_model") or metadata.get("active_model") or config.get("active_model") + routed = unified_model(alias, config) + if routed is None: + return None + model, billed = routed + session = obj(envelope.get("session")) + started = epoch_seconds(meta.get("start_time")) or epoch_seconds(session.get("createdAt")) + if started is None: + return None + usage = obj(session.get("tokenUsage")) + # Cached prompt tokens are part of inputTokens, as the Harness checks. + input_tokens = token_count(usage.get("inputTokens")) + cached = min(token_count(usage.get("cachedInputTokens")), input_tokens) + output_tokens = token_count(usage.get("outputTokens")) + owner = session_dir.name + prompts = 0 + if metadata.get("subagent_spawn_key"): + owner = store_name(metadata.get("root_session_id")) or store_name(session.get("rootSessionId")) or owner + else: + # The Harness leaves the turns Vibe injects out of the history, and marks + # what a fork copied as its own. + prompts = sum(1 for entry in entries if entry.get("type") == "message" and entry.get("role") == "user" + and entry.get("source") in ("turn_start", "turn_steer")) + if input_tokens + output_tokens <= 0 and prompts <= 0: + return None + owner = session_dir.parent / owner + return [started, model, input_tokens - cached, output_tokens, cached, prompts, str(owner), billed] + + +def unified_signature(session_dir, config_path): + journal_dir = session_dir / "journal" + journals = sorted(journal_dir.glob("*.jsonl")) if journal_dir.is_dir() else [] + return [ + file_signature(session_dir / "CURRENT"), + file_signature(session_dir / "meta.json"), + [[path.name, file_signature(path)] for path in journals], + file_signature(config_path), + ] + + +def load_index(path): + data = read_json(path) + if data.get("version") == INDEX_VERSION and isinstance(data.get("files"), dict): + return data["files"] + return {} + + +def harness_sources(): + return scan_sources() + + +def harness_match(): + """What harness_rows() decided owners under, for cached_scan(match=...): + the models and providers in config.toml, which say where a Unified Harness + session went, so a scan made before a route changed never stands in for + one after it, in this collector or in local.""" + config = vibe_config() + routes = json.dumps([config.get(key) for key in ("active_model", "models", "providers")], sort_keys=True, default=str) + return {"vibeRoutes": hashlib.sha256(routes.encode("utf-8", "surrogatepass")).hexdigest()[:16]} + + +def harness_rows(owner, index_file=False): + """Vibe sessions that `owner` bills, from the Unified Harness store and the + session folders older releases wrote, and whether every session was read, + so a partial scan is never cached. With no index file named, the one + beside this collector's scan cache keeps the parsed sessions. + + A session's subagents keep their own meta.json under its agents/ directory + with their own totals; they count toward the session that spawned them, and + their prompts are the parent's instructions rather than anything typed. + """ + if index_file is False: + try: + index_file = index_path(scan_sources()) + except Exception as exc: + warn(f"session index unavailable ({exc})") + index_file = None + save_dir, prefix = session_logging() + previous = load_index(index_file) if index_file else {} + current = {} + complete = True + rows = [] + try: + session_dirs = sorted(path for path in save_dir.glob(f"{prefix}_*") if path.is_dir()) + except OSError: + return rows, False + + for session_dir in session_dirs: + metas = [session_dir / "meta.json"] + agents = session_dir / "agents" + if agents.is_dir(): + metas.extend(sorted(agents.rglob("meta.json"))) + for meta_path in metas: + key = str(meta_path) + top_level = meta_path.parent == session_dir + try: + signature = [file_signature(meta_path), file_signature(meta_path.parent / "messages.jsonl") if top_level else None] + if signature[0] is None: + continue + held = previous.get(key) + if isinstance(held, dict) and held.get("sig") == signature: + entry = held + else: + entry = {"sig": signature, "row": parse_meta(meta_path, top_level)} + except (OSError, ValueError, UnicodeDecodeError) as exc: + # Mid-write or unreadable: skip it this run, uncached. + warn(f"skipping {meta_path}: {type(exc).__name__}") + complete = False + continue + current[key] = entry + row = entry.get("row") + if row is None: + continue + try: + started, model, input_tokens, output_tokens, cache_read, prompts, billed = row + day = start_day(float(started)) + counts = [int(input_tokens), int(output_tokens), int(cache_read), int(prompts)] + except (TypeError, ValueError, OverflowError, OSError): + continue + if billed == owner: + rows.append(UsageRow(day, str(session_dir), "", str(model), counts[0], counts[1], counts[2], 0, prompts=counts[3])) + + config = vibe_config() + config_path = vibe_home() / "config.toml" + unified = save_dir / UNIFIED_DIR + try: + unified_dirs = sorted(path for path in unified.iterdir() if path.is_dir()) if unified.is_dir() else [] + except OSError: + unified_dirs = [] + complete = False + for session_dir in unified_dirs: + key = str(session_dir / "CURRENT") + try: + signature = unified_signature(session_dir, config_path) + # A session the Harness hasn't published yet has nothing to count. + if signature[0] is None: + continue + held = previous.get(key) + if isinstance(held, dict) and held.get("sig") == signature: + entry = held + else: + entry = {"sig": signature, "row": parse_unified(session_dir, config)} + except (OSError, ValueError, UnicodeDecodeError) as exc: + # Mid-write or unreadable: skip it this run, uncached. + warn(f"skipping {session_dir}: {type(exc).__name__}") + complete = False + continue + current[key] = entry + row = entry.get("row") + if row is None: + continue + try: + started, model, input_tokens, output_tokens, cache_read, prompts, session, billed = row + day = start_day(float(started)) + counts = [int(input_tokens), int(output_tokens), int(cache_read), int(prompts)] + except (TypeError, ValueError, OverflowError, OSError): + continue + if billed == owner: + rows.append(UsageRow(day, str(session), "", str(model), counts[0], counts[1], counts[2], 0, prompts=counts[3])) + + if index_file and current != previous: + try: + write_json(index_file, {"version": INDEX_VERSION, "files": current}) + except Exception as exc: + warn(f"could not write session index ({exc})") + return rows, complete + + +def tier_label(): + """The plan from the newest whoami answer Vibe cached, named as Vibe's own + resolve_user_plan names it. A plan Vibe has no name for stays unnamed.""" + entries = read_json(vibe_home() / "whoami_cache.json") + newest = None + for entry in entries.values(): + entry = obj(entry) + stored = entry.get("stored_at_timestamp") + payload = obj(entry.get("payload")) + if isinstance(stored, bool) or not isinstance(stored, (int, float)) or not payload: + continue + if newest is None or stored > newest[0]: + newest = (stored, payload) + if newest is None: + return "" + kind = str(newest[1].get("plan_type") or "").strip().lower() + name = str(newest[1].get("plan_name") or "").strip().upper() + if kind == "chat": + return {"FREE": "Free", "INDIVIDUAL": "Pro", "EDU": "Student", "TEAM": "Team"}.get(name, "") + if kind == "mistral_code": + return {"F": "Free Codestral", "E": "Code Enterprise"}.get(name, "") + if kind == "api" and name: + return "Free API" if "FREE" in name else "PAYG API" + return "" + + +# ---------------------------------------------------- pi, omp, and opencode + + +# pi, omp, and opencode name the provider alike. Exact match: opencode +# provider ids are free-form, and a custom gateway named like Mistral is not +# Mistral's API. +def billed_here(provider, model, entry): + return provider == PROVIDER + + +def add_harness_rows(stats): + """Add pi/omp and opencode messages on the mistral provider. Returns + whether every source was read in full.""" + pi_rows, pi_complete = scan_pi_sessions(r'"provider"\s*:\s*"mistral"', billed_here) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode([PROVIDER + '"'], billed_here) + for row in pi_rows + db_rows: + model = model_label(row.model) + if model is not None: + stats.add_row(row, model) + return pi_complete and db_complete + + +# ---------------------------------------------------------------- cache + + +def scan_sources(): + # Every data path the scan reads, so another home or Vibe directory never + # reuses this one's numbers. + save_dir, prefix = session_logging() + return [Path.home(), save_dir, prefix, opencode_db(), *pi_session_roots()] + + +def index_path(sources): + # The per-session index lives beside the scan cache, under the same key. + cache_file, _ = scan_cache_paths(AGENT_ID, sources) + return cache_file.with_name(cache_file.name.replace(f"{AGENT_ID}-scan-", f"{AGENT_ID}-index-", 1)) + + +def run_scans(index_file): + stats = Stats() + rows, complete = harness_rows(AGENT_ID, index_file) + for row in rows: + stats.add_row(row, row.model) + complete = add_harness_rows(stats) and complete + return stats.snapshot(), complete + + +def cached_local_stats(max_age): + sources = scan_sources() + try: + index_file = index_path(sources) + except Exception as exc: + warn(f"session index unavailable ({exc})") + index_file = None + return cached_scan(AGENT_ID, sources, max_age, lambda: run_scans(index_file), match=harness_match()) + + +def main(): + parser = argparse.ArgumentParser(description="Print the Mistral Vibe usage record as JSON") + # --force rescans everything. --limits-only only promises fresh limits, and + # Mistral offers none, so it reuses a scan far longer than a normal run. + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + stats = cached_local_stats(scan_max_age(args.force, args.limits_only)) + tier = tier_label() + + # Mistral publishes no usage limits, so there is nothing to wait for or + # fail at: the status stays empty and local stats show alone. A machine + # with neither Mistral usage nor a cached plan prints a record the panel + # skips. + in_use = stats["totalSessions"] > 0 + record = base_record(AGENT_ID, AGENT_NAME, ready=in_use or bool(tier), hasLocalStats=in_use, tierLabel=tier) + record.update(stats) + emit(record) + 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..84770c67c64 --- /dev/null +++ b/bin/omarchy-agent-usage-zai @@ -0,0 +1,446 @@ +#!/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 shows local +stats alone, which is no trouble. + +Local stats come from pi/omp sessions and from opencode's message database +(both of its stores) that ran through a Z.ai or Zhipu 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 math +import os +import sys +import urllib.request +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +# The shared collector plumbing ships with the agents panel. +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import ( # noqa: E402 + HttpError, + LimitsCache, + Stats, + base_record as shared_base_record, + cached_scan, + config_home, + emit, + open_json, + opencode_db, + pi_session_roots, + read_json, + routed_match, + routed_sources, + scan_max_age, + scan_opencode, + scan_pi_sessions, + scan_routed, +) + +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. Update it in ~/.config/omarchy/agents/zai.json or ZAI_API_KEY." +) + +# Global (Z.ai) vs China (Zhipu / BigModel) split the same API under two hosts. +PLATFORMS = { + "zai": "https://api.z.ai", + "zhipu": "https://open.bigmodel.cn", +} +QUOTA_PATH = "/api/monitor/usage/quota/limit" +# Authentication failed (1000), no credential sent (1001), token expired +# (1003), two-factor authentication needed (1005). +AUTH_ERROR_CODES = {"1000", "1001", "1003", "1005"} + +# A key with no Coding Plan behind it (a pay-as-you-go account) gets the same +# answer every time, so it is asked again hourly rather than on every run; +# --force asks at once, for someone who just subscribed. +NO_PLAN_RECHECK_SECONDS = 3600 + +# The pi, omp and opencode provider ids for Z.ai's and Zhipu's APIs, whose +# turns this collector counts: pi's `zai` and `zai-coding-cn`, omp's `zai` +# and `zhipu-coding-plan`, and models.dev's (opencode's) `zai`, `zhipuai`, +# `zai-coding-plan` and `zhipuai-coding-plan`, plus `zhipu`. Exact names: a +# custom `zai-proxy` gateway is not this subscription. +PROVIDERS = ("zai", "zhipu", "zai-coding-cn", "zhipu-coding-plan", "zhipuai", "zai-coding-plan", "zhipuai-coding-plan") + +# 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)", +} + + +def base_record(**overrides: Any) -> dict[str, Any]: + return shared_base_record(AGENT_ID, AGENT_NAME, **overrides) + + +# -------------------------------------------------------------------------- +# 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, UnicodeError): + # Unreadable, or not UTF-8: no keys from it. + return {} + return values + + +def read_config() -> dict[str, Any]: + return read_json(config_home() / "omarchy" / "agents" / "zai.json") + + +def pi_agent_dir() -> Path: + configured = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + return Path(os.path.expanduser(configured)) if configured else Path.home() / ".pi" / "agent" + + +def auth_json_key(path: Path, provider: str) -> str: + # pi/omp keep a key they were given with /login in auth.json. Only a literal + # key is taken: a `$VAR` or `!command` entry is resolved by pi alone. + entry = read_json(path).get(provider) + if not isinstance(entry, dict): + return "" + key = str(entry.get("key") or "").strip() + return "" if key.startswith(("!", "$")) else key + + +def credentials() -> tuple[str, str]: + """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", pi_agent_dir() / ".env"): + values = read_env_file(env_path) + if values.get("ZAI_API_KEY"): + key, platform = values["ZAI_API_KEY"], platform or "zai" + 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 not key: + for auth_path in (pi_agent_dir() / "auth.json", Path.home() / ".omp" / "agent" / "auth.json"): + if found := auth_json_key(auth_path, "zai"): + key, platform = found, platform or "zai" + break + if found := auth_json_key(auth_path, "zhipu"): + key, platform = found, platform or "zhipu" + break + + if platform not in PLATFORMS: + platform = "zai" + return key, platform + + +# -------------------------------------------------------------------------- +# Coding-plan limits +# -------------------------------------------------------------------------- + +def iso_from_ms(value: Any) -> str: + if isinstance(value, bool): + return "" + try: + ts = float(value) + except (TypeError, ValueError): + return "" + if not math.isfinite(ts) or ts <= 0: + return "" + # The field is milliseconds, but a plain-seconds epoch (< year 33658) is + # accepted too rather than rendered as a date 1000x in the future. + 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 {"ok": True, "limits": [...], "tierLabel": str}. On a + miss returns {"ok": False, "detail": str} plus one of "transport" (reached + no server, so worth an early retry), "keyRejected" (the key itself is bad) + or "noPlan" (a working key without a Coding Plan). The endpoint answers + HTTP 200 even when the key has no plan, carrying + {"success": false, "msg": "…coding plan…"} instead of a data payload. + """ + url = PLATFORMS[platform] + QUOTA_PATH + 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: + body = open_json(request, timeout=15) + except HttpError as error: + if error.status in (401, 403): + return {"ok": False, "keyRejected": True, "detail": AUTH_HELP_BAD_KEY} + if error.status: + return {"ok": False, "detail": f"Z.ai's monitor endpoint returned HTTP {error.status}. Local stats are still shown."} + if error.transport: + # No answer from any server: no route, no DNS, a timeout, a reset. + return {"ok": False, "transport": True, "detail": "Couldn't reach Z.ai's monitor endpoint. Retrying shortly. Local stats are still shown."} + return {"ok": False, "detail": "Z.ai returned an invalid monitor response. Local stats are still shown."} + + data = body.get("data") if isinstance(body, dict) else None + if not isinstance(data, dict) or not isinstance(data.get("limits"), list): + msg = str(body.get("msg") or "").strip()[:200] if isinstance(body, dict) else "" + code = str(body.get("code") or "").strip() if isinstance(body, dict) else "" + # Z.ai names a refused credential in the body under HTTP 200, with the + # authentication codes of docs.z.ai/api-reference/api-code. + if code in AUTH_ERROR_CODES: + return {"ok": False, "keyRejected": True, "detail": AUTH_HELP_BAD_KEY} + if "coding plan" in msg.lower(): + return {"ok": False, "noPlan": True, "detail": ""} + return {"ok": False, "detail": (msg or "Z.ai's monitor endpoint returned no limits") + ". Local stats are still shown."} + + limits = [] + for item in data["limits"]: + if not isinstance(item, dict) or item.get("type") not in LIMIT_TYPES: + continue + try: + label = LIMIT_WINDOWS.get((item.get("unit"), item.get("number"))) + except TypeError: + # A list or an object for a tag can't name a window. + label = None + 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 + raw = item.get("percentage") + # A missing or unparseable percentage is skipped: a false 0% meter reads + # as a fresh allowance, which is worse than showing one fewer window. + if isinstance(raw, bool) or not isinstance(raw, (int, float)) or not math.isfinite(raw): + continue + limits.append({ + "label": label, + "percent": max(0.0, min(1.0, float(raw) / 100.0)), + "resetsAt": iso_from_ms(item.get("nextResetTime")), + }) + + if not limits: + return {"ok": False, "detail": "Z.ai's monitor endpoint returned no limits. Local stats are still shown."} + tier = str(data.get("level") or "").strip()[:64] + return {"ok": True, "limits": limits, "tierLabel": tier[:1].upper() + tier[1:].lower() if tier else ""} + + +def limits_cache_paths(api_key: str, platform: str) -> tuple[Path, Path]: + # Keyed by who the numbers belong to: another key, or the same key on the + # other host, never shows this one's allowance. + cache = LimitsCache(AGENT_ID, platform, api_key) + return cache.path, cache.lock_path + + +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 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. `live` says whether the limits are this check's (or a probe recent + enough to stand in for one), and `fetchedAtMs` when they were measured, so + the panel can tell kept limits from fresh ones. + """ + result = {"limits": [], "tierLabel": "", "usageStatusText": "", "authHelpText": "", "retryAdvised": False, "live": False, "fetchedAtMs": 0} + # Keyed by who the numbers belong to: another key, or the same key on the + # other host, never shows this one's allowance. + cache = LimitsCache(AGENT_ID, platform, api_key) + + def reuse() -> dict[str, Any]: + result.update(limits=cache.limits, tierLabel=str(cache.data.get("tierLabel") or ""), live=True, fetchedAtMs=cache.fetched_at if cache.limits else 0) + return result + + def no_plan() -> bool: + return cache.data.get("noWindows") is True + + def stands_in() -> bool: + # --force is a person asking for fresh numbers, so it skips the reuse + # window; the interval absorbs repeated panel opens, not a refresh press. + # A key without a plan is an answer too, and stands in for longer. + if no_plan() and not force and cache.younger_than(NO_PLAN_RECHECK_SECONDS): + return True + return cache.fresh(force, allow_empty=True) + + if stands_in(): + return reuse() + + with cache.lock(): + if not force: + # Another run may have probed while this one waited for the lock. + cache.reload() + if stands_in(): + return reuse() + + probe = fetch_limits(api_key, platform) + if probe["ok"]: + fetched_at = cache.save(probe["limits"], tierLabel=probe["tierLabel"]) + result.update(limits=probe["limits"], tierLabel=probe["tierLabel"], live=True, fetchedAtMs=fetched_at) + return result + + # A working key without a Coding Plan is a pay-as-you-go account: Z.ai + # has no limits to show for it, which is no trouble. The windows kept from + # a plan that has since lapsed go with it, so a later failed check can't + # bring them back. + if probe.get("noPlan"): + cache.save([], tierLabel="", noWindows=True) + result["live"] = True + return result + + # A transport failure reached no server — typically the seconds after login + # before the network is up — so ask the shell to retry sooner than its + # regular interval. + if probe.get("transport"): + result["retryAdvised"] = True + # A kept window whose reset has passed comes back at zero: the allowance it + # measured has been renewed, so it is untouched rather than unknown. + fallback, fetched_at = cache.kept() + result.update(limits=fallback, tierLabel=str(cache.data.get("tierLabel") or ""), fetchedAtMs=fetched_at) + if probe.get("keyRejected"): + result.update(usageStatusText="Sign-in expired", authHelpText=probe["detail"]) + elif fallback: + result["authHelpText"] = probe["detail"] + elif not no_plan(): + # A key that last said it has no plan has no limits to lose, so a + # failed recheck is no news. + result.update(usageStatusText="Z.ai limits unavailable", authHelpText=probe["detail"]) + return result + + +# -------------------------------------------------------------------------- +# Local session stats +# -------------------------------------------------------------------------- + +# pi, omp, and opencode name the provider alike. Exact names: a custom +# `zai-proxy` gateway is not this subscription. +def billed_here(provider: str, model: str, entry: dict) -> bool: + return provider.lower() in PROVIDERS + + +def scan_sessions(reuse: bool = True) -> tuple[dict[str, Any], bool]: + """Z.ai turns from pi/omp sessions, both of opencode's message stores, + Claude Code's transcripts (pointed at Z.ai, or answered by a GLM model) and + MiniMax Code's custom providers at Z.ai's hosts, and whether every source + was read (safe to cache).""" + stats = Stats() + pattern = r'"provider"\s*:\s*"(' + "|".join(PROVIDERS) + r')"' + pi_rows, pi_complete = scan_pi_sessions(pattern, billed_here) + # The closing quote makes the database's prefilter match the whole id. + db_rows, db_complete = scan_opencode([provider + '"' for provider in PROVIDERS], billed_here) + routed_rows, routed_complete = scan_routed(AGENT_ID, reuse) + for row in pi_rows + db_rows + routed_rows: + stats.add_row(row, row.model or "glm") + return {"hasLocalStats": bool(stats.model_usage), **stats.snapshot()}, pi_complete and db_complete and routed_complete + + +def cached_local_stats(max_age: float) -> dict[str, Any]: + # The digest covers every data path the scan reads: the pi/omp session + # roots, the opencode DB, and Claude Code's and MiniMax Code's sessions, + # whose owner depends on where they are sent. + sources = [*pi_session_roots(), opencode_db(), *routed_sources()] + # A forced scan reads every Codex rollout again too. + return cached_scan(AGENT_ID, sources, max_age, lambda: scan_sessions(reuse=max_age > 0), match=routed_match()) + + +# -------------------------------------------------------------------------- +# Assembly +# -------------------------------------------------------------------------- + +def scan(force: bool, limits_only: bool) -> dict[str, Any]: + record = base_record() + record.update(cached_local_stats(scan_max_age(force, limits_only))) + + api_key, platform = credentials() + if not api_key: + # Z.ai in use here without a key to ask for its limits; a machine that has + # never used it is simply skipped. + if record["hasLocalStats"]: + record.update(ready=True, usageStatusText="Waiting for auth", authHelpText=AUTH_HELP_NO_KEY) + return record + + limits = collect_limits(api_key, platform, force) + record.update( + ready=True, + limits=limits["limits"], + limitsStale=bool(limits["limits"]) and not limits["live"], + limitsFetchedAt=limits["fetchedAtMs"], + tierLabel=limits["tierLabel"], + usageStatusText=limits["usageStatusText"], + authHelpText=limits["authHelpText"], + ) + if limits["retryAdvised"]: + record["retryAdvised"] = True + 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 limits unavailable", authHelpText="Z.ai usage scan failed.") + print(f"omarchy-agent-usage-zai: {type(error).__name__}", file=sys.stderr) + emit(record) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bin/omarchy-default-agent b/bin/omarchy-default-agent index 189a23ef33e..69d4c229194 100755 --- a/bin/omarchy-default-agent +++ b/bin/omarchy-default-agent @@ -1,7 +1,7 @@ #!/bin/bash # omarchy:summary=Set and launch the default coding agent -# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi] +# omarchy:args=[pi|omp|opencode|ori|claude|codex|grok|openclaw|agy|hermes|copilot|crush|cursor-agent|muse|kimi|qwen|dsh|mcode|kiro-cli|freebuff|kilo|devin|amp|commandcode] # omarchy:examples=omarchy default agent | omarchy default agent codex | omarchy default agent claude installing=false @@ -39,13 +39,42 @@ copilot | github-copilot) agent="copilot"; name="GitHub Copilot" ;; muse | muse-code | musecode) agent="muse"; name="Muse Code" ;; cursor | cursor-agent) agent="cursor-agent"; name="Cursor CLI" ;; kimi | kimi-code) agent="kimi"; name="Kimi"; agent_package="npm:@moonshot-ai/kimi-code" ;; +qwen | qwen-code) agent="qwen"; name="Qwen Code" ;; +dsh | deepseek | deepseek-harness) + agent="dsh"; name="DeepSeek Harness"; agent_package="npm:@deepseek-ai/dsh" + # Every DSH release so far is a prerelease. A declared "latest" follows the + # npm latest tag, but `mise use` lists only stable versions unless told to + # include prereleases. + agent_prereleases=true + ;; +mcode | minimax-code | minimax) agent="mcode"; name="MiniMax Code"; agent_tool="npm:@minimax-ai/code"; agent_package="$agent_tool[allow_builds=better-sqlite3]" ;; +kiro-cli | kiro) + agent="kiro-cli"; name="Kiro CLI" + # Kiro publishes per-platform archives behind a manifest of the stable release. + agent_tool="http:kiro-cli" + agent_package="$agent_tool[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" + ;; +freebuff) agent="freebuff"; name="Freebuff"; agent_package="npm:freebuff" ;; +kilo | kilocode | kilo-code) agent="kilo"; name="Kilo"; agent_package="npm:@kilocode/cli" ;; +devin | devin-cli) + agent="devin"; name="Devin CLI" + # Cognition publishes per-platform archives behind a manifest of the current release. + agent_tool="http:devin" + agent_package="$agent_tool[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" + ;; +amp | ampcode) agent="amp"; name="Amp" ;; +commandcode | command-code | cmdc) agent="commandcode"; name="Command Code"; agent_package="npm:command-code" ;; *) - echo "Usage: omarchy-default-agent " + echo "Usage: omarchy-default-agent " exit 1 ;; esac agent_package=${agent_package:-$agent} +# Omarchy declares a tool that takes backend options under its bare name, with +# the options inside the declaration's table, because mise drops bracket options +# from a table-valued key. The bracketed package is only for the user's config. +agent_tool=${agent_tool:-$agent_package} # Hermes and OpenClaw are not mise tools: each is a self-updating runtime in # the user's home, Hermes the desktop app's and OpenClaw the one its Control @@ -60,23 +89,45 @@ if [[ -n ${agent_installer:-} ]]; then agent_present() { "$agent_installer" --check; } agent_install() { "$agent_installer" --now; } else + # omarchy-mise-install's wrappers run `mise use` and `mise x`, current ones on + # lines of their own and older ones after an inline variable; either way + # running one installs the tool it wraps. A symlink is the user's own install. + mise_wrapper() { + [[ ! -L $1 ]] && head -c 4096 "$1" | grep -qaE 'mise (use|x) ' + } + # Anything at the wrapper's path other than the wrapper is the user's own # install, such as the symlink Cursor's installer leaves. A mise copy would # only shadow it, since the mise shims precede ~/.local/bin on PATH. user_install() { - [[ -x $HOME/.local/bin/$agent ]] && - { [[ -L $HOME/.local/bin/$agent ]] || ! grep -q '^mise use -g' "$HOME/.local/bin/$agent"; } + [[ -x $HOME/.local/bin/$agent ]] && ! mise_wrapper "$HOME/.local/bin/$agent" } - agent_present() { user_install || mise where "$agent_package" &>/dev/null; } + agent_present() { user_install || mise where "$agent_tool" &>/dev/null; } # Omarchy declares the default agents as mise tools, so installing one is # enough and leaves the user's config alone. A user who removed the # preinstalls, like one without Omarchy's declarations, selects it instead. + # An npm-backed agent needs node, itself a lazy tool where Omarchy declares + # it, and mise won't install it as the agent's dependency while it's missing. + npm_backed() { + [[ $1 == npm:* || $(mise tool "${1%%\[*}" --backend 2>/dev/null) == npm:* ]] + } agent_install() { user_install && return - if omarchy-mise-default-tools declares "$agent_package"; then - mise install "$agent_package" + if omarchy-mise-default-tools declares "$agent_tool"; then + if npm_backed "$agent_tool"; then + mise install node "$agent_tool" + else + mise install "$agent_tool" + fi else - mise use -g "$agent_package" + if npm_backed "$agent_package" && ! mise where node &>/dev/null; then + mise install node || return 1 + fi + if [[ ${agent_prereleases:-} == "true" ]]; then + MISE_PRERELEASES=true mise use -g "$agent_package" + else + mise use -g "$agent_package" + fi fi } fi diff --git a/default/fonts/omarchy/README.md b/default/fonts/omarchy/README.md index 1bb4bfb8a34..df8a0d802ff 100644 --- a/default/fonts/omarchy/README.md +++ b/default/fonts/omarchy/README.md @@ -7,7 +7,7 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E902` — OpenCode, from - `U+E903` — omp, from - `U+E904` — Grok, from -- `U+E905` — Codex, from the OpenAI mark on +- `U+E905` — Codex, from , the Codex app's mark (the OpenAI blossom with `>_` cut in), since OpenAI publishes no SVG of it - `U+E906` — LM Studio, from - `U+E907` — Ollama, from - `U+E908` — T3 Code, traced from the app icon in , since upstream publishes no monochrome SVG @@ -18,6 +18,22 @@ The private-use glyphs in `omarchy.ttf` are: - `U+E90D` — Cursor, from - `U+E90E` — Claude, from - `U+E90F` — Tailscale, from +- `U+E910` — Qwen, from +- `U+E911` — DeepSeek, from +- `U+E912` — MiniMax, from +- `U+E913` — Goose, from +- `U+E914` — Kiro, from the ghost in , with its tile dropped +- `U+E915` — Mistral AI, from +- `U+E916` — Freebuff, from , its rects joined into one path +- `U+E917` — Kilo, from +- `U+E918` — Devin, from the official mark published by Cognition at +- `U+E919` — Amp, the wordmark from cut out of a rounded tile, laid out as in the app icon at +- `U+E91A` — Command Code, from , its traced paths joined into one +- `U+E91B` — Phantombot, redrawn as a monochrome ghost silhouette from the project logo at (`www/assets/img/logo.png`), since upstream publishes no monochrome SVG +- `U+E91C` — Antigravity, the arch from +- `U+E91D` — Kimi, from +- `U+E91E` — Crush, the heartbit from , its eyes and mouth cut out +- `U+E91F` — Muse Code, from : Meta publishes no separate Muse mark and dev.meta.ai uses this one The agent marks are monochrome so the menu can render them using the active theme's foreground and selection colors. diff --git a/default/fonts/omarchy/omarchy.ttf b/default/fonts/omarchy/omarchy.ttf index 80573d24bb1..10277b494ed 100644 Binary files a/default/fonts/omarchy/omarchy.ttf and b/default/fonts/omarchy/omarchy.ttf differ diff --git a/default/i18n/en/cli.json b/default/i18n/en/cli.json index 69f82393a34..d6a0b033fc5 100644 --- a/default/i18n/en/cli.json +++ b/default/i18n/en/cli.json @@ -11,6 +11,8 @@ "agents.newSessionsUse": "New %1 sessions now use %2", "agents.noAccountNamed": "No %1 account named %2.", "agents.offerSwitch": "%1 is at %2%. Click to switch new sessions to it.", + "agents.promptOnClipboard": "%1 takes no starting prompt", + "agents.promptOnClipboardBody": "Your prompt is on the clipboard. Paste it into the session.", "agents.runningStayOn": "Running sessions stay on %1.", "agents.sameAccount": "That's %1 (already added). Try again, and sign in as the other account in the private window.", "agents.sameAccountEmail": "That's %1 (%2). Try again, and sign in as the other account in the private window.", diff --git a/default/i18n/en/menu.json b/default/i18n/en/menu.json index 773918665e3..ada8c0bd21c 100644 --- a/default/i18n/en/menu.json +++ b/default/i18n/en/menu.json @@ -499,12 +499,18 @@ "setup.default.agent.agy": { "label": "Antigravity" }, + "setup.default.agent.amp": { + "label": "Amp" + }, "setup.default.agent.claude": { "label": "Claude" }, "setup.default.agent.codex": { "label": "Codex" }, + "setup.default.agent.commandcode": { + "label": "Command Code" + }, "setup.default.agent.copilot": { "label": "Copilot" }, @@ -514,15 +520,33 @@ "setup.default.agent.cursor-agent": { "label": "Cursor CLI" }, + "setup.default.agent.devin": { + "label": "Devin CLI" + }, + "setup.default.agent.dsh": { + "label": "DeepSeek Harness" + }, + "setup.default.agent.freebuff": { + "label": "Freebuff" + }, "setup.default.agent.grok": { "label": "Grok" }, "setup.default.agent.hermes": { "label": "Hermes" }, + "setup.default.agent.kilo": { + "label": "Kilo" + }, "setup.default.agent.kimi": { "label": "Kimi" }, + "setup.default.agent.kiro-cli": { + "label": "Kiro CLI" + }, + "setup.default.agent.mcode": { + "label": "MiniMax Code" + }, "setup.default.agent.muse": { "label": "Muse Code" }, @@ -541,6 +565,9 @@ "setup.default.agent.pi": { "label": "Pi" }, + "setup.default.agent.qwen": { + "label": "Qwen Code" + }, "setup.default.browser": { "label": "Browser", "title": "Default Browser" diff --git a/default/i18n/en/shell.json b/default/i18n/en/shell.json index 5fefef3ba0d..766d468a27c 100644 --- a/default/i18n/en/shell.json +++ b/default/i18n/en/shell.json @@ -47,6 +47,7 @@ "agents.modelWeekly": "%1 Weekly", "agents.mostlyModel": "Mostly %1", "agents.nameAccount": "Name this account. It signs in through a private window, so your browser's current account isn't picked up.", + "agents.outOfCredits": "Out of credits", "agents.pasteCode": "If the page shows a code instead of finishing, paste it here.", "agents.placeholderCode": "Code", "agents.placeholderWork": "Work", diff --git a/default/omarchy/omarchy-menu.jsonc b/default/omarchy/omarchy-menu.jsonc index a974316ad81..e21404c0cbd 100644 --- a/default/omarchy/omarchy-menu.jsonc +++ b/default/omarchy/omarchy-menu.jsonc @@ -195,21 +195,30 @@ "setup.network.qr": {"icon":"󰐲","label":"QR Code","aliases":["wifi-qr"],"when":"omarchy-network-wifi-connected","action":"omarchy-shell shell summon omarchy.wifiqr"}, "setup.default": {"icon":"","label":"Defaults","aliases":["default","defaults"]}, "setup.default.agent": {"icon":"󰚩","label":"Agent","title":"Default Agent"}, - "setup.default.agent.agy": {"icon":"󰫢","label":"Antigravity","checked":"[[ \"$(omarchy-default-agent)\" == \"agy\" ]]","action":"omarchy-default-agent agy"}, - "setup.default.agent.claude": {"icon":"󰛄","label":"Claude","checked":"[[ \"$(omarchy-default-agent)\" == \"claude\" ]]","action":"omarchy-default-agent claude"}, + "setup.default.agent.amp": {"icon":"","iconFont":"omarchy","label":"Amp","checked":"[[ \"$(omarchy-default-agent)\" == \"amp\" ]]","action":"omarchy-default-agent amp"}, + "setup.default.agent.agy": {"icon":"","iconFont":"omarchy","label":"Antigravity","checked":"[[ \"$(omarchy-default-agent)\" == \"agy\" ]]","action":"omarchy-default-agent agy"}, + "setup.default.agent.claude": {"icon":"","iconFont":"omarchy","label":"Claude","checked":"[[ \"$(omarchy-default-agent)\" == \"claude\" ]]","action":"omarchy-default-agent claude"}, "setup.default.agent.codex": {"icon":"","iconFont":"omarchy","label":"Codex","checked":"[[ \"$(omarchy-default-agent)\" == \"codex\" ]]","action":"omarchy-default-agent codex"}, + "setup.default.agent.commandcode": {"icon":"","iconFont":"omarchy","label":"Command Code","checked":"[[ \"$(omarchy-default-agent)\" == \"commandcode\" ]]","action":"omarchy-default-agent commandcode"}, "setup.default.agent.copilot": {"icon":"","label":"Copilot","checked":"[[ \"$(omarchy-default-agent)\" == \"copilot\" ]]","action":"omarchy-default-agent copilot"}, - "setup.default.agent.crush": {"icon":"󰋑","label":"Crush","checked":"[[ \"$(omarchy-default-agent)\" == \"crush\" ]]","action":"omarchy-default-agent crush"}, + "setup.default.agent.crush": {"icon":"","iconFont":"omarchy","label":"Crush","checked":"[[ \"$(omarchy-default-agent)\" == \"crush\" ]]","action":"omarchy-default-agent crush"}, "setup.default.agent.cursor-agent": {"icon":"","iconFont":"omarchy","label":"Cursor CLI","checked":"[[ \"$(omarchy-default-agent)\" == \"cursor-agent\" ]]","action":"omarchy-default-agent cursor-agent"}, + "setup.default.agent.dsh": {"icon":"","iconFont":"omarchy","label":"DeepSeek Harness","checked":"[[ \"$(omarchy-default-agent)\" == \"dsh\" ]]","action":"omarchy-default-agent dsh"}, + "setup.default.agent.devin": {"icon":"","iconFont":"omarchy","label":"Devin CLI","checked":"[[ \"$(omarchy-default-agent)\" == \"devin\" ]]","action":"omarchy-default-agent devin"}, + "setup.default.agent.freebuff": {"icon":"","iconFont":"omarchy","label":"Freebuff","checked":"[[ \"$(omarchy-default-agent)\" == \"freebuff\" ]]","action":"omarchy-default-agent freebuff"}, "setup.default.agent.grok": {"icon":"","iconFont":"omarchy","label":"Grok","checked":"[[ \"$(omarchy-default-agent)\" == \"grok\" ]]","action":"omarchy-default-agent grok"}, "setup.default.agent.hermes": {"icon":"","iconFont":"omarchy","label":"Hermes","checked":"[[ \"$(omarchy-default-agent)\" == \"hermes\" ]]","action":"omarchy-default-agent hermes"}, - "setup.default.agent.kimi": {"icon":"󰚩","label":"Kimi","checked":"[[ \"$(omarchy-default-agent)\" == \"kimi\" ]]","action":"omarchy-default-agent kimi"}, - "setup.default.agent.muse": {"icon":"󰛤","label":"Muse Code","checked":"[[ \"$(omarchy-default-agent)\" == \"muse\" ]]","action":"omarchy-default-agent muse"}, + "setup.default.agent.kilo": {"icon":"","iconFont":"omarchy","label":"Kilo","checked":"[[ \"$(omarchy-default-agent)\" == \"kilo\" ]]","action":"omarchy-default-agent kilo"}, + "setup.default.agent.kimi": {"icon":"","iconFont":"omarchy","label":"Kimi","checked":"[[ \"$(omarchy-default-agent)\" == \"kimi\" ]]","action":"omarchy-default-agent kimi"}, + "setup.default.agent.kiro-cli": {"icon":"","iconFont":"omarchy","label":"Kiro CLI","checked":"[[ \"$(omarchy-default-agent)\" == \"kiro-cli\" ]]","action":"omarchy-default-agent kiro-cli"}, + "setup.default.agent.mcode": {"icon":"","iconFont":"omarchy","label":"MiniMax Code","checked":"[[ \"$(omarchy-default-agent)\" == \"mcode\" ]]","action":"omarchy-default-agent mcode"}, + "setup.default.agent.muse": {"icon":"","iconFont":"omarchy","label":"Muse Code","checked":"[[ \"$(omarchy-default-agent)\" == \"muse\" ]]","action":"omarchy-default-agent muse"}, "setup.default.agent.omp": {"icon":"","iconFont":"omarchy","label":"omp","checked":"[[ \"$(omarchy-default-agent)\" == \"omp\" ]]","action":"omarchy-default-agent omp"}, "setup.default.agent.openclaw": {"icon":"","iconFont":"omarchy","label":"OpenClaw","checked":"[[ \"$(omarchy-default-agent)\" == \"openclaw\" ]]","action":"omarchy-default-agent openclaw"}, "setup.default.agent.opencode": {"icon":"","iconFont":"omarchy","label":"OpenCode","checked":"[[ \"$(omarchy-default-agent)\" == \"opencode\" ]]","action":"omarchy-default-agent opencode"}, "setup.default.agent.ori": {"icon":"","iconFont":"omarchy","label":"Ori","checked":"[[ \"$(omarchy-default-agent)\" == \"ori\" ]]","action":"omarchy-default-agent ori"}, "setup.default.agent.pi": {"icon":"","iconFont":"omarchy","label":"Pi","checked":"[[ \"$(omarchy-default-agent)\" == \"pi\" ]]","action":"omarchy-default-agent pi"}, + "setup.default.agent.qwen": {"icon":"","iconFont":"omarchy","label":"Qwen Code","checked":"[[ \"$(omarchy-default-agent)\" == \"qwen\" ]]","action":"omarchy-default-agent qwen"}, "setup.default.browser": {"icon":"","label":"Browser","title":"Default Browser"}, "setup.default.browser.chromium": {"icon":"","label":"Chromium","checked":"[[ \"$(omarchy-default-browser)\" == \"chromium\" ]]","action":"omarchy-default-browser chromium"}, "setup.default.browser.chrome": {"icon":"󰊯","label":"Chrome","checked":"[[ \"$(omarchy-default-browser)\" == \"chrome\" ]]","action":"omarchy-default-browser chrome"}, diff --git a/docs/agent-usage-collectors.md b/docs/agent-usage-collectors.md new file mode 100644 index 00000000000..387028042fc --- /dev/null +++ b/docs/agent-usage-collectors.md @@ -0,0 +1,112 @@ +# Agent Usage Collectors + +Every agent or provider in the agents panel is a collector: `bin/omarchy-agent-usage-`, a Python script that prints one JSON record. `omarchy-agent-usage-update` finds collectors by globbing `bin/omarchy-agent-usage-*`, runs them, and the panel (`shell/plugins/agents/Main.qml`, `Panel.qml`) reads the records. This is the contract every collector follows, so that each agent behaves the same in the panel. `bin/omarchy-agent-usage-claude` is the reference implementation. + +New collectors build on the shared module, [`shell/plugins/agents/collectors.py`](../shell/plugins/agents/collectors.py), which holds the parts of this contract every collector shares. A collector keeps only what is its own: its credentials, its endpoint and how to read it, its provider filter, and the specifics of its record. + +## The command + +- `#!/usr/bin/python3`, two-space indentation, and the header lines `# omarchy:summary=...`, `# omarchy:args=[--force] [--limits-only]`, and `# omarchy:hidden=true`. +- argparse accepts `--force` and `--limits-only`, and both mean something: + - A normal run reuses a local scan for `SCAN_REUSE_SECONDS = 20` and a limits probe for `PROBE_MIN_INTERVAL_SECONDS = 15`. + - `--limits-only` (the panel's near-limit refresh) reuses a local scan for `LIMITS_ONLY_REUSE_SECONDS = 900`, but asks for limits like a normal run. + - `--force` (a person pressing refresh) reuses nothing. +- The record is printed on stdout. Diagnostics go to stderr and never carry a credential. +- A collector never writes to an agent's own files. Credentials are only read, and a token is refreshed only where the agent's own CLI does the same under its lock, never by rewriting another harness's store. + +## The shared module + +The module ships with the agents panel in the omarchy package. A collector loads it from `$OMARCHY_PATH`, never from a path derived from its own file or `$HOME`: + +```python +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +from collectors import Stats, base_record, cached_scan, ... +``` + +Its main entry points, each documented where it is defined: + +- `base_record(id, name, **fields)` and `emit(record)`: the record with every field and the zeroed stats block, as an unused machine prints it, and printing it. +- `Stats(today=None)`: `add(day, session, model, input, output, cache_read, cache_write, prompts=1)` tokens and the prompts they answered (or `add_row()` one row from a reader), `add_prompts(day, session, count=1)` prompts a store books apart from their tokens, and `snapshot()` the record's stats fields. A store that books prompts on their own passes `prompts=0` with the tokens; one keeping running totals passes how many prompts they cover. A `day` of `None` counts toward the lifetime totals alone, a `session` of `None` toward no session, and zero tokens make no model bucket. `today` dates the record against another day (a `date` or an ISO day). +- `cached_scan(id, sources, max_age, scan, match=None, with_complete=False)` with `scan_max_age(force, limits_only)`: the per-day scan cache under `flock`, reused for 20, 900, or 0 seconds, never caching an incomplete scan. `with_complete=True` answers `(stats, complete)`, a cached scan being a complete one. +- `scan_pi_sessions(pattern, accept, roots=None)` and `scan_opencode(providers, accept)`: the pi/omp and opencode readers, handing back deduplicated `UsageRow`s for the replies the collector's `accept(provider, model, entry)` takes. `roots` reads other pi-format trees instead of `pi_session_roots()` (a harness's own, or pi's and omp's with another beside them), through `canonical_roots()`; a collector that passes them keys its scan cache by the same roots. +- `scan_claude_code(owner, projects=None, index_file=..., base_url=None)` (or `claude_code_usage()`, which also says whether every transcript was read, so a scan that missed one isn't cached) with `claude_code_owner(model, base_url)` and `claude_code_scan_match()`: Claude Code's transcripts, read through the per-file index every reader shares (one reader at a time brings it up to date under its `.lock`, and a reader finding it current for every transcript reads it without the lock), handing back the rows (replies, what a lineage's `cost-state` totals hold beyond its replies, typed prompts, sessions in use) of the one collector `claude_code_owner` names for each reply. A collector that reads them adds `claude_code_dir() / "projects"` to its scan sources and `claude_code_scan_match()` to its `match`, since where Claude Code is sent decides whose its replies are. +- `codex_provider_urls()`: every Codex `model_providers` id's `base_url`, from `/etc/codex/config.toml`, `config.toml` and the profiles, remembered per Codex home in `local-codex-providers.json` after the config drops one, and the built-in `openai` provider's `openai_base_url` under `openai`, from `/etc/codex/config.toml` and `config.toml` only, never remembered (a profile applies only with `--profile`, so its `openai_base_url` is kept under `openai@` for the model it sets, unless that is an OpenAI model, the one `config.toml` sets, or one two profiles send to different places); `codex_route_owner(provider, model, urls)` says which collector bills a rollout, and `codex_native_endpoint(urls, model)` where the built-in provider sent a turn on `model`. +- `codex_rollouts(cache_file, reuse=True, native=False)` with `codex_rollout_cache()`: the Codex rollouts of the last 30 days off the built-in provider, and on it (or naming no provider) too with `native`, as the Codex collector reads them and every collector while `openai_base_url` sends it elsewhere, each with its provider, the session it counts under (`codex_rollout_session`: a subagent's, guardian's or review's thread its parent's, a memory consolidation none) and per-day, per-model tokens (one `token_usage_record` per response where Codex 0.163 writes them, else each `token_count`'s last request; interactive Codex's thread-title requests are written to no rollout and are missing) and typed prompts, replayed from a per-file cache while a file's size and mtime stand, and read on from where the last read stopped when the same file only grew. The records hold every provider's turns, so every collector reading Codex shares one cache per Codex home (any `cache_file` stands for it, None reads uncached), brought up to date by one reader at a time under its `.lock`; a forced read (`reuse=False`) takes what another forced read of the same refresh wrote (`omarchy-agent-usage-update` exports when the refresh began as `OMARCHY_AGENT_USAGE_REFRESH_AT`; a lone `--force` begins its own), so the collectors one refresh forces read each rollout once between them, whenever each gets to them. +- `scan_routed(owner, reuse=True)` with `routed_sources()` and `routed_match()`: Claude Code's, MiniMax Code's and Codex's usage one collector bills (`scan_claude_code`, `scan_mcode` and `scan_codex_routed`), since all three send to whatever endpoint they are given; the local collector reads Codex with its own classifier instead; every collector a base URL can name (`PROVIDER_HOSTS`, and `local`) reads them, adding the sources and match to its scan cache. +- `scan_dsh_sessions(accept)` and `kimchi_sessions_root()`: DeepSeek Harness's billed attempts on the routes `accept(provider, model, request)` takes, a fork's inherited history left out, each carrying as its prompts the typed messages (a user message whose source is `user`) it answered, so a turn's tool steps and retries, and a turn dsh starts on its own, answer none; and where Kimchi keeps its pi-format sessions. +- `endpoint_owner(url)`, `self_hosted(url)` and `cloud_model(model)`: which collector bills a base URL (a provider host in `PROVIDER_HOSTS`, `local` for the user's own hardware, or nobody), and whether a model carries Ollama's cloud tag. +- `LimitsCache(id, *identity, kind="limits")`: one credential's last good answer at `--.json`, with elapsed windows zeroed, `fresh(force, seconds=15, allow_empty=False)` for the probe reuse, `kept()` for the stale fallback, and `save(limits, only_if_newer=False, **extra)`. `kind` keeps another answer beside the limits under the same key (credits, a balance, usage rows); `allow_empty=True` lets any answer stamped inside the age guard stand in, limits or none; `only_if_newer=True` leaves a file another run stamped later alone. `LimitsCache.at(path)` opens one under a name that predates the shared one. +- `open_json(request)` and `HttpError`: a request that refuses redirects and says whether a failure was a refusal (`status`, with its `headers`), a transport failure (`transport`, the only cause for `retryAdvised`), a reply cut short (`cut_short`), or an invalid answer. See [Limits](#limits) for where the line falls. +- `trusted_endpoint(url)`: whether an endpoint override a person set (`CURSOR_API_ENDPOINT`, `MUSE_KEY_ENDPOINT`, …) may carry a credential: https anywhere, plain http only to this machine. Every override that moves where a token goes passes through it. +- `cursor_token(auth_path=None, state_db=None)`, with `jwt_claims()`, `token_lapsed()`, `account_key()` and `finite()`: the Cursor account the Cursor and Grok Bot collectors both ask as (an exported `CURSOR_AUTH_TOKEN`; with `CURSOR_API_KEY` exported, only the token the CLI traded that key for, never the key itself; else the CLI's `auth.json` at `CURSOR_AUTH_PATH` or its default, then the editor's `state.vscdb`, a lapsed sign-in giving way to one that isn't), and reading any JWT's claims, expiry and subject. +- `registered_accounts(id)` and `account_entry(account, limits, plan, **extra)`: the accounts `omarchy agent account` registered, and one entry of the record's `accounts`. +- `local_day()`, `token_count()`, `usage_tokens()`, `number()`, `now_ms()`, `stamp()`, `age_below()`, `read_json()`, `write_json()`, `cache_root()`: reading untrusted values and files the same way everywhere. + +`test/shell.d/agent-usage-shared-test.sh` holds the module to this contract directly. Something goes in the module only once two collectors need it. + +## The record + +| Field | Meaning | +|---|---| +| `schemaVersion` | `1` | +| `id`, `name` | The id matches the filename; `name` is the product name | +| `updatedAt` | ISO time the record was made | +| `ready` | `true` once the agent is in use on this machine (a credential or local usage) | +| `tierLabel` | The plan as the service names it, or `""`; never a placeholder or a guess. `"Prepaid"` is special to the panel | +| `usageStatusText`, `authHelpText` | See [Status](#status) | +| `limits` | `[{label, percent, resetsAt, title?}]`, `percent` in 0..1, `resetsAt` ISO or `""`, never `null`. Add `title` (Session, Weekly, Monthly) when the label doesn't parse as one | +| `limitsStale`, `limitsFetchedAt` | Whether `limits` were kept past a failed check, and when they were measured (epoch ms, `0` when never) | +| `retryAdvised` | `true` only when the service couldn't be reached at all (DNS, connection, timeout), so the panel retries this agent in 30 seconds | +| `balance` | `{remaining, funded, spent, currency, estimated}` for prepaid accounts | +| `accounts`, `accountSwitch` | Only when two or more accounts are registered; each account carries `limits`, `stale`, `fetchedAt`, `usageStatusText`, `authHelpText` | +| `hasLocalStats` | Whether the stats below come from this machine | +| `hasPromptStats` | `false` when the source has no real prompt or session counts (an API that counts requests, not prompts) | +| `scope` | `"account"` when the stats are account-wide rather than this machine's | +| Stats | `todayPrompts`, `todaySessions`, `todayTotalTokens`, `todayTokensByModel`, `recentDays` (seven `{date, messageCount}`, where `messageCount` holds tokens), `totalPrompts`, `totalSessions`, `activeDays`, `activeDates` (every date), and `modelUsage` buckets of `inputTokens`, `outputTokens`, `cacheReadInputTokens`, `cacheCreationInputTokens` | + +The full stats block is always present, zeroed when there is nothing. An agent the machine doesn't use prints `ready: false` with empty stats, no limits and no balance, and the panel skips it. That's the normal answer for a missing credential on an agent with no local usage; it isn't an error. + +## Status + +The panel acts on a small vocabulary in `usageStatusText`, so the same situation reads the same for every agent: + +| Status | When | +|---|---| +| `""` | Healthy, or a failed check that kept the last good limits (`limitsStale` says the rest) | +| `"Waiting for auth"` | The agent is in use but has no credential for its limits | +| `"Sign-in expired"` | The credential was refused, or a refresh token lapsed; signing in again is needed | +| `"Limits paused"` | An access token lapsed that the agent's own CLI renews the next time it runs; signing in again is not needed | +| `" limits unavailable"` | A check failed with nothing kept to show | +| `" unavailable"` | The agent's stats source failed (an unreadable database, a billing API that answered nothing) with nothing kept to show | +| `"Balance unavailable"` | A prepaid balance couldn't be read and none was kept | +| `"Out of credits"` | A prepaid account has no credit left, or its gateway rate limits it for that reason; urgent, since a balance with no funded figure raises no low-balance alarm of its own | + +The detail goes in `authHelpText`: what failed and the one thing to do about it. It is `""` when healthy. A plan the service doesn't offer limits for (pay-as-you-go, no Go plan, a free tier) isn't trouble: the status stays `""` and local stats show alone. Any other status shows under the name in the urgent color, so collectors don't invent new ones. + +## Limits + +- Each credential or account caches its own last good answer at `cache_root() / "-limits-.json"`, `{fetchedAtMs, limits, ...}`, where `cache_root()` is `$XDG_CACHE_HOME/omarchy/agent-usage` and `` is a short hash of who the credential signs in as (a JWT subject, an email, an account id) or of the credential itself. Another sign-in never shows this one's numbers. An answer of another kind from the same credential (credits, a balance) sits beside it as `--.json`. +- An answer stands in for a probe for 15 seconds, even one with no windows (a plan without quotas, credits, a balance), so long as the collector keeps it. Refreshes can overlap, so an answer that only replaces one fetched before it never rolls a newer one back. +- Where a failed check falls decides `retryAdvised`: + - A connection that fails or drops before any reply (no DNS, no route, refused, reset, a timeout, a server hanging up without a word) is transport: nothing answered, so `retryAdvised: true` asks the panel to try again soon. + - A reply that arrived but was cut short or garbled on the way (`http.client.HTTPException`: an incomplete read, a bad status line) is an answer, `cut_short` and not `transport`, so there is no sooner retry. Neither is there for an error status or a body that isn't the expected JSON. +- The cache is read whatever its file time, so a clock that went back doesn't lose it, and every age check is `0 <= age < N`. Stamps are floored (`int(time.time() * 1000)`), never rounded: a stamp half a millisecond ahead reads as a clock that went back, and the next run probes again. +- A failed check keeps the cached limits with `limitsStale: true` and their `limitsFetchedAt`. The panel shows their age on hover. The one exception is a service that rate-limits its own usage endpoint (Claude's answers 429 readily): a 429 only says "ask later", so numbers measured within the last 15 minutes still read as current. +- A kept window whose `resetsAt` has passed comes back at `percent: 0` with `resetsAt: ""`, never dropped: a reset allowance is untouched, not unknown. A window with no reset time is kept as it is. Reset times are only ever the service's own, never computed. +- A meter for a balance or allowance is a real ratio of the service's numbers, never estimated from local usage unless its label says `(est.)`. + +## Local stats + +- Tokens are counted once, under the provider that bills them. pi, omp, and opencode sessions are split by provider into the matching collector; there is no harness tab that counts them again. Claude Code transcripts are split the same way, by `claude_code_owner()`, which every collector reading them applies. Claude Code records only the model that answered, not where it sent the request, so while `ANTHROPIC_BASE_URL` (its `settings.json` `env`, else the environment) names a provider's host in `PROVIDER_HOSTS`, every reply is that provider's, earlier ones included. Otherwise an Ollama cloud tag is `ollama`'s and a Claude model `claude`'s (through a proxy too), a self-hosted base URL makes any other model `local`'s, and with no override the model's name decides: `glm-*` is `zai`'s, `MiniMax-*` `minimax`'s, `deepseek-*` `deepseek`'s, `kimi-*` and `moonshot*` `kimi`'s, `hf:*` and `syn:*` `synthetic`'s. Any other model with no override (`qwen3-coder`), and any non-Claude model behind a host nothing can place (a gateway or proxy), counts nowhere. Codex rollouts are split by their `model_provider`, by `codex_route_owner()`: the built-in `openai` is `codex`'s until `openai_base_url` (the config key only; Codex ignores `OPENAI_BASE_URL`) sends it elsewhere, read the way Claude Code's base URL is: a provider's host bills every built-in rollout, an Ollama cloud tag is `ollama`'s, OpenAI's own models stay `codex`'s through any other URL, a self-hosted one makes any other model `local`'s, and a host nothing can place counts it nowhere; a `model_providers` entry whose `base_url` is a host in `PROVIDER_HOSTS` is that provider's whatever its id (an `ollama` entry at ollama.com included), `--oss` on the local daemon's `ollama` is `local`'s unless the model is cloud-tagged, which is `ollama`'s along with a provider named `ollama-cloud`, an entry on the user's own hardware is `local`'s whatever its id (one another collector claims included), and one at a host nothing can place counts nowhere; `codex_rollouts()` reads them for all of them, `codex` included, and books each typed prompt with the model of the first request after it (else the model the rollout was last set to), so a rollout that switches models between turns splits its prompts the way it splits its tokens; the model in effect is the last one a `turn_context` or a `thread_settings_applied` (written by `-m` on a resume or `/model`, ahead of the compaction a switch can send on the new model) names. Other harnesses that route by provider are split the same way: DeepSeek Harness's DeepSeek routes are `deepseek`'s and its pi-ai `openrouter` route `openrouter`'s (`scan_dsh_sessions`), Kimchi's gateway routes (`kimchi`, `kimchi-dev`, `kimchi-dev/`, `kimchi-experimental`) are `kimchi`'s in its harness sessions (pi format, under `kimchi_sessions_root()`) and in pi's, omp's and opencode's alike, and a harness session on `--provider openrouter` is `openrouter`'s; its gateway's hosts (`kimchi.dev`, every region) are in `PROVIDER_HOSTS`, so the Claude Code and Codex setups `kimchi claudecode` and `kimchi codex` write are `kimchi`'s whatever the model, while `kimchi claude` sets `ANTHROPIC_BASE_URL` only in the environment of the Claude Code it starts, which leaves its transcripts read like any with no override, and a self-hosted Kimchi gateway's own host is placed by nothing, and a MiniMax Code reply is `minimax`'s on its own routes and on a `custom_provider:` whatever collector bills that provider's `config.yaml` `baseURL` (`mcode_route_owner`, which gives a cloud-tagged model on a self-hosted route to `ollama`, and each answered prompt the model of its reply), nobody's for a host nothing can place. Kimi Code, Qwen Code, Hermes and Mistral Vibe keep stores of their own, and each one's collector exposes `harness_rows(owner)`, which reads its store and keeps the rows `owner` bills, a base URL read by `claude_code_owner()`: Kimi Code's model alias resolves through its `config.toml` to the provider's base URL (Kimi's endpoint is `kimi`'s, as is an alias the config no longer names when it is one of the plan's `kimi-code/` models or `KIMI_MODEL_NAME`'s; a removed provider's alias counts nowhere), a Qwen Code reply's auth type and its model's `modelProviders` `baseUrl` in `settings.json`, else the `OPENAI_BASE_URL` or `ANTHROPIC_BASE_URL` its auth type takes from the environment, the first of `QWEN_HOME/.env`, `~/.qwen/.env` and `~/.env` that sets it, `settings.json`'s `env`, or the `security.auth.baseUrl` `/auth` saved (Alibaba's hosts and Qwen OAuth are `qwen`'s; an OpenAI-compatible reply with no URL went to OpenAI's API, Qwen Code's default, and like Anthropic, Gemini and Vertex AI is nobody's), a Hermes session's and request's billing base URL (Nous Portal is `hermes`'s, as is a row from before Hermes kept the route), and a Vibe session's provider `api_base` (Mistral's API is `vibe`'s). Only Hermes and Vibe's older `session_*` folders record the route with each request; Kimi Code's aliases, Qwen Code's base URLs and the providers of Vibe's Unified Harness sessions resolve through the current config, so a route change moves earlier usage too, history included. Each prompt goes with the request that answered it; in Kimi Code and Qwen Code with the next request sent for it, answered or not, else with the alias its agent was bound to or the model `settings.json` selects. The harness's own collector counts its rows and `local` the self-hosted ones; rows another provider bills count nowhere, since no other collector reads these stores yet. A route nothing names stays uncounted rather than guessed. +- pi and omp sessions live under `~/.pi/agent/sessions` and `~/.omp/agent/sessions`, plus `profiles//agent/sessions` under each. `PI_CODING_AGENT_DIR`'s sessions, `PI_CODING_AGENT_SESSION_DIR`, and an absolute `sessionDir` in any of those agent trees' `settings.json` are read as well as those defaults, never instead of them (a relative `sessionDir`, a project's `.pi/settings.json`, and `--session-dir` can't be found and aren't read); roots are compared by realpath, and one nested inside another is dropped. They are found with `rg --no-ignore`, since a `$HOME` that is a git checkout would otherwise hide them. pi's `reasoning` and omp's `reasoningTokens` are the part of `output` that was thinking, so they are never added to it again (opencode's `tokens.reasoning` is kept apart from its output, and is added). Besides replies, the usage pi counts in a session's totals counts too: a compaction's or branch summary's `usage`, a `usage` entry (cache warming), and a tool result's nested usage, under the provider the entry names or else the model the session was on (the last model switch or reply before it in the file), answering no prompt; with neither, nothing says who billed it and it is skipped. omp writes no `usage` on its compaction entries (18.12 keeps only `tokensBefore` and `tokensAfter`), nor anywhere else in the session or its `agent.db`, so omp's compactions can't be counted and are missing from its tokens, as those of a Claude Code session with no `cost-state` line are (every session before 2.1.278, and many since). A reply is deduplicated by its provider, `responseId` and `timestamp` when it has a response id, since omp writes a reply that overflowed the model's declared context window twice under new ids, and any other message is deduplicated by `(id, timestamp)` when it has both, so a forked session doesn't count its parent's messages again, and a line without an id always counts. A typed prompt counts once by its own `(id, timestamp)` as well, under one collector whichever reply books it: the first reply that answered it in any copy (the oldest session's first), else the oldest session's booking of it unanswered, so one a session left unanswered and a fork copying it answered counts once, on the fork's provider even when that is another collector's. The copies are found by `parentSession`: the other files in the fork tree of each file a collector parses that hold the id of a prompt it booked are read for their bookings alone, and a `parentSession` naming a file that is gone (pi deletes a session from `/resume` and leaves its forks) is no parent, so the fork's copies count under it and the scan stays complete. A row with no time is today's; a row whose time is untrusted, zero included, counts nowhere. Each collector parses only the session files holding its own lines (one its `pattern` matches, or a model switch its `accept` takes, which says whose a summary's usage or an unanswered prompt there is), so another provider's history costs it no more than rg's search. +- opencode is read from both its `message` store and V2's `session_message`, V2 first, deduplicated by message id, leaving out the copies a fork makes (in V2, a `session_v2.fork_session_id` session's messages older than it; in the original store, a message whose `time.created` predates its `session` row, since a fork there is a plain new session holding its source's messages under new ids and their original times), with compactions counted like replies. The database is `OPENCODE_DB` when set (or the older `opencode-v2.db`), opened read-only. `scan_opencode` in the shared module is the reference reader, and `test/shell.d/agent-usage-partition-test.sh` holds every collector to it. +- A prompt is a typed user message, never a model request, so a tool step, a compaction or summary, and a subagent's session answer none. `UsageRow.prompts` says which rows book one: in pi and omp the first reply that spent tokens after a user message (omp's own `attribution: "agent"` messages aside; a file with no user messages at all keeps one per reply), in opencode's `message` store the first step answering a user message (they share its `parentID`) whose tokens and time the reader trusts, and in V2's `session_message` the first reply after a `user` row (a session without any keeps one per reply). In both stores that first step is chosen across every provider, so a prompt a tool loop answers on two providers counts once, under the collector that took the first step. A prompt no reply answered (its request failed, or none came) still counts once, as a row without tokens on the model the session was set to use for it, so the collector billing that model counts it and its session: in pi and omp the failed reply, which they write with its provider and no usage, before the next user message, else the last model switch or reply before the prompt, so a prompt whose agent was killed mid-request never rides on the next provider's reply (and a fork copying it books it once), while one omp refused before sending it (an error with its `errorId` and no `duration`, such as an invalid credential) counts nowhere; in opencode's `message` store the `model` its user row names, once a step answering it shows the request was made (a failed or cut-off step without tokens; a user row with no step at all was refused before sending, an unknown model or provider, and counts nowhere; a compaction's marker is answered by its summary, failed or not); and in V2's `session_message` the reply after the `user` row, spent or not. With no model known it counts nowhere, and like every prompt it is dated by when it was typed, an untrusted time counting nowhere. +- Timestamps and token fields from disk are untrusted: a zero, huge, negative, `NaN`, boolean or non-numeric value skips that field or row, never the record. +- A scan is cached per day (`scanDate`) and rebuilt under `flock`, and a scan that couldn't read all its sources isn't cached. + +## Assets, docs and tests + +- `shell/plugins/agents/assets/.svg`, plus `-light.svg` when the mark is white and needs a dark twin for light surfaces. +- A provider entry in `shell/plugins/agents/manifest.json`, a row in the README's collector table, and a `###` section in the README naming where its credentials and data come from. +- `test/shell.d/agent-usage--test.sh` (or several, `agent-usage---test.sh`) covering the collector's own sources, credentials and answer, and an adapter for the contract checks. The record an unused machine prints is held for every collector by `agent-usage-partition-test.sh`, and the shared reading mechanics (pi, omp and opencode stores, untrusted values, cache files and locks, HTTP failure classes) by `agent-usage-shared-test.sh`, so a collector's test doesn't repeat them. Its adapter hands `check_limits_contract` fresh limits, the probe reuse and `--force`, stale limits after a failed check, `retryAdvised` on an unreachable service, an elapsed kept window at zero, and a refused credential, and `check_scan_cache` the scan reuse windows; [`testing.md`](testing.md#agent-usage-collector-tests) describes the helpers and `agent-usage-grok-test.sh` is the worked example. diff --git a/docs/testing.md b/docs/testing.md index f367dc58671..0e53c6c962b 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -133,6 +133,33 @@ Roughly a quarter of the shell test files use this to test parsing, merging, and layout logic as pure functions, reserving compositor-gated tests for what only a live session can prove. +## Agent usage collector tests + +The collectors behind the agents panel (`bin/omarchy-agent-usage-*`, see [`agent-usage-collectors.md`](agent-usage-collectors.md)) share two layers of tests, and every collector's own test builds on two shared helpers rather than its own copies. + +The shared layer covers what every collector does the same way. `agent-usage-shared-test.sh` holds the shared module, `shell/plugins/agents/collectors.py`, to the contract directly: untrusted numbers and times, the record's shape, the scan and limits caches and their locks, HTTP failure classes, and the pi, omp, opencode, Claude Code and Cursor readers, with the rule that says who bills a Claude Code reply. `agent-usage-partition-test.sh` runs every collector against one machine holding usage on every provider, so it holds each one to the unused-machine record and to counting its own provider's tokens exactly once from every root and store, and each Claude Code reply once behind every endpoint Claude Code can be sent to. A collector's own test doesn't repeat either; it covers its own sources, its credentials, how it reads its service's answer, and its adapter for the contract checks below. + +`test/shell.d/agent-usage-helpers.sh` is sourced after `base-test.sh`. It is not named `*-test.sh`, so the runner never runs it on its own. + +- `agent_sandbox` creates `$SANDBOX` with one EXIT trap (which also stops processes handed to `track_pid ...`), exports `HOME` and every `XDG_*` directory under it, unsets every developer variable naming an agent credential or store (any `*_API_KEY`, and prefixes like `CLAUDE_`, `CURSOR_`, `PI_CODING_AGENT_`, `OPENCODE_`, `MISE_`; the bash and Python lists are checked to match), points HTTP proxies at a dead port so nothing reaches a real service, and puts the Python fixtures on `PYTHONPATH`. +- `new_home [name]` prints a fresh home under the sandbox with its XDG directories and a `bin/` for stubs. +- `run_collector [VAR=value ...] [-- args...]` runs a collector as a process under `env -i`, seeing only that home, `OMARCHY_PATH`, the caller's `TZ`, the dead proxy, the given variables, and a `PATH` of `/bin` plus python3, rg, jq and sqlite3. It runs with `--force` unless arguments follow `--`, prints the record, and leaves stderr in `/.run/.err`. +- `expect_jq "" '' '' "$json"` passes when `jq -c` of the filter prints exactly the expected text, and fails showing both and the JSON. +- `stub_cli ` writes `/bin/` from the body on stdin, and `fake_codex [mode]` a stand-in Codex app-server whose answer `CODEX_FAKE` picks (`ok`, `backend`, `auth`, `empty`, `died`, `crashed`, `silent`, `stall`), logging every start to `/codex-starts`. +- `age_scan_cache ` moves a scan cache's stamp into the past (or, negative, the future), and `edit_limits_cache ` rewrites every limits cache of a collector with `$now_ms`, `$past` and `$future` at hand, so no test sleeps for a window to pass. + +`test/shell.d/agent_usage_fixtures.py` is the Python half, imported from a heredoc with `from agent_usage_fixtures import *`. It finds files through `ROOT` and `OMARCHY_PATH`, never its own location or `HOME`, and writes no bytecode. + +- `check(description, condition, detail="")` and `expect(description, actual, expected, record=None)` print the same TAP lines as the bash assertions and end the test on the first failure. +- `sandbox(root=None)` points `HOME` and every XDG directory at a fresh directory and drops the developer variables; `home()`, `data_home()`, `config_home()`, `state_home()` and `cache_root()` name its places. +- `load_collector(id)` loads a collector as a fresh module, `run_main(module, *args)` runs its `main()` and returns the record it printed, and `collect(id, *args)` does both. +- `FakeHttp(routes, mode="up")` replaces `urllib.request.urlopen`, the one door every collector request goes through. Routes map a URL substring to a payload, raw bytes, an HTTP status, an exception, or a callable taking the request; `mode` fails every request at once (`down`, `timeout`, `reset`, `cut`, `garbage`, or a status), and `requests` logs each one. `stub_server(handler, tls=False)` serves a real loopback socket instead, carrying every `https://` request to it with `tls=True`. +- Builders for what collectors read: `iso_in`, `ms_in`, `today_stamp` and `jwt`; `pi_reply` and `write_pi_session`; `claude_code_session` (a Claude Code transcript, and where `ANTHROPIC_BASE_URL` sends it); `oc_v1`, `oc_v2`, `opencode_db` and `oc_add`; `cursor_state_db`, `pi_auth`, `opencode_auth`, `registry` and `write_json`; and `limits_cache_files`, `scan_cache_files`, `edit_limits_cache`, `age_scan_cache` and `edit_scan_cache` for the caches. + +A collector test describes its collector once, as an `Adapter`, and hands it to `check_limits_contract` and `check_scan_cache`. The adapter names the collector, a `plant()` that signs the sandbox in, and the healthy service's `answer` as FakeHttp routes; a collector that talks to a socket, a CLI or a file it reads instead passes `serve(mode)` and `count()`, and `skip` names any check that can't apply, with the reason beside it. `check_limits_contract` then drives `main()` through a fresh probe, the 15-second reuse for a normal and a `--limits-only` run, `--force`, a failed check kept stale, unreachable service with `retryAdvised`, a reply cut short and a garbage body kept stale with no retry (`cut` in `skip`, and `serve(mode)` gets `cut` and `garbage`), kept windows whose reset passed coming back at zero, and a refused credential, asserting on the printed record so every collector is held to the same field names, and returns the fresh record for the test's own reading of its answer. `check_scan_cache` grows the collector's local usage with the adapter's `grow()` between runs and rewrites the cache's stamps, holding it to the 20-second, 900-second and `--force` reuse windows and to missing a scan from another day or the future, and, given `pi_root`, to never caching a scan that couldn't read the pi sessions. `agent-usage-grok-test.sh` is the worked example. + +A test whose fixtures are dated today calls `freeze_today()` (and stamps rows with `fixture_now()`), so a run that crosses local midnight still sees the day it began; `check_scan_cache` freezes it on its own. Collectors run as processes under `run_collector` aren't reached by it, so a test that runs them that way and dates its rows today waits for the new day when it would start within two minutes of local midnight, as the partition test does. + ## Conventions worth copying - **Redirect background output.** Send background fixtures' stdout to a log or `/dev/null` and clean up the processes on exit. An inherited output pipe can hold the runner open after the test exits. diff --git a/etc/mise/conf.d/omarchy-tools.toml b/etc/mise/conf.d/omarchy-tools.toml index 63a271caf4f..b784d7198a9 100644 --- a/etc/mise/conf.d/omarchy-tools.toml +++ b/etc/mise/conf.d/omarchy-tools.toml @@ -6,7 +6,10 @@ system_installs_dir = "~/.local/share/mise/installs" system_shims_dir = "~/.local/share/mise/shims" # Agent CLIs update themselves before they run, at most once every four hours (mise auto_update). +# Backend options go inside a tool's table, never in a "tool[...]" key: mise drops bracket options from a +# key whose value is a table, so a tool declared that way installs without them. [tools] +amp = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["amp"], minimum_release_age = "0s" } antigravity-cli = { version = "latest", lazy = true, auto_update = "4h", minimum_release_age = "0s" } basecamp = { version = "latest", lazy = true, minimum_release_age = "0s" } claude = { version = "latest", lazy = true, auto_update = "4h", minimum_release_age = "0s" } @@ -18,14 +21,22 @@ gh = { version = "latest", lazy = true, minimum_release_age = "0s" } ghui = { version = "latest", lazy = true, minimum_release_age = "0s" } grok = { version = "latest", lazy = true, auto_update = "4h", minimum_release_age = "0s" } hey-cli = { version = "latest", lazy = true, minimum_release_age = "0s" } +"http:devin" = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["devin"], minimum_release_age = "0s", platforms_linux_x64_url = "https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz", platforms_linux_arm64_url = "https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz", bin_path = "bin", version_list_url = "https://static.devin.ai/cli/current/manifest.json", version_json_path = ".version" } +"http:kiro-cli" = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["kiro-cli"], minimum_release_age = "0s", platforms_linux_x64_url = "https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz", platforms_linux_arm64_url = "https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz", version_list_url = "https://prod.download.cli.kiro.dev/stable/latest/manifest.json", version_json_path = ".version" } hunk = { version = "latest", lazy = true, minimum_release_age = "0s" } muse = { version = "latest", lazy = true, auto_update = "4h", minimum_release_age = "0s" } node = { version = "latest", lazy = true, minimum_release_age = "0s" } +"npm:@deepseek-ai/dsh" = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["dsh"], minimum_release_age = "0s" } +"npm:@kilocode/cli" = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["kilo"], minimum_release_age = "0s" } +"npm:@minimax-ai/code" = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["mcode"], minimum_release_age = "0s", allow_builds = "better-sqlite3" } "npm:@moonshot-ai/kimi-code" = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["kimi"], minimum_release_age = "0s" } "npm:cf" = { version = "latest", lazy = true, lazy_bins = ["cf"], minimum_release_age = "0s" } +"npm:command-code" = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["commandcode"], minimum_release_age = "0s" } +"npm:freebuff" = { version = "latest", lazy = true, auto_update = "4h", lazy_bins = ["freebuff"], minimum_release_age = "0s" } oh-my-pi = { version = "latest", lazy = true, auto_update = "4h", minimum_release_age = "0s" } opencode = { version = "latest", lazy = true, auto_update = "4h", minimum_release_age = "0s" } ori = { version = "latest", lazy = true, minimum_release_age = "0s" } pi = { version = "latest", lazy = true, auto_update = "4h", minimum_release_age = "0s" } playwright = { version = "latest", lazy = true, minimum_release_age = "0s" } +qwen = { version = "latest", lazy = true, auto_update = "4h", minimum_release_age = "0s" } uv = { version = "latest", lazy = true, minimum_release_age = "0s" } diff --git a/manual/17-ai.md b/manual/17-ai.md index 18201410fe4..5da16aeab34 100644 --- a/manual/17-ai.md +++ b/manual/17-ai.md @@ -2,22 +2,31 @@ Omarchy includes launchers for the coding agents below without choosing a default for you. Run a command to install its CLI on first use, then authenticate when prompted. Hermes is the exception: nothing is installed for it on first use, and choosing it installs Hermes Desktop, as _Desktop apps_ below explains. -| Command | Agent | -|------------|------------------------------------------------------------------| -| `claude` | [Claude Code](https://code.claude.com/docs/en/overview) | -| `codex` | [OpenAI Codex](https://github.com/openai/codex) | -| `opencode` | [OpenCode](https://opencode.ai/) | -| `agy` | [Google Antigravity CLI](https://github.com/google-antigravity/antigravity-cli) | -| `copilot` | [GitHub Copilot CLI](https://github.com/github/copilot-cli) | -| `crush` | [Crush](https://github.com/charmbracelet/crush) | -| `grok` | [Grok Build](https://x.ai/build) from xAI | -| `kimi` | [Kimi Code](https://www.kimi.com/code/docs/en/) | -| `pi` | [Mario Zechner's Pi](https://github.com/badlogic/pi-mono) | -| `omp` | [Oh My Pi](https://github.com/can1357/oh-my-pi) | -| `ori` | [Ori](https://openrouter.ai/docs/guides/ori/harness), OpenRouter's harness | -| `hermes` | [Hermes](https://hermes-agent.nousresearch.com/), Nous Research's agent | -| `muse` | [Muse Code](https://dev.meta.ai), Meta's coding agent | -| `cursor-agent` | [Cursor CLI](https://cursor.com/cli) | +| Command | Agent | +|---|---| +| `claude` | [Claude Code](https://code.claude.com/docs/en/overview) | +| `codex` | [OpenAI Codex](https://github.com/openai/codex) | +| `opencode` | [OpenCode](https://opencode.ai/) | +| `agy` | [Google Antigravity CLI](https://github.com/google-antigravity/antigravity-cli) | +| `amp` | [Amp](https://ampcode.com/) | +| `commandcode` | [Command Code](https://commandcode.ai) | +| `copilot` | [GitHub Copilot CLI](https://github.com/github/copilot-cli) | +| `crush` | [Crush](https://github.com/charmbracelet/crush) | +| `cursor-agent` | [Cursor CLI](https://cursor.com/cli) | +| `devin` | [Devin CLI](https://devin.ai/cli) from Cognition | +| `dsh` | [DeepSeek Harness](https://github.com/deepseek-ai/deepseek-harness) | +| `freebuff` | [Freebuff](https://freebuff.com), Codebuff's free, ad-supported coding agent | +| `grok` | [Grok Build](https://x.ai/build) from xAI | +| `hermes` | [Hermes](https://hermes-agent.nousresearch.com/), Nous Research's agent | +| `kilo` | [Kilo](https://kilo.ai/docs/code-with-ai/platforms/cli), an OpenCode-based coding agent | +| `kimi` | [Kimi Code](https://www.kimi.com/code/docs/en/) | +| `kiro-cli` | [Kiro CLI](https://kiro.dev/docs/cli/) | +| `mcode` | [MiniMax Code](https://github.com/MiniMax-AI/minimax-code), MiniMax's coding agent | +| `muse` | [Muse Code](https://dev.meta.ai), Meta's coding agent | +| `omp` | [Oh My Pi](https://github.com/can1357/oh-my-pi) | +| `ori` | [Ori](https://openrouter.ai/docs/guides/ori/harness), OpenRouter's harness | +| `pi` | [Mario Zechner's Pi](https://github.com/badlogic/pi-mono) | +| `qwen` | [Qwen Code](https://github.com/QwenLM/qwen-code), Alibaba's coding agent | `ori` is the odd one out: it runs the other harnesses against OpenRouter's whole model catalog, so `ori claude`, `ori codex`, or `ori opencode` start those agents on whichever model you point them at, and `ori code` is Ori's own agent. @@ -31,19 +40,21 @@ To install all configured lazy tools before you need them, run `mise install --i `grok` is installed through mise's first-party `grok` package, the same way as `claude` and `codex`. Run `grok`, or pick it under _Setup > Defaults > Agent_. Skip the `curl | bash` installer on x.ai. That path puts a second copy in `~/.grok/bin`, prepends it on `PATH`, and overwrites Omarchy's wrapper. +`dsh` opens its browser UI on `127.0.0.1:3080`, with the launch directory as its default workspace. Its browser UI takes no starting prompt, so `omarchy agent prompt` opens it with the prompt on your clipboard to paste. Omarchy starts it with DeepSeek's no-approval permission preset. DeepSeek Harness uploads session logs to DeepSeek when you use its official model API; turn that off under _Settings > General > Upload Session Log_ in its browser UI. Feedback you send it is uploaded with the session too, unless you start it with `DSH_TELEMETRY_MODE=DISABLED`. + ### The default agent Pick your default agent with `omarchy default agent ` or under _Setup > Defaults > Agent_ in the Omarchy Menu (`Super + Space`). If the agent isn't installed yet, picking it installs it first. A fresh Omarchy will invite you to make this choice with a one-time notification. Choosing Claude also attempts to set up its browser extension for Chromium, Chrome, Brave, Brave Origin, and Edge. The extension setup applies to all users and may ask for your system password; cancelling or a failed extension install still selects and launches Claude. Restart your browser, enable the extension if prompted, and sign in to Claude to finish connecting it. Run `/chrome` in Claude to check the connection. Firefox and Zen do not support this extension. If Claude was already your default before this setup was added, select it again to install the extension. -Once you've chosen, `Super + Shift + Ctrl + A` launches the default agent in a dedicated terminal window (or brings up the picker if you haven't chosen yet). You can also launch it straight into a task with `omarchy agent prompt "Review this project"`. Agents launched this way run unattended in their respective don't-stop-to-ask modes, so be ready for them to actually do things! And since agents refuse to remember trust for your home directory, launches from `$HOME` start in `~/Work` instead. +Once you've chosen, `Super + Shift + Ctrl + A` launches the default agent in a dedicated terminal window (or brings up the picker if you haven't chosen yet). You can also launch it straight into a task with `omarchy agent prompt "Review this project"`. Freebuff and DeepSeek Harness take no starting prompt, so they open with the prompt on your clipboard to paste instead. Agents launched this way run unattended in their respective don't-stop-to-ask modes, so be ready for them to actually do things! MiniMax Code has no such mode for its interactive sessions, so it asks for permissions as usual. And since agents refuse to remember trust for your home directory, launches from `$HOME` start in `~/Work` instead. There are terminal shortcuts too: `a` runs the default agent inline in the current terminal, while `c`, `cx`, and `cy` start OpenCode, Claude Code, and Codex directly (again in their auto-approving modes). Theme changes sync to the agents as well: Claude Code, Pi, OpenCode, and Hermes (once Omarchy has installed it) all follow along when you switch the Omarchy theme. ### The agents panel -The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan, a meter for each of its 5-hour session and weekly limits with the time until it resets (hover for the exact percentage), or the remaining prepaid balance, under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, Fireworks, and Kimi are covered out of the box. Kimi currently reports account quota limits without local token or prompt statistics. +The top bar has an agents icon. Before you've set up an agent, its panel walks you through signing in to Claude, Codex, or Grok, and the first one becomes your default agent. After that, the panel puts every subscription on one page: each agent and account with its plan, a meter for each of its 5-hour session and weekly limits with the time until it resets (hover for the exact percentage), or the remaining prepaid balance, under a line that rotates through what you've used across all of them, like tokens this week and your most used model. Under _Make something_ it offers starter prompts for a new theme, plugin, or app that hand the task to your default agent. Apps are built the way Omarchy's own Hype, Monologue, and Omacut are, in C++ and Qt Quick, following the omarchy-app agent skill, and the compiler and Qt pieces they need come installed. Claude Code, Codex, Grok, Kimi, Cursor and its Grok Bot allowance, GitHub Copilot, Muse, Antigravity, Kimchi, and Devin are covered out of the box with their limits, and Hermes, Mistral Vibe, and Qwen Code show their sessions and tokens once you've used them. Provider plans and keys you use through pi, omp, or opencode show up as well: OpenCode Zen and Go, the Z.ai GLM Coding Plan, Ollama Cloud, the MiniMax Token Plan, and Synthetic with their limits; Fireworks, OpenRouter, and DeepSeek with their prepaid balances; and the models you run on your own hardware under Local models. Every agent is counted once, under whoever bills it, however many harnesses you drive it through. 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/migrations/1792504295.sh b/migrations/1792504295.sh new file mode 100644 index 00000000000..7c3eeed64b8 --- /dev/null +++ b/migrations/1792504295.sh @@ -0,0 +1,15 @@ +echo "Build the lazy shims for the newly declared default coding agents" + +# The settings package ships the declarations for Qwen Code, DeepSeek Harness, +# MiniMax Code, Kiro CLI, Freebuff, Kilo, Devin CLI, Amp, and Command Code. +# Until it has, stay pending so the shims are built once they arrive. +if ! omarchy-mise-default-tools list | grep -qxF qwen; then + echo "Waiting for the omarchy-settings package that declares the new coding agents" >&2 + exit 1 +fi + +# A user who removed the preinstalls skips Omarchy's declarations, so there is +# nothing to build for them. +if omarchy-mise-default-tools declares qwen; then + mise reshim --system +fi diff --git a/shell/plugins/agents/Panel.qml b/shell/plugins/agents/Panel.qml index 8d814e211e6..8f87615b922 100644 --- a/shell/plugins/agents/Panel.qml +++ b/shell/plugins/agents/Panel.qml @@ -419,9 +419,18 @@ Panel { return parts.join(" · ") } + // Only the agents `omarchy-agent-account-add` signs in get the link; any + // other agent's sign-in trouble stays text, with its own help to follow. + function supportsSignIn(p) { + for (var i = 0; i < addProviders.length; i++) + if (p && addProviders[i].providerId === p.providerId) return true + return false + } + // A lapsed or missing sign-in is something you can fix from here; any other // trouble is only reported. function needsSignIn(item) { + if (item && item.providerId && !supportsSignIn(item)) return false var status = String(item && item.usageStatusText || "") return status === "Sign-in expired" || status === "Waiting for auth" } @@ -438,7 +447,7 @@ Panel { // Sign an account that's already here in again, following along in the // add view just like adding one. function signInAgain(p, account) { - if (!p || addStage === "running") return + if (!supportsSignIn(p) || addStage === "running") return addProvider = p.providerId addPrivate = !!account && account.primary !== true addStatus = Commons.I18n.t("agents.starting") @@ -630,6 +639,7 @@ Panel { function statusDisplayText(status) { var value = String(status || "") if (value === "Balance unavailable") return Commons.I18n.t("agents.balanceUnavailable") + if (value === "Out of credits") return Commons.I18n.t("agents.outOfCredits") var limits = value.match(/^(.+) limits unavailable$/) if (limits) return Commons.I18n.t("agents.providerLimitsUnavailable").arg(limits[1]) var provider = value.match(/^(.+) unavailable$/) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 9578180d47c..581e3852243 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -19,10 +19,12 @@ Every subscription on one page, limits first. per limit window: its meter and the time until it resets (the exact percentage on hover). A model-scoped allowance on the same clock (Claude's Fable weekly limit) is a tick on that window's meter rather than a line of its own; the - row's tooltip names it. A lapsed or missing sign-in shows a _Sign-in required_ link that signs - that account in again from the panel; other endpoint trouble shows under the - name in the urgent color. Limits kept from an earlier check after a failed one dim, and their - tooltip says how old they are. + row's tooltip names it. A lapsed or missing Claude, Codex, or Grok sign-in + shows a _Sign-in required_ link that signs that account in again from the + panel; any other agent's sign-in trouble, and other endpoint trouble, shows + under the name in the urgent color, with what to do on hover. Limits kept + from an earlier check after a failed one look the same, and their tooltip + says how old they are. - **Accounts** — an agent with more than one subscription account (see `omarchy agent account`) lists each: name and plan on one line (the email on hover), and its own limit lines. An _ACTIVE_ label marks the account new sessions start as; the @@ -63,60 +65,196 @@ written by `omarchy-agent-usage-update`. That command runs one on its refresh timer and whenever you ask for a refresh, and picks up any record that lands in the directory regardless of who wrote it. -Adding an agent needs no QML changes: ship a collector that prints the record contract (see the `claude` and `codex` collectors in `bin/`), and the panel gains a section. An `assets/.svg` mark is optional, with an `assets/-light.svg` twin if the mark needs a dark variant for light surfaces; the bar glyph stands in when there is none. +Adding an agent needs no QML changes: ship a collector that follows the contract in [`docs/agent-usage-collectors.md`](../../../docs/agent-usage-collectors.md) (the `claude` collector is the reference), built on the plumbing the collectors share in `collectors.py` beside this README, and the panel gains a section. An `assets/.svg` mark is optional, with an `assets/-light.svg` twin if the mark needs a dark variant for light surfaces; the bar glyph stands in when there is none. | Collector | Limits | Local stats | |---|---|---| -| `claude` | Anthropic's OAuth usage endpoint (5-hour session + 7-day weekly) | `~/.claude/projects` transcripts, opencode sessions on an Anthropic provider, plus `stats-cache.json` and `history.jsonl` as fallback | -| `codex` | The Codex app-server RPC (rate windows, plus an individual spend cap and its reported reset time when present) | native Codex CLI session files on the built-in `openai` provider (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 | -| `kimi` | Kimi Code quotas (5-hour session, weekly, and monthly allowances) | none; quota only | - -When `~/.local/state/omarchy/agents/accounts/.json` -registers more than one account, the `claude`, `codex`, and `grok` records -also carry -`accounts: [{ id, label, email, plan, active, limits, stale, usageStatusText, -authHelpText }]`, each account probed with its own sign-in (Claude caches each -account's limits separately; Codex runs one app-server per account home), and -`accountSwitch: { mode, threshold }`. The record's top-level `limits` and -`tierLabel` keep describing the active account, and local stats stay one set, -since every account shares the primary home's history. After each run, -`omarchy-agent-usage-update` hands the fresh limits to -`omarchy-agent-account-state autoswitch`, which notifies or switches when the -active account crosses its threshold, and re-collects the record if the active -account changed. - -Codex CLI will front any OpenAI-compatible backend — `--oss`, or a custom `model_provider` in `config.toml` aimed at Ollama, LM Studio, or a gateway — and those rollouts sit in the same sessions directory as OpenAI-backed ones. The Codex collector skips a rollout whose first `session_meta.model_provider` names anything but the built-in `openai` provider. Rollouts written before Codex recorded that field carry no provider and still count. - -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 -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 -again. Fireworks reads -`FIREWORKS_API_KEY` and `FIREWORKS_ACCOUNT_ID` first, then -`~/.fireworks/auth.ini` (which `firectl set-api-key` creates), then the key -pi stores in `~/.pi/agent/auth.json` when Fireworks is signed in there -(honoring `PI_CODING_AGENT_DIR`, and pi's literal and `$ENV_VAR` key forms), -and finally the key opencode stores in `~/.local/share/opencode/auth.json`. - -Kimi reads `KIMI_API_KEY` first, then `credentials/kimi-code.json` under an explicit `KIMI_CODE_HOME`. Without that override it prefers `~/.kimi-code/credentials/kimi-code.json`, with the legacy `~/.kimi/credentials/kimi-code.json` as a compatibility fallback. An explicit home never borrows another home's sign-in. Its record reports account quotas, without local token or prompt statistics. - -Expired tokens are refreshed under the current CLI's renewable directory lock at `/oauth/kimi-code.lock`; legacy credentials use their sibling file lock. The rotated token pair is flushed and published atomically with private permissions. Credentials are reread after the exchange so a sign-in or logout completed during the request takes precedence. A failed lock prevents refresh, and credentials are sent only to trusted Kimi HTTPS hosts without following redirects. When sign-in expires, run `kimi login`; an invalid overriding `KIMI_API_KEY` must instead be updated or unset. - -Limits probes reuse a result for one minute unless `--force` is requested. A failed probe retains the last-known limits with stale metadata. Cache identity follows the token used by the request, and a refresh this collector performs carries the same account's previous limits forward. Different credentials or endpoints never share cached quota. - -### Fireworks balance - -The collector first asks the account's `:getBalance` endpoint for the real -prepaid ledger. That endpoint exists but is permission-gated, and as of -August 2026 no console-issued API key passes it — Fireworks appears to -reserve it for the dashboard session. The probe stays because it is cheap -and the live figure lights up automatically if Fireworks ever opens it to -keys. Until then the collector falls back to estimating the balance from -configuration in `~/.config/omarchy/agents/fireworks.json`: +| `claude` | Anthropic's OAuth usage endpoint (5-hour session, 7-day weekly, and any model-scoped windows) | `~/.claude/projects` transcripts, pi/omp sessions and opencode sessions on an Anthropic provider, plus `stats-cache.json` and `history.jsonl` as fallback | +| `codex` | The Codex app-server's `account/rateLimits/read`, every bucket (general windows plus any model's own), plus an individual spend cap and its reported reset time when present | native Codex CLI session files on the built-in `openai` provider, pi/omp sessions on `openai-codex`, and opencode sessions on the `openai` provider | +| `grok` | The credits endpoint behind Grok's `/usage` view (the billing period's included usage), plus the monthly included quota of a unified-billing account | Each session's `usage.json` (the ledger `grok usage` prints: tokens by model per finished turn, one turn per typed prompt, a fork's copied turns counted once), plus `summary.json` for sessions | +| `kimi` | Kimi Code quotas (5-hour session, weekly, and monthly allowances) | Kimi Code's per-agent `wire.jsonl` session logs under its home on Kimi's routes, pi, omp, and opencode sessions on the Kimi Code plan (`kimi-coding`, `kimi-code`, `kimi-code-plan-global`, `kimi-code-plan-cn`) or Moonshot's API (`moonshotai`, `moonshotai-cn`, `moonshot`), plus Claude Code turns sent to Kimi's or Moonshot's API or answered by a `kimi-*` or `moonshot*` model | +| `cursor` | Cursor's dashboard RPCs: the included total for the billing month, with the auto-model, named-model, and on-demand meters attached to it | the same dashboard RPCs, one call per day for the last week | +| `grokbot` | Grok Bot's weekly included usage from Cursor's `GetSandUsageStatus` RPC, asked as the Cursor sign-in | none; limit only | +| `copilot` | Account-wide AI credit allowance (GitHub quota endpoint, or estimated locally) | the Copilot CLI session store's `assistant_usage_events`, with `~/.copilot/session-state` transcripts as fallback | +| `muse` | The client's key endpoint (session + weekly percents), else local metering only — never estimated | native `muse` CLI session files, with prompt and session counts from its `session-index.db`, plus pi/omp and opencode sessions on the `meta` provider and omp's `muse-code` | +| `agy` | Google's Cloud Code API (`loadCodeAssist` for the plan, `retrieveUserQuotaSummary` for the 5-hour and weekly windows and any touched Claude/GPT allowance), asked with the sign-in agy, omp, pi, or opencode's Antigravity plugin keeps | `history.jsonl`, `conversation_summaries.db`, and each conversation's `transcript.jsonl` under `~/.gemini/antigravity-cli/` (`AGY_DIR`), plus pi, omp, and OpenClaw sessions on the `google-antigravity` provider and opencode's `antigravity-` models | +| `kimchi` | The plan, prepaid credits, and any spend budgets from the gateway's `/v1/credits` and `/v1/budget`, the endpoints Kimchi's own status line reads | Kimchi's pi session files under `~/.config/kimchi/harness/sessions` (`KIMCHI_CODING_AGENT_DIR`): tokens by model per assistant reply on Kimchi's own `kimchi` and `kimchi-dev` routes; a session `--provider openrouter` is OpenRouter's, and one on a custom provider counts nowhere. Also pi/omp sessions and opencode replies on those provider ids, and Claude Code turns and Codex rollouts sent to Kimchi's gateway (`llm.kimchi.dev`, `llm.eu.kimchi.dev`), as `kimchi claudecode` and `kimchi codex` set them up | +| `hermes` | None (Hermes has no subscription of its own) | `state.db` in Hermes' home and in each named profile under its root's `profiles/`, read-only, scanned once per refresh window: the sessions Nous Portal billed, and those from before Hermes kept the route | +| `devin` | The daily and weekly quota (plus credit buckets on credit plans) from the GetUserStatus answer the Devin CLI caches in `$XDG_CACHE_HOME/devin/cli/user_status.*.bin`; no network call | the CLI's `$XDG_DATA_HOME/devin/cli/sessions.db` (per-message metrics and `prompt_history`), plus pi/omp and opencode sessions on the `devin` provider | +| `vibe` | None (Mistral publishes no usage limits) | Vibe's Unified Harness store under `~/.vibe/logs/session/unified` and older releases' per-session `meta.json` under `~/.vibe/logs/session` (`VIBE_HOME`), only sessions on a Mistral-served model, plus pi/omp and opencode sessions on the `mistral` provider | +| `qwen` | None (neither Qwen Code nor Alibaba publishes a way to read a Coding Plan or Token Plan quota) | Session transcripts under `~/.qwen/projects/*/chats` (`QWEN_HOME`, `QWEN_RUNTIME_DIR`, `advanced.runtimeOutputDir`): billed API responses on Alibaba's routes and typed prompts, plus pi, omp, and opencode sessions on Alibaba's Coding Plan, Token Plan, and Qwen OAuth providers | +| `opencode` | OpenCode Go's usage endpoint (rolling + weekly + monthly), with an API key or the Console session; a key without a Go plan shows local usage alone | pi and omp sessions (profiles included) on the Zen and Go providers, plus opencode's own `message` and `session_message` stores | +| `zai` | Z.ai's coding-plan monitor endpoint (5-hour session + 7-day weekly), for keys that carry a GLM Coding Plan | pi, omp, and opencode sessions on Z.ai's and Zhipu's providers (`zai`, `zhipu`, pi's `zai-coding-cn`, omp's `zhipu-coding-plan`, opencode's `zhipuai`, `zai-coding-plan` and `zhipuai-coding-plan`; `glm-*` models), plus Claude Code turns sent to Z.ai or answered by a `glm-*` model | +| `ollama` | Ollama Cloud's balance endpoint (5-hour session + 7-day weekly on a legacy plan, or the monthly included credits), plus any purchased credits left as a balance | pi, omp, and opencode sessions on the `ollama-cloud` provider, plus cloud-tagged models reached through the local daemon's `ollama` provider (Codex `--oss` rollouts included), Codex rollouts on a `model_providers` entry at ollama.com, and Claude Code turns sent to ollama.com or answered by a cloud-tagged model | +| `minimax` | MiniMax Token Plan API (5-hour session + 7-day weekly) | MiniMax Code's session files under `~/.minimax/v2/sessions`, plus pi, omp, and opencode sessions on a MiniMax provider, and Claude Code turns sent to MiniMax or answered by a `MiniMax-*` model | +| `openrouter` | Prepaid balance from `/credits`, key-limit meter from `/key`, each kept stale past a failed check | opencode and pi/omp sessions on the OpenRouter provider, Kimchi sessions run with `--provider openrouter`, DeepSeek Harness's `openrouter` route, MiniMax Code custom providers at openrouter.ai, and Claude Code turns sent to OpenRouter, with or without a key, or the `/activity` billing API with a management key | +| `deepseek` | Live prepaid balance from `GET https://api.deepseek.com/user/balance`, kept stale past a failed check, with the optional `fundedAmount` estimate for the funded-versus-spent line | pi/omp session files (profiles included) whose assistant messages ran on the deepseek provider, opencode sessions on a deepseek provider, DeepSeek Harness (`dsh`) session logs on its DeepSeek routes, and Claude Code turns sent to DeepSeek or answered by a `deepseek-*` model | +| `fireworks` | Prepaid balance: the live ledger when the key can read it, else configured funding minus rated account costs; kept stale past a failed check | Fireworks billing API, grouped by day and model for the last 30 days, cached per key and account | +| `synthetic` | Synthetic's `GET /v2/quotas`: rolling five-hour requests, weekly credits, hourly search, free tool calls | pi/omp sessions and opencode messages on the `synthetic` provider; Claude Code, MiniMax Code and Codex sent to `api.synthetic.new`, and Claude Code replies by `hf:`/`syn:` models with no override | +| `local` | None (self-hosted work has no meter) | Codex rollouts, pi/omp sessions and opencode messages on a local runtime, a provider named in `local.json`, or a Codex provider whose `base_url` is on the LAN, never a provider id another collector counts, plus Claude Code turns while `ANTHROPIC_BASE_URL` is self-hosted (not Claude models or cloud-tagged ones), and Kimi Code, Qwen Code, Hermes and Mistral Vibe requests sent to a self-hosted base URL | + +Every collector follows the same rules, so the same situation reads the same for every agent: a recent answer is reused for 15 seconds unless refresh is pressed, each credential or account keeps its own last good answer, a failed check keeps it with its age on hover and windows that have since reset back at zero, an endpoint that can't be reached is retried soon, and an agent the machine doesn't use stays out of the panel. Tokens are counted once, under the provider that bills them: pi, omp, and opencode sessions are split by provider into the matching collector. A pi compaction's usage counts, but omp records only the context's size before and after one, not what its summary request used, so omp's compactions are missing from the tokens. pi and omp sessions are read from every profile, and from `PI_CODING_AGENT_DIR`, `PI_CODING_AGENT_SESSION_DIR`, and a `settings.json` `sessionDir` as well as the defaults when set, a forked session's copied messages once; opencode is read from both its `message` and V2 `session_message` stores in the database `OPENCODE_DB` names, a message migrated into both once and a fork's copied history not at all. A machine with none of an agent's home, CLI, or usage prints an empty record the panel skips; one that has them but no sign-in says it is waiting for one. The sections below cover where each agent's sign-in and data come from. + +When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records also carry `accounts: [{ id, label, email, plan, active, primary, limits, stale, fetchedAt, retryAdvised, usageStatusText, authHelpText }]`, each account probed with its own sign-in (each account caches its own last limits; Codex runs one app-server per account home), and `accountSwitch: { mode, threshold }`. The record's top-level `limits` and `tierLabel` keep describing the active account, and local stats stay one set, since every account shares the primary home's history. After each run, `omarchy-agent-usage-update` hands the fresh limits to `omarchy-agent-account-state autoswitch`, which notifies or switches when the active account crosses its threshold, and re-collects the record if the active account changed. + +### Claude Code + +Limits come from Anthropic's OAuth usage endpoint, asked with the access token the Claude Code CLI keeps in `.credentials.json` under `CLAUDE_CONFIG_DIR` (default `~/.claude`); the plan comes from the CLI's profile, else the credentials. Only the CLI refreshes that token, so one left to lapse reads "Limits paused" and keeps the last numbers until Claude Code runs again; a lapsed refresh token reads "Sign-in expired". Without credentials the panel says so and falls back to local stats only. + +Local stats come from the CLI's transcripts under `projects/`, falling back to `stats-cache.json` and `history.jsonl` (the cache's days are UTC dates that stop at yesterday, so its daily tokens are approximate and rarely cover today), plus pi and omp sessions (honouring `PI_CODING_AGENT_DIR` and `PI_CODING_AGENT_SESSION_DIR`) and opencode sessions on the `anthropic` provider. + +Only what Anthropic served counts here, and the collector that bills the rest counts it by the same rule (`claude_code_owner` in `collectors.py`), with the prompts it answered. Claude Code sends to whatever `ANTHROPIC_BASE_URL` names, and Z.ai, Ollama, MiniMax, DeepSeek, Kimi, OpenRouter, Alibaba's Model Studio ([Qwen](#qwen-code)), [Synthetic](#synthetic) and [Kimchi](#kimchi) (`kimchi claudecode`) document pointing it at their own APIs. A transcript doesn't record where it was sent, so while the `env` in `settings.json` (else the environment) points `ANTHROPIC_BASE_URL` at one of those providers' hosts, every reply is that provider's, earlier ones included. Otherwise the model that answered decides: a Claude model is Claude's (through a proxy or a local gateway too), an Ollama `:cloud` tag is Ollama Cloud's, a self-hosted base URL (Ollama's `localhost:11434`) makes any other model the `local` collector's, and with no override `glm-*` is Z.ai's, `MiniMax-*` MiniMax's, `deepseek-*` DeepSeek's, `kimi-*` and `moonshot*` Kimi's, and `hf:*` and `syn:*` Synthetic's. Any other model with no override (`qwen3-coder`), and any non-Claude model behind a host that can't be placed, counts nowhere. The `stats-cache.json` fallback drops non-Claude models the same way; the other collectors read only the transcripts. + +A reply counts once by its message and request ids, across the transcripts a resume or fork copies it into; a `/btw` side question's sidechain replay of its parent's message, a whole copy under a new request id, counts under the parent, while a subagent's own reply counts even when a gateway gave it its parent's message id. A prompt is one a person sent: a main-thread user line that isn't a tool's result, Claude Code's own meta line, a compaction summary, a task notification, a local command's output, or a built-in that runs locally (`/model`, `/clear`, `/compact`, `/exit`, any command Claude Code answers with a `local-command-stdout` or `-stderr` line), plus a prompt typed mid-turn, which Claude Code queues as a `queued_command` attachment; a custom slash command or skill asks the model, and counts. Each counts once by its uuid, which a resumed or forked transcript keeps, so a tool loop's API requests are tokens, not prompts. Some requests have no reply line: a compaction (a `compact_boundary` line records only the context's size), the session title, and side queries such as Haiku's. Some replies are understated: a subagent's reply can be written with only its `message_start` usage (`stop_reason` null, a few output tokens), and its real output never is. Claude Code 2.1.278 and later write `cost-state` lines in some sessions, not all, the running totals per model of a process's cost lineage (its `startTime`, which a resume or `--fork-session` carries on), and what those hold beyond the replies of the lineage's sessions (subagents included) counts too, once per lineage, under the collector that bills that model, on the day of the cost-state that first holds it, not the day of the requests it makes up for. A process still running counts that usage when it next saves its cost, and a session with no cost-state, as every one before 2.1.278 is, has its compactions and understated subagent replies missing from the tokens. + +Every collector reading the transcripts shares one index of them, which one brings up to date at a time while the others wait for it and then read it; an index that already holds every transcript as it is is read without waiting. + +### Codex + +Limits and the plan come from the Codex CLI's own app-server (`codex app-server`, `account/rateLimits/read`), started in each account's home (`CODEX_HOME`, default `~/.codex`) so it answers for that login and refreshes its tokens there. Every rate-limit bucket is read: the general windows, and any model's own allowance shown beside the window of the same span. An individual spend cap, such as a Business seat's, is its own "Spend limit" meter, the share of the cap spent with the reset time the app-server reports, since it names no period; it can come without any rate window. The plan reads as Codex's own `/status` names it (`pro` is Pro 200, `promax` Pro 500, `team` Business), and a plan newer than that table keeps its wire name. A home with no `auth.json` under the file store and no `CODEX_ACCESS_TOKEN` reads "Waiting for auth" without starting the app-server. An API-key login (`auth_mode` `apikey` in `auth.json`, or, with no `auth_mode`, an `OPENAI_API_KEY` there, ChatGPT tokens beside it or not, as Codex itself reads it) has no plan limits to read, so it starts no app-server either and shows local stats alone, with nothing to sign in to. + +Each login keeps its last good limits, keyed by the workspace and user in `auth.json` (or the access token); a keyring login can't be told apart and keeps none. A stalled app-server, or one that couldn't reach OpenAI, is retried soon. + +Local stats come from native session files on the built-in `openai` provider, pi and omp sessions on `openai-codex` (honouring `PI_CODING_AGENT_DIR` and `PI_CODING_AGENT_SESSION_DIR`), and opencode sessions on `openai`, from both of opencode's message stores and the database `OPENCODE_DB` names. A native session's tokens are its responses: one `token_usage_record` line each from Codex 0.163 on, a remote compaction's included, and before that each `token_count` event's last request. The thread-title requests interactive Codex sends are written to no rollout, so their (small) tokens are missing; `codex exec` sends none. A native session's prompts are the ones typed (a `UserMessage` item, or a legacy rollout's `user_message` event), not the requests that answered them, and a subagent's or guardian's thread types none and starts no session: its tokens count under the session that started it (its `session_id`, else its `parent_thread_id`). + +Codex CLI will front any OpenAI-compatible backend — `--oss`, or a custom `model_provider` in `config.toml` aimed at Ollama, LM Studio, or a gateway — and those rollouts sit in the same sessions directory as OpenAI-backed ones. The Codex collector skips a rollout whose first `session_meta.model_provider` names anything but the built-in `openai` provider. Rollouts written before Codex recorded that field carry no provider and still count. A skipped rollout on a `model_providers` entry whose `base_url` is a provider's API (Z.ai, OpenRouter, DeepSeek, Moonshot or Kimi, Alibaba's Model Studio, MiniMax, ollama.com, [Synthetic](#synthetic), [Kimchi](#kimchi)'s gateway as `kimchi codex` sets it up) counts under that provider's collector, whatever the entry's id; one on your own hardware under Local models; and one behind a gateway nothing can place nowhere. + +Codex also sends the built-in `openai` provider wherever `openai_base_url` in `/etc/codex/config.toml` or `config.toml` says (`OPENAI_BASE_URL` in the environment does nothing), and those rollouts still name only `openai`. While that `openai_base_url` names a provider's API, every built-in rollout counts under that provider's collector instead, earlier ones included. A profile's `.config.toml` applies only to a run given `--profile`, so its `openai_base_url` decides only for the model it sets, when that is neither an OpenAI model nor the one `config.toml` sets: a profile naming a provider's API moves only that model's rollouts. Behind any other base URL OpenAI's own models (`gpt-*` but `gpt-oss`, `*codex*`, `o3`) still count here, since a proxy may forward them to OpenAI; any other model counts under Local models when the URL is self-hosted and nowhere when it is a gateway nothing can place. The current config decides; once the key goes, every built-in rollout is Codex's again. + +### Grok + +The plan comes from the settings Grok caches in its home (`GROK_HOME`, default `~/.grok`), and the limit from the credits endpoint its own `/usage` view reads, asked with each account's sign-in from `auth.json`. A lapsed access token with a live refresh token is routine: Grok renews it when it starts, so until then the limits read "Limits paused" with the last credits kept stale. A refused sign-in, or one past Grok's 30 days, reads "Sign-in expired". Local stats come from each session's `usage.json` and `summary.json`: a turn is a typed prompt, except in a subagent's session, and the turns a fork copies from its source count once; sessions on disk count Grok as in use even without a sign-in. + +### Kimi + +Kimi reads `KIMI_API_KEY` first, then `credentials/kimi-code.json` under an explicit `KIMI_CODE_HOME`. Without that override it prefers `~/.kimi-code/credentials/kimi-code.json`, with the legacy kimi-cli's `credentials/kimi-code.json` under `KIMI_SHARE_DIR` (else `~/.kimi`) as a compatibility fallback. An explicit home never borrows another home's sign-in. Kimi counts as in use when it has a credential or a Kimi home; a home without a sign-in waits for one, and a machine with neither stays hidden. + +Local stats come from the log Kimi Code 2.x keeps for each agent of a session, subagents included, at `sessions///agents//wire.jsonl` under its home (`KIMI_CODE_HOME`, else `~/.kimi-code`): every `usage.record` is one billed model request, compactions included, and every `turn.prompt` or `turn.steer` someone typed, slash commands included, is a prompt. A fork copies its source's log before marking it forked, so the copies count once, under the source's session. A request counts here when Kimi billed it: its model alias resolves through `config.toml` (`[models.]`'s provider, its `base_url` or the `*_BASE_URL` in its `env` table, or the model's own `base_url`), and Kimi's or Moonshot's API, a `kimi` provider with no `base_url`, or an alias the config doesn't name is Kimi's. A request on a server you run counts under Local models, and one on Anthropic's, OpenAI's or Google's API or a gateway nowhere. Each prompt goes with the next request sent for it, answered or not, else with the alias its agent was bound to, so a prompt a server you run refused counts under Local models too. The log records only the alias, so the current `config.toml` decides for every earlier request too. The legacy kimi-cli's own session logs aren't read. pi, omp, and opencode sessions on the Kimi Code plan count too: pi's own `kimi-coding` provider and omp's `kimi-code` (the plan's `api.kimi.com/coding/v1`) and models.dev's `kimi-code-plan-global` and `kimi-code-plan-cn`; Moonshot's pay-as-you-go API (`moonshotai`, `moonshotai-cn`, omp's `moonshot`) counts here too, although these limits don't meter it. Claude Code turns sent to `api.kimi.com`, `api.kimi.ai` or `api.moonshot.*`, or answered by a `kimi-*` or `moonshot*` model with no `ANTHROPIC_BASE_URL`, count here too, by the [Claude Code](#claude-code) rule, and that usage alone puts Kimi in use. + +Expired tokens are refreshed under the current CLI's renewable directory lock at `/oauth/kimi-code.lock`; legacy credentials use their sibling file lock. The rotated token pair is flushed and published atomically with private permissions. Credentials are reread after the exchange so a sign-in or logout completed during the request takes precedence. A failed lock prevents refresh, and credentials are sent only to trusted Kimi HTTPS hosts without following redirects. A refused token or refresh reads as Sign-in expired: run `kimi login`, or update or unset an overriding `KIMI_API_KEY`. Cache identity follows the token used by the request, and a refresh this collector performs carries the same account's previous limits forward. Different credentials or endpoints never share cached quota. + +### Cursor + +Cursor keeps no token counts on disk, so everything comes from the dashboard RPCs the Cursor CLI itself calls, asked with `CURSOR_AUTH_TOKEN`, else, with `CURSOR_API_KEY` exported, the access token the CLI traded that key for (the key itself is never sent), else the sign-in `cursor-agent login` stores in `~/.config/cursor/auth.json`, else the editor's in `~/.config/Cursor/User/globalStorage/state.vscdb` (read-only); `CURSOR_API_ENDPOINT` moves the endpoint, over https only (plain http only to this machine). A sign-in that has lapsed gives way to the next one that hasn't; with all of them lapsed the limits read "Limits paused" until the CLI or editor renews its token. A machine with Cursor installed but signed in nowhere says it is waiting for a sign-in; one without Cursor stays out of the panel. + +The week of tokens is reused for 20 seconds (15 minutes on a near-limit refresh). A failed check keeps the account's last good meters, plan, and week of tokens. Tokens by day and by model cover the last seven local days, one aggregation call per day. + +The meters are Cursor's own percentages, so they always agree with the Cursor dashboard. The included total is the Monthly row; the auto-bucket and named model meters and the on-demand budget run on the same billing cycle and show in its tooltip. Spend is never divided into an allowance here: bonus usage the model providers hand out is spend with no allowance behind it, so dollars over the plan limit read far past 100% while Cursor's own meters do not. The on-demand meter is the exception and is a real ratio: dollars spent against the on-demand limit the account sets. + +### Grok Bot + +Grok Bot runs on a Cursor account, but its weekly allowance is metered apart from Cursor's plan: the dashboard's `GetSandUsageStatus` RPC says how much of this week's included usage is gone, when it resets, and the Grok plan. It is asked as the same account the Cursor collector reads, found the same way (`CURSOR_AUTH_TOKEN`, the token traded for `CURSOR_API_KEY`, then the CLI's `auth.json`, then the editor's `state.vscdb`, read-only). Grok Bot keeps nothing countable on disk, so the record carries the limit alone, and an account without a Grok Bot allowance, or a machine without a Cursor sign-in, stays out of the panel. + +### Copilot + +GitHub exposes no supported personal-account usage API — its Copilot REST endpoints are organization- or enterprise-scoped — so the token sections come from what the Copilot CLI keeps under `~/.copilot` (or `COPILOT_HOME`). Its session store (`session-store.db`) records every billed request in an `assistant_usage_events` table: day, model, session, the token split, and the cost in nano AI units (1e9 nano units = 1 AI credit ≈ $0.01). A row is a billed request rather than a prompt; a current CLI marks the ones a typed prompt started (`initiator` "user", the agent's own follow-up calls "agent"), and those are the prompts, while an older store's counts stay out of the panel's prompt totals. A request sent to a custom provider (BYOK, `COPILOT_PROVIDER_BASE_URL`) carries no GitHub usage report and isn't Copilot's to count; the store doesn't say where it went, so it counts nowhere. CLI versions whose store predates it are served by the session transcripts under `session-state` instead, whose `session.shutdown` events carry the same per-model token split and cost. + +The monthly meter is the account-wide allowance of AI credits from GitHub's internal quota endpoint (the same one the editor plugins ask), which counts every machine, IDE, and agent on the entitlement and names the plan and the reset date. It needs a token from `COPILOT_QUOTA_TOKEN`, `GH_TOKEN`, `GITHUB_TOKEN`, or `gh auth token`, and only runs where the Copilot CLI has been used, so a machine with `gh` signed in but no Copilot stays out of the panel. Each account keeps its own last answer, `gh`'s keyed by the user its `hosts.yml` names, so a reused answer doesn't run `gh` at all. Without any token the section says it is waiting for a sign-in. When the endpoint can't be asked at all, a budget in `~/.config/omarchy/agents/copilot.json` gives a meter labeled "(est.)" from this machine's CLI spend alone: + +```json +{ + "monthlyCredits": 1500, + "remote": true +} +``` + +`remote` set to `false` turns the quota probe off (and the sign-in reminder with it), and `true` runs it even without a Copilot CLI home. The estimated window resets on the 1st of each month at 00:00 UTC, GitHub's own boundary. + +### Muse + +Muse reads native CLI sessions from `$XDG_DATA_HOME/muse/sessions` (defaulting to `~/.local/share/muse/sessions`, and overridable via `MUSE_DATA_DIR`), with the real prompt and session counts from the CLI's `session-index.db` beside them (completions stand in when it is missing, and for a session it has no row for yet; since the index keeps each session's running prompt count, today's prompts are what each count grew by since the last scan before today), plus pi, omp, and opencode sessions on the `meta` provider and omp's `muse-code` (its Muse Code sign-in). Muse Spark through the OpenCode Zen gateway counts in the `opencode` record instead. The OAuth token comes from `MUSE_AUTH_PATH` when set, otherwise from `~/.config/muse/auth.json` (which `muse login` creates); the file is only ever read, never refreshed or rewritten. A Muse CLI set up on the machine without a login waits for one. + +A session's subagent logs roll up into it. Every billed call a child agent made counts too: most children keep a log of their own, and one that doesn't (a skill reminder) is counted from the `goal_usage_attribution` record the session's log keeps for it, on the session's model, answering no prompt. + +The collector asks the client's own key endpoint for the subscription's session and weekly percents with reset times — the same figures the TUI's `/usage` overlay shows. Minting is idempotent (the same Model API key comes back every call), so polling is safe; the token travels only in the `Authorization` header and neither it nor the returned key is persisted. `MUSE_KEY_ENDPOINT` moves the endpoint, over https only (plain http only to this machine). Pay-as-you-go accounts have no subscription windows and no plan name, so the tab shows measured local usage with no meters — never estimated, never zeroed. An active plan that returns no `subs_usage` (Muse Code Power Usage answers this way on every probe) behaves the same way, except the tab keeps the plan name. + +### Antigravity + +The record's id is `agy`, the name the default agent knows it by. The plan and limits are asked of Google's Cloud Code API with whichever Antigravity sign-in is on the machine: the one `agy` keeps in the Secret Service keyring (service `gemini`, username `antigravity`) or, without a keyring, in `antigravity-oauth-token` under `AGY_DIR`; omp's `google-antigravity` sign-in in `~/.omp/agent/agent.db`; pi's in `auth.json` (under `PI_CODING_AGENT_DIR` when set); and opencode's `google` sign-in, only when the Antigravity plugin's `antigravity-accounts.json` is there. The token refreshed most recently is asked first, and one Google refuses gives way to the next. Those stores are only ever read. A lapsed token shows the last limits as paused until the harness that owns it runs again; a lapsed omp sign-in is asked through `omp usage`, which renews it. Each Google account keeps its own last answer, so another account's limits never stand in. `CLOUD_CODE_URL` moves the endpoint, over https only. + +Local stats count `agy`'s history prompts, its conversations (a subagent's conversation counts toward the one that spawned it), and the tokens each transcript step reports, under the model the step or the last Model Selection names. Antigravity answers that pi, omp, OpenClaw (`OPENCLAW_STATE_DIR`), and opencode's plugin recorded count too, a forked copy once. No other collector counts these. + +### Kimchi + +Kimchi reads `KIMCHI_API_KEY` first, then the key in `~/.config/kimchi/config.json`, and asks the gateway its region names (or `KIMCHI_BASE_URL`), which for a self-hosted login is `/llm` (or `KIMCHI_SELF_HOSTED_URL`), and with no host saved is nowhere: that key is never sent to the public gateway; its sessions follow `KIMCHI_CODING_AGENT_DIR`. The credits answer says how much remains but not how much was bought, so the balance shows the remaining figure without a funded meter; a spend budget's cap is what may be spent, not what was bought, so it never stands in. Every budget is a limit line that resets when its period ends, titled Monthly or Weekly when the period is one, with each capped provider inside it a tick on that line (or a line of its own in a period of another length). Each endpoint falls back on its own: when one can't be reached, its last answer stays up beside the other's fresh one. A rejected key reads as Sign-in expired; sign in again through Kimchi's own CLI. Sessions on disk without a key wait for one. Only sessions on Kimchi's own gateway routes count here; one run with `--provider openrouter` counts under OpenRouter, and one on a custom provider in `harness/models.json` counts nowhere. Usage other agents send through Kimchi's gateway counts here too: pi and omp sessions and opencode replies on the `kimchi`, `kimchi-dev` or `kimchi-experimental` provider, and Claude Code and Codex as `kimchi claudecode` and `kimchi codex` configure them, whatever the model, Claude and Kimi ones included. `kimchi claude` points only the Claude Code it starts at the gateway, through its environment, so those transcripts can't be told apart and count by their model's name (a `kimi-*` reply under Kimi, say); run `kimchi claudecode` once to have them count here. A self-hosted Kimchi gateway's host isn't recognized as Kimchi's, so Claude Code or Codex sent there counts as through any other gateway. An account out of credits keeps its balance shown and reads as Out of credits in the urgent color, with its help saying the gateway is rate limiting it until it has credits again. + +### Hermes + +Hermes Agent keeps its sessions in `state.db` in its home (`HERMES_HOME` with `~` and `$VAR`s expanded, else `~/.hermes` plus `HERMES_DATA_DIR_SUFFIX`) and in each named profile under its root's `profiles/`. The root is the default home unless `HERMES_HOME` lies outside it: then a `/profiles/` home makes `` the root, and any other home is its own. The root is the default profile, so its own `state.db` counts even while `HERMES_HOME` points at a named profile. The default home `~/.hermes` is read as well, whole: its own `state.db` beside its `profiles/`. The collector reads them read-only: sessions, user prompts, and per-session, per-model token usage, including chats that arrived over messaging gateways. A session's tokens land on the day it started. Hermes records the provider and base URL each session and request was billed through, so the tab counts those Nous Portal billed, and those from before Hermes kept the route (a session with none goes where its first request did); a server you run counts under Local models, and any other provider, Hermes' `custom` gateways included, nowhere. Hermes has no subscription, so the tab shows local stats only. Without any `state.db` the tab stays hidden; a database that can't be read shows "Hermes unavailable" with the file and error on hover, while the other profiles still count. + +### Devin + +Devin's quota comes from the cache the Devin CLI writes each time it starts, `$XDG_CACHE_HOME/devin/cli/user_status.*.bin` (`DEVIN_CACHE_DIR`): the GetUserStatus answer it fetched, the same numbers its `/usage` command shows. The collector never makes a request of its own. Of the caches fetched since `credentials.toml` was last written, the newest wins, so a sign-in to another account never shows the last one's quota; until the CLI runs again after a sign-in, the panel says so. A cache more than fifteen minutes old still shows, marked stale with its age, and a window that reset since reads zero. The daily and weekly windows show unless the plan hides them, and a credit plan's prompt, flow and flex credits show as meters that reset when the plan period ends. + +Local stats come from the CLI's `$XDG_DATA_HOME/devin/cli/sessions.db` (`DEVIN_DATA_DIR`), read-only: tokens by model from each assistant message's metrics, prompts and sessions from `prompt_history`. Builds that kept `cli_sessions.db` or `sessions.db` straight under `$XDG_DATA_HOME/devin` are read only when the CLI's own database is absent; they keep one lifetime total per session, which lands on its last active day. pi, omp, and opencode sessions on the `devin` provider count too. + +### Mistral Vibe + +Vibe saves sessions under `~/.vibe/logs/session` (moved by `VIBE_HOME`, or by `session_logging.save_dir` in Vibe's `config.toml`). Current Vibe runs every session on its Unified Harness, which keeps each one under `unified//`: the running token totals are in the public state of the generation `CURRENT` names, with its journal folded in, and the session's `meta.json` names the model's alias, resolved against `config.toml`'s models and providers over Vibe's built-in default. Older releases wrote a `session_*` folder per session whose `meta.json` holds the running totals and the whole config it ran with, and those still count. Vibe drives other providers too, a local llama.cpp server out of the box, so a session only counts when the model it resolved to is served by Mistral's API: a provider with Vibe's `mistral` backend or a mistral.ai host. A session on a self-hosted `api_base` counts under Local models, and one on another provider's host or a gateway nowhere; a Unified Harness session's provider comes from the current `config.toml`, so it decides for earlier sessions too, while an older folder keeps the config it ran with. Totals are booked on the day a session started, a subagent's on the session that spawned it, and prompts are the user messages Vibe didn't inject itself (in a fork, not the history it copied). pi, omp, and opencode sessions on the `mistral` provider count here too. The plan is the one Vibe last cached from Mistral's whoami answer in `whoami_cache.json`, named as Vibe names it. Mistral publishes no usage limits, so the section shows local stats only. + +### Qwen Code + +Qwen Code writes each session to `projects//chats/.jsonl` under `~/.qwen` (`QWEN_HOME`), or under `QWEN_RUNTIME_DIR` or the `advanced.runtimeOutputDir` setting when either moves it. Every billed API response is a telemetry line there with the model and the provider's token counts, the same lines Qwen Code's own usage statistics are rebuilt from, read by the same rules: the prompt count includes its cached part, and thinking is added to output only where the provider counts it apart. Typed prompts come from the user lines, and a branched session's repeated lines count once. A reply counts here when Alibaba billed it: Qwen Code records only its auth type, so its model's `baseUrl` in `settings.json`'s `modelProviders` says where it went, else the base URL Qwen Code resolves from the environment (`OPENAI_BASE_URL`, or `ANTHROPIC_BASE_URL` for its Anthropic auth type, read from `QWEN_HOME/.env`, `~/.qwen/.env` or `~/.env` and `settings.json`'s `env`, a project's own `.env` and the `--openai-base-url` flag being out of reach), Alibaba's hosts (`*.aliyuncs.com`, `*.alibabacloud.com`, `qwen.ai`) and a Qwen OAuth reply being Qwen's. A reply on a server you run counts under Local models; one on Anthropic's, Gemini's or Vertex AI's API, another provider's host or a gateway counts nowhere, and so does an OpenAI-compatible reply with no base URL in reach, which Qwen Code sends to OpenAI's API by default: an `OPENAI_BASE_URL` exported only from a shell profile is out of reach, so set it in one of those files for its replies to count. Each prompt goes with the next request sent for it, answered or not, else with the model `settings.json` selects (`security.auth.selectedType` and `model.name`), so a prompt a server you run refused counts under Local models too. The current settings and `.env` files decide for every earlier reply too. + +pi, omp, and opencode turns that Alibaba's Qwen plans bill count here too, on exact provider ids: the Model Studio Coding Plan as `alibaba-coding-plan` (opencode and omp) and `alibaba-coding-plan-cn` (opencode's mainland-China one), the Token Plan as `alibaba-token-plan` (opencode and omp), `alibaba-token-plan-cn` (opencode), or `qwen-token-plan`, `qwen-token-plan-cn`, and `qwen-token-plan-individual` (pi), and omp's Qwen OAuth sign-in as `qwen-portal`. Alibaba's pay-as-you-go `alibaba` and `alibaba-cn` keys carry no plan and stay out. Claude Code sent to Model Studio's Anthropic-compatible endpoint (an Alibaba host in `ANTHROPIC_BASE_URL`) counts here as well, by the [Claude Code](#claude-code) rule. + +There is no meter. Alibaba shows a Coding Plan's 5-hour, weekly, and monthly request quotas only on its console page and documents no API to read them, and Qwen Code reads none of its own, so nothing waits on a credential and the tab shows sessions and tokens alone. + +### OpenCode + +One record covers the OpenCode subscription, Zen and Go together: usage on the `opencode` and `opencode-go` providers. opencode sessions on an Anthropic or OpenAI provider stay in the `claude` and `codex` records, so nothing counts twice. Both of opencode's message stores are read (`message` and V2's `session_message`, deduplicated by message id, leaving out the copies a fork makes), from the database `OPENCODE_DB` names when set. + +The Go meters read `OPENCODE_API_KEY` or `OPENCODE_GO_API_KEY` first, then the keys opencode keeps in its database and `auth.json`, then pi's (under `PI_CODING_AGENT_DIR` when set), then omp's, and last the OpenCode Console session, which asks Console's own endpoint with its workspace id. A refused credential hands over to the next one, and when every one is refused the record asks for a new sign-in. A rate-limited window reads as full. A key without a Go plan (Zen pay-as-you-go) is no trouble: the panel shows local token usage alone, and that key isn't asked again for an hour. Without any key, local usage shows the same way. + +### Z.ai + +Z.ai limits need a key that carries a GLM Coding Plan. A pay-as-you-go key reports "no coding plan", which is no trouble: the panel shows local stats alone, and that key is asked again only hourly, a failed recheck staying quiet. The key comes from `~/.config/omarchy/agents/zai.json`, then `ZAI_API_KEY` / `ZHIPU_API_KEY` in the environment, then pi/omp's `.env`, then the key pi or omp saved in `auth.json`. Set `platform` to `"zhipu"` in the config for the China (`open.bigmodel.cn`) host. Local stats count pi, omp, and opencode sessions on Z.ai's and Zhipu's providers (`zai`, `zhipu`, `zai-coding-cn`, `zhipu-coding-plan`, `zhipuai`, `zai-coding-plan`, `zhipuai-coding-plan`), and Claude Code turns sent to `api.z.ai` or `open.bigmodel.cn`, or answered by a `glm-*` model with no `ANTHROPIC_BASE_URL`, by the [Claude Code](#claude-code) rule, and MiniMax Code replies on a custom provider (its Z.AI preset) at those hosts. + +### Ollama Cloud + +Ollama Cloud reads `OLLAMA_API_KEY` first, then `apiKey` in `~/.config/omarchy/agents/ollama.json`, then the key pi or omp signed in with in `~/.pi/agent/auth.json` (or under `PI_CODING_AGENT_DIR`) or `~/.omp/agent/auth.json`. Each key keeps its own last good limits and plan. + +Ollama's balance endpoint (`/api/balance`) reports a legacy plan's session and weekly windows as the percent remaining with their reset times, and a current plan's included credits as a balance against the monthly allowance, shown as an Included credits meter that resets when the period ends. Credits bought on top of the plan (`purchased.balance_usd`, the unexpired ones left) show as a balance of what remains; the answer says nothing of how much was bought, so it raises no low-balance alarm. The limits cover the whole account — every model the key can reach — while the local stats count pi, omp, and opencode sessions on the `ollama-cloud` provider, and cloud-tagged models (`glm-4.6:cloud`, `gpt-oss:120b-cloud`) reached through the local daemon's `ollama` provider, which run on Ollama's servers too, from Codex's `--oss` rollouts as well (with the prompts typed in them, and any model on a Codex provider named `ollama-cloud`), and Claude Code turns sent to `ollama.com` or answered by a cloud-tagged model, by the [Claude Code](#claude-code) rule. 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` is a key of its own, used when `OLLAMA_API_KEY` is unset. + +### MiniMax + +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, pi, or omp stores when MiniMax is signed in there. A key is only ever read from MiniMax's own provider ids, never a look-alike such as a custom gateway. The host follows where the key came from: `mmx`'s own key uses the host `mmx` itself asks (its `base_url`, then its OAuth resource URL, then its region, which `MINIMAX_REGION` overrides), an opencode, pi, or omp key on a `minimax-cn` provider uses the China host, and anything else uses the global `https://api.minimax.io`; `MINIMAX_BASE_URL` overrides them all. The quota request never follows a redirect, so the key stays with MiniMax's host. The windows are the `general` quota's (or the first non-media model's), skipping any model mmx would call not in plan; a window MiniMax marks exhausted reads as full whatever percentage it still carries, and an unlimited one shows no meter. Local stats count MiniMax Code's own sessions under `MINIMAX_DATA_DIR`, `MAVIS_DATA_DIR`, or `~/.minimax` (user messages as prompts, assistant replies as tokens), only replies on MiniMax's own routes (`minimax` and `minimax_api`) and the prompts they answered; a reply on one of MiniMax Code's custom providers (`custom_provider:`, its Z.AI and OpenRouter presets included) is billed by whoever runs that provider's `options.baseURL` in `config.yaml`, so it counts here at MiniMax's own host, under Z.ai, Ollama Cloud, DeepSeek, Kimi or OpenRouter at theirs, under Local models on this machine or the LAN, and nowhere for a gateway that can't be placed or a provider gone from the config (the current config decides for every earlier reply too). pi, omp, and opencode sessions on a MiniMax provider count too. Without a key the panel still shows that local usage and waits for one to show the limits. + +### OpenRouter + +OpenRouter reads `OPENROUTER_API_KEY` first, then the key opencode stores in `~/.local/share/opencode/auth.json`, then the key pi (`~/.pi/agent/auth.json`, honoring `PI_CODING_AGENT_DIR`) or omp (`~/.omp/agent/auth.json`) saved when OpenRouter is signed in there, then `apiKey` in `~/.config/omarchy/agents/openrouter.json`. Without any key, the local pi, omp, and opencode usage on the OpenRouter provider still shows, and the panel waits for a key to show the balance and meter. 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 scan to the account-global `/activity` API; on its own it also reads the balance, with no key-limit meter. OpenRouter documents `/credits` as management-key only, so an inference key it answers with 403 shows its key-limit meter without a balance and is no sign-in to redo; with a management key as well, that key reads the balance instead. + +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. OpenRouter documents that endpoint as management-key only, so it is asked once per refresh, with the management key when there is one; an inference key alone asks until OpenRouter refuses it that way, which shows no balance rather than a refused key, and then isn't asked again for a day. 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. Each key keeps its own last good meter and balance. + +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), from pi/omp sessions whose assistant messages ran on the `openrouter` provider (and Kimchi's harness sessions run with `--provider openrouter`), from DeepSeek Harness attempts on its pi-ai `openrouter` route, from MiniMax Code replies on a custom provider at `openrouter.ai`, and from Claude Code turns while `ANTHROPIC_BASE_URL` points at `openrouter.ai`, by the [Claude Code](#claude-code) rule. 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. + +### DeepSeek + +DeepSeek reads `DEEPSEEK_API_KEY` first, then the key pi stores in `~/.pi/agent/auth.json` (under `PI_CODING_AGENT_DIR` when set), then omp's `~/.omp/agent/auth.json`, never a `!command` or `$VAR` entry, then the key opencode stores in its own `auth.json` when DeepSeek is signed in there, and last the key DeepSeek Harness keeps in `$DSH_HOME/.credentials.yaml` (or the `.env` beside it; `$DSH_HOME` defaults to `~/.dsh`). Each key keeps its own last balance. DeepSeek answers one balance per currency (CNY or USD), and the panel shows the one holding credit, or the first when none does. When DeepSeek says the balance no longer covers API calls (`is_available` false), or it is spent, the status reads Out of credits in the urgent color. Local usage also counts dsh's own session logs under `$DSH_HOME/sessions` on its DeepSeek routes (`deepseek-official`, `deepseek-account`, and its pi-ai adapter's `deepseek` route), leaving out the history a forked session inherits, and Claude Code turns sent to `api.deepseek.com` or answered by a `deepseek-*` model, by the [Claude Code](#claude-code) rule. dsh's other pi-ai routes are billed by that provider: its `openrouter` route counts under OpenRouter, and a hand-declared gateway counts nowhere. 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 + +Fireworks reads `FIREWORKS_API_KEY` and `FIREWORKS_ACCOUNT_ID` first, then `~/.fireworks/auth.ini` (which `firectl set-api-key` creates), then the key pi stores in `~/.pi/agent/auth.json` when Fireworks is signed in there (honoring `PI_CODING_AGENT_DIR`, and pi's literal and `$ENV_VAR` key forms), and finally the key opencode stores in `~/.local/share/opencode/auth.json`. + +The collector first asks the account's `:getBalance` endpoint for the real prepaid ledger. That endpoint exists but is permission-gated, and as of August 2026 no console-issued API key passes it — Fireworks appears to reserve it for the dashboard session. The probe stays because it is cheap and the live figure lights up automatically if Fireworks ever opens it to keys. Until then the collector falls back to estimating the balance from configuration in `~/.config/omarchy/agents/fireworks.json`: ```json { @@ -126,15 +264,25 @@ configuration in `~/.config/omarchy/agents/fireworks.json`: } ``` -Set `fundedAmount` to the credits purchased and optionally `fundedAt` to the -purchase date; with no date, the collector uses the account creation time. It -subtracts rated account costs and the panel labels the result as estimated. -For a later top-up, increase `fundedAmount` by the new credit while keeping -the original `fundedAt`, so both the funding and spend still cover the same -period. `accountId` only matters when one API key can access several -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. +Set `fundedAmount` to the credits purchased and optionally `fundedAt` to the purchase date; with no date, the collector uses the account creation time. It subtracts rated account costs and the panel labels the result as estimated. For a later top-up, increase `fundedAmount` by the new credit while keeping the original `fundedAt`, so both the funding and spend still cover the same period. `accountId` only matters when one API key can access several 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. + +Each API key keeps its account's last good usage rows and balance; another key or account never shows them. Usage is re-read at most every 20 seconds (15 minutes on the panel's near-limit refresh). A rejected key reads as an expired sign-in; with no key anywhere the tab stays hidden. Its stats come from the account's billing API, so they already include this machine's pi and opencode use, and no local scan adds them again. + +### Synthetic + +Synthetic reads `SYNTHETIC_API_KEY` first, then the key pi stores in `~/.pi/agent/auth.json` (honoring `PI_CODING_AGENT_DIR` and `$ENV_VAR` keys), then omp's `~/.omp/agent/auth.json`, then opencode's `auth.json`. The collector asks the documented `GET https://api.synthetic.new/v2/quotas`, which doesn't count against the quotas. The rolling five-hour request limit and the weekly credits regenerate a little at a time, so they show how much is used with no reset time; hourly search and free tool calls reset when Synthetic says they renew. An account without the credit system shows the documented subscription request bucket instead. Local stats count pi and omp sessions and opencode messages on the `synthetic` provider, and through `scan_routed()` the Claude Code, MiniMax Code and Codex usage sent to Synthetic's API at `api.synthetic.new` (its Anthropic-compatible `/anthropic` and OpenAI-compatible `/openai/v1`, as its guides set them up), plus Claude Code replies by an `hf:` or `syn:` model with no `ANTHROPIC_BASE_URL` override. + +### Local models + +The Local models section counts work that ran on hardware you own: Codex rollouts, pi and omp sessions, and opencode messages whose provider is a well-known local runtime (Ollama, LM Studio, llama.cpp, vLLM, ...), a provider named in `~/.config/omarchy/agents/local.json`, or a Codex `model_providers` entry whose `base_url` is on a loopback, private, link-local or LAN address: + +```json +{ "providers": ["homelab", "my-gateway"] } +``` + +A Codex rollout records its provider's id but not the URL it reached, so the collector remembers each `model_providers` endpoint it has seen in `/etc/codex/config.toml`, `config.toml` or a profile's `.config.toml` (in `~/.local/state/omarchy/agents/local-codex-providers.json`), and a provider's history stays local after it leaves the config. Every provider collector reads the same map, so an entry at a provider's API counts there instead, after it leaves the config too. A provider only ever given with `codex -c model_providers...` or in a project's `.codex/config.toml` is never seen there, so name it in `local.json`. + +It is the exact complement of the provider collectors: a provider id another collector counts is never counted here, whatever `local.json` says (that collector counts it, once), but for a Codex `model_providers` entry, whose `base_url` decides whatever its id, so one under a claimed id on your own hardware (a self-hosted Kimchi gateway's `kimchi codex` setup) counts here; and the Codex collector keeps every rollout on the built-in `openai` provider, but for the models of one that a self-hosted `openai_base_url` served (see [Codex](#codex)). Two providers are claimed only for some models, and only those rows are left out: Ollama's cloud models (a `:cloud` or `-cloud` tag) run on Ollama's servers even through the local daemon, so they count under Ollama Cloud instead, and opencode's Antigravity plugin records a `google` provider with `antigravity-` models, which the Antigravity collector counts. A Codex rollout's prompts are the ones typed, and a subagent's or guardian's thread starts no session, read as the Codex collector reads them. Claude Code counts here while `ANTHROPIC_BASE_URL` is self-hosted (Ollama's `http://localhost:11434`, a LAN box), for every model but a Claude one, which a local gateway forwards on to Anthropic, and a cloud-tagged one, by the [Claude Code](#claude-code) rule; so do MiniMax Code replies on a custom provider whose `baseURL` is self-hosted, and, by the same rule, Kimi Code, Qwen Code, Hermes and Mistral Vibe requests whose recorded base URL is self-hosted, read by each one's own collector. There are no limits or balance. ## Interactions @@ -186,7 +334,8 @@ edit `shell.json` directly): omarchy bar set omarchy.agents providers '{ "claude": { "enabled": true }, "codex": { "enabled": false }, - "fireworks": { "enabled": true } + "fireworks": { "enabled": true }, + "kimchi": { "enabled": true } }' --json ``` @@ -204,5 +353,5 @@ the same account synced from two machines is not counted twice. One caveat on "all-time": the Codex collector only reads native session files touched in the last 30 days, and Fireworks requests the last 30 days from its -billing API, so their totals and day counts cover that window. Claude's cover -every transcript still on disk. +billing API, so their totals and day counts cover that window. Claude's and +Muse's cover every session still on disk. diff --git a/shell/plugins/agents/assets/agy-light.svg b/shell/plugins/agents/assets/agy-light.svg new file mode 100644 index 00000000000..cf46f512c09 --- /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..4b3d3554217 --- /dev/null +++ b/shell/plugins/agents/assets/agy.svg @@ -0,0 +1 @@ +Antigravity diff --git a/shell/plugins/agents/assets/copilot-light.svg b/shell/plugins/agents/assets/copilot-light.svg new file mode 100644 index 00000000000..717e3d72142 --- /dev/null +++ b/shell/plugins/agents/assets/copilot-light.svg @@ -0,0 +1 @@ +GitHub Copilot \ No newline at end of file diff --git a/shell/plugins/agents/assets/copilot.svg b/shell/plugins/agents/assets/copilot.svg new file mode 100644 index 00000000000..4edbc713f39 --- /dev/null +++ b/shell/plugins/agents/assets/copilot.svg @@ -0,0 +1 @@ +GitHub Copilot \ No newline at end of file diff --git a/shell/plugins/agents/assets/cursor-light.svg b/shell/plugins/agents/assets/cursor-light.svg new file mode 100644 index 00000000000..ecf8c981852 --- /dev/null +++ b/shell/plugins/agents/assets/cursor-light.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/shell/plugins/agents/assets/cursor.svg b/shell/plugins/agents/assets/cursor.svg new file mode 100644 index 00000000000..8652a86f9c0 --- /dev/null +++ b/shell/plugins/agents/assets/cursor.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/shell/plugins/agents/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/devin-light.svg b/shell/plugins/agents/assets/devin-light.svg new file mode 100644 index 00000000000..17611b9db50 --- /dev/null +++ b/shell/plugins/agents/assets/devin-light.svg @@ -0,0 +1,3 @@ + + + diff --git a/shell/plugins/agents/assets/devin.svg b/shell/plugins/agents/assets/devin.svg new file mode 100644 index 00000000000..e2a93b3bcf5 --- /dev/null +++ b/shell/plugins/agents/assets/devin.svg @@ -0,0 +1,3 @@ + + + diff --git a/shell/plugins/agents/assets/fireworks.svg b/shell/plugins/agents/assets/fireworks.svg index a659bed9fa9..018709fe4f2 100644 --- a/shell/plugins/agents/assets/fireworks.svg +++ b/shell/plugins/agents/assets/fireworks.svg @@ -1,7 +1 @@ - - - - - - - +Fireworks AI diff --git a/shell/plugins/agents/assets/grokbot-light.svg b/shell/plugins/agents/assets/grokbot-light.svg new file mode 100644 index 00000000000..92d54f32957 --- /dev/null +++ b/shell/plugins/agents/assets/grokbot-light.svg @@ -0,0 +1 @@ +Grok Bot diff --git a/shell/plugins/agents/assets/grokbot.svg b/shell/plugins/agents/assets/grokbot.svg new file mode 100644 index 00000000000..80eb13cad8d --- /dev/null +++ b/shell/plugins/agents/assets/grokbot.svg @@ -0,0 +1 @@ +Grok Bot diff --git a/shell/plugins/agents/assets/hermes-light.svg b/shell/plugins/agents/assets/hermes-light.svg new file mode 100644 index 00000000000..46f08ad1fb5 --- /dev/null +++ b/shell/plugins/agents/assets/hermes-light.svg @@ -0,0 +1,4 @@ + + Hermes Agent + + diff --git a/shell/plugins/agents/assets/hermes.svg b/shell/plugins/agents/assets/hermes.svg new file mode 100644 index 00000000000..a7706742202 --- /dev/null +++ b/shell/plugins/agents/assets/hermes.svg @@ -0,0 +1,4 @@ + + Hermes Agent + + diff --git a/shell/plugins/agents/assets/kimchi.svg b/shell/plugins/agents/assets/kimchi.svg new file mode 100644 index 00000000000..f2fb29de6eb --- /dev/null +++ b/shell/plugins/agents/assets/kimchi.svg @@ -0,0 +1,4 @@ + + + + diff --git a/shell/plugins/agents/assets/kimi-light.svg b/shell/plugins/agents/assets/kimi-light.svg new file mode 100644 index 00000000000..c7bfa27f3a9 --- /dev/null +++ b/shell/plugins/agents/assets/kimi-light.svg @@ -0,0 +1 @@ +Kimi diff --git a/shell/plugins/agents/assets/kimi.svg b/shell/plugins/agents/assets/kimi.svg new file mode 100644 index 00000000000..4222df4ff97 --- /dev/null +++ b/shell/plugins/agents/assets/kimi.svg @@ -0,0 +1 @@ +Kimi diff --git a/shell/plugins/agents/assets/minimax-light.svg b/shell/plugins/agents/assets/minimax-light.svg new file mode 100644 index 00000000000..290d7c2411f --- /dev/null +++ b/shell/plugins/agents/assets/minimax-light.svg @@ -0,0 +1 @@ + diff --git a/shell/plugins/agents/assets/minimax.svg b/shell/plugins/agents/assets/minimax.svg new file mode 100644 index 00000000000..74a90280532 --- /dev/null +++ b/shell/plugins/agents/assets/minimax.svg @@ -0,0 +1 @@ + diff --git a/shell/plugins/agents/assets/muse-light.svg b/shell/plugins/agents/assets/muse-light.svg new file mode 100644 index 00000000000..0850837cf9c --- /dev/null +++ b/shell/plugins/agents/assets/muse-light.svg @@ -0,0 +1 @@ +Muse diff --git a/shell/plugins/agents/assets/muse.svg b/shell/plugins/agents/assets/muse.svg new file mode 100644 index 00000000000..daf5cc78f53 --- /dev/null +++ b/shell/plugins/agents/assets/muse.svg @@ -0,0 +1 @@ +Muse diff --git a/shell/plugins/agents/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/opencode-light.svg b/shell/plugins/agents/assets/opencode-light.svg new file mode 100644 index 00000000000..b79140a5070 --- /dev/null +++ b/shell/plugins/agents/assets/opencode-light.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/shell/plugins/agents/assets/opencode.svg b/shell/plugins/agents/assets/opencode.svg new file mode 100644 index 00000000000..b79c7332e20 --- /dev/null +++ b/shell/plugins/agents/assets/opencode.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/shell/plugins/agents/assets/openrouter-light.svg b/shell/plugins/agents/assets/openrouter-light.svg new file mode 100644 index 00000000000..34972d58fb4 --- /dev/null +++ b/shell/plugins/agents/assets/openrouter-light.svg @@ -0,0 +1 @@ +OpenRouter diff --git a/shell/plugins/agents/assets/openrouter.svg b/shell/plugins/agents/assets/openrouter.svg new file mode 100644 index 00000000000..47b7e653851 --- /dev/null +++ b/shell/plugins/agents/assets/openrouter.svg @@ -0,0 +1 @@ +OpenRouter diff --git a/shell/plugins/agents/assets/qwen.svg b/shell/plugins/agents/assets/qwen.svg new file mode 100644 index 00000000000..5f2a8899839 --- /dev/null +++ b/shell/plugins/agents/assets/qwen.svg @@ -0,0 +1 @@ +Qwen Code diff --git a/shell/plugins/agents/assets/synthetic-light.svg b/shell/plugins/agents/assets/synthetic-light.svg new file mode 100644 index 00000000000..5a8ae34ab38 --- /dev/null +++ b/shell/plugins/agents/assets/synthetic-light.svg @@ -0,0 +1,13 @@ + + + + + + + + + + + + + \ No newline at end of file diff --git a/shell/plugins/agents/assets/synthetic.svg b/shell/plugins/agents/assets/synthetic.svg new file mode 100644 index 00000000000..8a4eb744b7d --- /dev/null +++ b/shell/plugins/agents/assets/synthetic.svg @@ -0,0 +1,13 @@ + + + + + + + + + + + + + \ No newline at end of file diff --git a/shell/plugins/agents/assets/vibe.svg b/shell/plugins/agents/assets/vibe.svg new file mode 100644 index 00000000000..fd905ddb705 --- /dev/null +++ b/shell/plugins/agents/assets/vibe.svg @@ -0,0 +1 @@ +Mistral Vibe \ No newline at end of file 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/collectors.py b/shell/plugins/agents/collectors.py new file mode 100644 index 00000000000..15113e88fa9 --- /dev/null +++ b/shell/plugins/agents/collectors.py @@ -0,0 +1,3795 @@ +"""The plumbing every agent usage collector shares. + +Each `bin/omarchy-agent-usage-` keeps what is its own: where its +credentials live, the endpoint it asks and how to read the answer, which +providers' tokens it counts, and the quirks of its record. Everything the +collector contract (docs/agent-usage-collectors.md) asks of all of them lives +here, once: the cache files, how untrusted numbers and times are read, the +record's shape, the limits cache, and the readers for the stores several +collectors split between them by provider: pi, omp and opencode sessions, +Claude Code transcripts, Codex rollouts, DeepSeek Harness logs and MiniMax +Code sessions, with the one rule each split follows. + +It ships with the agents panel the records feed, and a collector loads it +from `$OMARCHY_PATH/shell/plugins/agents`. Nothing here prints a record, reads +a credential, or talks to a service on its own. +""" + +from __future__ import annotations + +import base64 +import contextlib +import fcntl +import hashlib +import http.client +import ipaddress +import json +import math +import os +import re +import shutil +import sqlite3 +import stat +import subprocess +import sys +import tempfile +import time +import tomllib +import urllib.error +import urllib.parse +import urllib.request +from collections import Counter +from datetime import date, datetime, timedelta, timezone +from pathlib import Path +from typing import Any, Callable, Iterable, Iterator, NamedTuple + +try: + from compression import zstd +except ImportError: + zstd = None + +# 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. +# --force reuses nothing. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +# A panel opened and closed repeatedly must not become one request per flick, +# so a limits answer this recent stands in for a probe; --force skips it. +PROBE_MIN_INTERVAL_SECONDS = 15 + +# The scan cache's envelope. A file of another version, or another shape, is +# a miss: rescanned and rewritten, never a crash or a garbage record. +SCAN_CACHE_VERSION = 2 +STATS_KEYS = ("todayPrompts", "todayTotalTokens", "recentDays", "activeDates", "modelUsage") + +# Counts and times read from disk or the network are untrusted: a boolean, +# NaN, negative, or absurdly large value reads as nothing, never as a number +# that would swamp the totals. The ceiling sits well above any epoch +# millisecond time. +NUMBER_CEILING = 10 ** 15 + +# Bodies larger than this are no answer any collector expects. +HTTP_BODY_LIMIT = 1_048_576 + + +# ------------------------------------------------------------- diagnostics + + +def program() -> str: + name = Path(sys.argv[0]).name if sys.argv and sys.argv[0] else "" + return name if name.startswith("omarchy-agent-usage-") else "omarchy-agent-usage" + + +def warn(message: str) -> None: + """A diagnostic on stderr, named for the collector. Never a credential.""" + print(f"{program()}: {message}", file=sys.stderr) + + +# ----------------------------------------------------------------- runtime + + +def runtime_env() -> dict[str, str]: + """The environment with the places agent CLIs and rg install to on PATH.""" + home = str(Path.home()) + path_parts = [ + os.environ.get("PATH", ""), + f"{home}/.local/bin", + f"{home}/.npm-global/bin", + f"{home}/.local/share/mise/shims", + ] + env = os.environ.copy() + env["PATH"] = os.pathsep.join(part for part in path_parts if part) + return env + + +def find_command(name: str) -> str | None: + return shutil.which(name, path=runtime_env().get("PATH")) + + +# --------------------------------------------------------- untrusted values + + +def obj(value: Any) -> dict[str, Any]: + """A nested field of the wrong shape (a list of tokens, a numeric time) + reads as empty rather than failing the row, or the whole scan with it.""" + return value if isinstance(value, dict) else {} + + +def number(value: Any) -> int: + """A count or time from disk or the network, 0 when it can't be trusted.""" + if isinstance(value, bool): + return 0 + try: + n = float(value or 0) + except (TypeError, ValueError, OverflowError): + return 0 + if not math.isfinite(n) or n < 0 or n >= NUMBER_CEILING: + return 0 + return round(n) + + +def token_count(value: Any) -> int: + """A token count from disk, 0 unless it is a plain JSON number number() + trusts: a count written as a string is as untrusted as a boolean.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + return number(value) + + +def usage_tokens(usage: dict[str, Any], *keys: str) -> int: + """pi, omp, and the harnesses that copy their format spell usage camelCase + or snake_case; the first spelling present wins.""" + for key in keys: + if key in usage: + return token_count(usage.get(key)) + return 0 + + +def date_string(value: datetime) -> str: + return value.strftime("%Y-%m-%d") + + +def today() -> str: + return date_string(datetime.now()) + + +def recent_dates() -> list[str]: + """The last seven local days, oldest first, today last.""" + now = datetime.now() + return [date_string(now - timedelta(days=offset)) for offset in range(6, -1, -1)] + + +def local_day(value: Any) -> str | None: + """The local day a row was spent on, or None when its time can't be + trusted, so the caller skips that row rather than the whole scan. + + A row with no time at all is today's. Epoch seconds and milliseconds, a + string of digits, and ISO times are read; zero, a huge or negative number, + NaN, a boolean, a date outside 2000-2100, and anything else are untrusted. + """ + if value is None: + return today() + # ASCII digits only, and few enough to be a time: isdigit() also passes + # other scripts' digits and int() refuses strings past 4300 digits. + if isinstance(value, str) and re.fullmatch(r"[0-9]{1,20}", value.strip()): + value = int(value.strip()) + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + # number() turns untrusted values (NaN, negative, huge) into 0; the time + # itself is floored, as stamps are, so the last half second of a day + # isn't read as the next one. + if number(value) <= 0: + return None + seconds = math.floor(value) + # pi, omp, and opencode write milliseconds. + if seconds > 10_000_000_000: + seconds = seconds / 1000 + try: + parsed = datetime.fromtimestamp(seconds) + except (OverflowError, OSError, ValueError): + return None + else: + try: + parsed = datetime.fromisoformat(str(value).strip().replace("Z", "+00:00")) + if parsed.tzinfo is not None: + parsed = parsed.astimezone() + except (OverflowError, OSError, ValueError): + return None + # No agent wrote a row before 2000 or after 2100: a time out there is a + # damaged stamp (a stray 1 reads as 1970), not usage to count. + if not 2000 <= parsed.year <= 2100: + return None + return date_string(parsed) + + +# -------------------------------------------------------------------- time + + +def now_ms() -> int: + """Now as an epoch-ms stamp, floored: a stamp rounded half a millisecond + ahead reads as a clock that went back, and the next run probes again.""" + return int(time.time() * 1000) + + +def stamp(value: Any) -> int: + """An epoch-ms stamp a collector wrote, read back without trusting it.""" + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return 0 + return max(0, int(value)) + + +def age_below(stamp_seconds: float, limit_seconds: float) -> bool: + """Whether a stamp is younger than the limit. A negative age is a stamp + from the future: the clock moved back since it was written, so it proves + nothing about freshness.""" + age = time.time() - stamp_seconds + return 0 <= age < limit_seconds + + +# ------------------------------------------------------------------- files + + +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 data_home() -> Path: + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + + +def config_home() -> Path: + return Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) + + +def read_json(path: Path) -> dict[str, Any]: + """A JSON object from disk, {} when missing, unreadable, or not an object. + + Read whatever the file's time: a clock that went back must not lose the + last good answer. Freshness is judged by the stamps stored inside. + """ + try: + parsed = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + return parsed if isinstance(parsed, dict) else {} + + +def write_json(path: Path, payload: Any, sort_keys: bool = False) -> None: + """Replace `path` atomically with `payload` as compact JSON.""" + # 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=sort_keys) + "\n") + # mkstemp opens at 0600; nothing in the cache is a credential, so open it + # up to the usual 0644. + 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. + + A lock file that isn't a regular file of ours, or one others may write, is + never trusted with the lock; the block still runs, since the lock only + dedups concurrent runs. + """ + 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) + + +# ------------------------------------------------------------------ record + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +def empty_stats() -> dict[str, Any]: + """The record's stats block with nothing in it: always present, zeroed.""" + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [{"date": day, "messageCount": 0} for day in recent_dates()], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + } + + +def base_record(agent_id: str, name: str, **overrides: Any) -> dict[str, Any]: + """A record with every field, empty stats, no limits and no balance. + + As it stands it is what a machine that doesn't use the agent prints: ready + false, so the panel skips it, and no status, since that is no trouble. + """ + record: dict[str, Any] = { + "schemaVersion": 1, + "id": agent_id, + "name": name, + "updatedAt": datetime.now(timezone.utc).isoformat(), + "ready": False, + "hasLocalStats": False, + "hasPromptStats": True, + "tierLabel": "", + "usageStatusText": "", + "authHelpText": "", + "limits": [], + "limitsStale": False, + "limitsFetchedAt": 0, + } + record.update(empty_stats()) + record.update(overrides) + return record + + +def _finite(value: Any) -> Any: + if isinstance(value, float) and not math.isfinite(value): + return 0 + if isinstance(value, dict): + return {key: _finite(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_finite(item) for item in value] + return value + + +def emit(record: dict[str, Any], sort_keys: bool = False) -> None: + """Print the record, the one thing on stdout. + + A number that overflowed to infinity (a huge balance read with float()) + prints as 0: Infinity and NaN aren't JSON, and the panel would drop the + whole record over one. + """ + print(json.dumps(_finite(record), separators=(",", ":"), sort_keys=sort_keys, allow_nan=False)) + + +class Stats: + """Per-day and per-model totals of every counted reply, turned into the + record's stats fields by snapshot(). + + `today` is the local day "today" means: today itself unless a caller hands + in another (a date or an ISO day), for rows summarized against a day of + their own. + """ + + def __init__(self, today: date | str | None = None) -> None: + if today is None: + day = datetime.now().date() + elif isinstance(today, str): + day = date.fromisoformat(today) + else: + day = today + self.today = date_string(day) + self.recent_dates = [date_string(day - timedelta(days=offset)) for offset in range(6, -1, -1)] + self.recent = {day: 0 for day in self.recent_dates} + self.today_prompts = 0 + self.today_tokens = 0 + self.today_sessions: set[str] = set() + self.today_by_model: dict[str, int] = {} + self.total_prompts = 0 + self.sessions: set[str] = set() + self.active_dates: set[str] = set() + self.model_usage: dict[str, dict[str, int]] = {} + + def add( + self, + day: str | None, + session: str | None, + model: str, + input_tokens: int, + output_tokens: int, + cache_read: int, + cache_write: int, + prompts: int = 1, + ) -> None: + """Tokens `model` spent on `day` in `session`, and the prompts they + answered: one reply is one prompt, the default. + + A store that books prompts apart from the tokens they spent passes + prompts=0 here and counts each prompt with add_prompts(); one that keeps + running totals passes how many prompts they cover. A day of None is a + row with no date to trust that still counts toward the lifetime totals, + and a session of None is tokens no session is known for. With no tokens + there is no model bucket, and with neither tokens nor prompts there is + nothing to count at all. + """ + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0 and prompts <= 0: + return + self.total_prompts += prompts + if session is not None: + self.sessions.add(session) + if day is not None: + self.active_dates.add(day) + if total > 0: + bucket = self.model_usage.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + if day in self.recent: + # recentDays' messageCount holds tokens, despite the name synced + # snapshots share. + self.recent[day] += total + if day is not None and day == self.today: + self.today_prompts += prompts + if session is not None: + self.today_sessions.add(session) + self.today_tokens += total + if total > 0: + self.today_by_model[model] = self.today_by_model.get(model, 0) + total + + def add_prompts(self, day: str | None, session: str | None, count: int = 1) -> None: + """`count` prompts typed on `day` in `session`, with no tokens of their + own: for a store that books what was asked apart from what it spent. + + A count of 0 still says the session was in use that day. A day or a + session of None reads as it does in add(). + """ + self.total_prompts += count + if session is not None: + self.sessions.add(session) + if day is not None: + self.active_dates.add(day) + if day == self.today: + self.today_prompts += count + if session is not None: + self.today_sessions.add(session) + + def add_row(self, row: UsageRow, model: str) -> None: + """A row from one of the shared readers, under the model name the + collector shows for it. A row with no tokens books its prompts alone, + and with none of those still says its session was in use that day.""" + if row.input_tokens + row.output_tokens + row.cache_read + row.cache_write <= 0: + self.add_prompts(row.day, row.session, row.prompts) + return + self.add(row.day, row.session, model, row.input_tokens, row.output_tokens, row.cache_read, row.cache_write, row.prompts) + + def snapshot(self) -> dict[str, Any]: + return { + "todayPrompts": self.today_prompts, + "todaySessions": len(self.today_sessions), + "todayTotalTokens": self.today_tokens, + "todayTokensByModel": self.today_by_model, + "recentDays": [{"date": day, "messageCount": self.recent[day]} for day in self.recent_dates], + "totalPrompts": self.total_prompts, + "totalSessions": len(self.sessions), + # Days with any recorded usage, for the all-time "N days" summary. The + # dates travel too: merging snapshots from several machines needs + # their union, which a count alone cannot give. + "activeDays": len(self.active_dates), + "activeDates": sorted(self.active_dates), + "modelUsage": self.model_usage, + } + + +# -------------------------------------------------------------- scan cache + + +def scan_max_age(force: bool, limits_only: bool) -> float: + """How old a cached scan this run may reuse.""" + if force: + return 0 + return LIMITS_ONLY_REUSE_SECONDS if limits_only else SCAN_REUSE_SECONDS + + +def scan_cache_paths(agent_id: str, sources: Iterable[Any]) -> tuple[Path, Path]: + """The scan cache and its lock, keyed by every data path the scan reads, so + moving a source is a miss.""" + identity = "\0".join(str(source) for source in sources) + digest = hashlib.sha1(identity.encode("utf-8", "surrogatepass")).hexdigest()[:16] + root = cache_root() + return root / f"{agent_id}-scan-{digest}.json", root / f"{agent_id}-scan-{digest}.lock" + + +def read_cached_scan(cache_file: Path, max_age: float, match: dict[str, Any] | None = None) -> dict[str, Any] | None: + """The stats a cached scan holds, or None when it can't stand in for one.""" + cached = read_json(cache_file) + if cached.get("schemaVersion") != SCAN_CACHE_VERSION: + return None + # today* fields only mean "today" on the day they were scanned. A cache + # from another local date (midnight passed, or the clock moved) is a miss, + # not merely old. + if cached.get("scanDate") != today(): + return None + # A scan only answers for the options it applied (a model filter, say). + if any(cached.get(key) != value for key, value in (match or {}).items()): + return None + if not age_below(stamp(cached.get("scannedAtMs")) / 1000, max_age): + return None + stats = cached.get("stats") + if not isinstance(stats, dict) or not all(key in stats for key in STATS_KEYS): + return None + return stats + + +def cached_scan( + agent_id: str, + sources: Iterable[Any], + max_age: float, + scan: Callable[[], tuple[dict[str, Any], bool]], + match: dict[str, Any] | None = None, + with_complete: bool = False, +) -> Any: + """Local stats, from a scan cached per day, or from `scan` under flock. + + `scan` returns (stats, complete). A scan that couldn't read all its sources + still serves this run, but caching it would hide the missing usage from + every reader until the cache expires, so only a complete one is kept. + `match` names options the scan applied, stored with it and required to + match on reuse. + + With `with_complete` the answer is (stats, complete), for a record that + says whether its stats cover every source; a cached scan was a complete + one. + + The cache is a pure optimization: any failure in it (an unwritable cache + root, a lock error, a full disk) degrades to a direct scan and a warning. + The JSON record is the contract; the cache is not. + """ + sources = list(sources) + match = dict(match or {}) + try: + cache_file, lock_file = scan_cache_paths(agent_id, sources) + cached = read_cached_scan(cache_file, max_age, match) + if cached is not None: + return (cached, True) if with_complete else cached + lock = exclusive_lock(lock_file) + except Exception as exc: + warn(f"cache unavailable ({exc}); scanning directly") + stats, complete = scan() + return (stats, complete) if with_complete else stats + + with lock: + # Another run may have scanned while this one waited for the lock. + cached = read_cached_scan(cache_file, max_age, match) + if cached is not None: + return (cached, True) if with_complete else cached + stats, complete = scan() + if complete: + try: + write_json(cache_file, {"schemaVersion": SCAN_CACHE_VERSION, "scanDate": today(), "scannedAtMs": now_ms(), **match, "stats": stats}) + except Exception as exc: + warn(f"could not write usage cache ({exc})") + return (stats, complete) if with_complete else stats + + +# ------------------------------------------------------- pi and omp sessions + + +class UsageRow(NamedTuple): + """One counted reply from a shared store, normalized. + + `session` is the session it counts toward: the realpath of a pi/omp + session file, or "opencode:". `model` is the model id as the + store wrote it, for the collector to name. `prompts` is how many prompts + it answered: tokens spent on a summary answered none. + """ + + day: str + session: str + provider: str + model: str + input_tokens: int + output_tokens: int + cache_read: int + cache_write: int + prompts: int = 1 + + +# accept(provider, model, entry) decides whether a reply bills this collector. +# `entry` is the pi/omp message object, or the opencode row's data, for a +# collector whose filter needs more than the provider and model. +Accept = Callable[[str, str, dict], bool] + + +def pi_session_roots() -> list[Path]: + """pi's and omp's session trees, each profile's too. + + `--profile=` moves the whole agent tree under + /profiles//. pi moves its agent tree with PI_CODING_AGENT_DIR + and its sessions alone with PI_CODING_AGENT_SESSION_DIR, both expanding a + leading ~; either is read as well as the default trees, which keep + whatever was written before the move. A `sessionDir` in an agent tree's + settings.json moves its sessions too and is read the same way; a relative + one resolves against the cwd pi started in, which can't be known here. + """ + agent_dirs: list[Path] = [] + for base in (Path.home() / ".pi", Path.home() / ".omp"): + agent_dirs.append(base / "agent") + try: + # Sorted so the scan order does not depend on directory order. + agent_dirs.extend(sorted(child / "agent" for child in (base / "profiles").iterdir() if child.is_dir())) + except OSError: + # No profiles directory, or it is unreadable: the default root stands. + pass + if os.environ.get("PI_CODING_AGENT_DIR", "").strip(): + agent_dirs.append(Path(os.environ["PI_CODING_AGENT_DIR"].strip()).expanduser()) + roots = [agent_dir / "sessions" for agent_dir in agent_dirs] + for agent_dir in agent_dirs: + configured = read_json(agent_dir / "settings.json").get("sessionDir") + if isinstance(configured, str) and configured.strip(): + configured_root = Path(configured.strip()).expanduser() + if configured_root.is_absolute(): + roots.append(configured_root) + if os.environ.get("PI_CODING_AGENT_SESSION_DIR", "").strip(): + roots.append(Path(os.environ["PI_CODING_AGENT_SESSION_DIR"].strip()).expanduser()) + return canonical_roots(roots) + + +def canonical_roots(roots: Iterable[Path]) -> list[Path]: + """Session roots by realpath, so an override that aliases another root + through a symlink is scanned under one spelling. A root inside another is + already read with it, so it is dropped rather than counted twice.""" + canonical = list(dict.fromkeys(Path(root).resolve() for root in roots)) + return [root for root in canonical if not any(other != root and root.is_relative_to(other) for other in canonical)] + + +def pi_parent_session(path: str) -> str | None: + """The session file a fork was cut from, as its header line names it.""" + try: + # A header can name any path; opening a FIFO or a device there would + # hang the scan, so only a regular file is read. + if not stat.S_ISREG(os.stat(path).st_mode): + return None + with open(path, encoding="utf-8", errors="replace") as handle: + header = json.loads(handle.readline()) + except (OSError, ValueError, RecursionError): + # RecursionError: a damaged header nested deeper than the parser goes. + return None + parent = header.get("parentSession") if isinstance(header, dict) else None + if not isinstance(parent, str) or not parent: + return None + # pi deletes a session from /resume, leaving its forks naming it. A parent + # that is gone still ties its forks together, so a prompt they both copied + # counts once; only reading it is skipped. + return os.path.realpath(os.path.expanduser(parent)) + + +def _rg_text(field: Any) -> str: + # rg hands back text that is not valid UTF-8 as base64 bytes. + field = obj(field) + if isinstance(field.get("text"), str): + return field["text"] + return os.fsdecode(base64.b64decode(field.get("bytes") or "")) + + +# Spent tokens with no provider of their own: pi writes a compaction's and a +# branch summary's on the entry, and a tool's nested model work on its +# result, all as the session's model ran them, and counts them in the +# session's totals. +_PI_SUMMARIES = ("compaction", "branch_summary") + + +def _pi_line(entry: Any) -> tuple | None: + """What one session line says, for the reader to fold in file order: + ("model", provider, model, reference) for a model switch, ("reply", ...) + for an assistant message, which also names the model the session is on, + ("spent", provider, model, entry, usage) for usage with its own provider, + ("summary", usage) for usage spent on the session's model, ("prompt",) + for a typed prompt, or None.""" + if not isinstance(entry, dict): + return None + kind = entry.get("type") + if kind == "message": + message = obj(entry.get("message")) + role = message.get("role") + if role == "assistant": + provider = str(message.get("provider") or "") + model = str(message.get("model") or "") + # Only what a collector's filter reads is kept, not the reply's text. + reference = {key: value for key, value in message.items() if key != "content"} + return "reply", provider, model, reference, obj(message.get("usage")) + if role == "toolResult" and isinstance(message.get("usage"), dict): + return "summary", message["usage"] + # omp marks the user messages it writes itself (a handoff, a compaction + # summary) as the agent's. + if role == "user" and message.get("attribution") != "agent" and message.get("synthetic") is not True: + return ("prompt",) + return None + if kind == "model_change": + # pi names the provider and model apart; omp joins them as provider/model. + provider = str(entry.get("provider") or "") + model = str(entry.get("modelId") or "") + if not provider and isinstance(entry.get("model"), str): + provider, _, model = entry["model"].partition("/") + return ("model", provider, model, {"provider": provider, "model": model}) if provider else None + if kind == "usage": + return "spent", str(entry.get("provider") or ""), str(entry.get("model") or ""), entry, obj(entry.get("usage")) + if kind in _PI_SUMMARIES and isinstance(entry.get("usage"), dict): + return "summary", entry["usage"] + return None + + +def _pi_unsent(message: dict) -> bool: + """Whether a failed omp reply never reached its provider. + + omp times every request it makes, a failed one included (`duration`, and + `errorStatus` for an HTTP answer), so an error it writes with its + `errorId` and neither, like a refused local credential, was raised before + sending. pi writes no `errorId`, and an aborted request may have gone out. + """ + return message.get("stopReason") == "error" and "errorId" in message and not any(key in message for key in ("duration", "errorStatus", "responseId")) + + +def _pi_key(entry: dict, timestamp: Any, path: str) -> tuple: + """The dedup key of one session line. + + A reply with a response id is that response, at the time it came. A + fork copies its parent's entries with their id and timestamp, so the + pair names a message across files; an id alone only within its file. A + line without an id has nothing tying it to a copy elsewhere, so it is its + own and always counts. + """ + message = obj(entry.get("message")) + response = message.get("responseId") + if message.get("role") == "assistant" and isinstance(response, str) and response and message.get("timestamp") is not None: + # omp writes a reply that overflowed the model's context window twice, + # as siblings with new ids; the response they repeat, and a fork's copy + # of it, is the one request the provider served. + return ("response", str(message.get("provider") or ""), response, str(message["timestamp"])) + message_id = str(entry.get("id") or "") + if message_id and timestamp: + return ("entry", message_id, str(timestamp)) + if message_id: + return ("path", path, message_id) + return ("line", path, object()) + + +def _pi_row(entry: dict, timestamp: Any, path: str, provider: str, model: str, usage: dict, prompts: int) -> tuple[tuple, UsageRow] | None: + """(dedup key, row) for usage one session line spent, or None.""" + input_tokens = usage_tokens(usage, "input", "inputTokens", "input_tokens") + # pi's `reasoning` and omp's `reasoningTokens` are the part of output that + # was thinking, already in it, so they are never added again. + output_tokens = usage_tokens(usage, "output", "outputTokens", "output_tokens") + cache_read = usage_tokens(usage, "cacheRead", "cacheReadTokens", "cache_read", "cache_read_input_tokens") + cache_write = usage_tokens(usage, "cacheWrite", "cacheWriteTokens", "cache_write", "cache_creation_input_tokens") + if not (input_tokens or output_tokens or cache_read or cache_write): + # Some writers keep only the total; it was still spent. + input_tokens = usage_tokens(usage, "totalTokens", "total_tokens") + if not input_tokens: + return None + # The first time present decides: none at all is today's, and an untrusted + # one, zero included, skips only this line. + day = local_day(timestamp) + if day is None: + return None + return _pi_key(entry, timestamp, path), UsageRow(day, path, provider, model, input_tokens, output_tokens, cache_read, cache_write, prompts) + + +def _pi_rows(path: str, lines: list[tuple[int, Any]], accept: Accept) -> Iterator[tuple[int, tuple, UsageRow, tuple, bool]]: + """(line number, dedup key, row, prompt keys, whether `accept` takes it) + for each (line number, entry) of one session file, in order, whoever's, + the prompt keys naming the typed prompts the row's `prompts` count. + + `lines` hold every model switch, user and assistant message in the file. + Usage with no provider of its own bills the model the session was on: the + last model switch or reply before it. With neither, nothing says who + billed it, and it is skipped. A typed prompt is booked on the first reply + that spent tokens answering it, so tool steps and summaries answer none, + and a prompt another provider answered stays with that one. One that no + reply answered before the next prompt or the file's end (its request + failed, or none came) is a row of its own without tokens, on the model + the session was set to use for it: the failed reply, which pi and omp + write with its provider and no usage, else the last model switch or reply + before it; with neither, or when omp refused the request before sending + it, it counts nowhere. A prompt counts on its own day, + not the reply's, and one whose time can't be trusted doesn't count. A file + with no user messages at all, as only a writer copying pi's format leaves, + keeps a prompt per reply. + """ + parsed = [(line_number, entry, _pi_line(entry)) for line_number, entry in lines] + typed = any(line == ("prompt",) for _, _, line in parsed) + # The model the session is on, from the last switch or reply. + context: tuple | None = None + # Each typed prompt no reply has answered yet: (line number, dedup key, + # day or None for an untrusted time, the model the session was on). + asked: list[tuple[int, tuple, str | None, tuple | None]] = [] + # The first reply since those prompts that spent nothing: a failed request, + # or () for one omp refused before sending it. + failed: tuple | None = None + + def takes(provider: str, model: str, reference: dict) -> bool: + try: + return bool(accept(provider, model, reference)) + except Exception: + return False + + def unanswered() -> Iterator[tuple[int, tuple, UsageRow, tuple, bool]]: + for line_number, key, day, before in asked: + target = before if failed is None else failed + if day is None or not target: + continue + provider, model, reference = target + yield line_number, ("asked", *key), UsageRow(day, path, provider, model, 0, 0, 0, 0, 1), (key,), takes(provider, model, reference) + + for line_number, entry, line in parsed: + try: + if line is None: + continue + kind = line[0] + if kind == "model": + context = tuple(line[1:4]) + continue + # The entry's own time, else the message's. + timestamp = entry.get("timestamp") + if timestamp is None: + timestamp = obj(entry.get("message")).get("timestamp") + if kind == "prompt": + # Only a prompt no reply answered is still asked: it is its own row, + # so a fork that copies it books it under the same key. One with no + # model known yet rides on the next reply, as before. + if asked: + yield from unanswered() + asked = [prompt for prompt in asked if failed is None and prompt[3] is None] + failed = None + asked.append((line_number, _pi_key(entry, timestamp, path), local_day(timestamp), context)) + continue + if kind == "summary": + if context is None: + continue + provider, model, reference = context + usage = line[1] + else: + provider, model, reference, usage = line[1:] + if kind == "reply": + context = (provider, model, reference) + counted = _pi_row(entry, timestamp, path, provider, model, obj(usage), int(kind == "reply" and not typed)) + answered: list[tuple] = [] + if kind == "reply": + if counted is not None: + answered, asked, failed = asked, [], None + elif asked and failed is None: + failed = () if _pi_unsent(reference) else (provider, model, reference) + except Exception: + # One odd line costs that line, never the scan. + continue + if counted is None: + continue + taken = takes(provider, model, reference) + key, row = counted + # Prompts typed on the reply's day ride on its row; one typed on another + # day is a row of its own there, keyed by the reply that answered it. + days: dict[str, list[tuple]] = {} + for _, prompt, day, _ in answered: + if day is not None: + days.setdefault(day, []).append(prompt) + asks = days.pop(row.day, []) + yield line_number, key, row._replace(prompts=row.prompts + len(asks)), tuple(asks), taken + for day, asks in sorted(days.items()): + yield line_number, ("prompt", day, *key), UsageRow(day, path, provider, model, 0, 0, 0, 0, len(asks)), tuple(asks), taken + yield from unanswered() + + +def scan_pi_sessions(pattern: str, accept: Accept, roots: Iterable[Path] | None = None) -> tuple[list[UsageRow], bool]: + """Every pi/omp reply `accept` takes, each counted once, and whether every + root was read in full. + + `pattern` is an rg regex every such line matches (the provider field, + say); it only keeps other lines and files out of the parse, and `accept` + stays the authority for the ones that reach it. A message a fork repeats is one row, + counted toward the oldest session holding it, so a fork is a session of + its own only for what it added, whichever file rg read first. + + `roots` are the session trees to read: pi_session_roots() unless the + collector names others in pi's format (a harness of its own, or pi's and + omp's with another beside them). They are read as canonical_roots() gives + them, so a message copied between any of them still counts once. A + collector that passes roots keys its scan cache by those same roots. + """ + roots = [root for root in canonical_roots(pi_session_roots() if roots is None else roots) if root.is_dir()] + if not roots: + return [], True + rg = find_command("rg") + if not rg: + warn("rg is missing; pi and omp sessions were not scanned") + return [], False + + # Usage that names no provider of its own, spent on the session's model, + # the prompts the replies answer, every reply, since one another provider + # gave still answers the prompt before it, and every model switch, which + # says whose a summary or a prompt no reply answered is. + switch = r'"type"\s*:\s*"model_change"' + context = r'"type"\s*:\s*"(compaction|branch_summary)"|"role"\s*:\s*"toolResult".*"usage"\s*:\s*\{|"role"\s*:\s*"(user|assistant)"|' + switch + + # (root, path, line) of every counted line, so the rows come back in one + # order however rg's threads interleave the files. + found: list[tuple[int, str, int, tuple, UsageRow, tuple]] = [] + complete = True + # Each file's lines, folded in order when rg finishes the file. + pending: dict[str, list[tuple[int, Any]]] = {} + # Every booking of each typed prompt, whichever provider took it, so one a + # fork answered on another provider books under one collector alone. + bookings: dict[tuple, list[tuple]] = {} + read: set[str] = set() + parents: dict[str, str | None] = {} + + def ancestors(path: str) -> set[str]: + found: set[str] = set() + parent = parents.setdefault(path, pi_parent_session(path)) + while parent and parent not in found and parent != path: + found.add(parent) + parent = parents.setdefault(parent, pi_parent_session(parent)) + return found + + def fold(index: int | None, path: str) -> None: + # `index` None reads a file for its bookings alone. + read.add(path) + depth = len(ancestors(path)) + for line_number, key, row, asks, taken in _pi_rows(path, sorted(pending.pop(path, []), key=lambda line: line[0]), accept): + for ask in asks: + # An answer before a prompt no reply answered, the oldest session's + # first, as a fork copies its parent's. + bookings.setdefault(ask, []).append((key[0] == "asked", depth, path, line_number, (key, row.provider, row.model))) + if taken and index is not None: + found.append((index, path, line_number, key, row, asks)) + + def run_rg(args: list[str], targets: list[bytes]) -> Iterator[bytes]: + """rg's output lines for `args` over `targets`, a batch of them at a time + so a long list of files stays under the argument limit.""" + nonlocal complete + batches: list[list[bytes]] = [[]] + size = 0 + for target in targets: + if batches[-1] and size + len(target) > 100_000: + batches.append([]) + size = 0 + batches[-1].append(target) + size += len(target) + 1 + for batch in batches: + if not batch: + continue + try: + # --no-ignore: these are data files, not a source tree. When $HOME is + # itself a git checkout (a common dotfiles setup), its ignore rules + # would otherwise hide every session file and count zero usage. + proc = subprocess.Popen([rg, "--no-ignore", "--hidden", *args, "--", *batch], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, env=runtime_env()) + except OSError: + complete = False + return + assert proc.stdout is not None + with proc.stdout: + yield from proc.stdout + try: + # rg exits 1 for no match and 2 when it couldn't read something. + if proc.wait(timeout=5) not in (0, 1): + complete = False + except subprocess.TimeoutExpired: + proc.kill() + complete = False + + def match_events(targets: list[bytes], patterns: list[str]) -> Iterator[tuple[str, str, int, Any]]: + """(event type, path, line number, entry) for each rg match and file end.""" + args = ["--json"] + [arg for pattern in patterns for arg in ("-e", pattern)] + for raw in run_rg(args, targets): + try: + event = json.loads(raw) + data = event["data"] + if event.get("type") == "end": + yield "end", os.path.realpath(_rg_text(data.get("path"))), 0, None + elif event.get("type") == "match": + yield "match", os.path.realpath(_rg_text(data.get("path"))), number(data.get("line_number")), json.loads(_rg_text(data.get("lines"))) + except Exception: + # One odd line costs that line, never the scan. + continue + + def matching_files(targets: list[bytes], patterns: list[str]) -> list[bytes]: + """The *.jsonl files under `targets` with a line matching any of `patterns`, + by realpath, as the scan names them.""" + args = ["--files-with-matches", "--null", "--glob", "*.jsonl"] + [arg for pattern in patterns for arg in ("-e", pattern)] + return [os.fsencode(os.path.realpath(os.fsdecode(name))) for name in b"".join(run_rg(args, targets)).split(b"\0") if name] + + for index, root in enumerate(roots): + # Only a file holding this collector's lines can hold its rows, so the + # files are picked first: those with a line `pattern` matches, and those + # with a model switch `accept` takes, which says whose a summary or an + # unanswered prompt there is (omp names the switch's provider/model, + # which `pattern` doesn't see). Every other file's prompts and summaries + # belong to another collector, and are never parsed here. + owned = matching_files([os.fsencode(root)], [pattern]) + switched = set(matching_files([os.fsencode(root)], [switch])) - set(owned) + for _, path, _, entry in match_events(sorted(switched), [switch]): + line = _pi_line(entry) + try: + if line is not None and line[0] == "model" and accept(*line[1:4]): + owned.append(os.fsencode(path)) + except Exception: + continue + for kind, path, line_number, entry in match_events(sorted(set(owned)), [pattern, context]): + if kind == "end": + fold(index, path) + else: + pending.setdefault(path, []).append((line_number, entry)) + for path in list(pending): + fold(index, path) + + # A prompt's copies live in the forks of the session it was typed in, so + # the other files of those fork trees holding a prompt this collector + # booked are read for their bookings alone: each collector then picks the + # same one among them. Only a prompt with an entry id has copies elsewhere. + ids = {ask[1] for *_, asks in found for ask in asks if ask[0] == "entry"} + children: dict[str, list[str]] = {} + for name in matching_files([os.fsencode(root) for root in roots], [r'"parentSession"']) if ids else (): + path = os.fsdecode(name) + parent = parents.setdefault(path, pi_parent_session(path)) + if parent: + children.setdefault(parent, []).append(path) + family: set[str] = set() + queue = list(read) if ids else [] + while queue: + path = queue.pop() + if path in family: + continue + family.add(path) + queue.extend(ancestors(path) | set(children.get(path, ()))) + others = sorted(os.fsencode(path) for path in family - read if os.path.isfile(path) and any(Path(path).is_relative_to(root) for root in roots)) + if others and all(re.fullmatch(r"[A-Za-z0-9_.:-]+", key) for key in ids): + # Such an id is written as itself, quoted, so a file holding none of + # them quoted holds no copy, and is never parsed. + try: + with tempfile.NamedTemporaryFile("w", suffix=".ids") as quoted: + quoted.write("".join(f'"{key}"\n' for key in sorted(ids))) + quoted.flush() + others = sorted({os.fsencode(os.path.realpath(os.fsdecode(name))) for name in b"".join(run_rg(["--files-with-matches", "--null", "--fixed-strings", "-f", quoted.name], others)).split(b"\0") if name}) + except OSError: + # With nowhere to write the ids, every file of the family is read. + pass + for kind, path, line_number, entry in match_events(others, [context]): + if kind == "end": + fold(None, path) + else: + pending.setdefault(path, []).append((line_number, entry)) + for path in list(pending): + fold(None, path) + owner = {ask: min(candidates, key=lambda candidate: candidate[:4])[-1] for ask, candidates in bookings.items()} + + first: dict[tuple, UsageRow] = {} + holders: dict[tuple, list[str]] = {} + for _, path, _, key, row, asks in sorted(found, key=lambda item: item[:3]): + if key not in first: + # A typed prompt is its own message, so it counts once by its own key + # however it was booked: a fork can copy one its parent left + # unanswered and answer it, on a reply the parent never had and on + # another provider too. The oldest session's answer books it, else + # its booking unanswered, whichever collector that is. + repeated = sum(1 for ask in asks if owner.get(ask) != (key, row.provider, row.model)) + row = row._replace(prompts=row.prompts - repeated) + first[key] = row + if path not in holders.setdefault(key, []): + holders[key].append(path) + + # Resolved after the scan, so rg's file order cannot decide which copy of + # a forked message counts toward which session. + rows = [] + for key, row in first.items(): + if not (row.prompts or row.input_tokens or row.output_tokens or row.cache_read or row.cache_write): + continue + paths = holders[key] + if len(paths) > 1: + roots_of_copy = [path for path in paths if not ancestors(path) & set(paths)] + row = row._replace(session=min(roots_of_copy or paths)) + rows.append(row) + return rows, complete + + +# ---------------------------------------------------------------- opencode + + +def opencode_db() -> Path: + """The database opencode itself reads. + + OPENCODE_DB points opencode at another database, as the CLI itself reads + it; a relative name lives beside the default one. Only a build whose + channel was literally "v2" wrote opencode-v2.db, so it is a fallback for + those installs alone. + """ + root = data_home() / "opencode" + configured = os.environ.get("OPENCODE_DB", "").strip() + if configured: + path = Path(configured).expanduser() + return path if path.is_absolute() else root / path + normal = root / "opencode.db" + legacy = root / "opencode-v2.db" + return legacy if not normal.is_file() and legacy.is_file() else normal + + +def _table_columns(conn: sqlite3.Connection, table: str) -> set[str]: + try: + return {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} + except sqlite3.Error: + return set() + + +def _opencode_tokens(value: Any) -> tuple[int, int, int, int]: + """(input, output, cache read, cache write) of an opencode `tokens` + object, each count one token_count() trusts.""" + tokens = obj(value) + cache = obj(tokens.get("cache")) + # opencode keeps thinking tokens out of output; both are generated. + return token_count(tokens.get("input")), token_count(tokens.get("output")) + token_count(tokens.get("reasoning")), token_count(cache.get("read")), token_count(cache.get("write")) + + +def _opencode_row(session_id: Any, entry: Any, row_type: str, created: Any, accept: Accept, empty: bool = False) -> UsageRow | None: + # Shared by both message stores. The original `message` table keeps the + # provider and model at the top level; the V2 `session_message` table nests + # them under `model` and keeps the role in its `type` column, with a + # compaction's tokens spent like any reply's. `empty` keeps a row that + # spent nothing, for the prompt it books; a user row is one such. + if not isinstance(entry, dict) or (entry.get("role") or row_type) not in (("assistant", "compaction", "user") if empty else ("assistant", "compaction")): + return None + nested = obj(entry.get("model")) + provider = str(entry.get("providerID") or nested.get("providerID") or "") + model = str(entry.get("modelID") or nested.get("id") or nested.get("modelID") or "") + if not accept(provider, model, entry): + return None + # A user row spent nothing, whatever it claims. + input_tokens, output_tokens, cache_read, cache_write = _opencode_tokens(entry.get("tokens") if entry.get("role") != "user" else {}) + if not (input_tokens or output_tokens or cache_read or cache_write or empty): + return None + # The message's own time, else the row's: none at all is today's, and an + # untrusted one, zero included, skips only this row. + created_at = obj(entry.get("time")).get("created") + day = local_day(created if created_at is None else created_at) + if day is None: + return None + # One key per session id across both stores: a migrated session is one. + return UsageRow(day, "opencode:" + str(session_id), provider, model, input_tokens, output_tokens, cache_read, cache_write) + + +def scan_opencode(providers: Iterable[str], accept: Accept) -> tuple[list[UsageRow], bool]: + """Every opencode reply `accept` takes, from both message stores, and + whether the scan ran to completion. + + `providers` are the provider ids the rows may carry, as LIKE prefixes: a + row whose JSON names none of them is skipped before any parse. End one + with `"` to match the whole id. They only keep other rows out of the + parse; `accept` stays the authority for the rows that reach it. + + V2's `session_message` is read first, then the original `message` store, + deduplicated by message id, leaving out the copies a fork makes. Read-only: + opencode may be writing right now. A scan cut short by a database error + still hands back what it read, but says it is incomplete, so it is never + cached as the whole story. + """ + db = opencode_db() + if not db.is_file(): + return [], True + try: + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return [], False + patterns = [f'%"providerID"%:%"{provider}%' for provider in providers] + rows: list[UsageRow] = [] + # The typed prompt each row answers, and where it falls among the steps + # answering the same one; the first step books the prompt. + turns: list[tuple[Any, tuple]] = [] + seen: set[str] = set() + # A subagent's session: the prompts in it are the model's, not typed. + children: set[str] = set() + # The original store's first step answering each user message, whichever + # provider took it, by (session, parentID). + firsts: dict[tuple[str, str], str] = {} + + def take(message_id: Any, session_id: Any, raw: Any, row_type: str, created: Any, since: Any = None, turn: Any = "", empty: bool = False, typed: str | None = None) -> None: + # `turn` is the prompt the row answers, None for none, or "" to read it + # from the row; `empty` keeps a row that spent nothing, for that prompt, + # dated `typed`, the day it was typed (None when that can't be trusted). + # One row that still fails must not take the scan down with it. + try: + entry = json.loads(raw) + # A copy a fork made predates the session holding it. + spent_at = number(obj(obj(entry).get("time")).get("created")) + if spent_at and spent_at < number(since): + return + row = _opencode_row(session_id, entry, row_type, created, accept, empty) + except Exception: + return + # A migrated message can sit in both stores under the same id. + if row is None or str(message_id) in seen: + return + seen.add(str(message_id)) + if empty and not (row.input_tokens or row.output_tokens or row.cache_read or row.cache_write): + # A failed reply is only its prompt, which counts on its own day. + if typed is None: + return + row = row._replace(day=typed) + if turn == "": + # The original store's steps name the user message they answer; a + # compaction's summary answers the marker opencode adds, no prompt. + parent = entry.get("parentID") + if entry.get("summary") is True or entry.get("mode") == "compaction": + turn = None + elif isinstance(parent, str) and parent: + turn = (str(session_id), parent) + # A later step, even one another provider's first, answers none. + if firsts.get(turn, str(message_id)) != str(message_id): + turn = None + else: + turn = ("message", str(message_id)) + if str(session_id) in children: + turn = None + rows.append(row) + turns.append((turn, (spent_at, str(message_id)))) + + def answered() -> list[UsageRow]: + first: dict[Any, int] = {} + for index, (turn, order) in enumerate(turns): + if turn is not None and (turn not in first or order < turns[first[turn]][1]): + first[turn] = index + return [row._replace(prompts=int(first.get(turn) == index)) for index, (row, (turn, _)) in enumerate(zip(rows, turns))] + + try: + conn.execute("PRAGMA query_only = ON") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + if "session" in tables and {"id", "parent_id"} <= _table_columns(conn, "session"): + children = {str(session_id) for (session_id,) in conn.execute("SELECT id FROM session WHERE parent_id IS NOT NULL")} + columns = _table_columns(conn, "session_message") if "session_message" in tables else set() + if {"id", "session_id", "type", "data"} <= columns and patterns: + # V2 keeps typed prompts as `user` rows without a model, so each books + # the first reply after it in its session, one whose request failed + # and spent nothing too; `synthetic` rows are the ones opencode adds + # itself. A session without any keeps a prompt per reply. + asked: set[str] = set() + # Each reply booking a prompt, and the day its first user row was typed. + replies: dict[str, str | None] = {} + order = "seq" if "seq" in columns else "time_created" if "time_created" in columns else "" + if order: + for session_id, reply_id, created_at, typed_at in conn.execute( + "SELECT u.session_id, (SELECT a.id FROM session_message a WHERE a.session_id = u.session_id AND a.type = 'assistant'" + f" AND (a.{order} > u.{order} OR (a.{order} = u.{order} AND a.id > u.id)) ORDER BY a.{order}, a.id LIMIT 1)," + " CASE WHEN json_valid(u.data) THEN json_extract(u.data, '$.time.created') END," + f" {'u.time_created' if 'time_created' in columns else 'NULL'}" + f" FROM session_message u WHERE u.type = 'user' ORDER BY u.{order}, u.id" + ): + asked.add(str(session_id)) + if reply_id is not None and str(reply_id) not in replies: + # The user row's own time, else the row's, as _opencode_row reads it. + replies[str(reply_id)] = local_day(typed_at if created_at is None else created_at) + created = "m.time_created" if "time_created" in columns else "NULL" + source = "session_message m" + fork_filter = "" + if "time_created" in columns and "session_v2" in tables and {"id", "fork_session_id", "time_created"} <= _table_columns(conn, "session_v2"): + # A fork copies its source session's messages under new ids; only + # what came after the fork was spent in it. + source += " LEFT JOIN session_v2 s ON s.id = m.session_id" + fork_filter = " AND (s.fork_session_id IS NULL OR m.time_created >= s.time_created)" + gate = " OR ".join("m.data LIKE ?" for _ in patterns) + for message_id, session_id, row_type, time_created, raw in conn.execute( + f"SELECT m.id, m.session_id, m.type, {created}, m.data FROM {source}" + f" WHERE m.type IN ('assistant', 'compaction') AND ({gate})" + fork_filter, + patterns, + ): + if row_type == "compaction" or (str(session_id) in asked and str(message_id) not in replies): + turn = None + else: + turn = ("message", str(message_id)) + take(message_id, session_id, raw, str(row_type or ""), time_created, turn=turn, empty=turn is not None and str(session_id) in asked, typed=replies.get(str(message_id))) + columns = _table_columns(conn, "message") if "message" in tables else set() + if {"session_id", "data"} <= columns and patterns: + row_id = "m.id" if "id" in columns else "m.rowid" + source = "message m" + session_created = "NULL" + if "session" in tables and {"id", "time_created"} <= _table_columns(conn, "session"): + # A fork here is a plain new session holding copies of its source's + # messages under new ids, each keeping its original time.created, so + # a message older than its session was spent in the source. + source += " LEFT JOIN session s ON s.id = m.session_id" + session_created = "s.time_created" + gate = " OR ".join("m.data LIKE ?" for _ in patterns) + # This collector's own steps and prompts, read before the first steps + # they need, which are looked up for their sessions alone. + own_steps = list(conn.execute( + f"SELECT {row_id}, m.session_id, m.data, {session_created} FROM {source}" + f" WHERE m.data LIKE '%\"role\"%:%\"assistant\"%' AND ({gate})", + patterns, + )) + own_prompts = list(conn.execute( + f"SELECT {row_id}, m.session_id, m.data, {session_created} FROM {source}" + f" WHERE m.data LIKE '%\"role\"%:%\"user\"%' AND ({gate})", + patterns, + )) + # The first step that spent tokens on each user message across every + # provider, so a tool loop that switches providers books its prompt + # once, on whichever collector took that step. Only valid JSON is read, + # so one odd row can't fail the query, and a step counts as first only + # if take() would keep it: its tokens and time pass the same checks, + # so a damaged step never takes the prompt from a valid one after it. + # A session's first steps are its own, so only the sessions above are + # read, through opencode's index on session_id, each id as stored. + sessions = list(dict.fromkeys(session_id for _, session_id, _, _ in own_steps + own_prompts)) + for session_id, parent, first_id, created, tokens, since in conn.execute( + "SELECT session_id, parent, id, created, tokens, since" + " FROM (SELECT id, session_id, j, since, json_extract(j, '$.parentID') AS parent, json_extract(j, '$.time.created') AS created," + " json_extract(j, '$.tokens') AS tokens" + f" FROM (SELECT {row_id} AS id, m.session_id, CASE WHEN json_valid(m.data) THEN m.data END AS j, {session_created} AS since" + f" FROM {source} WHERE m.session_id IN (SELECT value FROM json_each(?))" + " AND m.data LIKE '%\"parentID\"%' AND m.data LIKE '%\"role\"%:%\"assistant\"%'))" + " WHERE typeof(parent) = 'text' AND parent != '' AND typeof(tokens) = 'text'" + " AND json_extract(j, '$.summary') IS NOT 1 AND json_extract(j, '$.mode') IS NOT 'compaction'" + " ORDER BY session_id, parent, created, CAST(id AS TEXT)", + (json.dumps(sessions, default=str),), + ) if sessions else (): + turn = (str(session_id), str(parent)) + if turn in firsts: + continue + try: + spent_at = number(created) + if not any(_opencode_tokens(json.loads(tokens))) or local_day(created) is None or (spent_at and spent_at < number(since)): + continue + except Exception: + continue + firsts[turn] = str(first_id) + for message_id, session_id, raw, since in own_steps: + take(message_id, session_id, raw, "", None, since) + # A typed prompt no step answered (its request failed or was cut off) + # books itself, without tokens, on the model the user row names, the + # one opencode sent it to. A compaction's marker is answered by its + # summary, and one in a subagent's session isn't typed. + waiting: list[tuple[Any, str, dict]] = [] + for message_id, session_id, raw, since in own_prompts: + try: + entry = json.loads(raw) + if not isinstance(entry, dict) or entry.get("role") != "user" or str(session_id) in children or (str(session_id), str(message_id)) in firsts: + continue + typed_at = number(obj(entry.get("time")).get("created")) + if typed_at and typed_at < number(since): + continue + except Exception: + continue + waiting.append((message_id, str(session_id), entry)) + if waiting: + # The steps that did answer them: a compaction's summary, or a failed + # request, whose id V2 may have booked the prompt under already. + sessions = sorted({session_id for _, session_id, _ in waiting}) + steps: dict[tuple[str, str], str] = {} + compacted: set[tuple[str, str]] = set() + for start in range(0, len(sessions), 500): + batch = sessions[start:start + 500] + for step_id, session_id, raw in conn.execute( + f"SELECT {row_id}, m.session_id, m.data FROM message m WHERE m.session_id IN ({', '.join('?' for _ in batch)})" + " AND m.data LIKE '%\"parentID\"%' AND m.data LIKE '%\"role\"%:%\"assistant\"%'" + f" ORDER BY {'m.time_created, ' if 'time_created' in columns else ''}CAST({row_id} AS TEXT)", + batch, + ): + try: + step = json.loads(raw) + key = (str(session_id), str(step.get("parentID"))) + if step.get("summary") is True or step.get("mode") == "compaction": + compacted.add(key) + steps.setdefault(key, str(step_id)) + except Exception: + continue + for message_id, session_id, entry in waiting: + key = (session_id, str(message_id)) + # opencode writes the step before it sends the request, so a user + # row no step names was refused before sending (an unknown model or + # provider) and counts nowhere. + step_id = steps.get(key) + if step_id is None or key in compacted or step_id in seen: + continue + try: + row = _opencode_row(session_id, entry, "user", None, accept, empty=True) + except Exception: + continue + if row is None: + continue + seen.add(step_id) + rows.append(row) + turns.append((("asked", session_id, str(message_id)), (0, str(message_id)))) + except sqlite3.Error: + # Transient lock, schema migration, corruption: the numbers stop here, + # incomplete. + return answered(), False + finally: + conn.close() + return answered(), True + + +# --------------------------------------------------- who bills an endpoint +# +# Some harnesses send to whatever base URL they are given and record only the +# model that answered, or a route's name: Claude Code fronts any endpoint +# ANTHROPIC_BASE_URL names, and a MiniMax Code custom provider any baseURL. +# Z.ai, Ollama, MiniMax, DeepSeek, Kimi, Alibaba's Model Studio, OpenRouter, +# Synthetic and Kimchi each document pointing them at their own APIs, so the +# host says whose tokens those turns are. + +# The hosts a harness is pointed at in place of its own service, by the +# collector that bills them. A subdomain of one is the same service. +PROVIDER_HOSTS = { + "api.z.ai": "zai", + "open.bigmodel.cn": "zai", + "ollama.com": "ollama", + "api.minimax.io": "minimax", + "api.minimaxi.com": "minimax", + "api.deepseek.com": "deepseek", + "api.moonshot.ai": "kimi", + "api.moonshot.cn": "kimi", + "api.kimi.com": "kimi", + "api.kimi.ai": "kimi", + # Model Studio (DashScope, the Coding and Token Plans) and Qwen's own + # sign-in. + "aliyuncs.com": "qwen", + "alibabacloud.com": "qwen", + "qwen.ai": "qwen", + "openrouter.ai": "openrouter", + # Synthetic's Anthropic- and OpenAI-compatible APIs (/anthropic, /openai/v1). + "api.synthetic.new": "synthetic", + # Kimchi's gateway in each region (llm.kimchi.dev, llm.eu.kimchi.dev), at + # the /anthropic and /openai/v1 routes `kimchi claudecode`, `kimchi codex` + # and its opencode setup write. It bills every model it serves, Claude, + # Kimi, GLM and MiniMax ones too. + "kimchi.dev": "kimchi", +} + + +def url_host(url: Any) -> str: + """The lowercase host a URL names, or "" for none or a URL that isn't one.""" + try: + return (urllib.parse.urlsplit(str(url or "").strip()).hostname or "").lower().rstrip(".") + except ValueError: + return "" + + +def self_hosted(url: Any) -> bool: + """Whether a base URL points at hardware the user controls: this machine, + a private or link-local address, or a LAN name.""" + host = url_host(url) + if not host: + return False + if host == "localhost": + return True + try: + address = ipaddress.ip_address(host) + return address.is_loopback or address.is_private or address.is_link_local + except ValueError: + pass + # A bare hostname is a LAN name by construction; the rest are the suffixes + # home routers and mDNS hand out. + if "." not in host: + return True + return host.endswith((".home", ".lan", ".local", ".internal", ".localdomain", ".home.arpa")) + + +def endpoint_owner(url: Any) -> str: + """The collector that bills what a base URL serves: a provider host's, or + `local` for one the user runs, or "" when nothing here can place it (a + gateway or proxy that may forward anywhere).""" + host = url_host(url) + if not host: + return "" + owner = next((provider for name, provider in PROVIDER_HOSTS.items() if host == name or host.endswith("." + name)), "") + return owner or ("local" if self_hosted(url) else "") + + +def cloud_model(model: Any) -> bool: + """Whether a model is one of Ollama's hosted ones, which carry a cloud tag + ("glm-4.6:cloud", "gpt-oss:120b-cloud"): a local daemon only forwards + them, so Ollama's servers did the work.""" + text = str(model or "").lower() + tag = text.rpartition(":")[2] if ":" in text else "" + return tag == "cloud" or tag.endswith("-cloud") + + +# --------------------------------------------------------- Claude Code +# +# Claude Code writes one JSONL transcript per session under +# ~/.claude/projects (CLAUDE_CONFIG_DIR moves it). Each reply carries the +# model that answered and its usage, but not where it was sent: the claude +# collector counts the replies Anthropic served, and the collector of +# whichever provider ANTHROPIC_BASE_URL names, or whose model answered, +# counts the rest. claude_code_owner() is the one rule all of them apply, so +# every reply lands in exactly one record. + +# What a model's name says about who served it, when Claude Code is pointed +# at no other host: these are those providers' own models. +CLAUDE_CODE_MODEL_OWNERS = ( + ("glm-", "zai"), + ("minimax-", "minimax"), + ("deepseek-", "deepseek"), + ("kimi-", "kimi"), + ("moonshot", "kimi"), + # Synthetic's own naming: Hugging Face ids (hf:org/model) and its aliases + # (syn:large:text). + ("hf:", "synthetic"), + ("syn:", "synthetic"), +) + + +def claude_code_dir() -> Path: + return Path(os.path.expandvars(os.path.expanduser(os.environ.get("CLAUDE_CONFIG_DIR") or "~/.claude"))).resolve() + + +def claude_code_base_url(claude_dir: Path | None = None) -> str: + """Where ANTHROPIC_BASE_URL sends Claude Code now: its settings' env wins + over the environment, as it does in Claude Code.""" + env = read_json((claude_dir or claude_code_dir()) / "settings.json").get("env") + configured = env.get("ANTHROPIC_BASE_URL") if isinstance(env, dict) else None + return str(configured or os.environ.get("ANTHROPIC_BASE_URL") or "").strip() + + +def claude_code_endpoint(base_url: str) -> str: + """The kind of endpoint Claude Code is sent to, which is all that decides + who owns its replies: "" for Anthropic's own (no override), a provider + collector's id, `local`, or `other` for a host nothing here can place.""" + host = url_host(base_url) + if not host or host == "anthropic.com" or host.endswith(".anthropic.com"): + return "" + return endpoint_owner(base_url) or "other" + + +def claude_code_owner(model: Any, base_url: str = "") -> str | None: + """The collector that counts a Claude Code reply by `model` (None for a + prompt no reply has answered yet) while Claude Code is sent to `base_url`, + or None when nothing here can say who billed it. + + A transcript doesn't say where it was sent, so while ANTHROPIC_BASE_URL + names a provider's host every reply is read as that provider's, whatever + its model: a Claude or kimi-* reply through Kimchi's gateway is Kimchi's, + never Anthropic's or Kimi's. Otherwise an Ollama cloud tag is Ollama's, + wherever it was sent, and a Claude model + is Anthropic's, through a proxy that may forward to it too. Behind a + self-hosted base URL anything else ran on the user's own hardware. With no + override the model's name says whose it was (glm-* is Z.ai's, MiniMax-* + MiniMax's, deepseek-* DeepSeek's, kimi-* and moonshot* Kimi's, hf:* and + syn:* Synthetic's); behind a + gateway nothing here can place, it could be anyone's, and counts nowhere. + """ + endpoint = claude_code_endpoint(base_url) + if endpoint not in ("", "local", "other"): + return endpoint + name = str(model).lower() if model is not None else "" + if cloud_model(name): + return "ollama" + if model is None or "claude" in name: + return "claude" + if endpoint == "local": + return "local" + if endpoint == "other": + return None + return next((owner for prefix, owner in CLAUDE_CODE_MODEL_OWNERS if name.startswith(prefix)), None) + + +# Transcripts are append-only, and most of the corpus never changes between +# ticks, so the scan keeps a per-file index: the unique usage records already +# parsed out of each file plus the byte offset they end at. A file whose size +# and mtime match the index is not opened; a file that grew is read from the +# stored offset; anything else (truncated, rewritten, replaced, unknown) is +# read from the start. Every collector reading Claude Code shares the index, +# which holds every reply whoever owns it. +# +# Version 5 reads token counts and times the way every collector reads them +# (a count written as a string is no count, an untrusted time skips its line), +# so an index parsed the older way is read again. Version 6 stops reading a +# falsy time (0, false, "") as no time at all, which counted it as today. +# Version 7 keeps the typed prompts too, which are the record's prompts; a +# reply was counted as one before. Version 8 keys a reply by its request too, +# and keeps its bare message id and whether a sidechain wrote it. Version 9 +# reads the advisor iterations a reply carries. Version 10 keeps the model +# that answered each prompt. Version 11 stops counting a built-in local +# command (/model, /clear, /compact) as a prompt. Version 12 keeps each +# reply's time and the cost-state totals, which count the requests no reply +# records. +CLAUDE_CODE_INDEX_VERSION = 12 + +# How much of each end of what was read the index remembers a hash of, so a +# file replaced or rewritten with more than it had is not mistaken for an +# append. +CLAUDE_CODE_INDEX_HEAD_BYTES = 4096 + +# What Claude Code puts in a user line that nobody typed: a background task +# reporting back, a local command's output and its caveat, and the marker an +# interruption leaves. +CLAUDE_CODE_MACHINE_PROMPTS = ("", "", "", "", "[Request interrupted") +# What a local command's output starts with, which says the command it +# answers was no prompt. +CLAUDE_CODE_LOCAL_OUTPUT = ("", "") +# Who sent a prompt, where Claude Code records it (a user line's +# promptSource, a queued one's origin.kind): the system (a task notification) +# and another agent are not a person asking. +CLAUDE_CODE_MACHINE_SOURCES = ("system", "task-notification", "peer") + +# A line's own time, the last `timestamp` it holds (Claude Code writes it +# after the message), read without parsing the line. +CLAUDE_CODE_STAMP = re.compile(rb'"timestamp":\s*"([^"\\]{1,64})"') + +# Claude Code 2.1.308 and later write a `cost-state` line when a process +# saves its running cost: `startTime` names the process's cost lineage, +# which a resume or `--fork-session` carries on, and `modelUsage` holds what +# each model has used in it so far, the requests no reply line records +# (a compaction, the session title, a side query) included. Its model names +# carry the context window as a suffix the replies don't. +CLAUDE_CODE_COST_MODEL_SUFFIX = re.compile(r"\[1m\]$", re.IGNORECASE) + + +def claude_code_index_path(projects: Path) -> Path: + digest = hashlib.sha1(str(projects).encode("utf-8", "surrogatepass")).hexdigest()[:16] + return cache_root() / f"claude-index-{digest}.json" + + +# Records hold local days, so an index is only true in the timezone that +# wrote it; after a change every transcript is read again. +def _claude_code_index_zone() -> list[Any]: + return [os.environ.get("TZ"), list(time.tzname), time.timezone, time.altzone] + + +def _claude_code_load_index(index_file: Path) -> dict[str, Any]: + data = read_json(index_file) + if data.get("version") == CLAUDE_CODE_INDEX_VERSION and data.get("zone") == _claude_code_index_zone() and isinstance(data.get("files"), dict): + return data["files"] + return {} + + +# What was already read, by its first and last few kilobytes: a transcript +# rewritten in place rather than appended to changes one or the other. +def _claude_code_digest(path: Path, length: int) -> str: + digest = hashlib.sha1() + with path.open("rb") as handle: + digest.update(handle.read(min(length, CLAUDE_CODE_INDEX_HEAD_BYTES))) + handle.seek(max(0, length - CLAUDE_CODE_INDEX_HEAD_BYTES)) + digest.update(handle.read(min(length, CLAUDE_CODE_INDEX_HEAD_BYTES))) + return digest.hexdigest() + + +def claude_code_typed_prompt(line: dict[str, Any]) -> bool: + """Whether a transcript line is a prompt a person sent: a main-thread user + message that isn't a tool's result, Claude Code's own (isMeta), a + compaction summary, or a subagent's (isSidechain); or a prompt typed while + a turn was running, which Claude Code queues into the turn as a + queued_command attachment rather than a user line.""" + if line.get("isMeta") or line.get("isSidechain") or line.get("isCompactSummary"): + return False + if line.get("type") == "attachment": + queued = line.get("attachment") + if not isinstance(queued, dict) or queued.get("type") != "queued_command" or queued.get("commandMode") != "prompt" or queued.get("isMeta"): + return False + origin = queued.get("origin") + return not (isinstance(origin, dict) and origin.get("kind") in CLAUDE_CODE_MACHINE_SOURCES) + if line.get("type") != "user" or line.get("promptSource") in CLAUDE_CODE_MACHINE_SOURCES: + return False + text = _claude_code_user_text(line) + return text is not None and not text.lstrip().startswith(CLAUDE_CODE_MACHINE_PROMPTS) + + +def _claude_code_user_text(line: dict[str, Any]) -> str | None: + """The text of a user line, None for a tool's result or no content.""" + message = line.get("message") + content = message.get("content") if isinstance(message, dict) else None + if isinstance(content, str): + return content + if isinstance(content, list): + blocks = [block for block in content if isinstance(block, dict)] + if not blocks or any(block.get("type") == "tool_result" for block in blocks): + return None + return next((str(block.get("text") or "") for block in blocks if block.get("type") == "text"), "") + return None + + +def claude_code_local_command(line: dict[str, Any]) -> str | None: + """The uuid of the command a transcript line is the local output of. A + built-in that runs on its own (/model, /clear, /compact, /exit) answers its + `` line with a `` or `-stderr` child, + a user line or a `local_command` system line, and asks the model nothing; + a custom command or skill expands into a prompt instead, and has none.""" + if line.get("type") == "system": + text = line.get("content") if isinstance(line.get("content"), str) else None + elif line.get("type") == "user": + text = _claude_code_user_text(line) + else: + return None + parent = line.get("parentUuid") + if text is not None and isinstance(parent, str) and text.lstrip().startswith(CLAUDE_CODE_LOCAL_OUTPUT): + return parent + return None + + +def _claude_code_line_day(parsed: dict[str, Any], message: dict[str, Any]) -> tuple[bool, str | None]: + """(trusted, day) of a transcript line. A line with no time is today's, + kept as None so the index stays true past midnight; one whose time can't + be trusted counts nowhere. A time of 0, false or "" is present and + untrusted, not missing.""" + timestamp = parsed.get("timestamp") + if timestamp is None: + timestamp = message.get("timestamp") + if timestamp is None: + return True, None + day = local_day(timestamp) + return day is not None, day + + +def _claude_code_moment(value: Any) -> float | None: + """When a trusted transcript time was, as epoch seconds, to put a reply + before or after a cost-state line (None for no time).""" + if value is None or isinstance(value, bool): + return None + if isinstance(value, str) and re.fullmatch(r"[0-9]{1,20}", value.strip()): + value = int(value.strip()) + if isinstance(value, (int, float)): + seconds = math.floor(number(value)) + return seconds / 1000 if seconds > 10_000_000_000 else float(seconds) + try: + return datetime.fromisoformat(str(value).strip().replace("Z", "+00:00")).timestamp() + except (OverflowError, OSError, ValueError): + return None + + +def _claude_code_cost_state(parsed: dict[str, Any]) -> tuple[str, dict[str, list[int]]] | None: + """The cost lineage and per-model [input, output, cacheRead, cacheWrite] + totals of a cost-state line, None for one of another shape.""" + usage = parsed.get("modelUsage") + start = parsed.get("startTime") + if not isinstance(usage, dict) or start is None or isinstance(start, (bool, dict, list)) or str(start) == "": + return None + models: dict[str, list[int]] = {} + for name, counts in usage.items(): + if not isinstance(counts, dict): + continue + tokens = [ + usage_tokens(counts, "inputTokens", "input_tokens"), + usage_tokens(counts, "outputTokens", "output_tokens"), + usage_tokens(counts, "cacheReadInputTokens", "cache_read_input_tokens"), + usage_tokens(counts, "cacheCreationInputTokens", "cache_creation_input_tokens"), + ] + # One model under both names is two counters of the same model. + held = models.setdefault(CLAUDE_CODE_COST_MODEL_SUFFIX.sub("", str(name)), [0, 0, 0, 0]) + for kind, count in enumerate(tokens): + held[kind] += count + return str(start), models + + +def read_claude_code_transcript(path: Path, entry: dict[str, Any] | None) -> dict[str, Any]: + """Return the index entry for one transcript, reusing `entry` where the file + only grew. Records are [key, model, day, input, output, cacheRead, + cacheWrite, session, messageId, sidechain, moment]; prompts are [key, day, + session, model], the model of the first reply after it (None until one + is read); costs are [lineage, session, day, moment, {model: [input, + output, cacheRead, cacheWrite]}], each cost-state line's totals at the + time of the last line before it. A `day` is None, and a `moment` (epoch + seconds) too, when the line carries no timestamp.""" + info = path.stat() + size = info.st_size + mtime = info.st_mtime_ns + + if entry and entry.get("size") == size and entry.get("mtime") == mtime: + return entry + + # Only the same file, grown past what was read with what was read still in + # place, is an append. A transcript replaced by a larger one is a new file, + # and one rewritten in place has a different start or end. + if ( + entry + and size > number(entry.get("size")) + and entry.get("file") == [info.st_dev, info.st_ino] + and entry.get("head") == _claude_code_digest(path, number(entry.get("offset"))) + ): + records = list(entry.get("records") or []) + prompts = [list(prompt) for prompt in entry.get("prompts") or []] + costs = list(entry.get("costs") or []) + stamp = entry.get("stamp") + offset = number(entry.get("offset")) + line_number = number(entry.get("lines")) + else: + records = [] + prompts = [] + costs = [] + stamp = None + offset = 0 + line_number = 0 + + # A streamed response is several lines sharing one message id, and the + # first line's output_tokens is a placeholder while the last carries the + # real count. Keep each message's line with the highest output count (the + # last one read wins a tie), whole: a response that fell back to another + # model mid-stream has lines that are separate snapshots, and mixing fields + # across them over-counts cache tokens and credits the wrong model. + positions = {record[0]: index for index, record in enumerate(records)} + prompted = {prompt[0] for prompt in prompts} + # Prompts still waiting for their reply, which says who answered them. + unanswered = [prompt for prompt in prompts if prompt[3] is None] + stamped: tuple[bytes, int] | None = None + + def last_stamp() -> Any: + found = CLAUDE_CODE_STAMP.match(stamped[0], stamped[1]) if stamped else None + return found.group(1).decode("utf-8", errors="replace") if found else stamp + + with path.open("rb") as handle: + handle.seek(offset) + for raw in handle: + if not raw.endswith(b"\n"): + # A line still being written; pick it up whole next time. + break + offset += len(raw) + line_number += 1 + + # A cost-state line carries no time of its own: it was written when + # the line before it was, read only when one turns up. + at = raw.rfind(b'"timestamp":') + if at >= 0: + stamped = (raw, at) + + # Cheap pre-filter before JSON parsing keeps files with unrelated + # lines inexpensive. + if b'"usage":' not in raw and b'"user"' not in raw and b'"queued_command"' not in raw and b"`, no tokens). + if record_key == key and not sidechain: + for prompt in unanswered: + prompt[3] = model + unanswered.clear() + + held = positions.get(record_key) + if held is not None and output_tokens < records[held][4]: + continue + + record = [record_key, model, day, input_tokens, output_tokens, cache_read, cache_write, session_key, record_id, sidechain, moment] + if held is None: + positions[record_key] = len(records) + records.append(record) + else: + records[held] = record + + return { + "size": size, + "file": [info.st_dev, info.st_ino], + "mtime": mtime, + "head": _claude_code_digest(path, offset), + "offset": offset, + "lines": line_number, + "records": records, + "prompts": prompts, + "costs": costs, + "stamp": last_stamp(), + } + + +def _claude_code_index_stamp(index_file: Path) -> tuple[int, int, int] | None: + """Which index file is there now: a write replaces it whole.""" + try: + info = index_file.stat() + except OSError: + return None + return (info.st_ino, info.st_mtime_ns, info.st_size) + + +def _claude_code_index_fresh(index: dict[str, Any], files: list[Path]) -> bool: + """Whether `index` holds every transcript in `files` as it is now, and no + other, so nothing needs reading.""" + if len(index) != len(files): + return False + for path in files: + entry = index.get(str(path)) + try: + info = path.stat() + except OSError: + return False + if not entry or entry.get("size") != info.st_size or entry.get("mtime") != info.st_mtime_ns: + return False + return True + + +def _claude_code_read_transcripts(files: list[Path], previous: dict[str, Any]) -> tuple[dict[str, Any], list[str], bool]: + """Each transcript's index entry, reusing `previous`, the order they were + read in, and whether every one could be read.""" + current: dict[str, Any] = {} + order: list[str] = [] + unreadable_count = 0 + unreadable_examples: list[str] = [] + for path in files: + key = str(path) + try: + current[key] = read_claude_code_transcript(path, previous.get(key)) + except Exception as exc: + unreadable_count += 1 + if len(unreadable_examples) < 3: + unreadable_examples.append(f"{path}: {exc}") + # What was read before is still true of the file. + if key in previous: + current[key] = previous[key] + if key in current: + order.append(key) + + if unreadable_count: + examples = "; ".join(unreadable_examples) + print(f"Ignoring {unreadable_count} unreadable Claude project files (examples: {examples})", file=sys.stderr) + return current, order, not unreadable_count + + +_DEFAULT_INDEX = object() + + +def scan_claude_code( + owner: str, + projects: Path | None = None, + index_file: Any = _DEFAULT_INDEX, + base_url: str | None = None, + reread: bool = False, +) -> list[UsageRow]: + """claude_code_usage()'s rows alone.""" + return claude_code_usage(owner, projects, index_file, base_url, reread)[0] + + +def claude_code_usage( + owner: str, + projects: Path | None = None, + index_file: Any = _DEFAULT_INDEX, + base_url: str | None = None, + reread: bool = False, +) -> tuple[list[UsageRow], bool]: + """The Claude Code usage `owner` counts, and whether every transcript was + read: one that can't be read keeps what was read of it before, and a scan + missing the rest is not one to cache. The rows are each reply once (with + no prompts of its own), each typed prompt once (no tokens), and a tokenless + row for every session in use on a day, for Stats.add_row(). A prompt's row + names the model that answered it ("" for none yet), so a collector that + keeps some models alone can keep their prompts alone too. + + `projects` defaults to the Claude Code home's projects/, `index_file` to + the index every collector shares there (None reads without one), and + `base_url` to where Claude Code is sent now. `reread` reads every + transcript from the start and replaces the index, which other collectors + keep reading meanwhile. Every reply is deduplicated across transcripts + before claude_code_owner() splits them, so a reply a fork copied lands + once, under one owner. + """ + if projects is None: + projects = claude_code_dir() / "projects" + if base_url is None: + base_url = claude_code_base_url(projects.parent) + if index_file is _DEFAULT_INDEX: + try: + index_file = claude_code_index_path(projects) + except OSError as exc: + warn(f"scan index unavailable ({exc})") + index_file = None + # Every collector reading Claude Code runs at once (the update command + # backgrounds one per agent), so one of them brings the index up to date + # while the others wait on its lock and then find nothing left to read. An + # index that already answers for every transcript is read without waiting. + # Directory order, as before: when one API message shows up in several + # transcripts with the same output count, the last file visited wins, and + # the numbers must not move because the scan got faster. + files = list(projects.rglob("*.jsonl")) if projects.is_dir() else [] + if not index_file: + current, order, complete = _claude_code_read_transcripts(files, {}) + else: + lock_file = index_file.with_name(index_file.name + ".lock") + + def build(previous: dict[str, Any]) -> tuple[dict[str, Any], list[str], bool]: + current, order, complete = _claude_code_read_transcripts(files, previous) + if current != previous: + try: + write_json(index_file, {"version": CLAUDE_CODE_INDEX_VERSION, "zone": _claude_code_index_zone(), "files": current}, sort_keys=True) + except Exception as exc: + warn(f"could not write the Claude scan index {index_file} ({exc})") + return current, order, complete + + loaded = _claude_code_index_stamp(index_file) + previous = {} if reread else _claude_code_load_index(index_file) + built = None + if reread or not _claude_code_index_fresh(previous, files): + with exclusive_lock(lock_file): + if reread or _claude_code_index_stamp(index_file) == loaded: + built = build(previous) + # Whoever held the lock wrote the index meanwhile, most likely all + # this reader needs. + if built is None: + previous = _claude_code_load_index(index_file) + if not _claude_code_index_fresh(previous, files): + with exclusive_lock(lock_file): + built = build(_claude_code_load_index(index_file)) + # An index fresh for every transcript read them all. + current, order, complete = built or (previous, [str(path) for path in files], True) + + owners: dict[Any, str | None] = {} + + def ours(model: Any) -> bool: + if model not in owners: + owners[model] = claude_code_owner(model, base_url) + return owners[model] == owner + + # One API message can show up in several transcripts (a resumed session + # carries the earlier ones). As within a file, the record with the highest + # output count wins, the last one visited breaking a tie. Only the tokens + # are deduplicated: every transcript holding a reply on a day was a session + # in use that day, so a fork that copied a session's replies never takes + # that session out of the count. + seen: dict[str, list[Any]] = {} + # When a reply was first written, whichever copy of it counts. + earliest: dict[str, float] = {} + active: set[tuple[Any, str, str]] = set() + for key in order: + for record in current[key].get("records") or []: + if ours(record[1]): + active.add((record[2], record[7], record[1])) + if record[10] is not None and record[10] < earliest.get(record[0], math.inf): + earliest[record[0]] = record[10] + held = seen.get(record[0]) + if held is not None and record[4] < held[4]: + continue + seen[record[0]] = record + + # A sidechain (a /btw side question) replays its parent's messages whole + # under new request ids; the parent's own copy is the one that counts. A + # subagent's own reply is a request of its own, which a gateway reusing one + # message id can still give its parent's id, so only a whole copy is one. + parents = {(record[8], record[7], *record[3:7]) for record in seen.values() if record[8] and not record[9]} + scan_day = today() + rows: list[UsageRow] = [] + replies: list[list[Any]] = [] + for record in seen.values(): + _, model, day, input_tokens, output_tokens, cache_read, cache_write, session_key, message_id, sidechain, _ = record + if sidechain and (message_id, session_key, input_tokens, output_tokens, cache_read, cache_write) in parents: + continue + replies.append(record) + if ours(model): + rows.append(UsageRow(day or scan_day, session_key, owner, model, input_tokens, output_tokens, cache_read, cache_write, prompts=0)) + for day, session_key, model, counts in _claude_code_unrecorded(current, order, replies, earliest): + if ours(model): + rows.append(UsageRow(day or scan_day, session_key, owner, model, *counts, prompts=0)) + + # A prompt is what a person typed, not each API request a tool loop makes, + # so replies carry no prompts and the typed ones are counted on their own, + # once each across transcripts, under whoever answered them. + typed: dict[str, list[Any]] = {} + for key in order: + for prompt in current[key].get("prompts") or []: + if ours(prompt[3]): + active.add((prompt[1], prompt[2], prompt[3] or "")) + typed.setdefault(prompt[0], prompt) + for _, day, session_key, model in typed.values(): + if ours(model): + rows.append(UsageRow(day or scan_day, session_key, owner, model or "", 0, 0, 0, 0, prompts=1)) + for day, session_key, model in sorted(active, key=lambda item: (str(item[0]), item[1], item[2])): + rows.append(UsageRow(day or scan_day, session_key, owner, model, 0, 0, 0, 0, prompts=0)) + return rows, complete + + +def _claude_code_unrecorded( + current: dict[str, Any], order: list[str], replies: list[list[Any]], earliest: dict[str, float] +) -> list[tuple[str | None, str, str, list[int]]]: + """What the cost-state lines say was used beyond the replies, as (day, + session, model, [input, output, cacheRead, cacheWrite]): the requests no + reply line records, a compaction's, the session title's, a side query's. + + A cost lineage (one `startTime`) holds every session that carried it on, + forks included, so each is counted once: its cost-states are taken in + time order, each against every reply of its sessions written by then + (subagents' included, which run under their parent's session), and only + what one says beyond the most already counted is new, on that cost-state's + day. A fork's total includes its parent's, and the parent's own replies + are among those it is held against, so the parent is never counted twice; + two branches that both went on are held against each other's replies too, + which can only count too little. Nor can a model name the replies don't + use add anything: no more counts, kind by kind, than the whole cost-state + says beyond all the replies. + """ + lineages: dict[str, list[list[Any]]] = {} + for key in order: + for cost in current[key].get("costs") or []: + lineages.setdefault(cost[0], []).append(cost) + if not lineages: + return [] + + by_session: dict[str, list[tuple[float, str, list[int]]]] = {} + for record in replies: + # A reply with no time is held against every cost-state. + moment = earliest.get(record[0], -math.inf) + by_session.setdefault(record[7], []).append((moment, record[1], record[3:7])) + + unrecorded: dict[tuple[str | None, str, str], list[int]] = {} + for events in lineages.values(): + sessions = {event[1] for event in events} + lineage_replies = sorted((reply for session in sessions for reply in by_session.get(session, ())), key=lambda reply: reply[0]) + # A cost-state with no time is the last, held against every reply. + events.sort(key=lambda event: math.inf if event[3] is None else event[3]) + replied: dict[str, list[int]] = {} + replied_total = [0, 0, 0, 0] + counted: dict[str, list[int]] = {} + position = 0 + for _, session, day, moment, models in events: + until = math.inf if moment is None else moment + while position < len(lineage_replies) and lineage_replies[position][0] <= until: + _, model, counts = lineage_replies[position] + held = replied.setdefault(model, [0, 0, 0, 0]) + for kind in range(4): + held[kind] += counts[kind] + replied_total[kind] += counts[kind] + position += 1 + beyond = {model: [max(0, totals[kind] - replied.get(model, (0, 0, 0, 0))[kind]) for kind in range(4)] for model, totals in models.items()} + for kind in range(4): + room = max(0, sum(totals[kind] for totals in models.values()) - replied_total[kind]) + claimed = sum(extra[kind] for extra in beyond.values()) + if claimed > room: + for extra in beyond.values(): + extra[kind] = extra[kind] * room // claimed + for model, extra in beyond.items(): + done = counted.setdefault(model, [0, 0, 0, 0]) + new = [max(0, extra[kind] - done[kind]) for kind in range(4)] + if any(new): + counted[model] = [max(done[kind], extra[kind]) for kind in range(4)] + held = unrecorded.setdefault((day, session, model), [0, 0, 0, 0]) + for kind in range(4): + held[kind] += new[kind] + return [(day, session, model, counts) for (day, session, model), counts in unrecorded.items()] + + +def claude_code_scan_match(claude_dir: Path | None = None) -> dict[str, str]: + """What a scan of Claude Code's transcripts was made under, for a + collector's cached_scan(match=...): where Claude Code is sent decides + whose its replies are, so a scan only answers for the endpoint it saw.""" + return {"claudeCodeEndpoint": claude_code_endpoint(claude_code_base_url(claude_dir))} + + +# ------------------------------------------------------- Codex rollouts +# +# Codex writes one JSONL rollout per session under $CODEX_HOME/sessions (and +# archived_sessions), naming the model provider it ran on. Every rollout is +# read here, once for every collector that bills one: the Codex collector its +# own `openai` rollouts, and the rest (`--oss` on a local daemon, a custom +# model_providers entry) the local collector for a self-hosted provider, +# Ollama Cloud for a cloud-tagged model the local daemon forwarded, and a +# provider's collector for a model_providers entry whose base_url is that +# provider's host (codex_route_owner). While openai_base_url sends the +# built-in provider elsewhere, its models are split the same way. + +# The Codex collector counts rollouts on the built-in provider, and ones too +# old to name a provider at all. +CODEX_NATIVE_PROVIDER = "openai" + +# Codex rollouts are read back this far, matching the Codex collector. +CODEX_WINDOW_DAYS = 30 + +# Per-file totals of Codex rollouts survive between runs, so a rescan only +# reads the files that changed. Bumping it retires every older record; version +# 2 drops totals cached while an untrusted time still counted on the file's day, +# version 3 those that counted every request as a prompt, version 4 those +# that missed the responses only Codex 0.163's token_usage_record lines name, +# version 5 those with no session, version 6 those that left a built-in +# rollout unread without saying so, version 7 those that booked a +# switched-to model's first requests to the model before it, and version 8 +# those that read no turns from a rollout naming no session_meta. +CODEX_ROLLOUT_SCHEMA = 8 + +_CODEX_MARKERS = ( + '"token_count"', '"token_usage_record"', '"turn_context"', '"thread_settings_applied"', '"session_meta"', '"user_message"', '"UserMessage"', +) + +# What a token_usage_record's running total is compared by. +_CODEX_USAGE_FIELDS = ("input_tokens", "cached_input_tokens", "cache_write_input_tokens", "output_tokens", "reasoning_output_tokens", "total_tokens") + + +def codex_home() -> Path: + return Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + + +# Records hold local days, so they are only true in the timezone that wrote +# them. +def _codex_cache_zone() -> list[Any]: + return [os.environ.get("TZ"), list(time.tzname), time.timezone, time.altzone] + + +def _codex_record_valid(value: Any) -> bool: + try: + counts = [totals for models in value["days"].values() for totals in models.values()] + return ( + type(value["mtime"]) in (int, float) and type(value["size"]) is int and isinstance(value["provider"], str) + and (value["session"] is None or type(value["session"]) is str) + and all(len(totals) == 5 and all(type(n) is int for n in totals) for totals in counts) + ) + except Exception: + return False + + +def _codex_typed_prompt(entry: dict[str, Any]) -> bool: + """Whether a rollout line is a prompt the user typed, as the Codex collector + reads it: an item_completed UserMessage in paginated rollouts, a plain + user_message event in legacy ones.""" + if entry.get("type") != "event_msg": + return False + payload = obj(entry.get("payload")) + if payload.get("type") == "user_message": + return payload.get("kind") in (None, "plain") + return payload.get("type") == "item_completed" and obj(payload.get("item")).get("type") == "UserMessage" + + +def _codex_settings_model(entry: dict[str, Any]) -> str: + """The model a thread_settings_applied event switches the thread to, or "". + Codex 0.163 writes one when `-m` on a resume or /model picks another model, + ahead of any request on it: a compaction the switch sends goes out on the + new model before that turn's turn_context names it.""" + payload = obj(entry.get("payload")) + if entry.get("type") != "event_msg" or payload.get("type") != "thread_settings_applied": + return "" + settings = obj(payload.get("thread_settings")) + return str(settings.get("model") or obj(obj(settings.get("collaboration_mode")).get("settings")).get("model") or "") + + +def _codex_typed_session(meta: dict[str, Any]) -> bool: + # A subagent's thread, or one of Codex's own (a guardian review, memory + # consolidation), is prompted by Codex, not by anyone typing: its requests + # still spend tokens, but its messages are no prompts. + source = meta.get("source") + return not (isinstance(source, dict) and ("subagent" in source or "internal" in source)) + + +def codex_rollout_session(meta: dict[str, Any], typed: bool, path: Any) -> str | None: + """The session a rollout's tokens count under. A thread someone started is + its own, by its thread id. Codex's own threads (a review, a guardian check, + a spawned agent) name the root thread that started them in session_id, or + their parent in parent_thread_id, and count under it, as a Claude Code + sidechain counts under its parent's sessionId: they spend tokens but start + no session. One that names no thread but its own, such as a memory + consolidation, counts under none.""" + thread = meta.get("id") + if typed: + return str(thread or path) + for starter in (meta.get("session_id"), meta.get("parent_thread_id")): + if starter and starter != thread: + return str(starter) + return None + + +# A rollout that only grew is read on from where the last read stopped, +# with the reader's state as it was there, when the first and last few +# kilobytes of what was read are still in place. +CODEX_ROLLOUT_HEAD_BYTES = 4096 + +_CODEX_MARKER_BYTES = tuple(marker.encode() for marker in _CODEX_MARKERS) + +# The line ends a text-mode read splits a rollout on. +_CODEX_LINE_ENDS = re.compile(r"\r\n|\r|\n") + + +def _codex_digest(handle: Any, length: int) -> str: + digest = hashlib.sha1() + handle.seek(0) + digest.update(handle.read(min(length, CODEX_ROLLOUT_HEAD_BYTES))) + handle.seek(max(0, length - CODEX_ROLLOUT_HEAD_BYTES)) + digest.update(handle.read(min(length, CODEX_ROLLOUT_HEAD_BYTES))) + return digest.hexdigest() + + +def _codex_response_key(response: tuple[Any, ...]) -> str: + return hashlib.blake2b(json.dumps(response).encode(), digest_size=8).hexdigest() + + +def _codex_resume(handle: Any, previous: dict[str, Any]) -> dict[str, Any] | None: + """The reader state a cached record stopped in, when the open rollout is + that same file, grown past it with what was read still in place, else + None. The record's days have the prompts no request had answered yet + taken back out, since the read goes on to find who answers them.""" + try: + state = dict(previous["resume"]) + info = os.fstat(handle.fileno()) + offset = state["offset"] + if not ( + type(offset) is int and info.st_size > previous["size"] and info.st_size >= offset + and state["file"] == [info.st_dev, info.st_ino] + and isinstance(state["model"], str) and all(type(state[flag]) is bool for flag in ("meta", "typed", "records", "done")) + and all(isinstance(day, str) for day in state["pending"]) and all(isinstance(key, str) for key in state["responses"]) + and state["head"] == _codex_digest(handle, offset) + ): + return None + days = {day: {model: list(totals) for model, totals in models.items()} for day, models in previous["days"].items()} + for day in state["pending"]: + totals = days[day][state["model"]] + totals[4] -= 1 + if totals[4] < 0: + return None + # A prompt's bucket holds nothing else until a request lands in it. + if not any(totals): + del days[day][state["model"]] + if not days[day]: + del days[day] + state["days"] = days + return state + except Exception: + return None + + +def read_codex_rollout(path: Path, stat_result: os.stat_result, previous: dict[str, Any] | None = None, native: bool = False) -> dict[str, Any]: + """One rollout's provider, session (codex_rollout_session) and day -> + model -> [input, output, cacheRead, cacheWrite, prompts]. A rollout on the + built-in provider, or naming none, is left unread and marked `skipped`, + unless `native` asks for it, as the Codex collector does and + openai_base_url can. + + A prompt counts on its own day, under the model that answered it: the one + in effect at the next request, since a legacy rollout can write the prompt + before the turn_context naming its model. A partial read has an mtime of + None, so it is never cached as the file's whole story. + + `previous`, the file's last record, is read on from where it stopped when + the file only grew; the record's `resume` holds what that takes. A + rollout dated by its mtime, or ending in a line still being written, is + read whole next time. + """ + record: dict[str, Any] = {"mtime": stat_result.st_mtime, "size": stat_result.st_size, "provider": "", "session": str(path), "days": {}} + current_model = "local" + seen_meta = False + typed = True + pending_prompts: list[str] = [] + previous_total_usage = None + # Codex 0.163 writes a token_usage_record for every response it is billed + # for, a compaction's included, where token_count only follows a sampling + # request. Once one appears, the records are the usage and the token_count + # lines after it repeat them; those before it still count. + seen_records = False + seen_responses: set[str] = set() + offset = 0 + done = False + resumable = True + resume: dict[str, Any] = {} + + def bucket(day: str, model: str) -> list[int]: + return record["days"].setdefault(day, {}).setdefault(model, [0, 0, 0, 0, 0]) + + try: + with path.open("rb") as handle: + state = _codex_resume(handle, previous) if previous else None + if previous and state is not None: + record["provider"] = previous["provider"] + record["session"] = previous["session"] + record["days"] = state["days"] + current_model = state["model"] + seen_meta = state["meta"] + typed = state["typed"] + pending_prompts = list(state["pending"]) + previous_total_usage = state["total"] + seen_records = state["records"] + seen_responses = set(state["responses"]) + offset = state["offset"] + done = state["done"] + handle.seek(offset) + for raw in [] if done else handle: + # A line still being written is read again once it is whole. + whole = raw.endswith(b"\n") + if whole: + offset += len(raw) + if not any(marker in raw for marker in _CODEX_MARKER_BYTES): + continue + text = raw.decode("utf-8", errors="replace") + for line in _CODEX_LINE_ENDS.split(text) if "\r" in text else (text,): + if not any(marker in line for marker in _CODEX_MARKERS): + continue + try: + entry = json.loads(line) + except Exception: + # A line nested too deep to parse is a RecursionError, not a + # ValueError; either way only that line is lost. + continue + if not isinstance(entry, dict): + continue + if not whole: + resumable = False + if entry.get("type") == "session_meta": + # A forked rollout copies its parent's session_meta after its own, + # so only the first one, as Codex reads it, says who served this. + if seen_meta: + continue + seen_meta = True + meta = obj(entry.get("payload")) + record["provider"] = str(meta.get("model_provider") or "").strip().lower() + # The built-in provider, or none recorded, is the Codex + # collector's unless openai_base_url sends it elsewhere. + if record["provider"] in ("", CODEX_NATIVE_PROVIDER): + if not native: + record["skipped"] = True + done = True + break + # A turn naming no model is Codex's own, as the Codex + # collector reads it. + current_model = "codex" + typed = _codex_typed_session(meta) + record["session"] = codex_rollout_session(meta, typed, path) + continue + if not seen_meta: + # A rollout from before Codex wrote session_meta names no + # provider, so its turns are the built-in provider's. + if not native: + record["skipped"] = True + done = True + break + seen_meta = True + current_model = "codex" + if _codex_typed_prompt(entry): + # A prompt with no time is the file's; one whose time can't be + # trusted counts nowhere. + timestamp = entry.get("timestamp") + if timestamp is None: + resumable = False + day = local_day(stat_result.st_mtime if timestamp is None else timestamp) + if typed and day is not None: + pending_prompts.append(day) + continue + if entry.get("type") == "turn_context": + payload = obj(entry.get("payload")) + current_model = str(payload.get("model") or payload.get("model_slug") or current_model) + continue + switched = _codex_settings_model(entry) + if switched: + current_model = switched + continue + if entry.get("type") == "token_usage_record": + seen_records = True + payload = obj(entry.get("payload")) + # One response, with the thread's running total after it. Only the + # same response id under an unchanged total is a repeat: a backend + # may reuse an id from turn to turn. + thread_usage = obj(payload.get("thread_token_usage")) + response = _codex_response_key((str(payload.get("response_id") or ""), *(token_count(thread_usage.get(field)) for field in _CODEX_USAGE_FIELDS))) + if response in seen_responses: + continue + seen_responses.add(response) + usage = obj(payload.get("usage")) + total_usage = None + else: + payload = entry.get("payload") or entry + if entry.get("type") == "response_item" and isinstance(payload, dict): + payload = payload.get("payload") or payload + if not isinstance(payload, dict) or payload.get("type") != "token_count" or seen_records: + continue + info = obj(payload.get("info")) + # total_token_usage is cumulative for the session; the last turn is + # what this one spent. + usage = obj(info.get("last_token_usage")) + total_usage = info.get("total_token_usage") + cache_read = token_count(usage.get("cached_input_tokens")) + cache_write = token_count(usage.get("cache_write_input_tokens")) + # Cached tokens are inside input_tokens, and reasoning inside + # output_tokens. + input_tokens = max(0, token_count(usage.get("input_tokens")) - cache_read - cache_write) + output_tokens = token_count(usage.get("output_tokens")) + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + # Quota updates can repeat the previous request's last_token_usage. + # Only an unchanged cumulative snapshot proves this is a repeat: + # separate requests may spend identical counts. + if isinstance(total_usage, dict) and total_usage: + if total_usage == previous_total_usage: + continue + previous_total_usage = total_usage + # A turn with no time is dated by the file; one whose time can't be + # trusted counts nowhere, as in the Codex collector. + timestamp = entry.get("timestamp") + if timestamp is None: + resumable = False + day = local_day(stat_result.st_mtime if timestamp is None else timestamp) + if day is None: + continue + totals = bucket(day, current_model) + totals[0] += input_tokens + totals[1] += output_tokens + totals[2] += cache_read + totals[3] += cache_write + # One prompt can take many requests, one per tool call it makes; the + # requests count no prompts of their own. + for prompt_day in pending_prompts: + bucket(prompt_day, current_model)[4] += 1 + pending_prompts = [] + if done: + break + if resumable: + opened = os.fstat(handle.fileno()) + resume = { + "file": [opened.st_dev, opened.st_ino], "offset": offset, "head": _codex_digest(handle, offset), + "model": current_model, "meta": seen_meta, "typed": typed, "pending": list(pending_prompts), + "total": previous_total_usage, "records": seen_records, "responses": sorted(seen_responses), "done": done, + } + except OSError: + record["mtime"] = None + resumable = False + # A prompt no request answered yet, or ever, was put to the last model. + for prompt_day in pending_prompts: + bucket(prompt_day, current_model)[4] += 1 + if resumable: + record["resume"] = resume + return record + + +# How long a refresh's start, as omarchy-agent-usage-update hands it to the +# collectors it runs, speaks for a forced read: longer than any refresh takes, +# so a stamp inherited from elsewhere soon stops standing in for --force. +CODEX_REFRESH_SECONDS = 300 + + +def _codex_refresh_started(now: float) -> float: + """When the refresh this collector runs in began, or `now` outside one.""" + try: + started = float(os.environ.get("OMARCHY_AGENT_USAGE_REFRESH_AT", "")) + except ValueError: + return now + return started if now - CODEX_REFRESH_SECONDS <= started <= now else now + + +def codex_rollout_cache(agent_id: str | None = None) -> Path | None: + """Where the Codex rollout records of this Codex home are kept, or None + when the cache root is unusable: that costs the reuse, never the scan. + A record keeps every turn whatever its provider, so every collector shares + the one file; `agent_id` names nothing, for callers from before that.""" + digest = hashlib.sha1(str(codex_home()).encode("utf-8", "surrogatepass")).hexdigest()[:16] + try: + return cache_root() / f"codex-rollouts-{digest}.json" + except OSError: + return None + + +def codex_rollouts(cache_file: Path | None, reuse: bool = True, native: bool = False) -> tuple[dict[str, dict[str, Any]], bool]: + """Every Codex rollout of the last CODEX_WINDOW_DAYS, as + read_codex_rollout() reads it, keyed by path, and whether every one was + read in full. Those on the built-in provider are read only with `native`, + while openai_base_url sends it elsewhere (codex_native_endpoint), and a + record that skipped one is read again once they are. + + Rollouts are append-only, so a file whose mtime and size both match its + cached record cannot have grown a turn since, and one that grew is read on + from where its record stopped; `reuse=False` reads them all again, but + takes what another forced read wrote since this one or its refresh + (OMARCHY_AGENT_USAGE_REFRESH_AT) began, so the collectors one refresh + forces read every rollout once between them, whenever each starts. A + record keeps every turn whatever its provider, so the caller decides which + providers and models are its own, and every caller shares the cache + codex_rollout_cache() names: any `cache_file` stands for it, and None + reads with no cache at all. Every collector reading Codex runs at once, + so one of them brings the cache up to date under its `.lock` while the + others wait and then find nothing left to read; a cache that already + answers for every rollout is read without waiting. Only the files a scan + saw are written back, so deleted rollouts and ones past the window leave + the cache. + """ + shared = None if cache_file is None else codex_rollout_cache() + + def rollouts() -> list[tuple[str, Path, os.stat_result]]: + found = [] + cutoff = time.time() - CODEX_WINDOW_DAYS * 24 * 60 * 60 + for root in (codex_home() / "sessions", codex_home() / "archived_sessions"): + if not root.is_dir(): + continue + for path in root.rglob("*.jsonl"): + try: + stat_result = path.stat() + except OSError: + continue + if stat_result.st_mtime >= cutoff: + found.append((str(path), path, stat_result)) + return found + + started = _codex_refresh_started(time.time()) + # When a forced read last wrote the cache: every record in it was read + # afresh since, replayed from that read or read since. + forced_at: list[float | None] = [None] + + def load() -> dict[str, Any]: + forced_at[0] = None + if shared is None: + return {} + try: + payload = read_json(shared) + except OSError: + payload = {} + if payload.get("schemaVersion") != CODEX_ROLLOUT_SCHEMA or payload.get("zone") != _codex_cache_zone(): + return {} + stamp = payload.get("forcedAt") + stamp = stamp if type(stamp) in (int, float) else None + # A forced read reuses only what another forced read of the same refresh + # wrote, so a refresh rereads every rollout once, whichever collector gets + # to the lock first, and never replays an older cache. + if not reuse and (stamp is None or stamp < started): + return {} + forced_at[0] = stamp + return {path: record for path, record in obj(payload.get("files")).items() if _codex_record_valid(record)} + + def current(record: dict[str, Any] | None, stat_result: os.stat_result) -> bool: + return ( + record is not None and record["mtime"] == stat_result.st_mtime and record["size"] == stat_result.st_size + and not (native and record.get("skipped")) + ) + + def scan(files: list[tuple[str, Path, os.stat_result]], cached: dict[str, Any]) -> tuple[dict[str, dict[str, Any]], bool]: + records: dict[str, dict[str, Any]] = {} + kept: dict[str, dict[str, Any]] = {} + complete = True + for key, path, stat_result in files: + record = cached.get(key) + if not current(record, stat_result): + # A built-in rollout once skipped is read from the start. + record = read_codex_rollout(path, stat_result, None if record and record.get("skipped") else record, native) + if record["mtime"] is None: + complete = False + else: + kept[key] = record + records[key] = record + if shared is not None and kept != cached: + try: + stamp = forced_at[0] if reuse else time.time() + write_json(shared, {"schemaVersion": CODEX_ROLLOUT_SCHEMA, "zone": _codex_cache_zone(), "forcedAt": stamp, "files": kept}) + except OSError as exc: + warn(f"could not write file cache ({exc})") + return records, complete + + if shared is None: + return scan(rollouts(), {}) + files = rollouts() + cached = load() + if reuse and len(cached) == len(files) and all(current(cached.get(key), stat_result) for key, _, stat_result in files): + return {key: cached[key] for key, _, _ in files}, True + with exclusive_lock(shared.with_name(shared.name + ".lock")): + # Whoever held the lock meanwhile most likely read what this scan is after. + return scan(rollouts(), load()) + + +def _codex_config(path: Path) -> dict[str, Any]: + try: + with path.open("rb") as handle: + return tomllib.load(handle) + except (OSError, UnicodeError, tomllib.TOMLDecodeError): + return {} + + +def codex_native_key(pid: Any) -> bool: + """Whether a codex_config_urls() key is the built-in provider's + openai_base_url, plain or a profile's (`openai@`).""" + return pid == CODEX_NATIVE_PROVIDER or str(pid).startswith(CODEX_NATIVE_PROVIDER + "@") + + +def codex_config_urls() -> dict[str, str]: + """provider id -> base_url, from the Codex CLI's own config: the system + layer, the user's config.toml over it, and the profiles + (.config.toml) Codex layers over that when one is selected. The + built-in `openai` provider's is its openai_base_url. + + Codex 0.163 applies a profile only to a run given --profile, and its + rollouts still name only `openai`, so only the system layer and + config.toml say where the built-in provider goes. A profile's + openai_base_url is kept as `openai@` for the model it sets, the one + thing its rollouts can be told apart by: never an OpenAI model, which + plain `codex` runs too, nor the one the plain config sets, and none when + two profiles send that model to different places.""" + home = codex_home() + urls: dict[str, str] = {} + profiles: dict[str, str | None] = {} + plain_model = "" + for path in (Path("/etc/codex/config.toml"), home / "config.toml", *sorted(home.glob("*.config.toml"))): + config = _codex_config(path) + providers = config.get("model_providers") + if isinstance(providers, dict): + urls.update((str(name).lower(), str(entry.get("base_url") or "")) for name, entry in providers.items() if isinstance(entry, dict)) + model = config.get("model").strip().lower() if isinstance(config.get("model"), str) else "" + base = config.get("openai_base_url") + if path.name != "config.toml": + if model and isinstance(base, str) and not codex_openai_model(model): + profiles[model] = base if profiles.get(model, base) == base else None + continue + # Where the built-in provider goes in place of OpenAI's API. Codex 0.163 + # reads it from the config only, never from OPENAI_BASE_URL. + if isinstance(base, str): + urls[CODEX_NATIVE_PROVIDER] = base + plain_model = model or plain_model + for model, base in profiles.items(): + if base is not None and model != plain_model: + urls[f"{CODEX_NATIVE_PROVIDER}@{model}"] = base + return urls + + +def codex_providers_path() -> Path: + return Path(os.environ.get("XDG_STATE_HOME") or (Path.home() / ".local" / "state")) / "omarchy" / "agents" / "local-codex-providers.json" + + +def codex_provider_urls() -> dict[str, str]: + """provider id -> base_url, for every Codex provider config.toml has named. + + A rollout records only its provider's id, never the URL it reached, so a + custom provider that leaves config.toml would take its history with it. + Each run remembers the endpoints the config names, per Codex home, and a + provider the config no longer has keeps the last one it had. The current + config always wins. A provider only ever given on the command line (`-c + model_providers...`) or in a project's .codex/config.toml is never seen + here: local.json names those for the local collector. The built-in + provider's openai_base_url is never remembered: once the config drops it, + Codex is back on OpenAI's API, as Claude Code is once ANTHROPIC_BASE_URL + goes. + """ + path = codex_providers_path() + homes = read_json(path).get("homes") + homes = homes if isinstance(homes, dict) else {} + key = str(codex_home()) + known = homes.get(key) + known = {str(pid): url for pid, url in known.items() if isinstance(url, str) and url and not codex_native_key(pid)} if isinstance(known, dict) else {} + urls = {**known, **codex_config_urls()} + remembered = {pid: url for pid, url in urls.items() if url and not codex_native_key(pid)} + if remembered != known: + try: + path.parent.mkdir(parents=True, exist_ok=True) + write_json(path, {"homes": {**homes, key: remembered}}, sort_keys=True) + except OSError as exc: + # Remembering costs nothing when it fails: the current config still + # classifies. + warn(f"could not remember Codex providers ({exc})") + return urls + + +def codex_openai_model(model: Any) -> bool: + """Whether a model is one OpenAI serves Codex (gpt-5.1-codex, + codex-auto-review, o3), or none is named; gpt-oss is open weights anyone + runs.""" + name = str(model or "").strip().lower().rpartition("/")[2] + if name.startswith("gpt-oss"): + return False + return not name or name.startswith(("gpt-", "chatgpt-")) or "codex" in name or re.fullmatch(r"o[1-9](-.*)?", name) is not None + + +def codex_native_class(base: Any) -> str: + """Where an openai_base_url sends the built-in provider: "" for OpenAI's + own API (no override), a provider collector's id, `local`, or `other` for + a host nothing here can place.""" + host = url_host(base) + if not host or host in ("openai.com", "chatgpt.com") or host.endswith((".openai.com", ".chatgpt.com")): + return "" + return endpoint_owner(base) or "other" + + +def codex_native_endpoint(urls: dict[str, str], model: Any = None) -> str: + """codex_native_class() of the openai_base_url a built-in turn on `model` + went to: the profile's that sets that model, else the plain config's.""" + profile = f"{CODEX_NATIVE_PROVIDER}@{str(model or '').strip().lower()}" + return codex_native_class(urls[profile] if profile in urls else urls.get(CODEX_NATIVE_PROVIDER, "")) + + +def codex_native_routed(urls: dict[str, str]) -> bool: + """Whether any built-in rollout may be another collector's, so + codex_rollouts() has to read them (its `native`).""" + return any(codex_native_class(url) for pid, url in urls.items() if codex_native_key(pid)) + + +def codex_route_owner(provider: Any, model: Any, urls: dict[str, str]) -> str | None: + """The collector that bills a Codex rollout on `provider`, with the model + that answered, or None for one no provider's collector counts (the local + collector decides which of those `--oss` and local.json ran itself). + + The built-in provider is `codex`'s until openai_base_url (the profile's + for the model a profile sets, else the plain config's) sends it + elsewhere. A provider's host there bills all of it. Otherwise an Ollama + cloud tag is Ollama's, and OpenAI's own models stay `codex`'s through a + proxy that may forward to it; a self-hosted URL makes any other model + `local`'s, and a host nothing here can place counts it nowhere, as + claude_code_owner() reads ANTHROPIC_BASE_URL. + + A model_providers entry whose base_url is a provider's host is that + provider's, whatever its id; one at a host nothing here can place is no + provider's. Past that, `ollama-cloud`, and a cloud-tagged model through + the local daemon's `ollama`, are Ollama Cloud's, and any other model on + the user's own hardware is `local`'s, whatever the entry's id: the base + URL decides, not an id another collector claims elsewhere.""" + name = str(provider or "").strip().lower() + if name in ("", CODEX_NATIVE_PROVIDER): + endpoint = codex_native_endpoint(urls, model) + if endpoint not in ("", "local", "other"): + return endpoint + if endpoint and cloud_model(model): + return "ollama" + if not endpoint or codex_openai_model(model): + return "codex" + return "local" if endpoint == "local" else None + owner = endpoint_owner(urls.get(name)) + if owner and owner != "local": + return owner + if name == "ollama-cloud" or (name == "ollama" and cloud_model(model)): + return "ollama" + return "local" if owner == "local" and not cloud_model(model) else None + + +def codex_scan_match() -> dict[str, Any]: + """For a collector's cached_scan(match=...): which collector each Codex + provider's base_url bills, which the config can change without a rollout + changing. The built-in provider's is read as codex_route_owner() reads + it, which tells OpenAI's own API from a host nothing can place.""" + return {"codexRoutes": [ + [pid, codex_native_class(url) if codex_native_key(pid) else endpoint_owner(url)] for pid, url in sorted(codex_provider_urls().items()) + ]} + + +def scan_codex_routed(owner: str, reuse: bool = True) -> tuple[list[UsageRow], bool]: + """The Codex rollout totals codex_route_owner() says `owner` bills, as rows + (one per rollout, day and model, prompts counted), and whether every + rollout was read. Only Ollama Cloud bills a provider no base_url names, so + a collector no configured provider points at reads nothing.""" + urls = codex_provider_urls() + if owner != "ollama" and not any(endpoint_owner(url) == owner for url in urls.values()): + return [], True + records, complete = codex_rollouts(codex_rollout_cache(), reuse=reuse, native=codex_native_routed(urls)) + rows = [ + UsageRow(day, record["session"], owner, model, *totals[:4], prompts=totals[4]) + for record in records.values() + for day, models in record["days"].items() + for model, totals in models.items() + if codex_route_owner(record["provider"], model, urls) == owner + ] + return rows, complete + + +# ------------------------------------------------------------ Kimchi + + +def kimchi_sessions_root() -> Path: + """Where Kimchi's harness keeps its pi-format sessions: under + KIMCHI_CODING_AGENT_DIR when set, as the CLI itself reads it, else + ~/.config/kimchi/harness. Kimchi has no profiles that move it. Its own + gateway routes (kimchi, kimchi-dev, kimchi-dev/, + kimchi-experimental) are the kimchi collector's here and in pi's and omp's + trees alike; a session `--provider` sent elsewhere (OpenRouter) is that + provider's.""" + return Path(os.environ.get("KIMCHI_CODING_AGENT_DIR") or (Path.home() / ".config" / "kimchi" / "harness")) / "sessions" + + +# ---------------------------------------------------- DeepSeek Harness +# +# dsh keeps one append-only event log per session under +# $DSH_HOME/sessions///, Zstandard-compressed by default. +# Each billed model attempt settles as an assistant/message (or a log-only +# assistant/attempt) carrying the provider's usage; the request/header before +# it names the provider route and model. Its DeepSeek routes are the deepseek +# collector's, and a pi-ai route on another provider (OpenRouter, say) that +# provider's, so the reader hands each collector the routes it takes. + +# One canonical session log generation: session.jsonl, session.v4.jsonl, and +# their Zstandard twins. +DSH_LOG_NAME = re.compile(r"session(?:\.v([1-9][0-9]*))?\.jsonl(\.zstd)?") + +# The most one compressed session log may inflate to. A real log, even a +# long-running agent session's, is a few MiB of JSON lines, so 64 MiB is far +# past any of them; a log that inflates further is a corrupt or hostile +# frame, and the scan stops there rather than filling memory. +DSH_MAX_LOG_BYTES = 64 * 1024 * 1024 + + +def dsh_home() -> Path: + # An empty or blank DSH_HOME reads as unset, as dsh itself treats it. + configured = os.environ.get("DSH_HOME", "").strip() + return Path(configured).expanduser() if configured else Path.home() / ".dsh" + + +def dsh_session_logs(root: Path) -> tuple[list[Path], bool]: + """The current log of each session, and whether every directory could be + listed: a walk that couldn't read one has missed its logs.""" + # The newest generation in each session directory is the current log; + # older ones are kept only as migration history. + newest: dict[Path, tuple[int, Path]] = {} + unlisted: list[OSError] = [] + for directory, _, names in root.walk(on_error=unlisted.append): + for name in names: + match = DSH_LOG_NAME.fullmatch(name) + path = directory / name + if not match or not path.is_file(): + continue + version = int(match.group(1) or 0) + current = newest.get(directory) + if current is None or version > current[0]: + newest[directory] = (version, path) + return [path for _, path in sorted(newest.values(), key=lambda item: str(item[1]))], not unlisted + + +def dsh_decompress(data: bytes) -> bytes | None: + """Every complete Zstandard frame of a log, or None past DSH_MAX_LOG_BYTES. + + Decompressed a bounded piece at a time, so a corrupt or hostile frame that + inflates without end stops at the cap instead of filling memory. + """ + chunks: list[bytes] = [] + total = 0 + while data: + decompressor = zstd.ZstdDecompressor() + pending = data + try: + while True: + chunk = decompressor.decompress(pending, max_length=DSH_MAX_LOG_BYTES - total + 1) + pending = b"" + total += len(chunk) + if total > DSH_MAX_LOG_BYTES: + return None + chunks.append(chunk) + if decompressor.eof or decompressor.needs_input: + break + except zstd.ZstdError: + break + if not decompressor.eof: + break + data = decompressor.unused_data + return b"".join(chunks) + + +def dsh_log_lines(path: Path) -> tuple[list[bytes], bool]: + """The complete records of one log, and whether all of it could be read. + + A torn final line or frame from a crash mid-write is dropped, as dsh's own + reader does, without making the scan incomplete. + """ + data = path.read_bytes() + if path.name.endswith(".zstd"): + if zstd is None: + warn("this Python has no Zstandard support; DeepSeek Harness sessions were not scanned") + return [], False + data = dsh_decompress(data) + if data is None: + warn(f"a DeepSeek Harness session log inflates past {DSH_MAX_LOG_BYTES >> 20} MiB and was not scanned") + return [], False + lines = data.split(b"\n") + # Whatever follows the last newline is an unfinished record. + return lines[:-1], True + + +def dsh_usage_sample(event: dict[str, Any]) -> dict[str, Any] | None: + data = obj(event.get("data")) + if event.get("type") == "assistant/message" and isinstance(data.get("usage"), dict): + return data["usage"] + stream = data.get("stream") + if not isinstance(stream, list): + return None + for record in reversed(stream): + chunk = obj(obj(record).get("chunk")) + if obj(record).get("type") == "chunk" and chunk.get("type") == "usage": + return obj(chunk.get("usage")) + return None + + +def _dsh_rows(path: Path, accept: Accept) -> tuple[list[UsageRow], bool]: + lines, complete = dsh_log_lines(path) + if not lines: + return [], complete + try: + header = obj(json.loads(lines[0])) + except Exception: + # A line nested too deep to parse is a RecursionError, not a + # ValueError; either way only that line (here, its log) is lost. + return [], complete + if header.get("type") != "session": + return [], complete + events = [] + for line in lines[1:]: + try: + event = json.loads(line) + except Exception: + continue + if isinstance(event, dict): + events.append(event) + + # A fork's log opens with a copy of its parent's history; the last + # `session/end-seed` marked inherited is where its own spending begins. + start = 0 + if header.get("isSeeded") is True: + for index, event in enumerate(events): + if event.get("type") == "session/end-seed" and obj(event.get("data")).get("inherited") is True: + start = index + 1 + + session = "dsh:" + str(header.get("id") or path.parent.name) + provider, model, request = "", "", {} + # One slot per billed attempt. A settlement replaces the streaming sample + # of the same turn and step, until a retry starts a new attempt, the same + # fold dsh's own token meter runs. + slots: list[list[Any]] = [] + last: tuple[Any, Any] | None = None + # The messages someone typed (a user message whose source is `user`, as + # dsh's own session title reads them) since the last attempt: the next + # attempt answers them. A tool step, a retry, and a turn dsh starts on its + # own (a schedule, a goal round) answer none. + typed = 0 + for index, event in enumerate(events): + kind = event.get("type") + data = obj(event.get("data")) + if kind == "user/message": + if index >= start and obj(data.get("source")).get("kind") == "user": + typed += 1 + continue + if kind == "request/header": + request = obj(obj(data.get("header")).get("config")) + provider, model = str(request.get("provider") or ""), str(request.get("model") or "") + continue + if kind == "llm/retry-started": + if last == (data.get("turn"), data.get("step")): + last = None + continue + if kind not in ("assistant/message", "assistant/attempt"): + continue + usage = dsh_usage_sample(event) + if usage is None: + continue + slot = (data.get("turn"), data.get("step")) + sample = [ + number(usage.get("inputTokens")), + number(usage.get("outputTokens")), + number(usage.get("cacheReadTokens")), + number(usage.get("cacheWriteTokens")), + ] + entry = [index >= start and accept(provider, model, request), local_day(event.get("time")), provider, model, sample] + if last == slot and slots: + slots[-1] = entry + [slots[-1][5]] + else: + slots.append(entry + [typed]) + typed = 0 + last = slot + + rows = [] + for counted, day, slot_provider, slot_model, (input_tokens, output_tokens, cache_read, cache_write), prompts in slots: + if counted and day is not None and (input_tokens or output_tokens or cache_read or cache_write or prompts): + rows.append(UsageRow(day, session, slot_provider, slot_model, input_tokens, output_tokens, cache_read, cache_write, prompts)) + return rows, complete + + +def scan_dsh_sessions(accept: Accept) -> tuple[list[UsageRow], bool]: + """The billed attempts in every dsh session log on a route `accept` + (provider, model, request config) takes, a fork's inherited history left + out, and whether every log was read in full. Each attempt carries the + typed messages it answered as its prompts: the first attempt of a typed + turn answers them, its tool steps and retries none.""" + root = dsh_home() / "sessions" + if not root.is_dir(): + return [], True + try: + logs, complete = dsh_session_logs(root) + except OSError: + return [], False + rows: list[UsageRow] = [] + for path in logs: + try: + found, read = _dsh_rows(path, accept) + except OSError: + complete = False + continue + rows.extend(found) + complete = read and complete + return rows, complete + + + +# ------------------------------------------------------- MiniMax Code +# +# MiniMax Code keeps one directory per session under its data dir's +# v2/sessions, with a messages.jsonl whose assistant lines name the route +# they ran on: its own (`minimax`, `minimax_api`, `minimax_oauth`) bill +# MiniMax, and `custom_provider:` (a custom endpoint, or a models.dev +# preset such as Z.AI or OpenRouter) bills whoever runs that provider's +# baseURL in config.yaml. + +# MiniMax Code's own routes: `minimax` is its signed-in Token Plan (or a key +# set on it), `minimax_api` the MiniMax API key it reserves, and +# `minimax_oauth` the name its provider list gives the sign-in. +MCODE_PROVIDERS = ("minimax", "minimax_api", "minimax_oauth") + + +def mcode_data_dir() -> Path: + for name in ("MINIMAX_DATA_DIR", "MAVIS_DATA_DIR"): + value = os.environ.get(name, "").strip() + if value: + return Path(os.path.abspath(os.path.expanduser(value))) + return Path.home() / ".minimax" + + +def _yaml_blocks(text: str) -> dict[str, Any]: + """The block mappings of a YAML file, scalars as plain strings: enough to + read MiniMax Code's config.yaml without a YAML dependency. Lists, flow + values, anchors and multi-line scalars aren't read.""" + root: dict[str, Any] = {} + stack: list[tuple[int, dict[str, Any]]] = [(-1, root)] + for raw in text.splitlines(): + body = raw.rstrip() + stripped = body.lstrip(" ") + if not stripped or stripped.startswith(("#", "-")): + continue + indent = len(body) - len(stripped) + key, colon, value = stripped.partition(":") + if not colon or (value and not value.startswith(" ")): + continue + while stack[-1][0] >= indent: + stack.pop() + key = key.strip().strip("'\"") + value = value.split(" #", 1)[0].strip() + if value: + stack[-1][1][key] = value.strip("'\"") + else: + child: dict[str, Any] = {} + stack[-1][1][key] = child + stack.append((indent, child)) + return root + + +def mcode_routes(data_dir: Path | None = None) -> dict[str, str]: + """custom_provider name -> the collector that bills its baseURL ("" for + one nothing here can place), from MiniMax Code's config.yaml.""" + try: + text = ((data_dir or mcode_data_dir()) / "config.yaml").read_text(encoding="utf-8", errors="replace") + except OSError: + return {} + providers = _yaml_blocks(text).get("custom_provider") + if not isinstance(providers, dict): + return {} + routes = {} + for name, entry in providers.items(): + options = entry.get("options") if isinstance(entry, dict) else None + routes[str(name)] = endpoint_owner(options.get("baseURL")) if isinstance(options, dict) else "" + return routes + + +def mcode_route_owner(provider: Any, routes: dict[str, str], model: Any = None) -> str | None: + """The collector that bills a MiniMax Code reply by `model` on route + `provider`. A reply naming none predates third-party routes and ran on + MiniMax. A custom provider is its baseURL's (a provider's host, or the + user's own hardware, where a cloud-tagged model only passed through a + local Ollama daemon to Ollama's servers); one gone from the config, or + behind a gateway nothing here can place, is nobody's here.""" + if provider in (None, "", *MCODE_PROVIDERS): + return "minimax" + name = str(provider) + if name.startswith("custom_provider:"): + owner = routes.get(name.partition(":")[2]) or None + return "ollama" if owner == "local" and cloud_model(model) else owner + return None + + +def mcode_scan_match() -> dict[str, Any]: + """For a collector's cached_scan(match=...): which collector each custom + provider bills, which the config can change without a session changing.""" + return {"mcodeRoutes": [list(item) for item in sorted(mcode_routes().items())]} + + +def _mcode_tokens(usage: dict[str, Any]) -> tuple[int, int, int, int]: + input_tokens = token_count(usage.get("input")) + output_tokens = token_count(usage.get("output")) + cache_read = token_count(usage.get("cacheRead")) + cache_write = token_count(usage.get("cacheWrite")) + if not (input_tokens or output_tokens): + # Only a total: MiniMax's input excludes the cache tokens, which keep + # their own buckets. + input_tokens = max(0, token_count(usage.get("totalTokens")) - cache_read - cache_write) + return input_tokens, output_tokens, cache_read, cache_write + + +def _mcode_day(message: dict[str, Any], entry: dict[str, Any]) -> str | None: + """The day of a MiniMax Code line: the message's time, else the line's. + A line with neither, or with neither trusted, is skipped.""" + for value in (message.get("timestamp"), entry.get("timestamp")): + if value is not None and (day := local_day(value)): + return day + return None + + +def scan_mcode(owner: str) -> tuple[list[UsageRow], bool]: + """MiniMax Code's replies and typed prompts that `owner` bills, as rows + (replies with no prompts of their own, prompts with no tokens), and + whether every session was read. A prompt goes with the route and model of + the reply that answers it, so a collector keeping some models alone keeps + their prompts too; one left unanswered with the session's last route and + no model.""" + root = mcode_data_dir() / "v2" / "sessions" + if not root.is_dir(): + return [], True + routes = mcode_routes() + # A directory the walk can't list has sessions the scan missed. + unlisted: list[OSError] = [] + found = [directory / "messages.jsonl" for directory, _, names in root.walk(on_error=unlisted.append) if "messages.jsonl" in names] + complete = not unlisted + rows: list[UsageRow] = [] + for path in sorted(found): + # One directory per session; its messages carry no session id of their own. + session = "mcode:" + str(path.parent) + pending: list[str] = [] + route = "minimax" + try: + with path.open("r", encoding="utf-8", errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except Exception: + # A line nested too deep to parse is a RecursionError, not a + # ValueError; either way only that line is lost. + continue + message = entry.get("message") if isinstance(entry, dict) else None + if not isinstance(message, dict): + continue + day = _mcode_day(message, entry) + if day is None: + continue + # MiniMax Code books a typed prompt and the assistant's usage on + # lines of their own. + if message.get("role") == "user": + pending.append(day) + continue + if message.get("role") != "assistant": + continue + model = str(message.get("model") or entry.get("model") or "minimax").strip() or "minimax" + route = mcode_route_owner(message.get("provider"), routes, model) + if route == owner: + rows.extend(UsageRow(prompt_day, session, owner, model, 0, 0, 0, 0) for prompt_day in pending) + pending = [] + usage = message.get("usage") + if route != owner or not isinstance(usage, dict): + continue + input_tokens, output_tokens, cache_read, cache_write = _mcode_tokens(usage) + if not (input_tokens or output_tokens or cache_read or cache_write): + continue + rows.append(UsageRow(day, session, owner, model, input_tokens, output_tokens, cache_read, cache_write, prompts=0)) + if route == owner: + rows.extend(UsageRow(prompt_day, session, owner, "", 0, 0, 0, 0) for prompt_day in pending) + except OSError: + complete = False + return rows, complete + + +# -------------------------------------------- harnesses routed by base URL +# +# Claude Code, MiniMax Code and Codex send to whatever endpoint they are +# given, and record only a model or a route's name, so the collector of every +# provider they can be pointed at (and `local`, for the user's own hardware) +# reads them, each keeping what claude_code_owner(), mcode_route_owner() and +# codex_route_owner() say it bills. The local collector reads Codex with a +# rule of its own, since it also counts provider ids no base_url names. + + +def routed_sources() -> list[Path]: + """The stores scan_routed() reads, for a collector's scan sources.""" + return [claude_code_dir() / "projects", mcode_data_dir(), codex_home()] + + +def routed_match() -> dict[str, Any]: + """What scan_routed() decided owners under, for cached_scan(match=...).""" + return {**claude_code_scan_match(), **mcode_scan_match(), **codex_scan_match()} + + +def scan_routed(owner: str, reuse: bool = True) -> tuple[list[UsageRow], bool]: + """Claude Code's, MiniMax Code's and Codex's usage that `owner` bills, and + whether every session was read. `reuse=False` reads every Codex rollout + again rather than from its per-file cache.""" + claude_rows, claude_complete = claude_code_usage(owner) + mcode_rows, mcode_complete = scan_mcode(owner) + codex_rows, codex_complete = ([], True) if owner == "local" else scan_codex_routed(owner, reuse) + return claude_rows + mcode_rows + codex_rows, claude_complete and mcode_complete and codex_complete + + +# ------------------------------------------------------------------ limits + + +def limits_cache_path(agent_id: str, *identity: str, kind: str = "limits") -> Path: + """`--.json`, keyed by who the credential signs in as (or + the credential itself), so another sign-in never shows this one's numbers. + `kind` names an answer kept beside the limits (credits, a balance, usage + rows) under the same key.""" + # surrogatepass: an identity read from JSON can hold a lone surrogate, and + # that must not cost the record. + digest = hashlib.sha256("\0".join(identity).encode("utf-8", "surrogatepass")).hexdigest()[:16] + return cache_root() / f"{agent_id}-{kind}-{digest}.json" + + +def limit_window_open(entry: dict[str, Any], now: datetime) -> bool: + # A window with no reset time, or one that will not parse, is no reason to + # doubt a real number. + raw = str(entry.get("resetsAt") or "") + if raw == "": + return True + try: + resets_at = datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return True + if resets_at.tzinfo is None: + resets_at = resets_at.replace(tzinfo=timezone.utc) + return resets_at > now + + +def open_limits(entries: Any) -> list[dict[str, Any]]: + """Kept limit windows, as they read now. + + A window whose reset has passed comes back at percent 0 with resetsAt "": + the allowance it measured has been renewed, so it is untouched rather than + unknown, and never dropped. + """ + if not isinstance(entries, list): + return [] + now = datetime.now(timezone.utc) + return [ + entry if limit_window_open(entry, now) else dict(entry, percent=0.0, resetsAt="") + for entry in entries if isinstance(entry, dict) + ] + + +class LimitsCache: + """One credential's last good limits answer. + + Read whatever the file time, so a clock that went back keeps it. `limits` + are the kept windows as they read now, `fetched_at` when they were + measured (epoch ms, 0 when never), and `data` the whole file for anything + else a collector keeps beside them (a plan, a no-plan marker, a balance). + + `kind` keeps another answer of the same credential beside its limits + (`--.json`): credits, a balance, usage rows. at() opens a + cache under a name a collector already gave it. + """ + + def __init__(self, agent_id: str, *identity: str, kind: str = "limits") -> None: + self._open(limits_cache_path(agent_id, *identity, kind=kind)) + + @classmethod + def at(cls, path: Path) -> LimitsCache: + """The cache kept at `path`, for a name that predates the shared one.""" + cache = cls.__new__(cls) + cache._open(Path(path)) + return cache + + def _open(self, path: Path) -> None: + self.path = path + self.lock_path = path.with_suffix(".lock") + self.reload() + + def reload(self) -> None: + self.data = read_json(self.path) + self.fetched_at = stamp(self.data.get("fetchedAtMs")) + self.limits = open_limits(self.data.get("limits")) + + def younger_than(self, seconds: float) -> bool: + return age_below(self.fetched_at / 1000, seconds) + + def fresh(self, force: bool, seconds: float = PROBE_MIN_INTERVAL_SECONDS, allow_empty: bool = False) -> bool: + """Whether the kept answer stands in for a probe: recent, and no --force. + + Only kept limits count, unless `allow_empty`: then any answer stamped + inside the age guard does, limits or none, for a service whose answer + can hold no windows (a plan without quotas, credits, a balance). + """ + if force or not self.younger_than(seconds): + return False + return bool(self.limits) or (allow_empty and self.fetched_at > 0) + + def kept(self) -> tuple[list[dict[str, Any]], int]: + """The stale fallback: the kept limits and their limitsFetchedAt.""" + return self.limits, self.fetched_at if self.limits else 0 + + def save(self, limits: list[dict[str, Any]], only_if_newer: bool = False, **extra: Any) -> int: + """Keep a good answer, stamped now; returns the stamp. A cache that can't + be written costs only the reuse, never the answer. + + With `only_if_newer`, a file another run stamped later is left as it is, + so overlapping runs (the panel's own refresh and a sooner retry) never + roll a newer answer back. This run still uses its own answer. Only a + stamp up to five minutes ahead counts as newer: one further ahead was + written before the clock went back and would otherwise never be + replaced. + """ + fetched_at = now_ms() + payload = {"fetchedAtMs": fetched_at, "limits": limits, **extra} + try: + kept = stamp(read_json(self.path).get("fetchedAtMs")) if only_if_newer else 0 + if not fetched_at < kept <= fetched_at + 300_000: + write_json(self.path, payload) + except Exception as exc: + warn(f"could not write limits cache ({exc})") + self.data = payload + self.fetched_at = fetched_at + self.limits = open_limits(limits) + return fetched_at + + def lock(self) -> contextlib.AbstractContextManager[None]: + """flock for a probe, so concurrent runs ask the service once.""" + return exclusive_lock(self.lock_path) + + +# -------------------------------------------------------------------- http + + +class _RefuseRedirects(urllib.request.HTTPRedirectHandler): + # A redirect would carry the credential to whatever host it names, so it + # answers as the 3xx status it is. + def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: ANN001 + return None + + +# Every collector request goes through urllib.request.urlopen, which the +# collectors' tests replace, so the refusal is installed where urlopen looks. +urllib.request.install_opener(urllib.request.build_opener(_RefuseRedirects())) + + +class HttpError(Exception): + """A request that brought back no usable answer. + + `status` is the HTTP status of a refusal (a redirect included), 0 when + there is none, and `headers` the refusal's own (a 429's retry-after). + `transport` means the connection failed or dropped before any reply (DNS, + no route, a refused or reset connection, a timeout), which is the only + failure worth retryAdvised. Otherwise the server answered, with a body + that isn't JSON or is too large to be an answer (`invalid`), or with a + reply cut short or garbled on the way (`cut_short`, invalid too): an + answer all the same, so no reason to ask again sooner. + """ + + def __init__(self, status: int = 0, transport: bool = False, body: bytes = b"", headers: Any = None, cut_short: bool = False) -> None: + super().__init__(f"HTTP {status}" if status else ("unreachable" if transport else ("reply cut short" if cut_short else "invalid answer"))) + self.status = status + self.transport = transport + self.invalid = not status and not transport + self.cut_short = cut_short + self.body = body + self.headers = headers + + +class _TooSlow(Exception): + """A reply still arriving at the deadline.""" + + +def _read_by(response: Any, size: int, deadline: float) -> bytes: + """Up to `size` bytes of a reply, by `deadline`. The socket timeout only + bounds each wait, so a server dripping a byte at a time would otherwise + hold a collector forever.""" + if not isinstance(response, http.client.HTTPResponse): + return response.read(size) + chunks, got = [], 0 + while got < size: + if time.monotonic() > deadline: + raise _TooSlow() + chunk = response.read1(min(65536, size - got)) + if not chunk: + break + chunks.append(chunk) + got += len(chunk) + return b"".join(chunks) + + +def open_json(request: urllib.request.Request, timeout: float = 10, limit: int = HTTP_BODY_LIMIT) -> Any: + """The JSON a request answers with, or HttpError saying why there's none.""" + deadline = time.monotonic() + 3 * timeout + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + body = _read_by(response, limit + 1, deadline) + # A sized read hands back what arrived when the server stops short of + # its Content-Length, with the rest still owed. + owed = getattr(response, "length", None) + except urllib.error.HTTPError as error: + try: + body = error.read(limit) or b"" + except Exception: + body = b"" + raise HttpError(status=error.code, body=body, headers=error.headers) from error + except _TooSlow as error: + # The server answered, just never finished: a reply cut short, so no + # reason to ask again sooner. + raise HttpError(cut_short=True) from error + except OSError as error: + # URLError and every socket failure, a reset and a timeout included: the + # connection failed or dropped before a reply. RemoteDisconnected is + # both an OSError and an HTTPException, and lands here: a server that + # hung up without a word gave no answer. + raise HttpError(transport=True) from error + except http.client.HTTPException as error: + # A reply arrived but was cut short (IncompleteRead) or garbled (a bad + # status line, bad headers): the server answered. + raise HttpError(cut_short=True) from error + except ValueError: + # A header http.client won't send: a key with a line break or a + # character outside Latin-1. Not chained, as the error quotes the key. + raise HttpError() from None + if len(body) > limit: + raise HttpError() + if isinstance(owed, int) and owed > 0: + raise HttpError(cut_short=True) + try: + return json.loads(body.decode("utf-8", errors="replace")) + except (ValueError, RecursionError) as error: + # RecursionError: a body nested deeper than the parser will go. + raise HttpError() from error + + +def trusted_endpoint(url: str) -> bool: + """Whether a credential may go to `url`, an endpoint the environment can + move: HTTPS, or plain HTTP to loopback alone, and no credentials of its + own in the URL.""" + try: + parts = urllib.parse.urlsplit(url) + except ValueError: + return False + if parts.username is not None or parts.password is not None: + return False + loopback = (parts.hostname or "").lower() in ("localhost", "127.0.0.1", "::1") + return parts.scheme == "https" or (parts.scheme == "http" and loopback) + + +# ---------------------------------------------------------------- sign-ins + + +def finite(value: Any) -> float: + """A float from an answer or a token, 0.0 when it isn't a finite number.""" + if isinstance(value, bool): + return 0.0 + try: + parsed = float(value or 0) + except (TypeError, ValueError): + return 0.0 + return parsed if math.isfinite(parsed) else 0.0 + + +def jwt_claims(token: Any) -> dict[str, Any]: + """The claims a JWT carries, read without checking its signature (only + ever to tell whose it is, or when it lapses), {} for anything else.""" + if not isinstance(token, str): + return {} + parts = token.split(".") + if len(parts) != 3: + return {} + try: + claims = json.loads(base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))) + except (ValueError, UnicodeDecodeError): + return {} + return claims if isinstance(claims, dict) else {} + + +def token_lapsed(token: str) -> bool: + """A sign-in token whose expiry has passed. The agent renews its own the + next time it runs; an opaque API key never lapses here.""" + expires = finite(jwt_claims(token).get("exp")) + return 0 < expires <= datetime.now(timezone.utc).timestamp() + + +def account_key(token: str) -> str: + """Who a token signs in as (the JWT's subject), so a machine signed in to + someone else never shows the last one's numbers, and a refreshed token + keeps the same caches. An opaque key stands for itself; the caches hash + it.""" + return str(jwt_claims(token).get("sub") or "") or token + + +# The Cursor and Grok Bot collectors ask as the same Cursor account, found +# the same way here, so the two never describe different people. + + +def cursor_auth_path() -> Path: + """The Cursor CLI's auth.json: CURSOR_AUTH_PATH when set, else its home + under the config directory.""" + configured = os.environ.get("CURSOR_AUTH_PATH", "").strip() + return Path(configured).expanduser() if configured else config_home() / "cursor" / "auth.json" + + +def cursor_state_db() -> Path: + return config_home() / "Cursor" / "User" / "globalStorage" / "state.vscdb" + + +def _cursor_ide_token(state_db: Path) -> str: + # The sign-in the Cursor editor keeps, for people who never ran the CLI. + # Read-only: the editor may be writing its database right now. No + # immutable=1, so a live database with an active WAL still reads. + if not state_db.is_file(): + return "" + try: + conn = sqlite3.connect(state_db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + except sqlite3.Error: + return "" + try: + conn.execute("PRAGMA query_only = ON") + row = conn.execute("SELECT value FROM ItemTable WHERE key = 'cursorAuth/accessToken' LIMIT 1").fetchone() + except sqlite3.Error: + return "" + finally: + conn.close() + value = row[0] if row else None + if isinstance(value, bytes): + value = value.decode("utf-8", errors="replace") + return str(value or "").strip() + + +def cursor_token(auth_path: Path | None = None, state_db: Path | None = None) -> str: + """The token the user exported, the credential the Cursor CLI stores, or + the editor's sign-in, in that order. + + CURSOR_AUTH_TOKEN is a bearer as it stands. CURSOR_API_KEY is not: the CLI + trades it for an access token and keeps that beside the key in its auth + file, so the stored token stands for the key only while it was traded for + this one. The key itself never goes out as a bearer. + + A sign-in that has lapsed gives way to the next one that hasn't: someone + who ran the CLI once and now lives in the editor would otherwise wait on a + token nothing renews. With all of them lapsed, the first one decides. + """ + bearer = str(os.environ.get("CURSOR_AUTH_TOKEN") or "").strip() + if bearer: + return bearer + auth = read_json(auth_path or cursor_auth_path()) + key = str(os.environ.get("CURSOR_API_KEY") or "").strip() + traded = str(auth.get("accessToken") or "").strip() + if key: + # With no token traded for this key, any other sign-in is someone else's + # account than the one the CLI runs as. + return traded if str(auth.get("apiKey") or "").strip() == key else "" + cli = traded + if cli and not token_lapsed(cli): + return cli + ide = _cursor_ide_token(state_db or cursor_state_db()) + if ide and not token_lapsed(ide): + return ide + return cli or ide + + +# ---------------------------------------------------------------- accounts + + +def registered_accounts(agent_id: str) -> list[dict[str, Any]]: + """The accounts `omarchy agent account` registered for the agent, read + straight from its registry, each marked `active` and carrying the + registry's `switch` settings. Only a registry holding a second account + matters: with one, the record is exactly what it always was.""" + state = Path(os.environ.get("XDG_STATE_HOME") or (Path.home() / ".local" / "state")) + registry = read_json(state / "omarchy" / "agents" / "accounts" / f"{agent_id}.json") + accounts = [a for a in registry.get("accounts") or [] if isinstance(a, dict) and a.get("id")] + if len(accounts) < 2: + return [] + active_id = registry.get("active") or accounts[0]["id"] + if not any(a["id"] == active_id for a in accounts): + active_id = accounts[0]["id"] + # The threshold reads as `omarchy agent account` reads it: a percentage + # held to 50-100, 95 when it isn't one. + try: + threshold = max(50, min(100, int(registry.get("threshold")))) + except Exception: + threshold = 95 + for account in accounts: + account["active"] = account["id"] == active_id + account["switch"] = { + "mode": "auto" if registry.get("switch") == "auto" else "manual", + "threshold": threshold, + } + return accounts + + +def account_entry(account: dict[str, Any], limits: dict[str, Any], plan: str, **extra: Any) -> dict[str, Any]: + """One registered account in the record's `accounts`. + + `limits` is what the collector's limits check answered for it: `limits`, + `live`, `fetchedAtMs`, `usageStatusText`, `authHelpText`, and + `retryAdvised` when set. `plan` is the plan it names now, the registry's + standing in when it names none, and `extra` the collector's own fields, + last. + """ + return { + "id": account["id"], + "label": str(account.get("label") or account["id"]), + "email": str(account.get("email") or ""), + "plan": plan or str(account.get("plan") or ""), + "active": account["active"], + "primary": bool(account.get("primary")), + "limits": limits["limits"], + "stale": not limits["live"] and bool(limits["limits"]), + "fetchedAt": limits["fetchedAtMs"], + "usageStatusText": limits["usageStatusText"], + "authHelpText": limits["authHelpText"], + "retryAdvised": bool(limits.get("retryAdvised")), + **extra, + } diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 21dc959dcf2..25b11cf547c 100644 --- a/shell/plugins/agents/manifest.json +++ b/shell/plugins/agents/manifest.json @@ -5,7 +5,7 @@ "version": "1.0.0", "author": "Omarchy", "license": "MIT", - "description": "Claude Code, Codex, Grok, Fireworks, and Kimi 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,8 +22,26 @@ "claude": { "enabled": true }, "codex": { "enabled": true }, "grok": { "enabled": true }, + "kimi": { "enabled": true }, + "cursor": { "enabled": true }, + "grokbot": { "enabled": true }, + "copilot": { "enabled": true }, + "muse": { "enabled": true }, + "agy": { "enabled": true }, + "kimchi": { "enabled": true }, + "hermes": { "enabled": true }, + "devin": { "enabled": true }, + "vibe": { "enabled": true }, + "qwen": { "enabled": true }, + "opencode": { "enabled": true }, + "zai": { "enabled": true }, + "ollama": { "enabled": true }, + "minimax": { "enabled": true }, + "openrouter": { "enabled": true }, + "deepseek": { "enabled": true }, "fireworks": { "enabled": true }, - "kimi": { "enabled": true } + "synthetic": { "enabled": true }, + "local": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-account-test.sh b/test/shell.d/agent-account-test.sh index 237e502cd17..4de7a5795b5 100644 --- a/test/shell.d/agent-account-test.sh +++ b/test/shell.d/agent-account-test.sh @@ -324,6 +324,23 @@ shutil.rmtree(pending) PY pass "a failed registration rolls the new home back to pending" +# A settings cache whose settings aren't an object costs only the plan. +STATE="$ROOT/bin/omarchy-agent-account-state" python3 - <<'PY' || fail "a Grok settings cache of the wrong shape reads as no plan" +import importlib.machinery, importlib.util, json, os, tempfile +from pathlib import Path +loader = importlib.machinery.SourceFileLoader("state", os.environ["STATE"]) +spec = importlib.util.spec_from_loader(loader.name, loader) +state = importlib.util.module_from_spec(spec) +loader.exec_module(state) + +home = Path(tempfile.mkdtemp(dir=os.environ["HOME"])) +(home / "auth.json").write_text(json.dumps({"https://auth.x.ai::client": {"user_id": "u-odd", "email": "odd@example.com"}})) +(home / "settings_cache.json").write_text(json.dumps({"payload": json.dumps({"settings": "x"})})) +found = state.grok_identity(home) +assert found["plan"] == "" and found["email"] == "odd@example.com", found +PY +pass "a Grok settings cache of the wrong shape reads as no plan" + # ------------------------------------------------------------ panel add flow [[ $(omarchy-agent-account-add --check) == $'claude additional\ncodex additional\ngrok additional' ]] || diff --git a/test/shell.d/agent-usage-accounts-test.sh b/test/shell.d/agent-usage-accounts-test.sh index 40b20290fda..fb6e460319f 100644 --- a/test/shell.d/agent-usage-accounts-test.sh +++ b/test/shell.d/agent-usage-accounts-test.sh @@ -1,226 +1,139 @@ #!/bin/bash +# Several accounts registered with `omarchy agent account`: the Claude and +# Codex records list each one with its own plan and limits, asked with its +# own sign-in, while the record's own fields describe the active account. +# Grok's accounts are covered in its own test, and the registry's reading in +# the shared test. + set -euo pipefail source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" -require_command jq -require_command python3 +agent_sandbox -test_tmp=$(mktemp -d) -trap 'rm -rf "$test_tmp"' EXIT +# ------------------------------------------------------------------- claude -export HOME="$test_tmp/home" -export XDG_STATE_HOME="$test_tmp/state" -export XDG_CACHE_HOME="$test_tmp/cache" -export XDG_DATA_HOME="$test_tmp/data" -unset CLAUDE_CONFIG_DIR CODEX_HOME +python3 - <<'PY' +from agent_usage_fixtures import * -accounts="$XDG_STATE_HOME/omarchy/agents/accounts" -mkdir -p "$HOME/.claude/projects" "$HOME/.codex/sessions" "$accounts/claude/work" "$accounts/claude/old" "$accounts/codex/side" +sandbox() +accounts = state_home() / "omarchy" / "agents" / "accounts" -future_ms=$(( ($(date +%s) + 3600) * 1000 )) -open_at=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=3)).isoformat())') -credentials() { - printf '{"claudeAiOauth":{"accessToken":"%s","expiresAt":%s,"rateLimitTier":"%s","subscriptionType":"max"}}\n' "$1" "$2" "$3" -} -credentials token-main "$future_ms" default_claude_max_20x >"$HOME/.claude/.credentials.json" -credentials token-work "$future_ms" default_claude_max_5x >"$accounts/claude/work/.credentials.json" -credentials token-old 1000 default_claude_max_5x >"$accounts/claude/old/.credentials.json" +def credentials(directory, token, expires, tier): + write_json(Path(directory) / ".credentials.json", {"claudeAiOauth": {"accessToken": token, "expiresAt": expires, "rateLimitTier": tier, "subscriptionType": "max"}}) + + +credentials(home() / ".claude", "token-main", ms_in(3600), "default_claude_max_20x") +credentials(accounts / "claude" / "work", "token-work", ms_in(3600), "default_claude_max_5x") +credentials(accounts / "claude" / "old", "token-old", 1000, "default_claude_max_5x") +(home() / ".claude" / "projects").mkdir() # A parked account whose sign-in lapsed still has its last numbers on disk. -mkdir -p "$XDG_CACHE_HOME/omarchy/agent-usage" -jq -nc --arg open "$open_at" '{fetchedAtMs: 1, limits: [{label: "Weekly (7-day)", percent: 0.4, resetsAt: $open}]}' \ - >"$XDG_CACHE_HOME/omarchy/agent-usage/claude-limits-old.json" +write_json(cache_root() / "claude-limits-old.json", {"fetchedAtMs": 1, "limits": [{"label": "Weekly (7-day)", "percent": 0.4, "resetsAt": iso_in(3 * 3600)}]}) -cat >"$accounts/claude.json" <"$XDG_CACHE_HOME/omarchy/agent-usage/claude-limits-old.json" -rested=$(COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" python3 - <<'PY' -import importlib.machinery, importlib.util, io, json, os, sys - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) -collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(b'{"five_hour": {"utilization": 30.0}}') -collector.scan_pi_usage = lambda age: None -collector.scan_opencode_usage = lambda age: None -sys.argv = ["omarchy-agent-usage-claude", "--force"] -collector.main() -PY -) -[[ $(jq -c '.accounts[2] | {stale, limits: [.limits[] | {label, empty: (.percent == 0)}]}' <<<"$rested") == '{"stale":true,"limits":[{"label":"Session (5-hour)","empty":true}]}' ]] || - fail "a lapsed account whose windows all reset reads as untouched" "$rested" -pass "a lapsed account whose windows all reset reads as untouched" +ANY = {"api.anthropic.com": {"five_hour": {"utilization": 30.0}}} +old_cache = write_json(cache_root() / "claude-limits-old.json", {"fetchedAtMs": 1, "limits": [{"label": "Session (5-hour)", "percent": 0.99, "resetsAt": iso_in(-300)}]}) + + +def rested(record): + account = record["accounts"][2] + return {"stale": account["stale"], "limits": [{"label": entry["label"], "empty": entry["percent"] == 0} for entry in account["limits"]]} + + +record = claude(ANY) +expect("a lapsed account whose windows all reset reads as untouched", rested(record), {"stale": True, "limits": [{"label": "Session (5-hour)", "empty": True}]}, record) # A cache stamped ahead of a clock that later stepped back is still what the # account last saw, so its reset windows still read as untouched. -touch -d "@$(( $(date +%s) + 3600 ))" "$XDG_CACHE_HOME/omarchy/agent-usage/claude-limits-old.json" -rested_future=$(COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" python3 - <<'PY' -import importlib.machinery, importlib.util, io, json, os, sys - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) -collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(b'{"five_hour": {"utilization": 30.0}}') -collector.scan_pi_usage = lambda age: None -collector.scan_opencode_usage = lambda age: None -sys.argv = ["omarchy-agent-usage-claude", "--force"] -collector.main() -PY -) -[[ $(jq -c '.accounts[2] | {stale, limits: [.limits[] | {label, empty: (.percent == 0)}]}' <<<"$rested_future") == '{"stale":true,"limits":[{"label":"Session (5-hour)","empty":true}]}' ]] || - fail "a future-dated cache of reset windows still reads as untouched" "$rested_future" -pass "a future-dated cache of reset windows still reads as untouched" +future = time.time() + 3600 +os.utime(old_cache, (future, future)) +record = claude(ANY) +expect("a future-dated cache of reset windows still reads as untouched", rested(record), {"stale": True, "limits": [{"label": "Session (5-hour)", "empty": True}]}, record) # Signing the primary home in to another subscription must not inherit the # last one's numbers when the first probe for the new one fails. -printf '{"oauthAccount":{"accountUuid":"u-new"}}\n' >"$HOME/.claude.json" -now_ms=$(( $(date +%s) * 1000 )) -jq -nc --arg open "$open_at" --argjson now "$now_ms" '{fetchedAtMs: $now, limits: [{label: "Session (5-hour)", percent: 0.55, resetsAt: $open}]}' \ - >"$XDG_CACHE_HOME/omarchy/agent-usage/claude-limits.json" -resigned=$(COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" python3 - <<'PY' -import importlib.machinery, importlib.util, io, json, os, sys - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) - -def urlopen(request, timeout=None): - if request.get_header("Authorization").endswith("token-main"): - raise collector.urllib.error.URLError("down") - return io.BytesIO(b'{"five_hour": {"utilization": 12.0}}') - -collector.urllib.request.urlopen = urlopen -collector.scan_pi_usage = lambda age: None -collector.scan_opencode_usage = lambda age: None -sys.argv = ["omarchy-agent-usage-claude", "--force"] -collector.main() -PY -) -[[ $(jq -c '.accounts[0].limits | map(.percent)' <<<"$resigned") != *0.55* ]] || - fail "a primary home signed in to another subscription doesn't inherit its limits" "$resigned" -rm -f "$HOME/.claude.json" -pass "a primary home signed in to another subscription doesn't inherit its limits" +write_json(home() / ".claude.json", {"oauthAccount": {"accountUuid": "u-new"}}) +write_json(cache_root() / "claude-limits.json", {"fetchedAtMs": ms_in(0), "limits": [{"label": "Session (5-hour)", "percent": 0.55, "resetsAt": iso_in(3 * 3600)}]}) +USAGE.pop("token-main") +record = claude() +check("a primary home signed in to another subscription doesn't inherit its limits", + 0.55 not in [entry["percent"] for entry in record["accounts"][0]["limits"]], record) +(home() / ".claude.json").unlink() # A registry with one account changes nothing about the record. -jq '.accounts |= [.[0]] | .active = "main"' "$accounts/claude.json" >"$test_tmp/one.json" -mv "$test_tmp/one.json" "$accounts/claude.json" -single=$(COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" python3 - <<'PY' -import importlib.machinery, importlib.util, io, json, os, sys - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) -collector.urllib.request.urlopen = lambda request, timeout=None: io.BytesIO(b'{"five_hour": {"utilization": 30.0}}') -collector.scan_pi_usage = lambda age: None -collector.scan_opencode_usage = lambda age: None -sys.argv = ["omarchy-agent-usage-claude", "--force"] -collector.main() -PY -) -[[ $(jq 'has("accounts")' <<<"$single") == false && $(jq '.limits[0].percent' <<<"$single") == 0.3 ]] || - fail "a single account record carries no accounts list" "$single" -pass "a single account record is unchanged" +registry("claude", [{"id": "main", "label": "Main", "home": "", "primary": True, "email": "me@example.com"}], active="main", switch="auto", threshold=95) +record = claude(ANY) +expect("a single account record is unchanged, and limits checked just now aren't stale", + ["accounts" in record, record["limits"][0]["percent"], record["limitsStale"], record["limitsFetchedAt"] > 0], [False, 0.3, False, True], record) -[[ $(jq -c '{limitsStale, fresh: (.limitsFetchedAt > 0)}' <<<"$single") == '{"limitsStale":false,"fresh":true}' ]] || - fail "limits checked just now aren't stale" "$single" # A failed probe falls back on hour-old numbers, and the record says so. -hour_ago_ms=$(( ($(date +%s) - 3600) * 1000 )) -jq -nc --arg open "$open_at" --argjson at "$hour_ago_ms" '{fetchedAtMs: $at, limits: [{label: "Session (5-hour)", percent: 0.42, resetsAt: $open}]}' \ - >"$XDG_CACHE_HOME/omarchy/agent-usage/claude-limits.json" -kept=$(COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" python3 - <<'PY' -import importlib.machinery, importlib.util, os, sys - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) - -def urlopen(request, timeout=None): - raise collector.urllib.error.URLError("down") - -collector.urllib.request.urlopen = urlopen -collector.scan_pi_usage = lambda age: None -collector.scan_opencode_usage = lambda age: None -sys.argv = ["omarchy-agent-usage-claude", "--force"] -collector.main() +hour_ago = ms_in(-3600) +write_json(cache_root() / "claude-limits.json", {"fetchedAtMs": hour_ago, "limits": [{"label": "Session (5-hour)", "percent": 0.42, "resetsAt": iso_in(3 * 3600)}]}) +with FakeHttp(mode="down"): + record = collect("claude", "--force") +expect("the record says when its limits are kept from an earlier check", + [record["limitsStale"], record["limitsFetchedAt"], record["limits"][0]["percent"]], [True, hour_ago, 0.42], record) PY -) -[[ $(jq -c --argjson at "$hour_ago_ms" '{limitsStale, at: (.limitsFetchedAt == $at), first: .limits[0].percent}' <<<"$kept") == '{"limitsStale":true,"at":true,"first":0.42}' ]] || - fail "kept limits are marked stale with when they were measured" "$kept" -pass "the record says when its limits are kept from an earlier check" -# ------------------------------------------------------------------------ codex +# -------------------------------------------------------------------- codex + +home=$(new_home codex) +accounts="$home/.local/state/omarchy/agents/accounts" +mkdir -p "$home/.codex/sessions" "$accounts/codex/side" +touch "$home/.codex/auth.json" "$accounts/codex/side/auth.json" +cat >"$accounts/codex.json" <"$test_tmp/bin/codex" <<'PY' -#!/usr/bin/python3 +stub_cli "$home" codex >/dev/null <"$accounts/codex.json" <>"$HOME/omp-stdin"; exit 3; }\n' + 'cat "$HOME/omp-report.json" 2>/dev/null')): + (stubs / name).write_text(f"#!/bin/bash\n{body}\n") + (stubs / name).chmod(0o755) + os.environ["PATH"] = f"{stubs}:{REAL_PATH}" + + +def omp_login(token, expires, email): + database = home() / ".omp" / "agent" / "agent.db" + database.parent.mkdir(parents=True, exist_ok=True) + database.unlink(missing_ok=True) + conn = sqlite3.connect(database) + conn.execute("CREATE TABLE auth_credentials (id INTEGER PRIMARY KEY, provider TEXT, credential_type TEXT, data TEXT, disabled_cause TEXT)") + conn.execute("INSERT INTO auth_credentials (provider, credential_type, data) VALUES ('google-antigravity', 'oauth', ?)", + (json.dumps({"access": token, "refresh": "r", "expires": expires, "email": email}),)) + conn.execute("INSERT INTO auth_credentials (provider, credential_type, data, disabled_cause) VALUES ('google-antigravity', 'oauth', ?, 'revoked')", + (json.dumps({"access": "disabled-token", "refresh": "r", "expires": ms_in(99999), "email": email}),)) + conn.commit() + conn.close() + + +def pi_login(token, expires, email, directory=None): + write_json(Path(directory or home() / ".pi" / "agent") / "auth.json", + {"google-antigravity": {"type": "oauth", "access": token, "refresh": "r", "expires": expires, "email": email}}) + + +# Google answers every token in `accepted` and refuses the rest. +def agy(*accepted): + def answer(request): + token = request.get_header("Authorization").removeprefix("Bearer ") + if token not in accepted: + return 403 + return ANSWERS[request.full_url.rsplit(":", 1)[1]] + with FakeHttp({"v1internal:": answer}) as http: + record = collect("agy", "--force") + asked = sorted({r["headers"]["authorization"].removeprefix("Bearer ") for r in http.requests}) + return record, asked + + +# Every store the collector reads, so a test can tell none was written. +def stores(): + return sorted((str(path), hashlib.md5(path.read_bytes()).hexdigest()) + for root in (".omp", ".pi", ".local/share/opencode", ".gemini", "pi-elsewhere") + for path in (home() / root).rglob("*") if path.is_file()) + + +# ---------------------------------------------------------- the sign-ins + +machine() +omp_login("omp-token", ms_in(3600), "me@example.com") +record, asked = agy("omp-token") +expect("Antigravity collector reports limits from an omp sign-in without agy, never asking with a disabled one", + [record["ready"], record["tierLabel"], record["usageStatusText"], record["authHelpText"], record["limitsStale"], record["limits"][0]["percent"], asked], + [True, "Pro", "", "", False, 0.25, ["omp-token"]], record) + +pi_login("pi-token", ms_in(7200), "me@example.com") +_, asked = agy("omp-token", "pi-token") +expect("Antigravity collector asks with the most recently refreshed sign-in", asked, ["pi-token"]) + +record, asked = agy("omp-token") +expect("Antigravity collector moves on from a sign-in Google refuses", + [asked, record["limitsStale"], record["usageStatusText"]], [["omp-token", "pi-token"], False, ""], record) + +record, _ = agy() +expect("Antigravity collector reports every sign-in refused as expired, keeping the account's limits, and names the harness", + [record["usageStatusText"], record["limitsStale"], record["limits"][0]["percent"], record["authHelpText"]], + ["Sign-in expired", True, 0.25, "Google refused the Antigravity sign-in pi keeps. Start pi to sign in again."], record) + +# A lapsed token is renewed the next time its harness runs, so the limits are +# paused, not signed out, and the token goes nowhere. +shutil.rmtree(home() / ".pi") +omp_login("omp-token", ms_in(-60), "me@example.com") +before = stores() +record, asked = agy("omp-token") +check("Antigravity collector keeps the last limits while omp's token has lapsed, sending it nowhere", + asked == [] and record["ready"] is True and record["limitsStale"] is True and record["usageStatusText"] == "Limits paused" + and record["limits"][0]["percent"] == 0.25 and record["authHelpText"].endswith("Start omp to refresh it; signing in again is not needed."), record) +expect("Antigravity collector never writes omp's store", stores(), before) + +omp_login("other-token", ms_in(-60), "other@example.com") +record, _ = agy() +expect("Antigravity collector shows no other Google account's kept limits", [record["limits"], record["tierLabel"]], [[], ""], record) + +# omp renews its own sign-in when asked for usage, so a lapsed token in its +# database still gets live limits through `omp usage`. +write_json(home() / "omp-report.json", {"reports": [{"provider": "google-antigravity", "metadata": {"email": "other@example.com"}, "limits": [ + {"label": "Gemini", "amount": {"usedFraction": 0.6}, "window": {"id": "5h", "resetsAt": 32503680000000}}, + {"label": "Gemini", "amount": {"usedFraction": 0.3}, "window": {"id": "weekly", "resetsAt": 32503680000000}}, + {"label": "Claude & GPT (shared)", "amount": {"usedFraction": 0.1}, "window": {"id": "weekly", "resetsAt": 32503680000000}}, + {"label": "Claude & GPT (shared)", "amount": {"usedFraction": 0.1}, "window": {"id": "weekly", "resetsAt": 32503680000000}}, + {"label": "Gemini", "amount": {"usedFraction": "lots"}, "window": {"id": "daily"}}, +]}]}) +record, asked = agy() +expect("Antigravity collector takes omp's limits from omp usage, sending omp's lapsed token nowhere", + [asked, record["limitsStale"], record["usageStatusText"], record["authHelpText"], [(entry["title"], entry["percent"]) for entry in record["limits"]]], + [[], False, "", "", [("Session", 0.6), ("Weekly", 0.3), ("Claude/GPT Weekly", 0.1)]], record) + +check("omp usage runs with nothing on its stdin", not (home() / "omp-stdin").exists()) + +# An omp report of another shape is no answer: the kept limits stand, stale. +report = {"provider": "google-antigravity", "metadata": {"email": "other@example.com"}} +kept = [] +for shape in ({"reports": [dict(report, limits="x")]}, {"reports": [dict(report, limits=5)]}, {"reports": 5}): + write_json(home() / "omp-report.json", shape) + record, _ = agy() + kept.append([record["limitsStale"], [entry["percent"] for entry in record["limits"]]]) +expect("Antigravity collector keeps the last limits over an omp report of another shape", kept, [[True, [0.6, 0.3, 0.1]]] * 3, kept) + +# A window id of another shape still names one window, and a report naming no +# account, with several omp sign-ins, is shown but kept for none of them. +conn = sqlite3.connect(home() / ".omp" / "agent" / "agent.db") +conn.execute("INSERT INTO auth_credentials (provider, credential_type, data) VALUES ('google-antigravity', 'oauth', ?)", + (json.dumps({"access": "third-token", "refresh": "r3", "expires": ms_in(-60), "email": "third@example.com"}),)) +conn.commit() +conn.close() +for path in limits_cache_files("agy"): + path.unlink() +write_json(home() / "omp-report.json", {"reports": [{"provider": "google-antigravity", "limits": [ + {"label": "Gemini", "amount": {"usedFraction": 0.2}, "window": {"id": ["odd"], "resetsAt": 32503680000000}}, + {"label": "Gemini", "amount": {"usedFraction": 0.6}, "window": {"id": "5h", "resetsAt": 32503680000000}}, +]}]}) +record, asked = agy() +expect("an anonymous omp report with several sign-ins is shown, kept for none, past an odd window id", + [asked, record["limitsStale"], [entry["percent"] for entry in record["limits"]], limits_cache_files("agy")], [[], False, [0.6], []], record) + +# opencode files any Google sign-in under "google"; only the Antigravity +# plugin's accounts file makes it an Antigravity one. +machine() +write_json(data_home() / "opencode" / "auth.json", {"google": {"type": "oauth", "access": "opencode-token", "refresh": "r|p", "expires": ms_in(3600)}}) +record, asked = agy("opencode-token") +expect("Antigravity collector ignores opencode's Google sign-in without the Antigravity plugin", [record["ready"], record["limits"], asked], [False, [], []], record) +write_json(config_home() / "opencode" / "antigravity-accounts.json", {"version": 3, "accounts": []}) +record, asked = agy("opencode-token") +expect("Antigravity collector reads opencode's Antigravity plugin sign-in", [asked, record["limits"][0]["percent"]], [["opencode-token"], 0.25], record) + +# opencode records no email for its sign-in, so a replaced sign-in is told +# apart by its refresh token: the kept limits go rather than pass for the +# new account's. +write_json(data_home() / "opencode" / "auth.json", {"google": {"type": "oauth", "access": "switched-token", "refresh": "switched|p", "expires": ms_in(3600)}}) +record, _ = agy() +expect("Antigravity collector drops kept limits when the account cannot be matched", + [record["limits"], record["tierLabel"], record["usageStatusText"]], [[], "", "Sign-in expired"], record) + +# Without a keyring, agy keeps its sign-in in a file under its home. +machine() +write_json(home() / ".gemini" / "antigravity-cli" / "antigravity-oauth-token", + {"token": {"access_token": "agy-file-token", "refresh_token": "agy-refresh", "expiry": "2999-01-01T00:00:00Z"}}) +record, asked = agy("agy-file-token") +expect("Antigravity collector reads agy's sign-in file when there is no keyring", [asked, record["limits"][0]["percent"]], [["agy-file-token"], 0.25], record) + +# PI_CODING_AGENT_DIR moves pi's sign-in. +machine() +pi_login("moved-token", ms_in(3600), "me@example.com", home() / "pi-elsewhere") +os.environ["PI_CODING_AGENT_DIR"] = str(home() / "pi-elsewhere") +before = stores() +record, asked = agy("moved-token") +expect("Antigravity collector reads pi's sign-in where PI_CODING_AGENT_DIR puts it", [asked, record["limits"][0]["percent"]], [["moved-token"], 0.25], record) +expect("Antigravity collector never writes pi's store", stores(), before) + +# ---------------------------------------------------------- local stats + +# OpenClaw keeps pi's format, one directory per agent; its entries carry no +# id, so two alike answers are two answers. Gemini through any provider but +# Antigravity's counts nowhere: pi's google-gemini-cli, and opencode's +# "google" without the plugin's antigravity- model prefix. +machine() +openclaw = pi_reply("google-antigravity", 1, model="gemini-3-pro", usage={"input": 1, "output": 1}) +write_pi_session(home() / ".openclaw" / "agents" / "main" / "sessions" / "s.jsonl", [openclaw, openclaw]) +write_pi_session(home() / ".pi" / "agent" / "sessions" / "p" / "s.jsonl", + [pi_reply("google-gemini-cli", 999, id="g1", model="gemini-3-pro"), + pi_reply("google-antigravity", 10, id="a1", model="gemini-3-pro", usage={"input": 10, "output": 2})]) +opencode_db(message=[("m1", "s1", oc_v1("google", 999, "gemini-3-pro")), ("m2", "s2", oc_v1("google", 4, "antigravity-gemini-3-pro"))]) +record, _ = agy() +expect("Antigravity collector counts OpenClaw's answers and only Antigravity's provider", + [record["ready"], record["usageStatusText"], record["totalPrompts"], record["todayTotalTokens"], + {model: bucket["inputTokens"] for model, bucket in record["modelUsage"].items()}], + [True, "Waiting for auth", 4, 20, {"gemini-3-pro": 12, "antigravity-gemini-3-pro": 4}], record) +expect("waiting for auth with no limits kept is nothing stale", [record["limits"], record["limitsStale"]], [[], False], record) +PY 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..16fda997a5d --- /dev/null +++ b/test/shell.d/agent-usage-agy-limits-test.sh @@ -0,0 +1,224 @@ +#!/bin/bash + +# The Antigravity collector's limits with agy's own sign-in: how it reads +# Google's plan and quota answers, a lapsed or replaced sign-in, where the +# bearer may go, and its adapter for the contract checks. Sign-ins kept by +# other harnesses are in agy-harness; local stats in agy-scanner. The +# unused-machine record is held centrally by the partition test, and the +# limits cache's mechanics by the shared test. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +import re + +TIER = "v1internal:loadCodeAssist" +QUOTA = "v1internal:retrieveUserQuotaSummary" +ANSWER = { + TIER: {"paidTier": {"id": "g1-pro-tier", "name": "Google AI Pro"}}, + QUOTA: {"groups": [{"displayName": "Gemini Models", "buckets": [{"window": "5h", "remainingFraction": 0.75, "resetTime": iso_in(3 * 3600)}]}]}, +} +ME = jwt({"email": "me@example.com"}) +OTHER = jwt({"email": "other@example.com"}) +REAL_PATH = os.environ["PATH"] + + +# A machine where agy keeps its sign-in in the keyring: secret-tool answers +# with keyring.json, which sign_in() writes. +def machine(): + sandbox() + stub = home() / "bin" / "secret-tool" + stub.parent.mkdir(parents=True, exist_ok=True) + stub.write_text('#!/bin/bash\ncat "$HOME/keyring.json" 2>/dev/null\n') + stub.chmod(0o755) + os.environ["PATH"] = f"{stub.parent}:{REAL_PATH}" + os.environ["AGY_DIR"] = str(home() / "agy") + + +def sign_in(token="token-main", expiry=None, id_token=ME, refresh="refresh-main"): + write_json(home() / "keyring.json", {"auth_method": "consumer", "id_token": id_token, + "token": {"access_token": token, "refresh_token": refresh, "expiry": expiry or iso_in(3600)}}) + + +def agy(mode="up", routes=ANSWER, args=("--force",)): + with FakeHttp(routes, mode=mode) as http: + record = collect("agy", *args) + return record, http + + +def tier(payload): + return load_collector("agy").extract_tier(payload) + + +def limits(payload): + return load_collector("agy").extract_limits(payload) + + +# ------------------------------------------------------------- the answer + +sandbox() +expect("Antigravity collector extracts the plan from paidTier, else currentTier", [ + tier({"currentTier": {"id": "free-tier", "name": "Antigravity"}, "paidTier": {"id": "g1-ultra-tier", "name": "Google AI Ultra"}}), + tier({"currentTier": {"id": "free-tier", "name": "Antigravity"}, "paidTier": {"id": "g1-pro-tier", "name": "Google AI Pro"}}), + tier({"currentTier": {"id": "free-tier", "name": "Antigravity"}}), +], ["Ultra", "Pro", "Free"]) + +# Gemini's windows first, the session before the week, then the Claude and +# GPT allowance only where some of it is used. +read = limits({"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}, + ]}, +]}) +expect("Antigravity collector orders the windows, reads used from remaining, and omits an unused 3P window", + [(entry["label"], entry.get("title"), round(entry["percent"], 2)) for entry in read], + [("Session (5-hour)", "Session", 0.1), ("Weekly (7-day)", "Weekly", 0.05), ("Claude/GPT Weekly", "Claude/GPT Weekly", 0.02)], read) + +read = limits({"groups": [ + {"displayName": "Claude and GPT models", "buckets": [{"window": "5h", "remainingFraction": 1.0}]}, + {"displayName": "Gemini Models", "buckets": [{"window": "weekly", "remainingFraction": 0.5}]}, +]}) +expect("Antigravity collector finds the Gemini group wherever it is listed", [(entry["label"], entry["percent"]) for entry in read], [("Weekly (7-day)", 0.5)], read) + +read = limits({"buckets": [ + {"window": "5h", "remaining_fraction": 0.4, "reset_time": "2026-09-15T19:31:46Z"}, + {"window": "weekly", "remainingFraction": "plenty"}, + {"bucketId": "gemini-weekly", "remainingFraction": float("inf")}, +]}) +expect("Antigravity collector reads snake_case buckets listed without groups and skips unreadable fractions", + [(entry["label"], round(entry["percent"], 2), entry["resetsAt"]) for entry in read], [("Session (5-hour)", 0.6, "2026-09-15T19:31:46+00:00")], read) + +read = limits({"buckets": [{"window": "5h", "resetTime": "2026-09-15T19:31:46Z"}, {"window": "weekly"}]}) +expect("Antigravity collector reads a bucket whose zero fraction proto3 JSON left out as fully spent, and skips one with no reset either", + [(entry["label"], entry["percent"]) for entry in read], [("Session (5-hour)", 1.0)], read) + +# agy's QuotaSummaryBucket holds its remainder as a oneof of a fraction and an +# amount (an int64, a string in JSON), so a bucket on the amount arm has no +# fraction yet isn't spent unless the amount is zero; a disabled bucket is +# skipped, as omp's reader skips it. +read = limits({"groups": [ + {"displayName": "Gemini Models", "buckets": [ + {"window": "5h", "remainingAmount": "800", "resetTime": "2026-09-15T19:31:46Z"}, + {"window": "weekly", "remainingAmount": "0", "resetTime": "2026-09-17T18:19:42Z"}, + ]}, + {"displayName": "Claude and GPT models", "buckets": [ + {"window": "5h", "disabled": True, "resetTime": "2026-09-15T22:16:22Z"}, + ]}, +]}) +expect("Antigravity collector reads a bucket's amount arm, spent only at zero, and skips a disabled bucket", + [(entry["label"], entry["percent"]) for entry in read], [("Weekly (7-day)", 1.0)], read) + +# An answer of another shape is no answer, not a plan without quotas; another +# group's buckets of the wrong shape skip that group alone. +gemini = {"displayName": "Gemini Models", "buckets": [{"window": "5h", "remainingFraction": 0.5}]} +shapes = [None, [1], {"error": {"code": 500}}, {"groups": 5}, {"groups": "abc"}, {"buckets": 5}, + {"groups": [{"displayName": "Gemini Models", "buckets": 5}]}] +expect("Antigravity collector reads an answer of another shape as none", [limits(shape) for shape in shapes], [None] * len(shapes)) +read = limits({"groups": [gemini, {"displayName": "Claude and GPT models", "buckets": 5}]}) +expect("Antigravity collector skips another group's buckets of the wrong shape", [(entry["label"], entry["percent"]) for entry in read], [("Session (5-hour)", 0.5)], read) + +# ---------------------------------------------------------------- contract + +machine() + + +def add_prompt(): + with (home() / "agy" / "history.jsonl").open("a") as history: + history.write(json.dumps({"timestamp": iso_in(0), "conversation_id": f"c{time.time_ns()}"}) + "\n") + + +(home() / "agy").mkdir() +adapter = Adapter("agy", plant=sign_in, answer=ANSWER, requests_per_probe=2, grow=add_prompt, + pi_root=lambda: home() / ".pi" / "agent" / "sessions") +fresh = check_limits_contract(adapter) +expect("Antigravity collector reports its plan and live limits", [fresh["tierLabel"], fresh["limits"][0]["percent"]], ["Pro", 0.25], fresh) +check_scan_cache(adapter) + +record, http = agy() +expect("Antigravity collector asks with agy's token", {r["headers"].get("authorization") for r in http.requests}, {"Bearer token-main"}) + +# agy's RetrieveUserQuotaSummaryRequest carries the project, which omp sends +# from loadCodeAssist's cloudaicompanionProject, named as a string or {id}. +asked = [] +for project in ("proj-1", {"id": "proj-2", "name": "Default"}, None): + routes = dict(ANSWER, **{TIER: dict(ANSWER[TIER], cloudaicompanionProject=project)}) + _, http = agy(routes=routes) + asked += [json.loads(r["data"] or b"{}") for r in http.requests if QUOTA in r["url"]] +expect("Antigravity collector asks for the quota of the project loadCodeAssist names", asked, [{"project": "proj-1"}, {"project": "proj-2"}, {}]) +caches = [path.name for path in limits_cache_files("agy")] +check("Antigravity collector caches limits per account, named by a digest rather than the email", + len(caches) == 1 and re.fullmatch(r"agy-limits-[0-9a-f]{16}\.json", caches[0]) and "me@example.com" not in caches[0], caches) + +# A damaged entry in the kept answer is dropped, not shown. +edit_limits_cache("agy", resets_past=True) +for path in limits_cache_files("agy"): + data = json.loads(path.read_text()) + path.write_text(json.dumps(dict(data, limits=data["limits"] + ["damaged"]))) +record, _ = agy("down") +expect("Antigravity collector drops a damaged kept window and shows an elapsed one at zero", + record["limits"], [{"label": "Session (5-hour)", "title": "Session", "percent": 0.0, "resetsAt": ""}], record) +agy() + +# A quota answer of another shape is a failed check: the kept limits stand, +# stale, and stay kept for the next failure. +record, _ = agy(routes={TIER: ANSWER[TIER], QUOTA: {"groups": 5}}) +after, _ = agy("down") +expect("Antigravity collector keeps the last limits over a quota answer of another shape", + [[entry["percent"] for entry in record["limits"]], record["limitsStale"], record["retryAdvised"], "couldn't be read" in record["authHelpText"], + [entry["percent"] for entry in after["limits"]], after["limitsStale"]], + [[0.25], True, False, True, [0.25], True], [record, after]) + +# ------------------------------------------------------------ the sign-in + +# A lapsed access token is renewed the next time agy runs, so the kept limits +# stay up without a new sign-in; without a refresh token behind it, a new +# sign-in is needed. Neither is sent to Google. +sign_in(expiry=iso_in(-3600)) +record, http = agy() +check("Antigravity collector pauses limits while agy's access token has lapsed, asking nothing", + http.count == 0 and record["usageStatusText"] == "Limits paused" and record["limitsStale"] is True and len(record["limits"]) == 1 + and "signing in again is not needed" in record["authHelpText"], record) + +sign_in(expiry=iso_in(-3600), refresh="") +record, http = agy() +expect("Antigravity collector reports a lapsed sign-in without a refresh token as expired", [http.count, record["usageStatusText"]], [0, "Sign-in expired"], record) + +# Another account signed in here never wears the last one's limits. +sign_in("token-other", iso_in(-3600), OTHER) +record, _ = agy("down") +expect("Antigravity collector shows no other account's kept limits", + [record["limits"], record["tierLabel"], record["limitsFetchedAt"]], [[], "", 0], record) + +# A plan with no quota windows to report is healthy, with no meters. +sign_in("token-other", iso_in(3600), OTHER) +record, _ = agy(routes={TIER: ANSWER[TIER], QUOTA: {"groups": []}}) +expect("Antigravity collector treats a plan without quota windows as healthy", + [record["limits"], record["tierLabel"], record["usageStatusText"], record["authHelpText"], record["limitsStale"]], [[], "Pro", "", "", False], record) + +# The bearer token goes only to an https destination. +os.environ["CLOUD_CODE_URL"] = "http://evil.example" +_, http = agy() +expect("Antigravity collector ignores a plain-http destination", {r["url"].split("/")[2] for r in http.requests}, {"daily-cloudcode-pa.googleapis.com"}) +os.environ["CLOUD_CODE_URL"] = "https://mirror.example" +_, http = agy() +expect("Antigravity collector honors an https destination", {r["url"].split("/")[2] for r in http.requests}, {"mirror.example"}) +del os.environ["CLOUD_CODE_URL"] + +for path in limits_cache_files("agy"): + path.unlink() +record, _ = agy("down") +expect("Antigravity collector invents no limits when it has never reached Google", + [record["limits"], record["usageStatusText"], record["retryAdvised"]], [[], "Antigravity limits unavailable", True], record) +PY 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..7ae06f43b06 --- /dev/null +++ b/test/shell.d/agent-usage-agy-scanner-test.sh @@ -0,0 +1,90 @@ +#!/bin/bash + +# The Antigravity collector's local stats from agy's own home: prompts from +# history.jsonl, conversations from it and the conversation summaries, and +# tokens from each transcript in both shapes agy writes. The unused-machine +# record is held centrally by the partition test, and the scan cache by the +# contract checks in agy-limits. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +home=$(new_home) +agy_dir="$home/.gemini/antigravity-cli" +mkdir -p "$agy_dir/brain/test-conv/.system_generated/logs" + +# A keyring entry signed in without a token, so nothing is asked of Google. +stub_cli "$home" secret-tool >/dev/null <<'EOF' +if [[ "$*" == *"lookup service gemini username antigravity"* ]]; then + echo '{"auth_method":"consumer","id_token":"header.eyJlbWFpbCI6InRlc3RAZ21haWwuY29tIn0.signature"}' + exit 0 +fi +exit 1 +EOF + +timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)" +cat >"$agy_dir/history.jsonl" <"$agy_dir/brain/test-conv/.system_generated/logs/transcript.jsonl" <"$agy_dir/brain/test-conv/.system_generated/logs/transcript.jsonl" <"$agy_dir/brain/flat-conv/.system_generated/logs/transcript.jsonl" <"$agy_dir/history.jsonl" < 0 && end > start, 'agents panel exposes its limit-window helpers') -eval(source.slice(start, end).replaceAll('root.nowMs', 'panelRoot.nowMs')) -eval(source.slice(source.indexOf('function needsSignIn'), source.indexOf('// Sign an account'))) - -const paused = { usageStatusText: 'Limits paused', limits: [] } -assertEqual(pausedWithoutLimits(paused), true, 'paused usage without a cache shows a recovery placeholder') -assertEqual(otherTrouble(paused), '', 'paused usage never puts the confusing status in the header') -assertEqual(pausedWithoutLimits({ ...paused, limits: [{ label: 'Session (5-hour)', percent: 0, resetsAt: '' }] }), false, 'cached usage replaces the recovery placeholder') -assertEqual(pausedWithoutLimits({ usageStatusText: 'Sign-in expired', limits: [] }), false, 'expired sign-ins retain their sign-in action') -assertEqual(pausedWithoutLimits(null), false, 'missing records do not show a paused-usage placeholder') - -assertDeepEqual( - limitWindows({ limits: [ - { label: 'Session (5-hour)', percent: 0.78, resetsAt: '' }, - { label: 'Opus 5 (1M context) Weekly', title: 'Opus 5 (1M context) Weekly', percent: 0.42, resetsAt: '' } - ] }), - [ - { title: 'Session', percent: 0.78, resetAt: '' }, - { title: 'Opus 5 (1M context) Weekly', percent: 0.42, resetAt: '' } - ], - 'agents panel titles a limit off the collector when it states one' -) - -assertDeepEqual( - limitWindows({ limits: [{ label: 'Weekly (7-day)', percent: 0.12, resetsAt: '' }] }), - [{ title: 'Weekly', percent: 0.12, resetAt: '' }], - 'agents panel still reads a window out of a label that carries no title' -) - -const resetAt = new Date(panelRoot.nowMs + 1000).toISOString() -const measured = { limits: [{ label: 'Session (5-hour)', percent: 0.78, resetsAt: resetAt }] } -assertEqual(limitWindows(measured)[0].percent, 0.78, 'usage remains until the reset time') -panelRoot.nowMs += 1000 -assertDeepEqual(limitWindows(measured), [{ title: 'Session', percent: 0, resetAt: '' }], 'usage resets at the deadline while the panel stays open') -assertEqual(measured.limits[0].percent, 0.78, 'display reset leaves the measured record intact') - -// A model's weekly allowance rides under the Weekly row; one whose window has -// no row of its own keeps a row. -assertDeepEqual( - displayWindows({ limits: [ - { label: 'Session (5-hour)', percent: 0.94, resetsAt: '' }, - { label: 'Weekly (7-day)', percent: 0.25, resetsAt: 'w' }, - { label: 'Fable Weekly', title: 'Fable Weekly', percent: 0.09, resetsAt: 'w' }, - { label: 'Opus Monthly', title: 'Opus Monthly', percent: 0.5, resetsAt: '' } - ] }), - [ - { title: 'Session', percent: 0.94, resetAt: '', scoped: [] }, - { title: 'Weekly', percent: 0.25, resetAt: 'w', scoped: [{ title: 'Fable', percent: 0.09, resetAt: 'w' }] }, - { title: 'Opus Monthly', percent: 0.5, resetAt: '', scoped: [] } - ], - 'agents panel shows a model-scoped window under the window it runs on' -) -JS - -# Anthropic rate-limits its usage endpoint readily. A refused re-check of -# numbers measured minutes ago still describes the account; only older numbers -# count as stale, and either way the record says when they were measured. -rate_limited() { - COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" FETCHED_AGO="$1" XDG_CACHE_HOME="$CACHE_HOME" python3 - <<'PY' -import datetime as dt, importlib.machinery, importlib.util, json, os, time, 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) - -reset = (dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=3)).isoformat() -fetched = round((time.time() - float(os.environ["FETCHED_AGO"])) * 1000) -(collector.cache_root() / "claude-limits-rate.json").write_text(json.dumps( - {"fetchedAtMs": fetched, "limits": [{"label": "Session (5-hour)", "percent": 0.4, "resetsAt": reset}]})) - -def refused(request, timeout=None): - raise urllib.error.HTTPError(request.full_url, 429, "Too Many Requests", {}, None) - -collector.urllib.request.urlopen = refused -result = collector.collect_limits("token", int((time.time() + 3600) * 1000), 0, True, "claude-limits-rate.json") -print(json.dumps({"live": result["live"], "percent": result["limits"][0]["percent"], "fetched": result["fetchedAtMs"] == fetched})) +(home() / ".claude" / ".credentials.json").unlink() +record, _ = claude() +check("Claude collector serves open cached windows without a token", + record["usageStatusText"] == "Waiting for auth" and labels(record) == ["Session (5-hour)", "Weekly (7-day)"], record) PY -} - -[[ $(rate_limited 60) == '{"live": true, "percent": 0.4, "fetched": true}' ]] || - fail "a refused re-check of numbers a minute old still counts as current" "$(rate_limited 60)" -[[ $(rate_limited 1800) == '{"live": false, "percent": 0.4, "fetched": true}' ]] || - fail "a refused re-check of half-hour-old numbers counts as stale" "$(rate_limited 1800)" -pass "a rate-limited check only goes stale once the numbers are old" diff --git a/test/shell.d/agent-usage-claude-plan-test.sh b/test/shell.d/agent-usage-claude-plan-test.sh index 8daf62a12ea..858dcc7fd2a 100644 --- a/test/shell.d/agent-usage-claude-plan-test.sh +++ b/test/shell.d/agent-usage-claude-plan-test.sh @@ -1,118 +1,101 @@ #!/bin/bash -source "$(dirname "$0")/base-test.sh" - -require_command python3 - -# oauth_login reads two files that only ever exist side by side, so the plan it -# settles on is exercised over a planted config directory: credentials for the -# token, the CLI's profile for the account it belongs to. The profile sits in -# $HOME by default and inside the config directory when CLAUDE_CONFIG_DIR moved -# the whole config elsewhere, so the caller says which shape to plant, and may -# leave a stale copy in the place the CLI does not read. -SANDBOX=$(mktemp -d) -trap 'rm -rf "$SANDBOX"' EXIT - -read_plan() { - local credentials="$1" profile="$2" where="${3:-beside}" stray="${4:-}" - local home dir live other relocated plan - home=$(mktemp -d "$SANDBOX/home.XXXXXX") - dir="$home/.claude" - mkdir "$dir" - printf '%s' "$credentials" >"$dir/.credentials.json" - if [[ $where == "inside" ]]; then - live="$dir/.claude.json" other="$home/.claude.json" relocated="$dir" - elif [[ $where == "legacy" ]]; then - live="$dir/.config.json" other="$home/.claude.json" relocated="" - else - live="$home/.claude.json" other="$dir/.claude.json" relocated="" - fi - [[ -n $profile ]] && printf '%s' "$profile" >"$live" - [[ -n $stray ]] && printf '%s' "$stray" >"$other" - - plan=$(env -u CLAUDE_CONFIG_DIR ${relocated:+CLAUDE_CONFIG_DIR="$relocated"} HOME="$home" \ - COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" CLAUDE_DIR="$dir" python3 - <<'PY' -import importlib.machinery, importlib.util, os, pathlib - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) - -print(collector.oauth_login(pathlib.Path(os.environ["CLAUDE_DIR"]))[-1]) -PY - ) +# The plan label the Claude collector settles on from the two files Claude +# Code keeps side by side: the credentials for the token, and the CLI's +# profile for the account it belongs to. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +CREDENTIALS = {"claudeAiOauth": {"accessToken": "token", "expiresAt": 1, "subscriptionType": "max", "rateLimitTier": "default_claude_max_5x"}} +UPGRADED = {"oauthAccount": {"organizationRateLimitTier": "default_claude_max_20x", "userRateLimitTier": None}} +STALE = {"oauthAccount": {"organizationRateLimitTier": "default_claude_max_5x", "userRateLimitTier": None}} - printf '%s' "$plan" -} -credentials='{"claudeAiOauth":{"accessToken":"token","expiresAt":1,"subscriptionType":"max","rateLimitTier":"default_claude_max_5x"}}' -upgraded='{"oauthAccount":{"organizationRateLimitTier":"default_claude_max_20x","userRateLimitTier":null}}' +def plan_of(directory): + return load_collector("claude").oauth_login(Path(directory))[-1] + + +# The profile sits in $HOME by default, inside the config directory when +# CLAUDE_CONFIG_DIR moved the whole config elsewhere, or as a legacy +# .config.json inside it; `where` says which shape to plant, and `stray` +# leaves a copy in the place the CLI does not read. +def read_plan(credentials, profile=None, where="beside", stray=None): + sandbox() + directory = home() / ".claude" + write_json(directory / ".credentials.json", credentials) + live, other = { + "beside": (home() / ".claude.json", directory / ".claude.json"), + "inside": (directory / ".claude.json", home() / ".claude.json"), + "legacy": (directory / ".config.json", home() / ".claude.json"), + }[where] + if where == "inside": + os.environ["CLAUDE_CONFIG_DIR"] = str(directory) + if profile is not None: + write_json(live, profile) + if stray is not None: + write_json(other, stray) + return plan_of(directory) + + +def credentials(tier, subscription): + return {"claudeAiOauth": {"accessToken": "token", "expiresAt": 1, "rateLimitTier": tier, "subscriptionType": subscription}} + # The tier in the credentials is the one the token was minted with. An upgraded # account keeps it there, so the profile the CLI refreshes decides the label. -plan=$(read_plan "$credentials" "$upgraded") -[[ $plan == "Max 20x" ]] || - fail "Claude collector labels the plan from the refreshed profile" "$plan" -pass "Claude collector labels the plan from the refreshed profile" +expect("Claude collector labels the plan from the refreshed profile", read_plan(CREDENTIALS, UPGRADED), "Max 20x") # CLAUDE_CONFIG_DIR takes the profile along with the rest of the config. -plan=$(read_plan "$credentials" "$upgraded" inside) -[[ $plan == "Max 20x" ]] || - fail "Claude collector finds the profile inside a relocated config directory" "$plan" -pass "Claude collector finds the profile inside a relocated config directory" +expect("Claude collector finds the profile inside a relocated config directory", read_plan(CREDENTIALS, UPGRADED, "inside"), "Max 20x") # A copy of the profile in the place the CLI does not read is never refreshed, # so it must not outrank the live one, whichever way the config is laid out. -stale='{"oauthAccount":{"organizationRateLimitTier":"default_claude_max_5x","userRateLimitTier":null}}' -for where in beside inside legacy; do - plan=$(read_plan "$credentials" "$upgraded" "$where" "$stale") - [[ $plan == "Max 20x" ]] || - fail "Claude collector reads only the profile the CLI keeps" "$where -> $plan" -done -pass "Claude collector reads only the profile the CLI keeps" +expect("Claude collector reads only the profile the CLI keeps", + {where: read_plan(CREDENTIALS, UPGRADED, where, STALE) for where in ("beside", "inside", "legacy")}, + {"beside": "Max 20x", "inside": "Max 20x", "legacy": "Max 20x"}) # A seat with a tier of its own is limited by that tier, not by its org's. -plan=$(read_plan "$credentials" '{"oauthAccount":{"organizationRateLimitTier":"default_claude_max_20x","userRateLimitTier":"default_claude_max_5x"}}') -[[ $plan == "Max 5x" ]] || - fail "Claude collector prefers the seat's own tier over the organization's" "$plan" -pass "Claude collector prefers the seat's own tier over the organization's" - -# Without a profile — or with one that states no tier, or a tier the label -# cannot read — the credentials are all there is, and they still answer. -for profile in '' '{}' '{"oauthAccount":{"organizationRateLimitTier":null}}' '{"oauthAccount":{"organizationRateLimitTier":"default_claude_unknown"}}'; do - plan=$(read_plan "$credentials" "$profile") - [[ $plan == "Max 5x" ]] || - fail "Claude collector falls back to the credentials when the profile names no readable tier" "$profile -> $plan" -done -pass "Claude collector falls back to the credentials when the profile names no readable tier" +expect("Claude collector prefers the seat's own tier over the organization's", + read_plan(CREDENTIALS, {"oauthAccount": {"organizationRateLimitTier": "default_claude_max_20x", "userRateLimitTier": "default_claude_max_5x"}}), "Max 5x") + +# Without a profile, or with one that states no tier, or a tier the label +# cannot read, the credentials are all there is, and they still answer. +profiles = [None, {}, {"oauthAccount": {"organizationRateLimitTier": None}}, {"oauthAccount": {"organizationRateLimitTier": "default_claude_unknown"}}] +expect("Claude collector falls back to the credentials when the profile names no readable tier", + [read_plan(CREDENTIALS, profile) for profile in profiles], ["Max 5x"] * len(profiles)) # A Team seat runs on a Max tier; the refreshed profile changes the multiplier # but never turns the seat into Max. -team='{"claudeAiOauth":{"accessToken":"token","expiresAt":1,"subscriptionType":"team","rateLimitTier":"default_claude_max_5x"}}' -plan=$(read_plan "$team" "$upgraded") -[[ $plan == "Team 20x" ]] || - fail "Claude collector keeps a Team seat's label when the profile names its tier" "$plan" -pass "Claude collector keeps a Team seat's label when the profile names its tier" +expect("Claude collector keeps a Team seat's label when the profile names its tier", + read_plan(credentials("default_claude_max_5x", "team"), UPGRADED), "Team 20x") # A secondary account's home has its own profile; ~/.claude.json describes the # primary account and must never lend it its tier. -second_home=$(mktemp -d "$SANDBOX/home.XXXXXX") -mkdir -p "$second_home/.claude" "$second_home/accounts/work" -printf '%s' "$upgraded" >"$second_home/.claude.json" -printf '%s' "$credentials" >"$second_home/accounts/work/.credentials.json" -plan=$(env -u CLAUDE_CONFIG_DIR HOME="$second_home" COLLECTOR="$ROOT/bin/omarchy-agent-usage-claude" \ - CLAUDE_DIR="$second_home/accounts/work" python3 - <<'PY' -import importlib.machinery, importlib.util, os, pathlib - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) - -print(collector.oauth_login(pathlib.Path(os.environ["CLAUDE_DIR"]))[-1]) +sandbox() +write_json(home() / ".claude.json", UPGRADED) +write_json(home() / "accounts" / "work" / ".credentials.json", CREDENTIALS) +(home() / ".claude").mkdir() +expect("Claude collector never labels a secondary account from the primary's profile", plan_of(home() / "accounts" / "work"), "Max 5x") + +# The credentials carry two plan facts: the subscription and the rate-limit +# tier it runs under. With no profile beside them, they settle the label. A +# Team premium seat runs on the Max 5x rate-limit tier, and the tier alone +# used to label it "Max 5x", the plan the seat is not on. +expect("Claude collector labels a Team seat by its subscription, not its rate-limit tier", + read_plan(credentials("default_claude_max_5x", "team")), "Team 5x") + +# Max keeps its multiplier, with or without a subscription type alongside the +# tier, and a plan whose tier names no multiplier is just the plan. +expect("Claude collector keeps Max multipliers and plain plan names", [ + read_plan(credentials("default_claude_max_20x", "max")), read_plan(credentials("default_claude_max_5x", "")), + read_plan(credentials("default_claude_pro", "pro")), read_plan(credentials("", "team")), +], ["Max 20x", "Max 5x", "Pro", "Team"]) PY -) -[[ $plan == "Max 5x" ]] || - fail "Claude collector never labels a secondary account from the primary's profile" "$plan" -pass "Claude collector never labels a secondary account from the primary's profile" diff --git a/test/shell.d/agent-usage-claude-scanner-test.sh b/test/shell.d/agent-usage-claude-scanner-test.sh index 533c2ad6ab1..1812dc77f19 100644 --- a/test/shell.d/agent-usage-claude-scanner-test.sh +++ b/test/shell.d/agent-usage-claude-scanner-test.sh @@ -1,446 +1,737 @@ #!/bin/bash -source "$(dirname "$0")/base-test.sh" +# The Claude collector's local stats: Claude Code's own transcripts, read +# through a per-file index that only reads what was appended, its +# history.jsonl fallback, and its own filter over pi and opencode sessions +# (the anthropic provider). How the shared pi and opencode readers work, and +# the unused-machine record, are held centrally by the shared and partition +# tests, and the scan cache by the contract check. Claude's limits and plan +# have tests of their own. -require_command jq -require_command python3 +set -euo pipefail -# Every fixture home lives under one scratch directory, cleaned up at exit. -SCRATCH=$(mktemp -d) -trap 'rm -rf "$SCRATCH"' EXIT -TEST_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" -projects="$TEST_HOME/.claude/projects/example" -mkdir -p "$projects" +agent_sandbox +require_command rg -timestamp="$(date +%Y-%m-%d)T12:00:00Z" -cat >"$projects/session.jsonl" <>"$projects/session.jsonl" </dev/null | head -1) -[[ -n $index && $(jq -r '.files | to_entries[0].value.records | length' "$index") == "3" ]] || - fail "Claude collector keeps one record per API message in the scan index" "$(cat "$index" 2>/dev/null)" -pass "Claude collector keeps one record per API message in the scan index" +def claude(*args): + return collect("claude", *(args or ("--force",))) -# An unchanged file is not opened again: its records come from the index. -chmod 000 "$projects/session.jsonl" -result=$(HOME="$TEST_HOME" XDG_CACHE_HOME="$TEST_HOME/.cache" XDG_DATA_HOME="$TEST_HOME/.local/share" \ - "$ROOT/bin/omarchy-agent-usage-claude" --cache-seconds 0 2>/dev/null) -chmod 644 "$projects/session.jsonl" -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "59000" ]] || - fail "Claude collector serves unchanged transcripts from the scan index" "$result" -pass "Claude collector serves unchanged transcripts from the scan index" +def rescan(): + return collect("claude", "--cache-seconds", "0") -# A file that shrank was rewritten, not appended to; it is read from the start. -head -n 1 "$projects/session.jsonl" >"$projects/session.jsonl.new" -mv "$projects/session.jsonl.new" "$projects/session.jsonl" -result=$(HOME="$TEST_HOME" XDG_CACHE_HOME="$TEST_HOME/.cache" XDG_DATA_HOME="$TEST_HOME/.local/share" \ - "$ROOT/bin/omarchy-agent-usage-claude" --cache-seconds 0) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "29090" ]] || - fail "Claude collector rescans a transcript that was rewritten" "$result" -pass "Claude collector rescans a transcript that was rewritten" - -# A transcript replaced by a larger one, or rewritten in place with more than -# it had, is a new file rather than an append: nothing of the old one stays, -# and the new one is read from its start. -INDEX_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -index_projects="$INDEX_HOME/.claude/projects/example" -mkdir -p "$index_projects" -index_line() { - printf '{"timestamp":"%s","type":"assistant","sessionId":"s","message":{"id":"%s","role":"assistant","model":"claude-test","usage":{"input_tokens":%s,"output_tokens":0}}}\n' "$timestamp" "$1" "$2" -} -index_scan() { - HOME="$INDEX_HOME" XDG_CACHE_HOME="$INDEX_HOME/.cache" XDG_DATA_HOME="$INDEX_HOME/.local/share" \ - "$ROOT/bin/omarchy-agent-usage-claude" --cache-seconds 0 "$@" -} - -index_line old-1 1000 >"$index_projects/session.jsonl" -index_scan >/dev/null -{ index_line new-1 20; index_line new-2 30; } >"$index_projects/session.jsonl.new" -mv "$index_projects/session.jsonl.new" "$index_projects/session.jsonl" -[[ $(index_scan | jq -c '[.todayTotalTokens, .totalPrompts]') == '[50,2]' ]] || - fail "Claude collector reads a transcript replaced by a larger one from its start" "$(index_scan)" - -{ index_line other-1 7; index_line other-2 8; index_line other-3 9; } >"$index_projects/session.jsonl" -[[ $(index_scan | jq -c '[.todayTotalTokens, .totalPrompts]') == '[24,3]' ]] || - fail "Claude collector reads a transcript rewritten in place from its start" "$(index_scan)" -pass "Claude collector rereads a transcript that was replaced or rewritten" - -# A rewrite in place that grows the file and changes only its end keeps the -# first kilobytes, so the end of what was read is checked too. -for i in $(seq 1 60); do index_line "long-$i" 100; done >"$index_projects/session.jsonl" -index_scan >/dev/null -{ for i in $(seq 1 59); do index_line "long-$i" 100; done; index_line long-60 5000; index_line long-61 1; } >"$index_projects/session.jsonl" -[[ $(index_scan | jq -r '.todayTotalTokens') == "10901" ]] || - fail "Claude collector rereads a transcript whose end was rewritten in place" "$(index_scan)" -pass "Claude collector rereads a transcript whose end was rewritten in place" - -# A larger replacement that matches both ends of what was read but differs in -# between is still a different file, so it is read from its start. -for i in $(seq 1 200); do index_line "mid-$i" 100; done >"$index_projects/session.jsonl" -index_scan >/dev/null -{ for i in $(seq 1 200); do if (( i == 100 )); then index_line "mid-$i" 900; else index_line "mid-$i" 100; fi; done; index_line mid-201 1; } >"$index_projects/session.jsonl.new" -mv "$index_projects/session.jsonl.new" "$index_projects/session.jsonl" -[[ $(index_scan | jq -r '.todayTotalTokens') == "20801" ]] || - fail "Claude collector rereads a replacement whose ends match" "$(index_scan)" -pass "Claude collector rereads a replacement whose ends match" -# The index holds local days, so a new timezone reads every transcript again -# rather than keep the days another timezone gave them. -zone_timestamp="$(date -u +%Y-%m-%d)T01:00:00Z" -printf '{"timestamp":"%s","type":"assistant","sessionId":"s","message":{"id":"zone-1","role":"assistant","model":"claude-test","usage":{"input_tokens":5,"output_tokens":0}}}\n' "$zone_timestamp" >"$index_projects/session.jsonl" -TZ=UTC index_scan >/dev/null -zone_dates=$(TZ=America/Los_Angeles index_scan | jq -r '.activeDates | join(",")') -[[ $zone_dates == "$(TZ=America/Los_Angeles date -d "$zone_timestamp" +%Y-%m-%d)" ]] || - fail "Claude collector recomputes indexed days after a timezone change" "$zone_dates" -pass "Claude collector recomputes indexed days after a timezone change" - -# A streamed response is several lines sharing one message id. The first -# line's output_tokens is a placeholder and the last line has the real count, -# so the message is counted from the line with the highest output. -STREAM_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -stream_projects="$STREAM_HOME/.claude/projects/example" -mkdir -p "$stream_projects" -cat >"$stream_projects/session.jsonl" <>"$stream_projects/session.jsonl" <>"$projects/session.jsonl" <"$HISTORY_HOME/.claude/history.jsonl" <"$PI_HOME/.pi/agent/sessions/project/pi.jsonl" <"$PI_HOME/.omp/agent/sessions/project/omp.jsonl" <` moves the whole agent tree under profiles//, so a -# subscription spent entirely through a profile leaves the default root empty. -cat >"$PI_HOME/.omp/profiles/work/agent/sessions/project/omp-profile.jsonl" <\nb1\n", promptSource="system"), + tool_reply("msg_3", 53, "text"), + user("meta-1", "context", isMeta=True), + user("summary-1", "This session is being continued from a previous conversation.", isCompactSummary=True, isVisibleInTranscriptOnly=True), + user("command-1", "/model"), + user("stdout-1", "Set model"), + user("interrupted-1", [{"type": "text", "text": "[Request interrupted by user]"}]), + user("image-1", [{"type": "image", "source": {}}, {"type": "text", "text": "what is this"}], promptSource="typed"), + tool_reply("msg_4", 54, "text"), + queued("queued-1", "also check the tests", {"kind": "human"}, humanTurn=True), + queued("queued-2", "done", {"kind": "peer", "from": "a1"}, isMeta=True), + {"type": "attachment", "uuid": "notified-2", "sessionId": "session-1", "timestamp": STAMP, "isSidechain": False, + "attachment": {"type": "queued_command", "commandMode": "task-notification", "prompt": ""}}, + name="parent.jsonl", ) +transcript({**user("side-1", "explore the repo"), "isSidechain": True}, name="parent/subagents/agent-a1.jsonl") +fork = [{**line, "sessionId": "session-2"} for line in loop] +transcript(*fork, user("typed-2", "forked", session="session-2", promptSource="typed"), tool_reply("msg_5", 55, "tool_use"), tool_reply("msg_6", 56, "text"), name="fork.jsonl") +record = claude() +expect("Claude counts the prompts a person typed, once each, not the API requests they made", + [record["todayPrompts"], record["totalPrompts"], record["todayTotalTokens"]], [5, 5, 6321], record) + +# A built-in that runs locally asks the model nothing, so it is no prompt: +# Claude Code 2.1.308 answers its `` line with a +# `` child, a user line (/compact, /model) or a +# `local_command` system line (/clear). A custom slash command expands into a +# prompt the model answers, and counts. An output read after its command, in +# a later scan, still takes the command back. +sandbox() + + +def command(uuid, name, parent=None): + return user(uuid, f"/{name}\n {name}\n ", parentUuid=parent) + + +def stdout(uuid, parent, text=""): + return user(uuid, f"{text}", parentUuid=parent) + + +transcript( + user("typed-1", "hello", promptSource="typed"), tool_reply("msg_1", 51, "text"), + command("compact-1", "compact", "typed-1"), stdout("compact-out", "compact-1", "Compacted (ctrl+o to see full summary)"), + {**user("summary-1", "This session is being continued from a previous conversation.", isCompactSummary=True), "parentUuid": "compact-out"}, + command("clear-1", "clear"), + {"type": "system", "subtype": "local_command", "uuid": "clear-out", "parentUuid": "clear-1", "isSidechain": False, "sessionId": "session-1", + "timestamp": STAMP, "content": "", "isMeta": False, "commandRun": {"command": "clear", "args": ""}}, + user("review-1", "review is running…\n/review"), + user("review-meta", "Review the current diff.", isMeta=True, parentUuid="review-1"), tool_reply("msg_2", 52, "text"), + command("model-1", "model"), +) +record = rescan() +transcript(stdout("model-out", "model-1", "Set model to Opus 5"), mode="a") +grown = rescan() +expect("Claude counts no built-in local command as a prompt, a custom command still", + [record["todayPrompts"], grown["todayPrompts"], grown["todayTotalTokens"]], [3, 2, 2103], grown) + +# A transcript replaced by a larger one, rewritten in place with more than it +# had, or rewritten at its end or in its middle with both ends matching, is a +# new file rather than an append, read from its start. +sandbox() + + +def replaced(*lines): + transcript(*lines, name="session.jsonl.new") + (projects() / "session.jsonl.new").replace(projects() / "session.jsonl") + + +transcript(message("old-1", 1000)) +rescan() +replaced(message("new-1", 20), message("new-2", 30)) +record = rescan() +expect("Claude reads a transcript replaced by a larger one from its start", record["todayTotalTokens"], 50, record) +transcript(message("other-1", 7), message("other-2", 8), message("other-3", 9)) +record = rescan() +expect("Claude reads a transcript rewritten in place from its start", record["todayTotalTokens"], 24, record) + +transcript(*(message(f"long-{i}", 100) for i in range(1, 61))) +rescan() +transcript(*(message(f"long-{i}", 100) for i in range(1, 60)), message("long-60", 5000), message("long-61", 1)) +record = rescan() +expect("Claude rereads a transcript whose end was rewritten in place", record["todayTotalTokens"], 10901, record) + +transcript(*(message(f"mid-{i}", 100) for i in range(1, 201))) +rescan() +replaced(*(message(f"mid-{i}", 900 if i == 100 else 100) for i in range(1, 201)), message("mid-201", 1)) +record = rescan() +expect("Claude rereads a larger replacement whose ends match", record["todayTotalTokens"], 20801, record) + +# The index holds local days, so a new timezone reads every transcript again +# rather than keep the days another one gave them. +zone_stamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT01:00:00Z") +transcript(message("zone-1", 5, timestamp=zone_stamp)) +saved_tz = os.environ.get("TZ") +try: + os.environ["TZ"] = "UTC" + time.tzset() + rescan() + os.environ["TZ"] = "America/Los_Angeles" + time.tzset() + record = rescan() + expected_day = datetime.fromisoformat(zone_stamp.replace("Z", "+00:00")).astimezone().date().isoformat() +finally: + if saved_tz is None: + os.environ.pop("TZ", None) + else: + os.environ["TZ"] = saved_tz + time.tzset() +expect("Claude recomputes indexed days after a timezone change", record["activeDates"], [expected_day], record) + +# Every collector reading Claude Code runs at once, and the index is built +# once: one reads the transcripts while the others wait for it, then take +# what it wrote. A transcript grown since is the only one read again, and an +# index current for every transcript is read as it is. +sandbox() +for n in range(3): + transcript(message(f"busy-{n}", 100 + n, 1), name=f"busy-{n}.jsonl") +import collectors +reads = home() / "reads.log" +totals = home() / "totals.log" +start = home() / "start" +real_read = collectors.read_claude_code_transcript + + +def logged_read(path, entry): + result = real_read(path, entry) + if result is not entry: + with reads.open("a") as handle: + handle.write(f"{path.name}\n") + time.sleep(0.1) + return result + + +def scan_at_once(count=4): + for path in (reads, totals, start): + path.unlink(missing_ok=True) + children = [] + for _ in range(count): + pid = os.fork() + if pid == 0: + try: + collectors.read_claude_code_transcript = logged_read + while not start.exists(): + time.sleep(0.005) + rows = collectors.scan_claude_code("claude", projects().parent, base_url="") + with totals.open("a") as handle: + handle.write(f"{sum(row.input_tokens for row in rows)}\n") + finally: + os._exit(0) + children.append(pid) + start.touch() + for pid in children: + os.waitpid(pid, 0) + lines = lambda path: sorted(path.read_text().split()) if path.exists() else [] + return lines(reads), lines(totals) + + +cold = scan_at_once() +transcript(message("busy-3", 1000, 1), name="busy-0.jsonl", mode="a") +grown = scan_at_once() +warm = scan_at_once() +expect("Claude Code's index is built once while every collector reads it at once", + [cold, grown, warm], + [(["busy-0.jsonl", "busy-1.jsonl", "busy-2.jsonl"], ["303"] * 4), (["busy-0.jsonl"], ["1303"] * 4), ([], ["1303"] * 4)]) + +# A streamed response is several lines under one message id: the first line's +# output is a placeholder, so a message counts from its highest output. One +# that fell back to another model mid-stream keeps its final line whole rather +# than a maximum per field. +sandbox() +transcript( + message("message-1", 8, 1, cache_write=100, cache_read=2000, event="event-1"), + message("message-1", 8, 1, cache_write=100, cache_read=2000, event="event-2"), + message("message-1", 8, 260, cache_write=100, cache_read=2000, event="event-3"), + message("message-2", 10, 1, cache_write=300, cache_read=4000, event="event-4"), + message("message-2", 10, 40, cache_read=3500, model="claude-fallback", event="event-5"), +) +record = claude() +expect("Claude counts a streamed message from its final usage line, whole after a fallback", { + "test": record["modelUsage"]["claude-test"], "fallback": record["modelUsage"]["claude-fallback"], "today": record["todayTotalTokens"], +}, { + "test": {"cacheCreationInputTokens": 100, "cacheReadInputTokens": 2000, "inputTokens": 8, "outputTokens": 260}, + "fallback": {"cacheCreationInputTokens": 0, "cacheReadInputTokens": 3500, "inputTokens": 10, "outputTokens": 40}, + "today": 5918, +}, record) + +# A stream still being written when the index was taken finishes in lines +# appended later, and the indexed message takes the higher count from them. +rescan() +transcript(message("message-2", 10, 90, cache_read=3500, model="claude-fallback", event="event-6"), mode="a") +record = rescan() +expect("Claude takes a streamed message's final count from appended lines", + [record["modelUsage"]["claude-fallback"]["outputTokens"], record["todayTotalTokens"]], [90, 5968], record) + +# A line of JSON that isn't an object skips that line, never the transcript. +sandbox() +transcript(message("listed-1", 10), [message("listed-x", 99)], message("listed-2", 20)) +record = claude() +expect("Claude skips a transcript line that isn't an object and keeps the rest", + record["todayTotalTokens"], 30, record) + +# Only a missing time is today's: a time of 0, false or "" is untrusted and +# counts nowhere. +sandbox() +transcript(message("stamped", 10), *(message(f"falsy-{n}", 99, timestamp=stamp) for n, stamp in enumerate((0, False, "")))) +record = claude() +expect("Claude counts a transcript line with a falsy time nowhere", + [record["todayTotalTokens"], sum(bucket["inputTokens"] for bucket in record["modelUsage"].values())], [10, 10], record) + +# `--fork-session` copies the parent's replies into a new transcript under the +# fork's session id. Their tokens count once, but the parent was still a +# session in use, whichever transcript the scan visits last. +sandbox() +transcript(*({**message(f"msg_{n}", 1000, 50 + n), "sessionId": "parent"} for n in (1, 2)), name="parent.jsonl") +transcript(*({**message(f"msg_{n}", 1000, 50 + n), "sessionId": "fork"} for n in (1, 2, 3)), name="fork.jsonl") +module = load_collector("claude") +real_rglob = Path.rglob +for reverse in (False, True): + module.Path.rglob = lambda path, pattern, reverse=reverse: iter(sorted(real_rglob(path, pattern), reverse=reverse)) + try: + forked = module.scan_projects(projects().parent).snapshot() + finally: + module.Path.rglob = real_rglob + expect(f"Claude counts a forked session's parent and its tokens once ({'parent' if reverse else 'fork'} read first)", + [forked["todaySessions"], forked["totalSessions"], forked["todayTotalTokens"]], [2, 2, 3156], forked) + +# A response is its message id and its request id. A gateway that reuses one +# message id for every response still served each of them; a copy in a fork +# keeps both ids and counts once; and a /btw side question's sidechain replays +# its parent's message under a new request id, which the parent's copy covers. +sandbox() + + +def served(message_id, request_id, output, session="session-1", **extra): + return {**message(message_id, 1000, output), "requestId": request_id, "sessionId": session, **extra} + + +transcript(served("msg_gateway", "req_1", 10), served("msg_gateway", "req_2", 20), served("msg_parent", "req_3", 30), name="main.jsonl") +transcript(served("msg_gateway", "req_1", 10, session="fork"), served("msg_gateway", "req_4", 40, session="fork"), name="fork.jsonl") +transcript(served("msg_parent", "req_5", 30, isSidechain=True), served("msg_aside", "req_6", 50, isSidechain=True), name="main/subagents/agent-aside.jsonl") +record = claude() +expect("Claude counts a response by message and request id, once across forks and sidechain replays", + record["todayTotalTokens"], 5000 + 10 + 20 + 30 + 40 + 50, record) + +# A gateway reusing one message id gives a subagent's reply its parent's id +# too, as Claude Code 2.1.308 wrote it behind one: the subagent's request is +# its own, not a replay of the parent's, and counts. +sandbox() + + +def gateway(request_id, n, **extra): + return {**message("msg_same", 5000234 + 1000 * n, 50 + n, cache_write=30 + n, cache_read=70 + n), "requestId": request_id, **extra} + + +transcript(gateway("req_0002", 1), gateway("req_0004", 3), gateway("req_0005", 4), name="main.jsonl") +transcript(gateway("req_0003", 2, isSidechain=True, agentId="a1"), name="main/subagents/agent-a1.jsonl") +record = claude() +expect("Claude counts a subagent's reply that shares its parent's message id", record["todayTotalTokens"], 20011566, record) + +# An advisor the model consulted is billed apart under its own model, from +# the reply's `usage.iterations`; the top-level usage is the main model's +# alone, and no other iteration is added again. ccusage's own fixture line, +# with Claude models. +sandbox() +transcript({"timestamp": STAMP, "version": "1.2.3", "sessionId": "session-a", "type": "assistant", "requestId": "req-parent", "costUSD": 1.23, + "message": {"id": "msg-parent", "model": "claude-main", "usage": {"input_tokens": 1, "output_tokens": 2, "iterations": [ + {"type": "message", "input_tokens": 1, "output_tokens": 2, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, + {"type": "advisor_message", "model": "claude-advisor", "input_tokens": 10, "output_tokens": 2, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, + ]}}}) +record = claude() +expect("Claude counts an advisor iteration under the advisor's model", + [record["todayTokensByModel"], record["modelUsage"].get("claude-advisor")], + [{"claude-main": 3, "claude-advisor": 12}, {"inputTokens": 10, "outputTokens": 2, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}], record) + +# Claude Code fronts whatever ANTHROPIC_BASE_URL names, and Z.ai and Ollama +# document pointing it at their own APIs; their transcripts look like any +# other, with the model that answered on each reply, as a real Claude Code +# 2.1 run against them wrote: glm-4.6 through Z.ai, then `--continue` back on +# Anthropic in the same session, and qwen3-coder and glm-4.6:cloud through +# Ollama. Only what Anthropic served is Claude's, prompts included. A reply +# Claude Code made up itself (``, no tokens) answers nothing. +sandbox() +transcript(user("zai-1", "zaihello", session="zai", promptSource="sdk"), {**message("msg_z1", 501234, 51, model="glm-4.6"), "sessionId": "zai"}, + user("zai-2", "continued", session="zai", promptSource="sdk"), {**message("msg_z2", 1000, 52, model="claude-sonnet-4-5"), "sessionId": "zai"}, + name="zai.jsonl") +transcript(user("local-1", "ollamaclaude", session="local", promptSource="sdk"), {**message("msg_o1", 21000, 53, model="qwen3-coder"), "sessionId": "local"}, + name="local.jsonl") +transcript(user("cloud-1", "ollamacloud", session="cloud", promptSource="sdk"), {**message("msg_c1", 31000, 54, model="glm-4.6:cloud"), "sessionId": "cloud"}, + name="cloud.jsonl") +transcript(user("synthetic-1", "hello", session="synthetic", promptSource="typed"), {**message("msg_s1", model=""), "sessionId": "synthetic"}, + name="synthetic.jsonl") +record = claude() +expect("Claude leaves out the turns Claude Code sent to another provider's model", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"claude-sonnet-4-5": 1052}, 2, 2], record) +# What Claude leaves out, the provider's collector counts by the same rule: +# the GLM reply is Z.ai's, once. +zai_record = collect("zai", "--force") +expect("Z.ai counts the Claude Code reply Claude leaves out", zai_record["todayTokensByModel"], {"glm-4.6": 501285}, zai_record) + +# A transcript doesn't say where it was sent, so while Claude Code is pointed +# at another provider (its settings' env, else the environment), a Claude +# model's replies are that provider's too. A host this can't place, such as a +# proxy that may forward to Anthropic, leaves them Claude's. +for env, settings, claude_tokens in ( + ({}, {"ANTHROPIC_BASE_URL": "https://api.z.ai/api/anthropic", "ANTHROPIC_AUTH_TOKEN": "zai-key"}, 0), + ({"ANTHROPIC_BASE_URL": "https://ollama.com"}, {}, 0), + ({"ANTHROPIC_BASE_URL": "https://ollama.com"}, {"ANTHROPIC_BASE_URL": "http://127.0.0.1:4000"}, 1052), + ({}, {"ANTHROPIC_BASE_URL": "https://llm-proxy.example.com/anthropic"}, 1052), +): + write_json(home() / ".claude" / "settings.json", {"env": settings}) + os.environ.update(env) + try: + record = claude() + finally: + for name in env: + os.environ.pop(name) + expect(f"Claude reads Claude model replies as {'Claude' if claude_tokens else 'not'} Claude's with the endpoint at {settings or env}", + [record["todayTotalTokens"], record["todayPrompts"]], [claude_tokens, 2 if claude_tokens else 0], record) + +# Unreadable transcripts are expected when an agent ran as another user, so +# they make one bounded warning rather than a line each. +sandbox() +for n in range(1, 6): + transcript(name=f"unreadable-{n}.jsonl") +module = load_collector("claude") +real_open = Path.open -[[ $(jq -c '.failures' <<<"$race_output") == "[]" ]] || - fail "Claude collector survives concurrent writes to one cache file" "$race_output" -[[ $(jq -r '.payload.writer != null and (.leftovers | length) == 0' <<<"$race_output") == "true" ]] || - fail "Claude collector leaves one intact cache file and no temp files" "$race_output" -[[ $(jq -r '.mode' <<<"$race_output") == "0o644" ]] || - fail "Claude collector keeps cache files readable" "$race_output" -pass "Claude collector survives concurrent writes to one cache file" - -# Unreadable transcript files are expected when an agent has run as another -# user, so report one bounded summary rather than one log line per file. -UNREADABLE_DIR=$(mktemp -d "$SCRATCH/home.XXXXXX") -mkdir -p "$UNREADABLE_DIR/project" -touch "$UNREADABLE_DIR/project/unreadable-1.jsonl" \ - "$UNREADABLE_DIR/project/unreadable-2.jsonl" \ - "$UNREADABLE_DIR/project/unreadable-3.jsonl" \ - "$UNREADABLE_DIR/project/unreadable-4.jsonl" \ - "$UNREADABLE_DIR/project/unreadable-5.jsonl" - -warning=$(python3 - "$ROOT/bin/omarchy-agent-usage-claude" "$UNREADABLE_DIR" <<'PY' -import contextlib -import importlib.util -import io -import sys -from importlib.machinery import SourceFileLoader -from pathlib import Path - -spec = importlib.util.spec_from_loader("collector", SourceFileLoader("collector", sys.argv[1])) -collector = importlib.util.module_from_spec(spec) -spec.loader.exec_module(collector) - -original_open = Path.open def deny_unreadable(path, *args, **kwargs): if path.name.startswith("unreadable-"): raise PermissionError("permission denied") - return original_open(path, *args, **kwargs) + return real_open(path, *args, **kwargs) -collector.Path.open = deny_unreadable -with contextlib.redirect_stderr(io.StringIO()) as stderr: - collector.scan_projects(Path(sys.argv[2])) -print(stderr.getvalue(), end="") -PY -) -[[ $(grep -c '^Ignoring ' <<<"$warning") == 1 ]] || - fail "Claude collector emits one bounded warning for unreadable transcripts" "$warning" -grep -q '^Ignoring 5 unreadable Claude project files' <<<"$warning" || - fail "Claude collector reports the unreadable transcript count" "$warning" -pass "Claude collector bounds unreadable transcript warnings" +module.Path.open = deny_unreadable +try: + with contextlib.redirect_stderr(io.StringIO()) as warnings: + module.scan_projects(projects().parent) +finally: + module.Path.open = real_open +lines = [line for line in warnings.getvalue().splitlines() if line.startswith("Ignoring ")] +check("Claude bounds unreadable transcripts to one warning with their count", + len(lines) == 1 and lines[0].startswith("Ignoring 5 unreadable Claude project files"), warnings.getvalue()) + +# ---------------------------------------------------------------- sources + +# A machine with no transcripts and no stats-cache still gets today's counts +# from history.jsonl. +sandbox() +(home() / ".claude").mkdir() +(home() / ".claude" / "history.jsonl").write_text("".join(json.dumps(line) + "\n" for line in ( + {"timestamp": 86400000, "sessionId": "old", "display": "ancient"}, + {"timestamp": ms_in(0), "sessionId": "s1", "display": "one"}, + {"timestamp": ms_in(0), "sessionId": "s2", "display": "two"}, +))) +record = claude() +expect("Claude falls back on history.jsonl without a stats-cache", [record["todayPrompts"], record["todaySessions"]], [2, 2], record) + +# A stats-cache shaped as Claude Code writes it gives the record's own +# shapes: seven local days of tokens from dailyModelTokens (dailyActivity +# counts messages), and per model the four token kinds alone, each a trusted +# count. A model another provider served (glm-4.6 through Z.ai) is not +# Claude's here either. +sandbox() +(home() / ".claude").mkdir() +week = [(datetime.now() - timedelta(days=offset)).date().isoformat() for offset in range(6, -1, -1)] +write_json(home() / ".claude" / "stats-cache.json", { + "totalMessages": 120, "totalSessions": 3, + "dailyActivity": [{"date": "2026-01-02", "messageCount": 120, "sessionCount": 3, "toolCallCount": 40}], + "dailyModelTokens": [ + {"date": "2026-01-02", "tokensByModel": {"claude-opus": 999}}, + {"date": week[4], "tokensByModel": {"claude-opus": 300, "claude-haiku": 20, "glm-4.6": 4000}}, + {"date": week[6], "tokensByModel": {"claude-opus": 50}}, + ], + "modelUsage": { + "claude-opus": {"inputTokens": 10, "outputTokens": -5, "cacheReadInputTokens": 7, "cacheCreationInputTokens": "x", "costUSD": 1.5, "webSearchRequests": 2}, + "claude-odd": [1, 2], + "glm-4.6": {"inputTokens": 4000, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, + }, +}) +record = claude() +expect("Claude's stats-cache fallback gives seven days of tokens and trusted model buckets", [record["recentDays"], record["modelUsage"]], [ + [{"date": day, "messageCount": {4: 320, 6: 50}.get(n, 0)} for n, day in enumerate(week)], + {"claude-opus": {"inputTokens": 10, "outputTokens": 0, "cacheReadInputTokens": 7, "cacheCreationInputTokens": 0}}, +], record) + +# history.jsonl as interactive Claude Code 2.1.308 wrote it: two typed +# prompts, /model and /exit (built-ins that ask the model nothing, with and +# without arguments), a custom command, and a prompt typed while a turn ran. +def history_line(display, session): + return {"display": display, "pastedContents": {}, "timestamp": ms_in(0), "project": "/home/me/proj", "sessionId": session} + + +HISTORY = [history_line("tuihello", "s1"), history_line("tuisecond", "s1"), history_line("/model opus", "s1"), history_line("/exit", "s1"), + history_line("/mycmd", "s2"), history_line("queuedtwo", "s2"), history_line("/exit", "s2")] + +# Without transcripts, history.jsonl's built-in commands are no prompts. +sandbox() +(home() / ".claude").mkdir() +(home() / ".claude" / "history.jsonl").write_text("".join(json.dumps(line) + "\n" for line in HISTORY)) +record = claude() +expect("Claude's history.jsonl fallback counts no built-in local command", [record["todayPrompts"], record["todaySessions"]], [4, 2], record) + +# Sent to Z.ai, those prompts were Z.ai's, and neither file says otherwise. +write_json(home() / ".claude" / "settings.json", {"env": {"ANTHROPIC_BASE_URL": "https://api.z.ai/api/anthropic"}}) +write_json(home() / ".claude" / "stats-cache.json", {"totalMessages": 40, "totalSessions": 3, "dailyActivity": [{"date": datetime.now().date().isoformat(), "messageCount": 4}]}) +record = claude() +expect("Claude sent to Z.ai takes no prompts or sessions from its fallbacks", + [record["todayPrompts"], record["todaySessions"], record["totalPrompts"], record["totalSessions"], record["activeDays"]], [0, 0, 0, 0, 0], record) + +# Transcripts whose every reply another provider served were counted by that +# provider's collector, prompts included: the fallbacks are not read again. +(home() / ".claude" / "settings.json").unlink() +transcript( + {"type": "user", "uuid": "u1", "sessionId": "s1", "timestamp": STAMP, "promptSource": "typed", "message": {"role": "user", "content": "tuihello"}}, + message("glm-1", 10, 5, model="glm-5"), +) +for where, settings in (("named by its model", {}), ("sent to Z.ai", {"env": {"ANTHROPIC_BASE_URL": "https://api.z.ai/api/anthropic"}})): + write_json(home() / ".claude" / "settings.json", settings) + record = claude() + expect(f"Claude with only another provider's transcripts ({where}) reads no fallback", + [record["ready"], record["todayPrompts"], record["todaySessions"], record["totalPrompts"], record["totalSessions"], record["activeDays"]], + [False, 0, 0, 0, 0, 0], record) + +# A stats-cache or history line of the wrong shape costs that field or line, +# never the record. +sandbox() +(home() / ".claude").mkdir() +for cache in ([1, 2], {"dailyModelTokens": 5, "dailyActivity": 7, "modelUsage": [1]}, + {"dailyModelTokens": [{"date": datetime.now().date().isoformat(), "tokensByModel": ["x"]}]}): + write_json(home() / ".claude" / "stats-cache.json", cache) + (home() / ".claude" / "history.jsonl").write_text(json.dumps([1]) + "\n" + json.dumps({"timestamp": ms_in(0), "sessionId": "s1"}) + "\n") + record = claude() + expect(f"a stats-cache shaped {json.dumps(cache)[:40]} and a list history line read as nothing", + [record["todayPrompts"], record["todayTokensByModel"], record["modelUsage"]], [1, {}, {}], record) + +# pi and opencode replies count on the anthropic provider alone: not another +# provider driving Anthropic's API, and not a prompt. +sandbox() +write_pi_session(home() / ".pi" / "agent" / "sessions" / "project" / "pi.jsonl", [ + pi_reply("anthropic", 10, id="pi-1", model="claude-pi", api="anthropic-messages"), + pi_reply("kimi-coding", 999, id="kimi-1", model="k3", api="anthropic-messages"), +]) +opencode_db(message=[ + ("msg_1", "ses_1", oc_v1("anthropic", 100, model="claude-opus-5")), + ("msg_2", "ses_1", oc_v1("anthropic", 999, model="claude-opus-5", role="user")), +]) +record = claude() +expect("Claude counts pi and opencode replies on the anthropic provider alone", + [record["todayTotalTokens"], sorted(record["modelUsage"])], [110, ["claude-opus-5", "claude-pi"]], record) + +# A Claude Code home without a sign-in is in use, and the sign-in is missing. +sandbox() +(home() / ".claude").mkdir() +record = claude() +expect("Claude asks a signed-out Claude Code home to sign in", [record["ready"], record["usageStatusText"]], [False, "Waiting for auth"], record) + +# ------------------------------------------------------------- cost-state + +# Claude Code 2.1.278 and later write a cost-state line when a process saves its +# cost: the lineage's startTime, carried through a resume and a +# --fork-session, and what each model has used in it so far, the requests no +# reply line records (a /compact, the session title, a side query) +# included. Shapes as Claude Code wrote them against a mock API that served +# each request once. +sandbox() +YESTERDAY = (fixture_now().astimezone() - timedelta(days=1)).replace(hour=12, minute=0, second=0, microsecond=0).isoformat() + + +def reply(session, message_id, counts, timestamp, model="claude-sonnet-4-5", sidechain=False): + return {"parentUuid": None, "isSidechain": sidechain, "type": "assistant", "sessionId": session, "timestamp": timestamp, + "requestId": f"req_{message_id}", "uuid": f"uuid-{session}-{message_id}", + "message": {"model": model, "id": message_id, "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello"}], + "usage": dict(zip(("input_tokens", "output_tokens", "cache_read_input_tokens", "cache_creation_input_tokens"), counts))}} + + +def typed(session, uuid, timestamp, text="Say hello"): + return {"parentUuid": None, "isSidechain": False, "type": "user", "promptSource": "sdk", "sessionId": session, "timestamp": timestamp, + "uuid": uuid, "message": {"role": "user", "content": text}} + + +def cost_state(session, start, usage): + return {"type": "cost-state", "sessionId": session, "totalCostUSD": 0.0046, "totalAPIDuration": 9, "startTime": start, + "modelUsage": {model: {"inputTokens": counts[0], "outputTokens": counts[1], "thinkingTokens": 0, "cacheReadInputTokens": counts[2], + "cacheCreationInputTokens": counts[3], "webSearchRequests": 0, "costUSD": 0.0046} + for model, counts in usage.items()}, + "hasUnknownModelCost": False} + + +def claude_code_days(owner="claude"): + import collectors + days = {} + for row in collectors.scan_claude_code(owner, projects().parent, base_url=""): + by_model = days.setdefault(row.day, {}) + by_model[row.model] = by_model.get(row.model, 0) + row.input_tokens + row.output_tokens + row.cache_read + row.cache_write + return days + + +yesterday, today_day = YESTERDAY[:10], fixture_now().astimezone().date().isoformat() +# A one-shot run asks for a title too, and is resumed the next day for a +# /compact, whose request no reply line records. +compacted = [ + typed("compact", "u-compact-1", YESTERDAY), + reply("compact", "msg_c1", (1234, 51, 71, 31), YESTERDAY), + {"type": "last-prompt", "lastPrompt": "Say hello", "sessionId": "compact"}, + cost_state("compact", 1792669864166, {"claude-sonnet-4-5[1m]": (1234, 51, 71, 31), "claude-haiku-4-5-20251001": (893, 24, 0, 0)}), + {"type": "queue-operation", "operation": "enqueue", "timestamp": STAMP, "sessionId": "compact", "content": "/compact"}, + {"type": "system", "subtype": "compact_boundary", "content": "Conversation compacted", "timestamp": STAMP, "sessionId": "compact", + "compactMetadata": {"trigger": "manual", "preTokens": 1387}}, + {"type": "user", "isCompactSummary": True, "sessionId": "compact", "timestamp": STAMP, "uuid": "u-summary", "message": {"role": "user", "content": "Summary"}}, + cost_state("compact", 1792669864166, {"claude-sonnet-4-5[1m]": (3002468, 102, 142, 62), "claude-haiku-4-5-20251001": (893, 24, 0, 0)}), +] +transcript(*compacted[:4], name="compact.jsonl") +expect("Claude counts the title query a cost-state holds beyond the replies", claude_code_days(), + {yesterday: {"claude-sonnet-4-5": 1387, "claude-haiku-4-5-20251001": 917}}) +transcript(*compacted[4:], name="compact.jsonl", mode="a") +expect("Claude counts a /compact's request, read from the cost-state appended after it, on the day it ran", claude_code_days(), + {yesterday: {"claude-sonnet-4-5": 1387, "claude-haiku-4-5-20251001": 917}, today_day: {"claude-sonnet-4-5": 3001387}}) +index = json.loads(next(cache_root().glob("claude-index-*.json")).read_text()) +(projects() / "compact.jsonl").touch() +expect("Claude reads a cost-state appended to a transcript as it reads the whole transcript again", + [claude_code_days(), [entry["costs"] for entry in json.loads(next(cache_root().glob("claude-index-*.json")).read_text())["files"].values()]], + [{yesterday: {"claude-sonnet-4-5": 1387, "claude-haiku-4-5-20251001": 917}, today_day: {"claude-sonnet-4-5": 3001387}}, + [entry["costs"] for entry in index["files"].values()]]) + +# A --fork-session carries the lineage on: its cost-state includes the +# parent's, compaction and all, and its transcript copies the parent's +# replies under its own session. A subagent runs under its parent's session +# in a file of its own, and its replies are in the cost-state too. A +# process still running has replies past its last cost-state. +sandbox() +parent = [typed("parent", "u-p1", STAMP), reply("parent", "msg_p1", (101234, 51, 71, 31), STAMP)] +transcript(*parent, cost_state("parent", 1792669919352, {"claude-sonnet-4-5": (203468, 103, 143, 63)}), name="parent.jsonl") +transcript(*[{**line, "sessionId": "fork"} for line in parent], typed("fork", "u-f1", STAMP), reply("fork", "msg_f1", (701234, 51, 71, 31), STAMP), + cost_state("fork", 1792669919352, {"claude-sonnet-4-5": (904702, 154, 214, 94)}), name="fork.jsonl") +transcript(typed("main", "u-m1", STAMP), reply("main", "msg_m1", (100, 1, 0, 0), STAMP), + cost_state("main", 1792681739702, {"claude-sonnet-4-5": (5600, 6, 0, 0), "claude-haiku-4-5-20251001": (300, 3, 0, 0)}), + reply("main", "msg_m2", (90000, 9, 0, 0), today_stamp(13)), name="main.jsonl") +transcript(reply("main", "msg_s1", (5000, 5, 0, 0), STAMP, sidechain=True), name="main/subagents/agent-a1.jsonl") +expect("Claude counts a fork's lineage once, a subagent's replies once, and keeps what a running process's cost-state held beyond its replies", + claude_code_days(), + {today_day: {"claude-sonnet-4-5": 905164 + 101 + 5005 + 500 + 90009, "claude-haiku-4-5-20251001": 303}}) + +# A gateway can name a model one way in the request and another in the +# reply: no more counts than the cost-state says beyond every reply. +sandbox() +transcript(typed("gw", "u-g1", STAMP), reply("gw", "msg_g1", (500, 5, 0, 0), STAMP, model="claude-sonnet-4-5-20250929"), + cost_state("gw", 1792681376828, {"anthropic/claude-sonnet-4.5": (500, 5, 0, 0)}), name="gw.jsonl") +expect("Claude counts nothing more for a cost-state that names the replies' model another way", claude_code_days(), + {today_day: {"claude-sonnet-4-5-20250929": 505}}) + +# Claude Code 2.1.282 can write a subagent's reply with only its +# message_start usage (stop_reason null, a few output tokens) and never its +# real output. The cost-state's totals make up the difference, on the day of +# the cost-state, not the day the subagent ran. +sandbox() +snapshot = {**reply("snap", "msg_s1", (3, 10, 40000, 2000), YESTERDAY, model="claude-opus-4-5", sidechain=True)} +snapshot["message"] = {**snapshot["message"], "stop_reason": None, "content": [{"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {}}]} +transcript(snapshot, name="snap/subagents/agent-a1.jsonl") +transcript(typed("snap", "u-n1", YESTERDAY), reply("snap", "msg_n1", (100, 50, 0, 0), YESTERDAY, model="claude-opus-4-5"), + typed("snap", "u-n2", STAMP), reply("snap", "msg_n2", (10, 1, 0, 0), STAMP, model="claude-opus-4-5"), + cost_state("snap", 1790338485827, {"claude-opus-4-5": (113, 2461, 40000, 2000)}), name="snap.jsonl") +expect("Claude counts the output a cost-state holds beyond a subagent's message_start snapshot, on the cost-state's day", claude_code_days(), + {yesterday: {"claude-opus-4-5": 42163}, today_day: {"claude-opus-4-5": 2411}}) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_message(): + add_message.count = getattr(add_message, "count", 0) + 1 + transcript(message(f"grown-{add_message.count}", 10), mode="a") + + +check_scan_cache(Adapter("claude", grow=add_message, pi_root=lambda: home() / ".pi" / "agent" / "sessions")) + +# A scan that couldn't read one of the transcripts isn't kept for the next +# run, so the missing tokens show as soon as the file can be read. +if os.geteuid() != 0: + sandbox() + claude_code_session("seen", [("claude-sonnet-4-5", 300)]) + locked = claude_code_session("late", [("claude-sonnet-4-5", 40)]) + locked.chmod(0) + try: + claude() + finally: + locked.chmod(0o644) + record = claude("--limits-only") + expect("Claude does not cache a scan that missed a transcript", record["todayTokensByModel"], {"claude-sonnet-4-5": 340}, record) +else: + print("ok - Claude does not cache a scan that missed a transcript # SKIP root reads anything") +PY diff --git a/test/shell.d/agent-usage-codex-limits-test.sh b/test/shell.d/agent-usage-codex-limits-test.sh new file mode 100644 index 00000000000..4bc4ccdf72b --- /dev/null +++ b/test/shell.d/agent-usage-codex-limits-test.sh @@ -0,0 +1,364 @@ +#!/bin/bash + +# The Codex collector's limits: what its app-server's rate-limit answer says +# (single and multi-bucket views, named models, spend caps, reset credits), +# how each way the app-server can fail reads, and whose limits a home's +# cache may show. The unused-machine record is held centrally by the +# partition test, the limits cache by the contract checks, run here against +# a stand-in app-server, and Codex's local stats by its scanner test. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox +fake_codex "$HOME" +export PATH="$HOME/bin:$PATH" + +python3 - <<'PY' +from agent_usage_fixtures import * + +MODES = {"up": "ok", 500: "backend", "down": "died", 401: "auth"} + + +# The stand-in app-server answers as CODEX_FAKE says. +@contextlib.contextmanager +def serve(mode): + os.environ["CODEX_FAKE"] = MODES.get(mode, mode) + yield + + +def starts(): + try: + return len((home() / "codex-starts").read_text().splitlines()) + except FileNotFoundError: + return 0 + + +def codex(mode, *args, **env): + saved = {name: os.environ.get(name) for name in env} + os.environ.update(env) + try: + with serve(mode): + return collect("codex", *(args or ("--force",))) + finally: + for name, value in saved.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +# Signs ~/.codex in as a user of a workspace, the way Codex writes it: the +# workspace as account_id, the user in the ID token's claims. +def login(workspace, user): + claims = {"https://api.openai.com/auth": {"chatgpt_account_id": workspace, "chatgpt_user_id": user}} + write_json(home() / ".codex" / "auth.json", {"tokens": {"id_token": jwt(claims), "account_id": workspace}}) + + +# ---------------------------------------------------------------- contract + +login("ws-one", "user-a") +# "cut": the app-server talks JSON-RPC over a pipe, not HTTP, so there is no +# reply to cut short; its own failures are covered below. +fresh = check_limits_contract(Adapter("codex", serve=serve, count=starts, refused_status="Waiting for auth", retry_modes=("down",), + skip=("cut",))) +expect("fresh Codex limits read both windows and the plan", + [[entry["percent"] for entry in fresh["limits"]], fresh["tierLabel"]], [[0.4, 0.7], "Pro 200"], fresh) + +# ------------------------------------------------------- a failed probe + +# However the app-server fails, the last limits stand, stale, and the help +# says what went wrong. Only a network that isn't there yet asks the panel to +# retry soon: a stalled app-server, or one that died unable to reach OpenAI. +# A refusal, or a crash of its own, would only fail again; a sign-out says so. +good = codex("ok") +for mode, retry, status in (("backend", False, ""), ("died", True, ""), ("crashed", False, ""), ("empty", False, ""), ("auth", False, "Waiting for auth")): + kept = codex(mode) + expect(f"a failed Codex probe ({mode}) keeps the last limits, stale, and says what went wrong", { + "limits": kept["limits"] == good["limits"], "stale": kept["limitsStale"], "at": kept["limitsFetchedAt"] == good["limitsFetchedAt"], + "tier": kept["tierLabel"], "status": kept["usageStatusText"], "help": bool(kept["authHelpText"]), "retry": kept.get("retryAdvised", False), + }, {"limits": True, "stale": True, "at": True, "tier": "Pro 200", "status": status, "help": True, "retry": retry}, kept) + +# A plan without windows is an answer too: a later failed probe has no limits +# to lose, so it is no trouble, whatever went wrong. A sign-out still says so. +no_windows = codex("ok", CODEX_RATE_LIMITS=json.dumps({"planType": "free"})) +expect("a Codex plan without windows is a healthy answer", [no_windows["limits"], no_windows["usageStatusText"], no_windows["tierLabel"]], [[], "", "Free"], no_windows) +for mode, retry, status in (("backend", False, ""), ("died", True, ""), ("crashed", False, ""), ("auth", False, "Waiting for auth")): + failed = codex(mode) + expect(f"a failed Codex probe ({mode}) after a plan without windows is no trouble", { + "limits": failed["limits"], "tier": failed["tierLabel"], "status": failed["usageStatusText"], "help": bool(failed["authHelpText"]), + "retry": failed.get("retryAdvised", False), + }, {"limits": [], "tier": "Free", "status": status, "help": bool(status), "retry": retry}, failed) +good = codex("ok") + +# A stall costs the app-server's whole 8-second answer timeout, so this run's +# collector waits half a second instead. (The collector reads its +# environment when it loads, so it loads inside the mode.) +with serve("stall"): + collector = load_collector("codex") + rpc_request = collector.rpc_request + collector.rpc_request = lambda proc, request_id, method, params=None, timeout=8: rpc_request(proc, request_id, method, params, min(timeout, 0.5)) + stalled = run_main(collector, "--force") +expect("a stalled Codex app-server advises a retry beside the last limits", + [stalled.get("retryAdvised"), stalled["limitsStale"], stalled["limits"] == good["limits"]], [True, True, True], stalled) + +# --------------------------------------------------------- whose limits + +for workspace, user in (("ws-two", "user-a"), ("ws-one", "user-b")): + login(workspace, user) + other = codex("backend") + expect(f"a home signed in to someone else never shows the last one's limits ({workspace} {user})", + [other["limits"], other["usageStatusText"]], [[], "Codex limits unavailable"], other) +login("ws-one", "user-a") + +# The login auth.json names only answers when nothing overrides it. +token = codex("backend", CODEX_ACCESS_TOKEN="another") +expect("an access token in the environment never shows auth.json's account's limits", token["limits"], [], token) + +# The token signs in the primary home alone: a registered home without a login +# of its own waits for one rather than answering for the token's account. +side = state_home() / "omarchy" / "agents" / "accounts" / "codex" / "side" +side.mkdir(parents=True) +registry("codex", [ + {"id": "main", "label": "Main", "home": "", "primary": True}, + {"id": "side", "label": "Side", "home": str(side), "primary": False}, +], active="main") +before_starts = starts() +tokened = codex("ok", CODEX_ACCESS_TOKEN="another") +expect("CODEX_ACCESS_TOKEN signs in only the primary of several Codex homes", + [[(account["id"], account["usageStatusText"], bool(account["limits"])) for account in tokened["accounts"]], starts() - before_starts], + [[("main", "", True), ("side", "Waiting for auth", False)], 1], tokened) +for path in (state_home() / "omarchy" / "agents" / "accounts").glob("codex.json"): + path.unlink() + +(home() / ".codex" / "config.toml").write_text('cli_auth_credentials_store = "keyring"\n') +keyring = codex("backend") +(home() / ".codex" / "config.toml").unlink() +expect("a keyring login, which can't be told apart, keeps no limits", keyring["limits"], [], keyring) + +# An API-key login has no plan limits, so no app-server is started to ask and +# there is no sign-in to ask for: Codex's own auth_mode, or a key without +# ChatGPT tokens. A ChatGPT auth_mode beside a key still asks Codex. +for shape in ({"auth_mode": "apikey", "OPENAI_API_KEY": "sk-test"}, {"OPENAI_API_KEY": "sk-test"}, {"OPENAI_API_KEY": "sk-test", "tokens": None}, + {"OPENAI_API_KEY": "sk-test", "tokens": {"id_token": "x", "access_token": "y", "refresh_token": "z"}}): + write_json(home() / ".codex" / "auth.json", shape) + before_starts = starts() + keyed = [codex("ok"), codex("ok", "--limits-only")] + expect(f"an API-key Codex login starts no app-server and waits for no auth ({json.dumps(sorted(shape))})", + [[(run["limits"], run["usageStatusText"], run["authHelpText"], run["limitsStale"]) for run in keyed], starts() - before_starts], + [[([], "", "Signed in with an API key, which has no plan limits.", False)] * 2, 0], keyed) +write_json(home() / ".codex" / "auth.json", {"auth_mode": "chatgpt", "OPENAI_API_KEY": "sk-test"}) +before_starts = starts() +chatgpt = codex("ok") +expect("a ChatGPT auth_mode beside an API key still asks the app-server", [bool(chatgpt["limits"]), starts() - before_starts], [True, 1], chatgpt) +login("ws-one", "user-a") + +own = codex("backend") +expect("the file store's own account keeps its limits", [own["limitsStale"], own["limits"] == good["limits"]], [True, True], own) + +# Signed in to someone else while the probe ran: the answer is shown, but kept +# for neither login. It reads 55% where the kept answer read 40%. +os.environ["CODEX_RATE_LIMITS"] = json.dumps({"planType": "pro", "primary": {"usedPercent": 55, "windowDurationMins": 300, "resetsAt": int(time.time()) + 3600}}) +with serve("up"): + collector = load_collector("codex") + fetch = collector.fetch_codex_rpc + + def fetch_across_login_change(home_dir=None, primary=True): + answer = fetch(home_dir, primary) + login("ws-three", "user-c") + return answer + + collector.fetch_codex_rpc = fetch_across_login_change + switched = run_main(collector, "--force") +del os.environ["CODEX_RATE_LIMITS"] +expect("a probe across a login change still reports what it got", [switched["limitsStale"], switched["limits"][0]["percent"]], [False, 0.55], switched) +after = codex("backend") +expect("limits from a probe across a login change aren't kept for the new login", after["limits"], [], after) +login("ws-one", "user-a") +before = codex("backend") +expect("limits from a probe across a login change don't replace the old login's", [before["limits"] == good["limits"], before["limitsFetchedAt"] == good["limitsFetchedAt"]], [True, True], before) + +# A user id that isn't valid text, a token list, or a missing ID token: the +# collector still prints a record, with nothing kept to show. +login("ws-one", "\udced\udca0\udc80") +for shape in ("surrogate user", {"tokens": ["unexpected"]}, {"tokens": {"account_id": "ws-one"}}): + if shape != "surrogate user": + write_json(home() / ".codex" / "auth.json", shape) + odd = codex("backend") + expect(f"an auth.json of another shape still gives a record ({json.dumps(shape)})", odd["usageStatusText"], "Codex limits unavailable", odd) + +# ---------------------------------------------------- a plan without windows + +# A plan the app-server gives no windows for is an answer, reused like any +# other rather than starting an app-server every run. +login("ws-five", "user-e") +bare = codex("ok", CODEX_RATE_LIMITS=json.dumps({"planType": "free"})) +before_starts = starts() +reused = codex("ok", "--limits-only", CODEX_RATE_LIMITS=json.dumps({"planType": "free"})) +expect("an answer without windows stands in for the next probe", + [bare["limits"], bare["tierLabel"], reused["limits"], reused["tierLabel"], reused["limitsStale"], starts() - before_starts], + [[], "Free", [], "Free", False, 0], reused) + +# ------------------------------------------------------------ a spend cap + +# A spend cap is kept like any window: stale past a failed check, and back at +# zero once the reset it named has passed (a cache rewrite, no sleeping). +login("ws-four", "user-d") +spend = {"planType": "business", "individualLimit": {"limit": "6000", "used": "3000", "remainingPercent": 50, "resetsAt": int(time.time()) + 3600}} +spent = codex("ok", CODEX_RATE_LIMITS=json.dumps(spend)) +expect("a Codex spend cap reads as a titled meter with its reset time", + [spent["limitsStale"], [[entry["title"], entry["percent"], bool(entry["resetsAt"])] for entry in spent["limits"]]], + [False, [["Spend limit", 0.5, True]]], spent) +edit_limits_cache("codex", resets_past=True) +kept = codex("backend") +expect("a kept Codex spend cap whose reset passed comes back at zero, stale", + [kept["limitsStale"], kept["limitsFetchedAt"] == spent["limitsFetchedAt"], [[entry["title"], entry["percent"], entry["resetsAt"]] for entry in kept["limits"]]], + [True, True, [["Spend limit", 0.0, ""]]], kept) + +# ------------------------------------------------------------- no CLI + +# A Codex home puts Codex in use, so a missing CLI is then worth saying. +sandbox() +(home() / ".codex" / "sessions").mkdir(parents=True) +path = os.environ["PATH"] +os.environ["PATH"] = str(home() / "empty-bin") +try: + missing = collect("codex", "--force") +finally: + os.environ["PATH"] = path +expect("a Codex home without the CLI says the CLI is missing", + [missing["usageStatusText"], missing["authHelpText"]], ["Codex limits unavailable", "codex not found in PATH"], missing) + +# --------------------------------------------------------------- the answer + +# The app-server's rate-limit answer as fetch_codex_rpc reads it, from a +# stand-in that answers RPC_LIMITS whole and logs every account/read. +sandbox() +codex_home = home() / ".codex" +write_json(codex_home / "auth.json", {}) +stub = home() / "rpc-bin" / "codex" +stub.parent.mkdir() +stub.write_text("#!" + sys.executable + "\n" + ''' +import json, os, sys +from pathlib import Path + +for line in sys.stdin: + request = json.loads(line) + if "id" not in request: + continue + method = request["method"] + if method == "account/read": + with Path(os.environ["RPC_CALLS"]).open("a") as log: + log.write(method + "\\n") + reply = {"account": {"planType": "plus"}} + elif method == "account/rateLimits/read": + reply = json.loads(os.environ["RPC_LIMITS"]) + else: + reply = {} + print(json.dumps({"id": request["id"], "result": reply}), flush=True) +''') +stub.chmod(0o755) +calls = home() / "calls" +os.environ.update(PATH=str(stub.parent) + os.pathsep + path, RPC_CALLS=str(calls)) + + +def window(percent, minutes=300): + return {"usedPercent": percent, "windowDurationMins": minutes} + + +def entry(percent, label="5h window", title=None): + result = {"label": label, "percent": percent / 100, "resetsAt": ""} + if title: + result["title"] = title + return result + + +def reads(name, payload, expected, plan="Pro 200", account_read=False, reset_credits=None): + calls.write_text("") + os.environ["RPC_LIMITS"] = json.dumps(payload) + result = load_collector("codex").fetch_codex_rpc(codex_home) + expect(name, { + "status": result["usageStatusText"], "help": result["authHelpText"], "plan": result["tierLabel"], "limits": result["limits"], + "accountRead": bool(calls.read_text()), "resetCredits": result.get("resetCredits"), + }, {"status": "", "help": "", "plan": plan, "limits": expected, "accountRead": account_read, "resetCredits": reset_credits}, result) + + +legacy = {"planType": "pro", "primary": window(12)} +reads("legacy single-bucket limits keep their display shape", {"rateLimits": legacy}, [entry(12)]) +reads("map-only limits supply both windows and the plan", { + "rateLimitsByLimitId": {"codex": {"planType": "pro", "primary": window(25), "secondary": window(42, 10080)}} +}, [entry(25), entry(42, "Weekly (7-day)")]) +reads("multi-bucket limits supersede the legacy view without duplicates", { + "rateLimits": legacy, "rateLimitsByLimitId": {"codex": {"planType": "pro", "primary": window(36)}} +}, [entry(36)]) +reads("named buckets are titled for their model and window, general Codex limits first", { + "rateLimits": {"planType": "plus", "primary": window(99)}, + "rateLimitsByLimitId": { + "codex_other": {"planType": "plus", "limitName": "Other model", "primary": window(80, 60)}, + "codex": {"planType": "pro", "primary": window(25)}} +}, [entry(25), entry(80, "1h window", "Other model Session")]) +reads("a single named bucket keeps its model label", { + "rateLimitsByLimitId": {"codex_other": {"planType": "pro", "limitName": "Other model", "primary": window(80)}} +}, [entry(80, title="Other model Session")]) +reads("unnamed model buckets are titled with their identifiers", { + "rateLimitsByLimitId": {"codex": legacy, "codex_other": {"primary": window(80)}} +}, [entry(12), entry(80, title="codex_other Session")]) +for unavailable in (None, {}, {"codex": None}, []): + reads("unavailable multi-bucket view falls back to legacy: " + repr(unavailable), { + "rateLimits": legacy, "rateLimitsByLimitId": unavailable + }, [entry(12)]) +reads("invalid map entries do not hide valid buckets", {"rateLimitsByLimitId": {"invalid": None, "codex": legacy}}, [entry(12)]) +reads("legacy plan metadata is retained when map buckets omit it", { + "rateLimits": legacy, "rateLimitsByLimitId": {"codex": {"primary": window(33)}} +}, [entry(33)]) +reads("account metadata remains an optional plan fallback", { + "rateLimitsByLimitId": {"codex": {"primary": window(33)}} +}, [entry(33)], plan="Plus", account_read=True) +reads("malformed legacy metadata does not erase valid map limits", { + "rateLimits": [], "rateLimitsByLimitId": {"codex": {"primary": window(33)}} +}, [entry(33)], plan="Plus", account_read=True) +reads("a weekly model bucket attaches to the weekly window", { + "rateLimitsByLimitId": {"codex": legacy, "codex_spark": {"limitName": "Spark", "secondary": window(30, 10080)}} +}, [entry(12), entry(30, "Weekly (7-day)", "Spark Weekly")]) +reads("an untrusted percentage or reset time costs only that window or field", { + "rateLimits": {"planType": "pro", "primary": {"usedPercent": "lots", "windowDurationMins": 300}, + "secondary": {"usedPercent": 250, "windowDurationMins": 10080, "resetsAt": -5}} +}, [entry(100, "Weekly (7-day)")]) +reads("multi-bucket limits retain earned reset credits", { + "rateLimitsByLimitId": {"codex": legacy}, + "rateLimitResetCredits": {"credits": [{"status": "available", "expiresAt": 1900000000}]} +}, [entry(12)], reset_credits={"available": 1, "nextExpiresAt": "2030-03-17T17:46:40+00:00"}) + +cap = {"limit": "6000", "used": "1500", "remainingPercent": 75, "resetsAt": 1793491200} + + +def spend_entry(share, resets_at="2026-11-01T00:00:00+00:00"): + return {"label": "Spend limit", "percent": share, "resetsAt": resets_at, "title": "Spend limit"} + + +reads("a spend cap follows the windows once, however many views carry it", { + "rateLimits": dict(legacy, individualLimit=cap), + "rateLimitsByLimitId": {"codex": dict(legacy, individualLimit=cap), + "codex_spark": {"limitName": "Spark", "secondary": window(30, 10080), "individualLimit": cap}} +}, [entry(12), entry(30, "Weekly (7-day)", "Spark Weekly"), spend_entry(0.25)]) +reads("a spend cap only the backward-compatible view carries still counts", { + "rateLimits": dict(legacy, individualLimit=cap), "rateLimitsByLimitId": {"codex": legacy} +}, [entry(12), spend_entry(0.25)]) +reads("a spend cap without usable amounts reads its reported percentage", { + "rateLimits": dict(legacy, individualLimit={"limit": True, "used": "NaN", "remainingPercent": 75, "resetsAt": -5}) +}, [entry(12), spend_entry(0.25, "")]) +reads("an untrusted spend cap costs only its own meter", { + "rateLimits": dict(legacy, individualLimit={"limit": "6000", "used": "lots", "remainingPercent": float("nan")}) +}, [entry(12)]) + +# The plan reads as Codex 0.163's /status names it (SubscriptionDisplay::Status), +# not as its wire planType, and one newer than that keeps its wire name. +for wire, shown in (("free", "Free"), ("go", "Go"), ("plus", "Plus"), ("pro", "Pro 200"), ("prolite", "Pro 100"), ("promax", "Pro 500"), + ("team", "Business"), ("self_serve_business_usage_based", "Business"), ("self_serve_business_prolite", "Business Premium"), + ("business", "Enterprise"), ("enterprise_cbp_automation", "Enterprise (Automation)"), ("enterprise_cbp_usage_based", "Enterprise"), + ("ent26", "Enterprise"), ("enterprise", "Enterprise"), ("edu", "Edu"), ("edu_plus", "Edu Plus"), ("edu_pro", "Edu Pro"), + ("ultra_max", "ultra max")): + reads(f"the {wire} plan reads as {shown!r}", {"rateLimits": dict(legacy, planType=wire)}, [entry(12)], plan=shown) +PY diff --git a/test/shell.d/agent-usage-codex-notifications-test.sh b/test/shell.d/agent-usage-codex-notifications-test.sh index ceeecaecd23..ee9c4055a09 100644 --- a/test/shell.d/agent-usage-codex-notifications-test.sh +++ b/test/shell.d/agent-usage-codex-notifications-test.sh @@ -1,22 +1,26 @@ #!/bin/bash +# Codex writes a token_count notification after every request, and repeats +# it whenever its quota or timestamp changes. The Codex collector counts each +# request once from the cumulative counters, within each rollout, without +# losing requests that carry no counters or follow a counter reset. + set -euo pipefail source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" -require_command python3 +agent_sandbox +# A stand-in Codex on PATH, so the collector never reaches for a real one. +# Nothing here is signed in, so it is never started either. +fake_codex "$HOME" +rm "$HOME/.codex/auth.json" +export PATH="$HOME/bin:$PATH" -python3 - "$ROOT/bin/omarchy-agent-usage-codex" <<'PY' -import json -import os -import runpy -import sys -import tempfile -from datetime import datetime, timedelta -from pathlib import Path +python3 - <<'PY' +from agent_usage_fixtures import * -collector = sys.argv[1] -now = datetime.now().astimezone().replace(hour=12, minute=0, second=0, microsecond=0) +noon = datetime.fromisoformat(today_stamp()) def tokens(input_tokens=100, output_tokens=10): @@ -26,55 +30,48 @@ def tokens(input_tokens=100, output_tokens=10): def event(total, last=None, offset=0): - return {"type": "event_msg", "timestamp": (now + timedelta(seconds=offset)).isoformat(), + return {"type": "event_msg", "timestamp": (noon + timedelta(seconds=offset)).isoformat(), "payload": {"type": "token_count", "info": { "total_token_usage": total, "last_token_usage": last or tokens()}, "rate_limits": {"primary": {"used_percent": offset}}}} -def check(name, rollouts, expected_tokens, expected_prompts): - with tempfile.TemporaryDirectory() as temporary: - root = Path(temporary) - os.environ["CODEX_HOME"] = str(root) - for relative_path, events in rollouts.items(): - path = root / relative_path - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text("\n".join(json.dumps(item) for item in events) + "\n") - loaded = runpy.run_path(collector, run_name="notification_test") - loaded["scan_native_codex_sessions"]({}, {}) - stats = loaded["local_stats"]() - actual = sum(sum(bucket.values()) for bucket in stats["modelUsage"].values()) - assert actual == expected_tokens, (name, actual, expected_tokens) - assert stats["totalPrompts"] == expected_prompts, (name, stats) - assert sum(day["messageCount"] for day in stats["recentDays"]) == expected_tokens, (name, stats) - print("ok - " + name) - - -check("native repeated snapshots ignore timestamp and quota changes", { +def counts(name, rollouts, expected_tokens): + sandbox() + for relative_path, events in rollouts.items(): + write_pi_session(home() / ".codex" / relative_path, events) + record = collect("codex", "--force") + counted = sum(sum(bucket.values()) for bucket in record["modelUsage"].values()) + # Requests are no prompts: none here was typed. + expect(name, [counted, record["totalPrompts"], sum(day["messageCount"] for day in record["recentDays"])], + [expected_tokens, 0, expected_tokens], record) + + +counts("repeated snapshots ignore timestamp and quota changes", { "sessions/main.jsonl": [event(tokens()), event(tokens(), offset=1), event(tokens(), offset=2)] -}, 110, 1) +}, 110) -check("native equal request sizes count when cumulative usage advances", { +counts("equal request sizes count when cumulative usage advances", { "sessions/main.jsonl": [event(tokens()), event(tokens(200, 20), offset=1)] -}, 220, 2) +}, 220) -check("native notification deduplication stays within each rollout", { +counts("deduplication stays within each rollout", { "sessions/first.jsonl": [event(tokens())], "archived_sessions/second.jsonl": [event(tokens()), event(tokens(), offset=1)] -}, 220, 2) +}, 220) -check("native requests without cumulative counters remain countable", { +counts("requests without cumulative counters remain countable", { "sessions/main.jsonl": [event(None), event(None, offset=1), event({}), event({}, offset=1)] -}, 440, 4) +}, 440) -check("native counter resets do not discard new requests", { +counts("counter resets do not discard new requests", { "sessions/main.jsonl": [event(tokens()), event(tokens(200, 20)), event(tokens())] -}, 330, 3) +}, 330) -check("native quota-only events and model context do not reset deduplication", { +counts("quota-only events and model context do not reset deduplication", { "sessions/main.jsonl": [event(tokens()), {"type": "event_msg", "payload": {"type": "token_count", "info": None}}, {"type": "turn_context", "payload": {"model": "another-model"}}, event(tokens(), offset=1)] -}, 110, 1) +}, 110) PY diff --git a/test/shell.d/agent-usage-codex-rollouts-test.sh b/test/shell.d/agent-usage-codex-rollouts-test.sh new file mode 100755 index 00000000000..5f1f9973fd2 --- /dev/null +++ b/test/shell.d/agent-usage-codex-rollouts-test.sh @@ -0,0 +1,303 @@ +#!/bin/bash + +# The Codex rollout records every collector billing a custom provider reads +# (codex_rollouts): one cache per Codex home that they all share, and a +# rollout that only grew read on from where the last read stopped, with the +# same totals a whole read gives. The lines are shaped as Codex 0.163 wrote +# them in a real run on a Z.ai model_providers entry. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +import json +import os +import threading +import time +from pathlib import Path + +from agent_usage_fixtures import * +import collectors + +THREAD = "01a1644c-0d61-7072-be1b-fe5b114d8a73" +TURN = "01a1644c-0d83-7963-a2ff-038fc748bbae" +codex = home() / ".codex" +sessions = codex / "sessions" / "2026" / "10" / "22" +sessions.mkdir(parents=True) +(codex / "config.toml").write_text( + 'model = "glm-4.6"\n' + '[model_providers.zai]\nname = "zai"\nbase_url = "https://api.z.ai/api/coding/paas/v4"\n' + '[model_providers.mm]\nname = "mm"\nbase_url = "https://api.minimax.io/v1"\n' +) + + +def stamp(n): + return f"2026-10-22T15:47:{n:02d}.229Z" + + +def meta(provider): + return {"timestamp": stamp(0), "type": "session_meta", "payload": { + "session_id": THREAD, "id": THREAD, "timestamp": stamp(0), "cwd": "/home/u/proj", "originator": "codex_exec", + "cli_version": "0.163.0", "source": "exec", "thread_source": "user", "model_provider": provider, + "base_instructions": "You are Codex.", "history_mode": "paginated"}} + + +def turn_context(n, model): + return {"timestamp": stamp(n), "type": "turn_context", "payload": {"turn_id": TURN, "root_turn_id": TURN, "cwd": "/home/u/proj", + "approval_policy": "never", "sandbox_policy": {"type": "danger-full-access"}, "model": model, + "collaboration_mode": {"mode": "default", "settings": {"model": model}}, "summary": "none"}} + + +def prompt(n, text): + return {"timestamp": stamp(n), "type": "event_msg", "payload": {"type": "item_completed", "thread_id": THREAD, "turn_id": TURN, + "item": {"type": "UserMessage", "id": f"msg{n}", "content": [{"type": "text", "text": text, "text_elements": []}]}}} + + +def usage(input_tokens, cached, written, output): + return {"input_tokens": input_tokens, "cached_input_tokens": cached, "cache_write_input_tokens": written, + "output_tokens": output, "reasoning_output_tokens": 0, "total_tokens": input_tokens + output} + + +def response(n, rid, spent, thread): + return {"timestamp": stamp(n), "type": "token_usage_record", "payload": {"thread_id": THREAD, "turn_id": TURN, "session_id": THREAD, + "root_turn_id": TURN, "response_id": rid, "usage": usage(*spent), "turn_token_usage": usage(*thread), "thread_token_usage": usage(*thread)}} + + +def token_count(n, last, total): + return {"timestamp": stamp(n), "type": "event_msg", "payload": {"type": "token_count", + "info": {"total_token_usage": usage(*total), "last_token_usage": usage(*last), "model_context_window": 258400}}} + + +def noise(n): + return {"timestamp": stamp(n), "type": "response_item", "payload": {"type": "message", "role": "assistant", + "content": [{"type": "output_text", "text": "x" * 4000}]}} + + +def lines(*entries): + return "".join(json.dumps(entry) + "\n" for entry in entries) + + +FIRST = (10234, 71, 31, 51) +SECOND = (11234, 72, 32, 52) +THIRD = (12234, 73, 33, 53) + + +def total(*spent): + return tuple(map(sum, zip(*spent))) + + +# A session's first turn, then a prompt typed while nothing has answered it. +HEAD = lines( + meta("zai"), turn_context(1, "glm-4.6"), noise(2), prompt(3, "zcompone"), + response(4, "resp_0001", FIRST, FIRST), token_count(5, FIRST, FIRST), noise(6), + prompt(7, "and again"), +) +# Codex answers it on another model, writes the earlier response's record a +# second time, a token_count repeating it, and one more response. +TAIL = lines( + turn_context(8, "glm-4.6-air"), response(9, "resp_0002", SECOND, total(FIRST, SECOND)), + response(10, "resp_0002", SECOND, total(FIRST, SECOND)), token_count(11, SECOND, total(FIRST, SECOND)), + noise(12), response(13, "resp_0003", THIRD, total(FIRST, SECOND, THIRD)), +) + +zai = sessions / f"rollout-2026-10-22T11-47-48-{THREAD}.jsonl" +mm = sessions / "rollout-2026-10-22T11-50-48-01a1644e-cbd2-72d3-8d17-8c9af79b8b7d.jsonl" +zai.write_text(HEAD) +mm.write_text(lines(meta("mm"), turn_context(1, "MiniMax-M2"), prompt(2, "hi"), response(3, "resp_m1", FIRST, FIRST))) + + +def tokens(spent): + input_tokens, cached, written, output = spent + return [input_tokens - cached - written, output, cached, written] + + +def days(record): + return {model: totals for models in record["days"].values() for model, totals in models.items()} + + +def read_counting(scan): + """What `scan` returns, and how many rollout lines it parsed.""" + real = json.loads + parsed = [] + + def counting(text, *args, **kwargs): + if isinstance(text, str) and '"payload"' in text: + parsed.append(text) + return real(text, *args, **kwargs) + + json.loads = counting + try: + return scan(), len(parsed) + finally: + json.loads = real + + +def grow(path, text): + with path.open("a") as handle: + handle.write(text) + + +records, complete = collectors.codex_rollouts(collectors.codex_rollout_cache("zai")) +expect("the waiting prompt counts on the model last in effect until something answers it", + days(records[str(zai)]), {"glm-4.6": [*tokens(FIRST), 2]}) + +# --------------------------------------------------------------- resuming + +grow(zai, TAIL) +(records, complete), parsed = read_counting(lambda: collectors.codex_rollouts(collectors.codex_rollout_cache("zai"))) +whole = collectors.read_codex_rollout(zai, zai.stat()) +expect("a rollout that grew counts what a whole read counts, the prompt moving to the model that answered it", + days(records[str(zai)]), days(whole)) +expect("the response written twice and the token_count repeating it count once", + days(records[str(zai)]), {"glm-4.6": [*tokens(FIRST), 1], "glm-4.6-air": [*map(sum, zip(tokens(SECOND), tokens(THIRD))), 1]}) +grown = sum(any(marker in line for marker in collectors._CODEX_MARKERS) for line in TAIL.splitlines()) +check("a rollout that grew is read on from where the last read stopped, not again from its start", + parsed == grown, f"parsed {parsed} lines, the grown part has {grown} worth reading") + +# ---------------------------------------------------- one shared cache + +before = records[str(mm)] +grow(mm, lines(response(4, "resp_m2", SECOND, total(FIRST, SECOND)))) +collectors.scan_codex_routed("zai") +reads = [] +real_reader = collectors.read_codex_rollout + + +def reader(*args, **kwargs): + reads.append(args[0]) + return real_reader(*args, **kwargs) + + +collectors.read_codex_rollout = reader +try: + minimax_rows, _ = collectors.scan_codex_routed("minimax") + collectors.scan_codex_routed("ollama") + collectors.codex_rollouts(collectors.codex_rollout_cache("local")) +finally: + collectors.read_codex_rollout = real_reader +check("once one collector has read a rollout that grew, the others billing Codex providers reuse its read", + not reads, f"read again: {[str(path) for path in reads]}") +expect("the collector a provider's base_url names counts that rollout from the shared read", + sorted((row.model, row.input_tokens, row.output_tokens) for row in minimax_rows), + [("MiniMax-M2", *tokens(total(FIRST, SECOND))[:2])]) +expect("every collector keeps its Codex records in one cache file per Codex home", + sorted(path.name for path in cache_root().glob("*codex*.json")), [collectors.codex_rollout_cache().name]) + +# ------------------------------------------ what a resume can't trust + +# A rollout replaced by a longer copy that differs in what was read, as a +# restore from a backup leaves it, is read whole again. +copy = zai.with_suffix(".new") +copy.write_text(zai.read_text().replace('"output_tokens": 51', '"output_tokens": 91', 1) + lines(noise(14))) +copy.replace(zai) +records, _ = collectors.codex_rollouts(collectors.codex_rollout_cache()) +expect("a rollout replaced by a longer one is read whole again", + days(records[str(zai)]), days(collectors.read_codex_rollout(zai, zai.stat()))) +check("the rewritten count is the one counted", days(records[str(zai)])["glm-4.6"][1] == FIRST[3] + 40, days(records[str(zai)])) + +# A last line still being written counts as a whole read counts it, and +# the read once it is whole counts it once. +partial = sessions / "rollout-2026-10-22T12-00-00-01a16450-0000-7000-8000-000000000000.jsonl" +partial.write_text(lines(meta("zai"), turn_context(1, "glm-4.6"), prompt(2, "p")) + json.dumps(response(3, "resp_p1", FIRST, FIRST))) +records, _ = collectors.codex_rollouts(collectors.codex_rollout_cache()) +expect("a rollout ending in a line with no newline yet counts that line", days(records[str(partial)]), {"glm-4.6": [*tokens(FIRST), 1]}) +grow(partial, "\n" + lines(response(4, "resp_p2", SECOND, total(FIRST, SECOND)))) +records, _ = collectors.codex_rollouts(collectors.codex_rollout_cache()) +expect("once that line is whole and more follow, it still counts once", + days(records[str(partial)]), {"glm-4.6": [*map(sum, zip(tokens(FIRST), tokens(SECOND))), 1]}) + +# The same totals whichever way the collectors got there. +resumed = sorted(collectors.scan_codex_routed("zai")[0]) +expect("Z.ai's Codex rows are the same read on as read whole", resumed, sorted(collectors.scan_codex_routed("zai", reuse=False)[0])) + +# A guardian thread counts under the session that started it, as Codex 0.160 +# writes one: its session_id and parent_thread_id name the thread it checks. +GUARDIAN = "01a1284c-6fa4-7c13-8d29-e2d8b566ae36" +guardian = sessions / f"rollout-2026-10-22T12-05-00-{GUARDIAN}.jsonl" +guardian_meta = meta("zai") +guardian_meta["payload"].update({"id": GUARDIAN, "parent_thread_id": THREAD, "originator": "codex-tui", + "source": {"subagent": {"other": "guardian"}}, "thread_source": "guardian_review"}) +guardian.write_text(lines(guardian_meta, turn_context(1, "codex-auto-review"), response(2, "resp_g1", FIRST, FIRST))) +rows = collectors.scan_codex_routed("zai")[0] +expect("a guardian thread's rows count under the session it checks", + sorted({row.session for row in rows if row.model == "codex-auto-review"}), [THREAD]) +# Every Z.ai rollout here belongs to that one thread, so it is one session. +stats = collectors.Stats() +for row in rows: + stats.add_row(row, row.model) +expect("so Z.ai counts one session, not one per rollout file", stats.snapshot()["totalSessions"], 1) + +# A built-in rollout is left unread, and marked so, until openai_base_url +# sends the built-in provider elsewhere; then the cached stub is read whole, +# session and all, and the full record stays once nothing needs it. +native = sessions / "rollout-2026-10-22T12-10-00-01a16451-0000-7000-8000-000000000000.jsonl" +native.write_text(lines(meta("openai"), turn_context(1, "hf:zai-org/GLM-4.6"), prompt(2, "p"), response(3, "resp_n1", FIRST, FIRST))) +records, _ = collectors.codex_rollouts(collectors.codex_rollout_cache()) +check("a built-in rollout is skipped while nothing sends it elsewhere", records[str(native)].get("skipped") and not records[str(native)]["days"], records[str(native)]) +records, _ = collectors.codex_rollouts(collectors.codex_rollout_cache(), native=True) +expect("once something does, the skipped built-in rollout is read whole", days(records[str(native)]), {"hf:zai-org/GLM-4.6": [*tokens(FIRST), 1]}) +expect("under the thread that started it", records[str(native)]["session"], THREAD) +records, _ = collectors.codex_rollouts(collectors.codex_rollout_cache()) +expect("a built-in rollout read whole stays read", days(records[str(native)]), {"hf:zai-org/GLM-4.6": [*tokens(FIRST), 1]}) + +# ------------------------------------------------- one forced refresh + +# A refresh forces every collector at once. The ones that wait for the lock +# take what the first forced read wrote meanwhile, built-in rollouts read +# whole included, so each rollout is parsed once per refresh, not once per +# collector, while a forced read after that one parses them all again. +files = sorted(str(path) for path in sessions.glob("*.jsonl")) +lock = collectors.codex_rollout_cache().with_name(collectors.codex_rollout_cache().name + ".lock") +reads = [] +forced = [] +collectors.read_codex_rollout = reader +try: + # Every one of them is waiting for the lock before any gets it. + with collectors.exclusive_lock(lock): + threads = [threading.Thread(target=lambda: forced.append(collectors.codex_rollouts(collectors.codex_rollout_cache(), reuse=False, native=True))) for _ in range(3)] + for thread in threads: + thread.start() + time.sleep(0.5) + for thread in threads: + thread.join() + once = sorted(str(path) for path in reads) + reads.clear() + collectors.codex_rollouts(collectors.codex_rollout_cache(), reuse=False, native=True) + again = sorted(str(path) for path in reads) +finally: + collectors.read_codex_rollout = real_reader +expect("collectors forced together parse each rollout once between them", once, files) +check("and each gets every rollout, read whole", len(forced) == 3 and all(sorted(records) == files and not any(r.get("skipped") for r in records.values()) for records, _ in forced), forced) +expect("a forced read after theirs parses every rollout again", again, files) + +# The collectors of one refresh start together but get to the rollouts when +# each does, so one forced read can be done before the next begins. Told when +# the refresh began, as omarchy-agent-usage-update tells them, the later ones +# take what the earlier one read: only the built-in rollout it skipped is +# read again, for the one that needs it. A refresh stamp from long ago forces +# as --force alone does. +cache = collectors.codex_rollout_cache() + + +def forced_reads(native=False): + reads.clear() + collectors.codex_rollouts(cache, reuse=False, native=native) + return sorted(str(path) for path in reads) + + +collectors.read_codex_rollout = reader +try: + os.environ["OMARCHY_AGENT_USAGE_REFRESH_AT"] = str(time.time()) + first, second, built_in = forced_reads(), forced_reads(), forced_reads(native=True) + os.environ["OMARCHY_AGENT_USAGE_REFRESH_AT"] = str(time.time() - collectors.CODEX_REFRESH_SECONDS - 1) + stale = forced_reads() +finally: + collectors.read_codex_rollout = real_reader + del os.environ["OMARCHY_AGENT_USAGE_REFRESH_AT"] +expect("one refresh's forced reads, one after another, parse each rollout once between them", [first, second, built_in], [files, [], [str(native)]]) +expect("a refresh stamp from long ago parses every rollout again", stale, files) +PY diff --git a/test/shell.d/agent-usage-codex-scanner-test.sh b/test/shell.d/agent-usage-codex-scanner-test.sh index f05a2eebfeb..2ba49f41621 100644 --- a/test/shell.d/agent-usage-codex-scanner-test.sh +++ b/test/shell.d/agent-usage-codex-scanner-test.sh @@ -1,1271 +1,672 @@ #!/bin/bash -source "$(dirname "$0")/base-test.sh" +# The Codex collector's own sources: native Codex rollouts and their per-file +# cache, its filter over pi and opencode sessions, and how it finds, starts and +# reads the Codex app-server for its limits. How the shared pi and opencode +# readers work, and the unused-machine record, are held centrally by the shared +# and partition tests, and the scan cache by the contract check. The limits +# cache has a test of its own. -require_command jq -require_command python3 -require_command git +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox require_command rg +require_command jq -# Every fixture home lives under one scratch directory, cleaned up at exit. -SCRATCH=$(mktemp -d) -trap 'rm -rf "$SCRATCH"' EXIT - -# Run a scratch copy of the collector with its 8s RPC timeout shortened to -# 0.5s; the stall scenarios only need it to give up, not wait the full bound. -collector="$SCRATCH/omarchy-agent-usage-codex" -sed 's/timeout=8/timeout=0.5/g' "$ROOT/bin/omarchy-agent-usage-codex" >"$collector" -chmod +x "$collector" -TEST_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") - -# A fixture home signed in to Codex, with an empty bin/ for its CLI. -signed_in_home() { - local home - home=$(mktemp -d "$SCRATCH/home.XXXXXX") - mkdir -p "$home/bin" "$home/.codex" - touch "$home/.codex/auth.json" - printf '%s\n' "$home" -} +# The stand-in app-server every signed-in home below copies. +fake_codex "$HOME" -mkdir -p "$TEST_HOME/.codex/sessions/$(date +%Y/%m/%d)" "$TEST_HOME/bin" -touch "$TEST_HOME/.codex/auth.json" +python3 - <<'PY' +import shutil +import tempfile -cat >"$TEST_HOME/bin/codex" <<'EOF' -#!/bin/bash +from agent_usage_fixtures import * + +FAKE_CODEX = home() / "bin" / "codex" +FAKE_AUTH = home() / ".codex" / "auth.json" +BASE_PATH = os.environ["PATH"] +# Only the interpreter and file tools, so no codex or mise on this machine can +# answer for the stubs. +TOOLS = Path(os.environ["SANDBOX"]) / "tools" +STAMP = today_stamp() + + +def new_home(codex=FAKE_CODEX, auth=True, path=None): + """A sandbox home with `codex` (a stub's path or its bash body) in its bin/, + signed in unless `auth` is False ("empty" leaves an empty auth.json).""" + sandbox() + bin_dir = home() / "bin" + bin_dir.mkdir() + if isinstance(codex, Path): + shutil.copy(codex, bin_dir / "codex") + elif codex is not None: + stub(bin_dir / "codex", codex) + (home() / ".codex").mkdir() + if auth == "empty": + (home() / ".codex" / "auth.json").touch() + elif auth: + shutil.copy(FAKE_AUTH, home() / ".codex" / "auth.json") + os.environ["PATH"] = f"{bin_dir}:{path or BASE_PATH}" + return home() + + +def stub(path, body): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(body if body.startswith("#!") else "#!/bin/bash\n" + body) + path.chmod(0o755) + return path + + +def codex(*args, env=None, quick=False): + """The collector's record, run with `env` added. `quick` cuts the 8-second + app-server timeout to half a second, for the cases that only need it to + give up.""" + env = {key: str(value) for key, value in (env or {}).items()} + saved = {key: os.environ.get(key) for key in env} + os.environ.update(env) + try: + module = load_collector("codex") + if quick: + real = module.rpc_request + module.rpc_request = lambda proc, request_id, method, params=None, timeout=8: real(proc, request_id, method, params, min(timeout, 0.5)) + with contextlib.redirect_stderr(io.StringIO()): + return run_main(module, *(args or ("--force",))) + finally: + for key, value in saved.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def sessions(): + return home() / ".codex" / "sessions" / datetime.now().strftime("%Y/%m/%d") + + +def rollout(name, *lines, stamp=STAMP, mode="w"): + path = sessions() / name + path.parent.mkdir(parents=True, exist_ok=True) + with path.open(mode) as handle: + for line in lines: + if isinstance(line, str): + handle.write(line + "\n") + else: + handle.write(json.dumps({"timestamp": stamp, **line}) + "\n") + return path + + +def context(model): + return {"type": "turn_context", "payload": {"model": model}} + + +def meta(**payload): + return {"type": "session_meta", "payload": payload} + + +def tokens(input, output, cached=0, total=None, kind="event_msg", **extra): + info = {"last_token_usage": {"input_tokens": input, "cached_input_tokens": cached, "output_tokens": output, **extra}} + if total is not None: + info["total_token_usage"] = total + return {"type": kind, "payload": {"type": "token_count", "info": info}} + + +def typed(text="hello", **extra): + # A prompt as Codex 0.163 writes it in a paginated rollout. + item = {"type": "UserMessage", "id": "u1", "content": [{"type": "text", "text": text, "text_elements": []}]} + return {"type": "event_msg", **extra, "payload": {"type": "item_completed", "thread_id": "t1", "turn_id": "turn1", "item": item}} + + +# ---------------------------------------------------------- native rollouts + +# A request counts once: a quota notification repeats the last request's usage +# under an unchanged running total. Cached input and reasoning are parts of +# input and output, not more of them, and the two requests one typed prompt +# took are still one prompt. +new_home() +first = {"input_tokens": 100, "cached_input_tokens": 60, "output_tokens": 20, "reasoning_output_tokens": 5, "total_tokens": 120} +second = {"input_tokens": 180, "cached_input_tokens": 110, "output_tokens": 30, "reasoning_output_tokens": 8, "total_tokens": 210} +rollout("rollout.jsonl", context("gpt-test"), typed(), + tokens(100, 20, 60, first, reasoning_output_tokens=5, total_tokens=120), + tokens(80, 10, 50, second, reasoning_output_tokens=3, total_tokens=90), + {**tokens(80, 10, 50, second, reasoning_output_tokens=3, total_tokens=90), "rate_limits": {"primary": {"used_percent": 10}}}) +record = codex("--force", env={"CODEX_ARGS_FILE": home() / "codex-args"}) +expect("Codex counts each request once, without a quota notification or cache and reasoning twice, and a prompt once", { + "today": record["todayTotalTokens"], "prompts": record["todayPrompts"], "usage": record["modelUsage"]["gpt-test"], "help": record["authHelpText"], +}, { + "today": 210, "prompts": 1, "usage": {"inputTokens": 70, "outputTokens": 30, "cacheReadInputTokens": 110, "cacheCreationInputTokens": 0}, "help": "", +}, record) +# Separate arguments, not "-a on-request" as one, which is no policy Codex +# accepts. +expect("Codex starts its app-server read-only with a supported approval policy", + (home() / "codex-args").read_text().split("\0"), ["-s", "read-only", "-a", "on-request", "app-server"]) + +# Codex 0.163 writes a token_usage_record for every response it is billed for. +# A remote compaction is one only a record names: the token_count after it +# spends nothing and leaves the running total alone, and a resumed turn may +# write a record and no token_count at all. Each record counts once, a line +# repeated under the same id and total being the same response, while a mock +# backend reusing an id in a later turn is a new one; the token_count lines +# after the first record repeat them, and one before it, from an older Codex, +# still counts. +new_home() + + +def usage(input, output, cached=0, written=0): + return {"input_tokens": input, "cached_input_tokens": cached, "cache_write_input_tokens": written, "output_tokens": output, + "reasoning_output_tokens": 0, "total_tokens": input + output} + + +def usage_record(turn, response, spent, thread): + return {"type": "token_usage_record", "payload": {"thread_id": "t1", "turn_id": turn, "session_id": "t1", "root_turn_id": turn, + "response_id": response, "usage": spent, "turn_token_usage": spent, "thread_token_usage": thread}} + + +def snapshot(last, total): + return {"type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": total, "last_token_usage": last, "model_context_window": 258400}}} + + +first, compaction, reply, resumed = usage(31234, 51, 71, 31), usage(61234, 51, 71, 31), usage(62234, 52, 72, 32), usage(1000, 10) +after_reply = usage(93468, 103, 143, 63) +rollout("rollout-compacted.jsonl", meta(model_provider="openai", source="exec", thread_source="user", id="t1", session_id="t1", cli_version="0.163.0"), + context("gpt-test"), typed("first"), usage_record("turn1", "resp_mock_0001", first, first), snapshot(first, first), + usage_record("turn2", "resp_mock_0001", compaction, usage(92468, 102, 142, 62)), + {"type": "compacted", "payload": {"message": "", "replacement_history": []}}, context("gpt-test"), + snapshot(usage(0, 0), first), typed("second"), + usage_record("turn2", "resp_mock_0002", reply, usage(154702, 154, 214, 94)), snapshot(reply, after_reply), + usage_record("turn2", "resp_mock_0002", reply, usage(154702, 154, 214, 94)), + typed("third"), usage_record("turn3", "resp_mock_0001", resumed, usage(155702, 164, 214, 94))) +rollout("rollout-upgraded.jsonl", meta(model_provider="openai", source="cli"), context("gpt-older"), typed(), + snapshot(usage(100, 10), usage(100, 10)), typed(), usage_record("turn2", "resp_2", usage(200, 20), usage(300, 30)), + snapshot(usage(200, 20), usage(300, 30))) +record = codex() +expect("Codex counts every response a token_usage_record names, a compaction's and a resumed turn's included, each once", { + "today": record["todayTotalTokens"], "byModel": record["todayTokensByModel"], "prompts": record["todayPrompts"], "usage": record["modelUsage"]["gpt-test"], +}, { + "today": 31285 + 61285 + 62286 + 1010 + 330, "byModel": {"gpt-test": 155866, "gpt-older": 330}, "prompts": 5, + "usage": {"inputTokens": 155702 - 214 - 94, "outputTokens": 164, "cacheReadInputTokens": 214, "cacheCreationInputTokens": 94}, +}, record) -if [[ -n ${CODEX_ARGS_FILE:-} ]]; then - printf '%s\0' "$@" >"$CODEX_ARGS_FILE" -fi +# Codex CLI can front any OpenAI-compatible backend (`--oss`, or a custom +# model_provider), whose rollouts land beside the subscription's but spend a +# local box or a third party. A rollout from before Codex recorded the +# provider still counts, and a fork's own session_meta, first, decides. +new_home() +for name, provider, model, input, output in (("native", "openai", "gpt-native", 100, 10), ("ollama", "ollama", "qwen3-coder:30b", 9000, 900), + ("openrouter", "openrouter", "claude-test", 5000, 500)): + rollout(f"rollout-{name}.jsonl", meta(model_provider=provider), context(model), tokens(input, output, cache_write_input_tokens=0, reasoning_output_tokens=0)) +rollout("rollout-legacy.jsonl", meta(id="legacy"), context("gpt-legacy"), tokens(40, 0)) +rollout("rollout-fork.jsonl", meta(model_provider="openai"), meta(model_provider="ollama"), context("gpt-fork"), tokens(3, 0)) +record = codex() +expect("Codex counts only subscription-backed rollouts, legacy ones included", + [record["todayTotalTokens"], sorted(record["modelUsage"]), record["todaySessions"]], [153, ["gpt-fork", "gpt-legacy", "gpt-native"], 3], record) + +# Lines that aren't usage, or aren't JSON, are skipped, and turn_context +# switches the model the following turns count under. +new_home() +rollout("session.jsonl", meta(id="session-123"), "this is plain non-json text that should be skipped by the prefilter", + context("gpt-4o"), {"type": "user_message", "payload": {"content": "Write a sorting algorithm"}}, + {"type": "agent_reasoning", "payload": {"content": "Thinking about quicksort..."}}, + tokens(50, 15, 20, kind="response_item", total_tokens=65), {"type": "agent_message", "payload": {"content": "Here is quicksort..."}}, + context("o3-mini"), tokens(30, 5, 10, total_tokens=35)) +record = codex() +expect("Codex skips lines that aren't usage and tracks the model across turn_context", { + "today": record["todayTotalTokens"], "byModel": record["todayTokensByModel"], "4o": record["modelUsage"]["gpt-4o"], "o3": record["modelUsage"]["o3-mini"], +}, { + "today": 100, "byModel": {"gpt-4o": 65, "o3-mini": 35}, + "4o": {"inputTokens": 30, "outputTokens": 15, "cacheReadInputTokens": 20, "cacheCreationInputTokens": 0}, + "o3": {"inputTokens": 20, "outputTokens": 5, "cacheReadInputTokens": 10, "cacheCreationInputTokens": 0}, +}, record) + +# A line of the wrong shape, or a payload, info or usage that isn't an +# object, loses that line alone: the turns after it still count. +new_home() +rollout("shaped.jsonl", '[{"type": "session_meta"}]', {"type": "session_meta", "payload": [1]}, + {"type": "turn_context", "payload": ["gpt-x"]}, context("gpt-test"), typed(), tokens(10, 1), + {"type": "event_msg", "payload": {"type": "token_count", "info": [1]}}, + {"type": "event_msg", "payload": {"type": "item_completed", "item": [1]}}, + {"type": "event_msg", "payload": {"type": "token_count", "info": {"last_token_usage": [1]}}}, typed(), tokens(20, 2)) +record = codex() +expect("Codex skips a wrong-shaped rollout line and keeps reading the file", + [record["todayTotalTokens"], record["todayPrompts"]], [33, 2], record) + +# Only a line with no time at all takes the file's; a time of 0, false, or "" +# is one that can't be trusted, so it counts nowhere. +new_home() +rollout("falsy.jsonl", context("gpt-test"), *(json.dumps({"timestamp": falsy, **tokens(100, 0)}) for falsy in (0, False, "")), + *(json.dumps({"timestamp": falsy, **typed()}) for falsy in (0, False, "")), typed(), tokens(10, 1)) +record = codex() +expect("Codex counts a rollout line with a time of 0, false, or empty nowhere, not on the file's day", + [record["todayTotalTokens"], record["todayPrompts"]], [11, 1], record) + +# Prompts are the ones typed, as Codex 0.163 records them: one prompt whose +# tool call took two requests is one, as is a legacy rollout's user_message +# event; the context Codex adds itself, and a subagent's or guardian's +# messages, are none, though the tokens they spent count. Those threads start +# no session either, counting under the one that started them: its root in +# session_id, as Codex 0.163 writes it, or else its parent_thread_id. A thread +# of Codex's own that names no other, such as a memory consolidation, counts +# under none. +new_home() +rollout("rollout-paginated.jsonl", meta(model_provider="openai", source="exec", thread_source="user", id="t-exec", session_id="t-exec"), + {"type": "event_msg", "payload": {"type": "task_started", "turn_id": "turn1"}}, context("gpt-test"), + {"type": "response_item", "payload": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": ""}]}}, + typed("run ls"), tokens(10, 1, total={"input_tokens": 10}), + {"type": "event_msg", "payload": {"type": "item_completed", "item": {"type": "CommandExecution", "id": "c1"}}}, + tokens(20, 2, total={"input_tokens": 30}), {"type": "event_msg", "payload": {"type": "task_complete", "turn_id": "turn1"}}) +rollout("rollout-legacy.jsonl", meta(model_provider="openai", source="cli", id="t-cli"), context("gpt-test"), + {"type": "event_msg", "payload": {"type": "user_message", "message": "", "kind": "user_instructions"}}, + {"type": "event_msg", "payload": {"type": "user_message", "message": "fix it", "kind": "plain"}}, + {"type": "event_msg", "payload": {"type": "user_message", "message": "and test it"}}, tokens(40, 4)) +rollout("rollout-review.jsonl", meta(model_provider="openai", source={"subagent": "review"}, thread_source="subagent", id="t-review", session_id="t-exec", + parent_thread_id="t-exec"), context("gpt-test"), typed(), tokens(100, 10)) +rollout("rollout-memory.jsonl", meta(model_provider="openai", source={"internal": "memory_consolidation"}, thread_source="memory_consolidation", id="t-memory", + session_id="t-memory"), context("gpt-test"), tokens(300, 30)) +rollout("rollout-guardian.jsonl", meta(model_provider="openai", source={"subagent": {"other": "guardian"}}, thread_source="guardian_review", id="t-guardian", + parent_thread_id="t-cli"), + context("gpt-test"), typed(), tokens(200, 20)) +record = codex() +expect("Codex counts the prompts typed, not the requests answering them or a subagent's, and a subagent's thread as no session", + [record["todayTotalTokens"], record["todayPrompts"], record["todaySessions"], record["totalSessions"]], [737, 3, 2, 2], record) + +# Session history only grows, so a refresh reads the files that changed and +# replays its per-file totals for the rest. pi sessions stay out of that +# cache: a fork repeats its parent's messages, so they are only deduplicated +# across the whole tree at once. +new_home() +early = today_stamp(9) +native = rollout("rollout.jsonl", context("gpt-test"), tokens(10, 5), stamp=early) +pi = write_pi_session(home() / ".pi" / "agent" / "sessions" / "session.jsonl", [pi_reply("openai-codex", 0, id="m1", model="gpt-test", timestamp=early, usage={"input": 4, "output": 1})]) + + +def file_cache(): + return json.loads(next(cache_root().glob("codex-rollouts-*.json")).read_text()) + + +def incremental(*args): + # The aggregate scan cache would answer on its own; the rescan beneath it + # is what this covers. + for path in scan_cache_files("codex"): + age_scan_cache(path, 7200) + return codex(*(args or ("--limits-only",))) + + +record = codex("--limits-only") +cached = file_cache()["files"] +expect("Codex counts both sources and records the native file, not pi's, per file", + [record["todayTotalTokens"], len(cached.get(str(native), {}).get("days", {})), str(pi) in cached], [20, 1, False], record) + +if os.geteuid() != 0: + # Unreadable, but unchanged in mtime and size: the same totals can only + # have been replayed. + native.chmod(0) + replayed = incremental() + forced = incremental("--force") + native.chmod(0o644) + expect("Codex replays an unchanged rollout, and --force rereads the history", + [replayed["todayTotalTokens"], forced["todayTotalTokens"]], [20, 5]) + # The forced run couldn't read the rollout, so cache it again first. + incremental() + +rollout("rollout.jsonl", tokens(30, 2, 10), stamp=early, mode="a") +record = incremental() +expect("Codex rereads a rollout that grew", record["todayTotalTokens"], 52, record) + +# A long session's rollout that grew is read on from where the last read +# stopped, through the reader every collector shares, not again from the +# start: a request rewritten in the middle of what was read, the same size, +# still counts as it was read. +padding = {"type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "x" * 6000}]}} +grown = rollout("rollout-grown.jsonl", meta(model_provider="openai", source="exec", id="t-grown"), context("gpt-grown"), typed(), tokens(10, 5), + padding, padding, tokens(100, 0), padding, padding, tokens(1, 0), stamp=early) +incremental() +grown.write_text(grown.read_text().replace('"input_tokens": 100,', '"input_tokens": 900,')) +rollout("rollout-grown.jsonl", typed(), tokens(30, 2), stamp=early, mode="a") +record = incremental() +expect("Codex reads a grown rollout on from where it stopped", record["todayTokensByModel"].get("gpt-grown"), 148, record) + +# Per-file totals are bucketed by local day, so a new timezone reads them again. +zone_stamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT02:00:00Z") +zone = home() / ".codex" / "sessions" / "zone.jsonl" +zone.write_text(json.dumps({"timestamp": zone_stamp, **tokens(3, 1)}) + "\n") +saved_tz = os.environ.get("TZ") +try: + os.environ["TZ"] = "UTC" + time.tzset() + incremental() + os.environ["TZ"] = "America/Los_Angeles" + time.tzset() + incremental() + expected_day = datetime.fromisoformat(zone_stamp.replace("Z", "+00:00")).astimezone().date().isoformat() +finally: + if saved_tz is None: + os.environ.pop("TZ", None) + else: + os.environ["TZ"] = saved_tz + time.tzset() +expect("Codex rereads per-file totals after a timezone change", list(file_cache()["files"][str(zone)]["days"]), [expected_day]) + +# A line that stops the read keeps the usage read before it. +(home() / ".codex" / "sessions" / "broken.jsonl").write_text( + json.dumps({"timestamp": early, **context("gpt-broken")}) + "\n" + json.dumps({"timestamp": early, **tokens(6, 0)}) + "\nnull\n") +record = incremental() +expect("Codex keeps the usage read before a line that stops the read", record["modelUsage"].get("gpt-broken", {}).get("inputTokens"), 6, record) + +# A scan that couldn't read a rollout serves its own run, but isn't cached as +# the whole story: the next run reads it again. +if os.geteuid() != 0: + new_home() + unreadable = rollout("unreadable.jsonl", context("gpt-test"), tokens(7, 0)) + unreadable.chmod(0) + partial = codex("--force") + unreadable.chmod(0o644) + record = codex("--limits-only") + expect("Codex never caches a scan that couldn't read a rollout", + [partial["todayTotalTokens"], record["todayTotalTokens"]], [0, 7], record) + +# ------------------------------------------------------- pi and opencode + +# pi and omp replies count on the openai-codex provider or the Responses API +# pi drives for it, written compact as pi writes them, and no other provider. +# Two messages sharing an id at different times are two messages, and a fork +# naming its parent through a symlinked spelling of pi's tree still repeats it. +new_home() +root = home() / ".pi" / "agent" / "sessions" / "project" +root.mkdir(parents=True) +(home() / "pi-alias").symlink_to(home() / ".pi" / "agent") + + +def compact(id, provider, input, timestamp=STAMP, api=None): + message = {"role": "assistant", "provider": provider, "model": "gpt-pi", "usage": {"input": input, "output": 0}} + if api: + message["api"] = api + return json.dumps({"type": "message", "id": id, "timestamp": timestamp, "message": message}, separators=(",", ":")) + + +write_pi_session(root / "pi.jsonl", [ + compact("deadbeef", "openai-codex", 10), + compact("deadbeef", "openai-codex", 8, timestamp=STAMP.replace(":00:00", ":00:01")), + compact("api", "openai", 20, api="openai-codex-responses"), + compact("other", "anthropic", 999), +]) +write_pi_session(root / "fork.jsonl", [compact("deadbeef", "openai-codex", 10)], parent=home() / "pi-alias" / "sessions" / "project" / "pi.jsonl") +opencode_db(message=[("c1", "ses_1", json.dumps(json.loads(oc_v1("openai", 5, model="gpt-5.2-codex")), separators=(",", ":")))]) +record = codex() +expect("Codex counts compact pi replies on its provider or API, an id reused at another time, a fork once, and compact opencode rows", + [record["todayTotalTokens"], record["todayPrompts"], record["totalSessions"]], [43, 4, 2], record) + +# ------------------------------------------------------------- scan cache + +new_home() +opencode_db(message=[("c1", "ses_1", oc_v1("openai", 5, model="gpt-5.2-codex"))]) +codex() +# A scan cached before native notifications were deduplicated must not bring +# its inflated counts back, even for --limits-only. +edit_scan_cache("codex", lambda data: data.update(schemaVersion=1, stats={**data["stats"], "todayTotalTokens": 999})) +record = codex("--limits-only") +expect("Codex rescans a scan cached before notifications were deduplicated", + [record["todayTotalTokens"], json.loads(scan_cache_files("codex")[0].read_text())["schemaVersion"]], [5, 2], record) +# A cache file of the wrong shape is a miss, written again. +for path in scan_cache_files("codex"): + path.write_text("[]") +record = codex("--limits-only") +expect("Codex rescans past a scan cache of the wrong shape, and writes it again", + [record["todayTotalTokens"], json.loads(scan_cache_files("codex")[0].read_text()).get("schemaVersion")], [5, 2], record) +# A cache that can't be written still leaves a record. +(home() / "not-a-dir").write_text("") +record = codex("--limits-only", env={"XDG_CACHE_HOME": home() / "not-a-dir"}) +expect("Codex still prints its record when the cache can't be written", record["todayTotalTokens"], 5, record) + +# ----------------------------------------------------------- the app-server + +# Rate limits with optional free resets, and an account/read that can hang +# for good, as Codex 0.158's does. +ANSWERS = r''' while read -r request; do id=$(jq -r '.id // empty' <<<"$request") - method=$(jq -r '.method // empty' <<<"$request") - - case "$method" in - initialize) - jq -cn --argjson id "$id" '{id: $id, result: {}}' - ;; + case "$(jq -r '.method // empty' <<<"$request")" in + initialize) jq -cn --argjson id "$id" '{id: $id, result: {}}' ;; account/read) - # Codex 0.158 can leave this one unanswered for good. [[ -n ${CODEX_ACCOUNT_READ_HANGS:-} ]] || - jq -cn --argjson id "$id" '{id: $id, result: {account: {}}}' - ;; + jq -cn --argjson id "$id" --argjson result "${CODEX_ACCOUNT_RESULT:-null}" '{id: $id, result: ($result // {account: {}})}' ;; account/rateLimits/read) - if [[ -n ${CODEX_LIMITS_ERROR:-} ]]; then - jq -cn --argjson id "$id" --arg message "$CODEX_LIMITS_ERROR" '{id: $id, error: {code: -32600, message: $message}}' - continue - fi - jq -cn --argjson id "$id" --argjson limits "${CODEX_RATE_LIMITS:-{\}}" --argjson credits "${CODEX_RESET_CREDITS:-null}" \ - '{id: $id, result: {rateLimits: $limits, rateLimitResetCredits: $credits}}' - ;; + jq -cn --argjson id "$id" --argjson limits "$CODEX_RATE_LIMITS" --argjson credits "${CODEX_RESET_CREDITS:-null}" \ + '{id: $id, result: {rateLimits: $limits, rateLimitResetCredits: $credits}}' ;; esac done -EOF -chmod +x "$TEST_HOME/bin/codex" - -timestamp="$(date +%Y-%m-%d)T12:00:00Z" -collision_timestamp="$(date +%Y-%m-%d)T12:00:01Z" -session="$TEST_HOME/.codex/sessions/$(date +%Y/%m/%d)/rollout.jsonl" -cat >"$session" <"$PI_HOME/bin/rg" <"$PI_HOME/.pi/agent/sessions/project/pi.jsonl" <"$PI_HOME/.pi/agent/sessions/project/pi-fork.jsonl" <"$PI_HOME/.pi/agent/sessions/project/pi-id-collision.jsonl" <"$PI_HOME/.omp/agent/sessions/project/omp.jsonl" <` moves the whole agent tree under profiles//, so a -# subscription spent entirely through a profile leaves the default root empty. -cat >"$PI_HOME/.omp/profiles/codex/agent/sessions/project/omp-profile.jsonl" <"$GIT_HOME/.pi/agent/sessions/project/pi.jsonl" <"$GIT_HOME/.gitignore" - -result=$(HOME="$GIT_HOME" CODEX_HOME="$GIT_HOME/.codex" XDG_DATA_HOME="$GIT_HOME/.local/share" \ - PATH="$GIT_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "8" ]] || - fail "Codex collector counts pi sessions when HOME is a git checkout" "$result" -pass "Codex collector counts pi sessions when HOME is a git checkout" - -# A subscription burned entirely through opencode has no native session files; -# usage must come from opencode's message database, filtered to OpenAI. -OPENCODE_HOME=$(signed_in_home) -cp "$TEST_HOME/bin/codex" "$OPENCODE_HOME/bin/codex" - -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, role="assistant", input=0, output=0, reasoning=0, read=0, write=0): - return (id, "ses_1", now_ms, now_ms, json.dumps({ - "role": role, - "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("msg_1", "openai", "gpt-5.2-codex", input=80, output=40, reasoning=5, read=30), - message("msg_2", "anthropic", "claude-opus-5", input=999, output=999), - message("msg_3", "openai", "gpt-5.2-codex", role="user"), - message("msg_4", "openai-local", "gpt-5.2-codex", input=999, output=999), -]) -conn.execute("INSERT INTO message VALUES ('msg_5', 'ses_1', ?, ?, '[\"not\",\"an\",\"object\"]')", (now_ms, now_ms)) -conn.commit() -conn.close() -PY - -result=$(HOME="$OPENCODE_HOME" CODEX_HOME="$OPENCODE_HOME/.codex" XDG_DATA_HOME="$OPENCODE_HOME/.local/share" \ - PATH="$OPENCODE_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "155" ]] || - fail "Codex collector counts OpenAI usage, reasoning included, from opencode sessions" "$result" -pass "Codex collector counts OpenAI usage, reasoning included, from opencode sessions" - -[[ $(jq -c '.modelUsage' <<<"$result") == '{"gpt-5.2-codex":{"inputTokens":80,"outputTokens":45,"cacheReadInputTokens":30,"cacheCreationInputTokens":0}}' ]] || - fail "Codex collector ignores prefix-colliding providers, user messages, and malformed rows" "$result" -pass "Codex collector ignores prefix-colliding providers, user messages, and malformed rows" - -# A warm cache makes --limits-only cheap: local stats come from the last scan -# instead of another walk over the opencode database, and --force bypasses it. -CACHE_HOME=$(signed_in_home) -cp "$TEST_HOME/bin/codex" "$CACHE_HOME/bin/codex" - -python3 - "$CACHE_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, role="assistant", input=0, output=0, reasoning=0, read=0, write=0): - return (id, "ses_1", now_ms, now_ms, json.dumps({ - "role": role, - "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("c_1", "openai", "gpt-5.2-codex", input=5), -]) -conn.commit() -conn.close() -PY - -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "5" ]] || - fail "Codex collector writes a fresh local-stats cache on first scan" "$result" -cache_file=$(ls "$CACHE_HOME/.cache/omarchy/agent-usage/"/codex-scan-*.json 2>/dev/null | head -n 1) -[[ -n $cache_file && -s $cache_file ]] || - fail "Codex collector leaves a cache file behind" "$result" -[[ $(stat -c %a "$cache_file") == "644" ]] || - fail "Codex collector keeps cache files readable" "$result" -[[ $(jq -r '.schemaVersion' "$cache_file") == "2" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "5" ]] || - fail "Codex collector writes a versioned cache envelope" "$result" -pass "Codex collector writes a local-stats cache on first scan" - -# A still-fresh cache from before native notification deduplication must not -# restore inflated counts, even when only quota limits were requested. -jq '.schemaVersion = 1 | .stats.todayTotalTokens = 999' "$cache_file" >"$CACHE_HOME/old-cache.json" -mv "$CACHE_HOME/old-cache.json" "$cache_file" -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector" --limits-only) -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "5" && $(jq -r '.schemaVersion' "$cache_file") == "2" ]] || - fail "Codex collector invalidates pre-deduplication cached totals" "$result" -pass "Codex collector invalidates pre-deduplication cached totals" - -# A corrupt-but-parseable cache (wrong shape) is a cache miss: rescan and -# rewrite instead of emitting a garbage record. -printf '[]' >"$cache_file" -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "5" ]] || - fail "Codex collector recovers from a corrupt cache file" "$result" -[[ $(jq -r '.schemaVersion' "$cache_file") == "2" ]] || - fail "Codex collector rewrites the cache after a corrupt read" "$result" -pass "Codex collector recovers from a corrupt cache file" - -# A new opencode message changes what a scan would find; a --limits-only run -# must reuse the cached stats instead of rescanning. -python3 - "$CACHE_HOME/.local/share/opencode/opencode.db" <<'PY' -import json -import sqlite3 -import sys -import time -from pathlib import Path - -db = Path(sys.argv[1]) -conn = sqlite3.connect(db) -now_ms = int(time.time() * 1000) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ( - "c_2", "ses_1", now_ms, now_ms, json.dumps({ - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": 10, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - }), -)) -conn.commit() -conn.close() -PY - -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "5" ]] || - fail "Codex collector --limits-only reuses cached local stats" "$result" -[[ $(jq -c '.modelUsage' <<<"$result") == '{"gpt-5.2-codex":{"inputTokens":5,"outputTokens":0,"cacheReadInputTokens":0,"cacheCreationInputTokens":0}}' ]] || - fail "Codex collector --limits-only emits a complete record from cache" "$result" -pass "Codex collector --limits-only reuses cached local stats" - -# --force must ignore the cache and pick up the new message. -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector" --force) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "15" ]] || - fail "Codex collector --force rescans past the cache" "$result" -pass "Codex collector --force rescans past the cache" - -# The forced scan refreshed the cache, so a following --limits-only sees it. -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "15" ]] || - fail "Codex collector --limits-only sees a refreshed cache after --force" "$result" -pass "Codex collector --limits-only sees a refreshed cache after --force" - -# An expired cache makes --limits-only rescan too: stale today* stats must -# never be served under a fresh updatedAt. -python3 - "$CACHE_HOME/.local/share/opencode/opencode.db" <<'PY' -import json -import sqlite3 -import sys -import time -from pathlib import Path - -db = Path(sys.argv[1]) -conn = sqlite3.connect(db) -now_ms = int(time.time() * 1000) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ( - "c_3", "ses_1", now_ms, now_ms, json.dumps({ - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": 10, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - }), -)) -conn.commit() -conn.close() -PY -touch -d "2 hours ago" "$cache_file" -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "25" ]] || - fail "Codex collector --limits-only rescans when the cache is stale" "$result" -pass "Codex collector --limits-only rescans when the cache is stale" - -# The 15-minute reuse window belongs to --limits-only alone. A no-flag run -# (the widget's periodic refresh) reuses a scan only while it is young enough -# to be a concurrent collector run; past that it rescans, so stats stay as -# fresh as refreshIntervalSec, however low the user sets it. -python3 - "$CACHE_HOME/.local/share/opencode/opencode.db" <<'PY' -import json -import sqlite3 -import sys -import time -from pathlib import Path - -db = Path(sys.argv[1]) -conn = sqlite3.connect(db) -now_ms = int(time.time() * 1000) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ( - "c_4", "ses_1", now_ms, now_ms, json.dumps({ - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": 10, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - }), -)) -conn.commit() -conn.close() -PY - -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "25" ]] || - fail "Codex collector no-flag reuses a seconds-old cache" "$result" - -# 30 seconds is the lowest refreshIntervalSec the widget supports, so a -# cache that old must already be past the no-flag reuse window. -touch -d "30 seconds ago" "$cache_file" -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "35" ]] || - fail "Codex collector no-flag rescans past the concurrent-run window" "$result" -pass "Codex collector no-flag mode rescans instead of serving a stale cache" - -# The same age from the other side: a cache far past the no-flag window but -# well inside 15 minutes is still good enough for --limits-only. -python3 - "$CACHE_HOME/.local/share/opencode/opencode.db" <<'PY' -import json -import sqlite3 -import sys -import time -from pathlib import Path - -db = Path(sys.argv[1]) -conn = sqlite3.connect(db) -now_ms = int(time.time() * 1000) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ( - "c_5", "ses_1", now_ms, now_ms, json.dumps({ - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": 10, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - }), -)) -conn.commit() -conn.close() -PY -touch -d "10 minutes ago" "$cache_file" -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "35" ]] || - fail "Codex collector --limits-only reuses a scan the no-flag mode would refresh" "$result" -pass "Codex collector --limits-only reuses a scan the no-flag mode would refresh" - -# A cache written on another local date holds another day's today* stats even -# under a fresh mtime (midnight passed, or the clock moved backwards): the -# envelope's scanDate must turn it into a miss. -jq -c '.scanDate = "1999-01-01"' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "45" ]] || - fail "Codex collector treats a cache from another day as a miss" "$result" -[[ $(jq -r '.scanDate' "$cache_file") == "$(date +%Y-%m-%d)" ]] || - fail "Codex collector stamps the rewritten cache with the scan date" "$result" -pass "Codex collector treats a cache from another day as a miss" - -# A cache stamped in the future (the clock was set backwards after the write) -# has no trustworthy age: it must be a miss, not fresh until the clock -# catches up. -python3 - "$CACHE_HOME/.local/share/opencode/opencode.db" <<'PY' -import json -import sqlite3 -import sys -import time -from pathlib import Path - -db = Path(sys.argv[1]) -conn = sqlite3.connect(db) -now_ms = int(time.time() * 1000) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ( - "c_6", "ses_1", now_ms, now_ms, json.dumps({ - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": 10, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - }), -)) -conn.commit() -conn.close() -PY -touch -d "@$(( $(date +%s) + 3600 ))" "$cache_file" -result=$(HOME="$CACHE_HOME" CODEX_HOME="$CACHE_HOME/.codex" XDG_CACHE_HOME="$CACHE_HOME/.cache" XDG_DATA_HOME="$CACHE_HOME/.local/share" \ - PATH="$CACHE_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "55" ]] || - fail "Codex collector treats a future-dated cache as a miss" "$result" -pass "Codex collector treats a future-dated cache as a miss" - -# First --limits-only on a machine with no cache falls back to a full scan. -FRESH_HOME=$(signed_in_home) -cp "$TEST_HOME/bin/codex" "$FRESH_HOME/bin/codex" - -python3 - "$FRESH_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) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ( - "f_1", "ses_1", now_ms, now_ms, json.dumps({ - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": 7, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - }), -)) -conn.commit() -conn.close() -PY - -result=$(HOME="$FRESH_HOME" CODEX_HOME="$FRESH_HOME/.codex" XDG_CACHE_HOME="$FRESH_HOME/.cache" XDG_DATA_HOME="$FRESH_HOME/.local/share" \ - PATH="$FRESH_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "7" ]] || - fail "Codex collector --limits-only falls back to a full scan without a cache" "$result" -pass "Codex collector --limits-only falls back to a full scan without a cache" - -# A malformed opencode row must not abort the scan: json_valid() guards the -# parse, so the good rows are still counted. Real opencode data also stores -# compact JSON, so one row is serialized compactly here on purpose. -MALFORMED_HOME=$(signed_in_home) -cp "$TEST_HOME/bin/codex" "$MALFORMED_HOME/bin/codex" - -python3 - "$MALFORMED_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, input=0, compact=False): - payload = { - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": input, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - } - if compact: - return (id, "ses_1", now_ms, now_ms, json.dumps(payload, separators=(",", ":"))) - return (id, "ses_1", now_ms, now_ms, json.dumps(payload)) - -conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ - message("mm_1", input=5, compact=True), - message("mm_2", input=7), -]) -# Valid JSON followed by trailing garbage: without json_valid() this row -# makes json_extract() raise and aborts the whole scan. -good = json.dumps({"role": "assistant", "providerID": "openai", "modelID": "gpt-5.2-codex", - "tokens": {"input": 999, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}}) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ("mm_3", "ses_1", now_ms, now_ms, good + " trailing-garbage")) -# Completely broken row: not JSON at all. -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ("mm_4", "ses_1", now_ms, now_ms, "this is not json")) -conn.commit() -conn.close() -PY - -result=$(HOME="$MALFORMED_HOME" CODEX_HOME="$MALFORMED_HOME/.codex" XDG_CACHE_HOME="$MALFORMED_HOME/.cache" XDG_DATA_HOME="$MALFORMED_HOME/.local/share" \ - PATH="$MALFORMED_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "12" ]] || - fail "Codex collector counts good opencode rows past malformed ones" "$result" -pass "Codex collector counts good opencode rows past malformed ones" - -# An unwritable cache must not kill the collector: the record is the contract. -UNWRITABLE_HOME=$(signed_in_home) -cp "$TEST_HOME/bin/codex" "$UNWRITABLE_HOME/bin/codex" - -python3 - "$UNWRITABLE_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) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ( - "u_1", "ses_1", now_ms, now_ms, json.dumps({ - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": 3, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - }), -)) -conn.commit() -conn.close() -PY - -# XDG_CACHE_HOME points at a regular file, so mkdir inside cache_root fails. -touch "$UNWRITABLE_HOME/not-a-dir" -result=$(HOME="$UNWRITABLE_HOME" CODEX_HOME="$UNWRITABLE_HOME/.codex" XDG_CACHE_HOME="$UNWRITABLE_HOME/not-a-dir" XDG_DATA_HOME="$UNWRITABLE_HOME/.local/share" \ - PATH="$UNWRITABLE_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "3" ]] || - fail "Codex collector still prints a complete record when the cache is unwritable" "$result" -pass "Codex collector still prints a complete record when the cache is unwritable" - -# A scan cut short by a database error (schema migration, transient lock, -# corruption) must not be cached as the whole story, or the missing usage -# would be suppressed for every reader until the cache expires. -INTERRUPTED_HOME=$(signed_in_home) -cp "$TEST_HOME/bin/codex" "$INTERRUPTED_HOME/bin/codex" - -# A database without the message table makes the scan fail mid-flight. -python3 - "$INTERRUPTED_HOME/.local/share/opencode/opencode.db" <<'PY' -import sqlite3 -import sys -from pathlib import Path - -db = Path(sys.argv[1]) -db.parent.mkdir(parents=True, exist_ok=True) -conn = sqlite3.connect(db) -conn.execute("CREATE TABLE unrelated (id text PRIMARY KEY)") -conn.commit() -conn.close() -PY - -result=$(HOME="$INTERRUPTED_HOME" CODEX_HOME="$INTERRUPTED_HOME/.codex" XDG_CACHE_HOME="$INTERRUPTED_HOME/.cache" XDG_DATA_HOME="$INTERRUPTED_HOME/.local/share" \ - PATH="$INTERRUPTED_HOME/bin:$PATH" "$collector") - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "0" ]] || - fail "Codex collector reports what it could read from a broken database" "$result" -[[ -z $(ls "$INTERRUPTED_HOME/.cache/omarchy/agent-usage/"codex-scan-*.json 2>/dev/null) ]] || - fail "Codex collector must not cache an interrupted scan" "$result" - -# Once the database is whole again, the very next --limits-only run scans it -# instead of reusing a zero snapshot. -python3 - "$INTERRUPTED_HOME/.local/share/opencode/opencode.db" <<'PY' -import json -import sqlite3 -import sys -import time -from pathlib import Path - -db = Path(sys.argv[1]) -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) -conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ( - "i_1", "ses_1", now_ms, now_ms, json.dumps({ - "role": "assistant", - "providerID": "openai", - "modelID": "gpt-5.2-codex", - "tokens": {"input": 9, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, - "time": {"created": now_ms}, - }), -)) -conn.commit() -conn.close() -PY - -result=$(HOME="$INTERRUPTED_HOME" CODEX_HOME="$INTERRUPTED_HOME/.codex" XDG_CACHE_HOME="$INTERRUPTED_HOME/.cache" XDG_DATA_HOME="$INTERRUPTED_HOME/.local/share" \ - PATH="$INTERRUPTED_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "9" ]] || - fail "Codex collector does not reuse a snapshot from an interrupted scan" "$result" -pass "Codex collector does not cache an interrupted opencode scan" +''' # The limits name the plan themselves, so an account/read that never answers -# costs nothing: the limits still arrive, and quickly. -started=$(date +%s) -result=$(HOME="$TEST_HOME" CODEX_HOME="$TEST_HOME/.codex" XDG_DATA_HOME="$TEST_HOME/.local/share" PATH="$TEST_HOME/bin:$PATH" \ - CODEX_ACCOUNT_READ_HANGS=1 CODEX_RATE_LIMITS='{"planType":"pro","primary":{"usedPercent":36,"windowDurationMins":10080}}' \ - "$collector" --limits-only) -(( $(date +%s) - started < 4 )) || fail "Codex collector doesn't wait on account/read when the limits name the plan" -[[ $(jq -c '{tierLabel, usageStatusText, limits: [.limits[] | {label, percent}]}' <<<"$result") == '{"tierLabel":"pro","usageStatusText":"","limits":[{"label":"Weekly (7-day)","percent":0.36}]}' ]] || - fail "Codex collector reads limits even when account/read never answers" "$result" -pass "Codex collector reads limits even when account/read never answers" - -# Free full resets ride along with the limits: only available ones count, and -# the soonest to lapse is the one worth mentioning. -result=$(HOME="$TEST_HOME" CODEX_HOME="$TEST_HOME/.codex" XDG_DATA_HOME="$TEST_HOME/.local/share" PATH="$TEST_HOME/bin:$PATH" \ - CODEX_RATE_LIMITS='{"planType":"pro","primary":{"usedPercent":42,"windowDurationMins":10080}}' \ - CODEX_RESET_CREDITS='{"availableCount":2,"credits":[{"status":"available","expiresAt":2000000000},{"status":"available","expiresAt":1900000000},{"status":"used","expiresAt":1800000000}]}' \ - "$collector" --limits-only) -[[ $(jq -c '.resetCredits' <<<"$result") == '{"available":2,"nextExpiresAt":"2030-03-17T17:46:40+00:00"}' ]] || - fail "Codex collector reports its available free resets" "$result" -pass "Codex collector reports its available free resets" +# costs nothing: it isn't asked. +new_home(ANSWERS) +started = time.monotonic() +record = codex(env={"CODEX_ACCOUNT_READ_HANGS": 1, "CODEX_RATE_LIMITS": '{"planType":"pro","primary":{"usedPercent":36,"windowDurationMins":10080}}'}) +elapsed = time.monotonic() - started +expect("Codex reads limits naming their plan without waiting on account/read", { + "tier": record["tierLabel"], "status": record["usageStatusText"], "limits": [[entry["label"], entry["percent"]] for entry in record["limits"]], "quick": elapsed < 1.8, +}, {"tier": "Pro 200", "status": "", "limits": [["Weekly (7-day)", 0.36]], "quick": True}, record) + +# Free full resets ride along: only available ones count, and the soonest to +# lapse is the one worth mentioning. +record = codex(env={"CODEX_RATE_LIMITS": '{"planType":"pro","primary":{"usedPercent":42,"windowDurationMins":10080}}', + "CODEX_RESET_CREDITS": json.dumps({"availableCount": 2, "credits": [ + {"status": "available", "expiresAt": 2000000000}, {"status": "available", "expiresAt": 1900000000}, {"status": "used", "expiresAt": 1800000000}]})}) +expect("Codex reports its available free resets", record.get("resetCredits"), {"available": 2, "nextExpiresAt": "2030-03-17T17:46:40+00:00"}, record) + +# Credits and account/read are extras: a malformed one loses itself, never +# the limits already read, which stay live. +for name, credits in (("milliseconds", {"credits": [{"status": "available", "expiresAt": 1900000000000}]}), ("a list", [1]), + ("a number", {"credits": 3}), ("a string", "lots")): + record = codex(env={"CODEX_RATE_LIMITS": '{"planType":"pro","primary":{"usedPercent":42,"windowDurationMins":10080}}', + "CODEX_RESET_CREDITS": json.dumps(credits)}) + expect(f"Codex keeps its limits live when the reset credits are {name}", { + "status": record["usageStatusText"], "stale": record["limitsStale"], "limits": [[entry["label"], entry["percent"]] for entry in record["limits"]], + }, {"status": "", "stale": False, "limits": [["Weekly (7-day)", 0.42]]}, record) +for name, answer in (("a list", [1]), ("a string account", {"account": "pro"})): + record = codex(env={"CODEX_RATE_LIMITS": '{"primary":{"usedPercent":42,"windowDurationMins":10080}}', "CODEX_ACCOUNT_RESULT": json.dumps(answer)}) + expect(f"Codex keeps its limits when account/read answers {name}", { + "tier": record["tierLabel"], "status": record["usageStatusText"], "stale": record["limitsStale"], + "limits": [[entry["label"], entry["percent"]] for entry in record["limits"]], + }, {"tier": "", "status": "", "stale": False, "limits": [["Weekly (7-day)", 0.42]]}, record) # A Business seat with a spend cap and no rate windows reports the cap alone, -# with no inferred frequency and a share taken from the amounts. -result=$(HOME="$TEST_HOME" CODEX_HOME="$TEST_HOME/.codex" XDG_DATA_HOME="$TEST_HOME/.local/share" PATH="$TEST_HOME/bin:$PATH" \ - CODEX_RATE_LIMITS='{"planType":"business","primary":null,"secondary":null,"individualLimit":{"limit":"6000","used":"1037.07","remainingPercent":83,"resetsAt":1793491200}}' \ - "$collector" --limits-only) -jq -e ' - (.limits | length) == 1 - and .limits[0].label == "Spend limit" - and ((.limits[0].percent - (1037.07 / 6000)) | fabs) < 0.000001 - and .limits[0].resetsAt == "2026-11-01T00:00:00+00:00" -' <<<"$result" >/dev/null || - fail "Codex collector reports a Business spend cap with its supplied reset time" "$result" -pass "Codex collector reports a Business spend cap with its supplied reset time" - -# A zero cap cannot be divided into usage. Use the reported percentage -# instead: 40% remaining means 60% used. This is a synthetic edge case. -result=$(HOME="$TEST_HOME" CODEX_HOME="$TEST_HOME/.codex" XDG_DATA_HOME="$TEST_HOME/.local/share" PATH="$TEST_HOME/bin:$PATH" \ - CODEX_RATE_LIMITS='{"planType":"business","individualLimit":{"limit":"0","used":"0","remainingPercent":40,"resetsAt":1793491200}}' \ - "$collector" --limits-only) -jq -e '[.limits[] | select(.label == "Spend limit") | .percent] == [0.6]' <<<"$result" >/dev/null || - fail "Codex collector uses the reported percentage when the spend cap is zero" "$result" -pass "Codex collector uses the reported percentage when the spend cap is zero" +# its share taken from the amounts; a zero cap takes the reported percentage. +new_home() +record = codex(env={"CODEX_RATE_LIMITS": '{"planType":"business","primary":null,"secondary":null,"individualLimit":{"limit":"6000","used":"1037.07","remainingPercent":83,"resetsAt":1793491200}}'}) +check("Codex reports a Business spend cap with its own reset time", + [entry["label"] for entry in record["limits"]] == ["Spend limit"] and abs(record["limits"][0]["percent"] - 1037.07 / 6000) < 1e-6 + and record["limits"][0]["resetsAt"] == "2026-11-01T00:00:00+00:00", record) +record = codex(env={"CODEX_RATE_LIMITS": '{"planType":"business","individualLimit":{"limit":"0","used":"0","remainingPercent":40,"resetsAt":1793491200}}'}) +expect("Codex takes the reported percentage when the spend cap is zero", + [entry["percent"] for entry in record["limits"] if entry["label"] == "Spend limit"], [0.6], record) # A home nobody is signed in to answers with an error, which reads as a # sign-in to restore rather than as missing numbers. -result=$(HOME="$TEST_HOME" CODEX_HOME="$TEST_HOME/.codex" XDG_DATA_HOME="$TEST_HOME/.local/share" PATH="$TEST_HOME/bin:$PATH" \ - CODEX_LIMITS_ERROR="codex account authentication required to read rate limits" \ - "$collector" --limits-only) -[[ $(jq -r '.usageStatusText' <<<"$result") == "Waiting for auth" ]] || - fail "Codex collector reports a missing sign-in as one" "$result" -pass "Codex collector reports a missing sign-in as one" - -# The app-server batches notifications with replies in one write. A reply -# that shares a write with a notification must not be stranded in a read -# buffer, and bytes left over from one request must carry into the next. -BATCHED_HOME=$(signed_in_home) -cat >"$BATCHED_HOME/bin/codex" <<'EOF' -#!/bin/bash +record = codex(env={"CODEX_FAKE": "auth"}) +expect("Codex reports a missing sign-in as one", record["usageStatusText"], "Waiting for auth", record) +# The app-server batches notifications with replies in one write: a reply +# sharing a write must not be stranded in a read buffer, and bytes left from +# one request carry into the next. +new_home(r''' while read -r request; do id=$(jq -r '.id // empty' <<<"$request") - method=$(jq -r '.method // empty' <<<"$request") - - case "$method" in + case "$(jq -r '.method // empty' <<<"$request")" in initialize) - # One write: this reply and a trailing notification. - printf '%s\n%s\n' \ - "$(jq -cn --argjson id "$id" '{id: $id, result: {}}')" \ - '{"method":"remoteControl/status/changed","params":{"status":"disabled"}}' - ;; + printf '%s\n%s\n' "$(jq -cn --argjson id "$id" '{id: $id, result: {}}')" '{"method":"remoteControl/status/changed","params":{"status":"disabled"}}' ;; account/rateLimits/read) - # One write: a notification ahead of this reply. - printf '%s\n%s\n' \ - '{"method":"account/updated","params":{"authMode":"chatgpt","planType":"plus"}}' \ - "$(jq -cn --argjson id "$id" '{id: $id, result: {rateLimits: {planType: "plus", primary: {usedPercent: 5, windowDurationMins: 300, resetsAt: 1790659194}}}}')" - ;; + printf '%s\n%s\n' '{"method":"account/updated","params":{"authMode":"chatgpt","planType":"plus"}}' \ + "$(jq -cn --argjson id "$id" '{id: $id, result: {rateLimits: {planType: "plus", primary: {usedPercent: 5, windowDurationMins: 300, resetsAt: 1790659194}}}}')" ;; esac done -EOF -chmod +x "$BATCHED_HOME/bin/codex" +''') +record = codex() +expect("Codex reads replies batched with notifications", { + "tier": record["tierLabel"], "status": record["usageStatusText"], "limits": [[entry["label"], entry["percent"]] for entry in record["limits"]], +}, {"tier": "Plus", "status": "", "limits": [["5h window", 0.05]]}, record) -result=$(HOME="$BATCHED_HOME" CODEX_HOME="$BATCHED_HOME/.codex" XDG_CACHE_HOME="$BATCHED_HOME/.cache" XDG_DATA_HOME="$BATCHED_HOME/.local/share" \ - PATH="$BATCHED_HOME/bin:$PATH" "$collector") +# Without temporary space the probe still runs, only without Codex's error +# text. +new_home() +real_temporary_file = tempfile.TemporaryFile -[[ $(jq -c '{tierLabel, usageStatusText, limits: [.limits[] | {label, percent}]}' <<<"$result") == '{"tierLabel":"plus","usageStatusText":"","limits":[{"label":"5h window","percent":0.05}]}' ]] || - fail "Codex collector reads replies batched with notifications" "$result" -pass "Codex collector reads replies batched with notifications" - -# Without temporary space the probe still runs; it only loses Codex's error -# text, so stderr goes nowhere instead of failing the probe. -NO_TEMP_PYTHON="$BATCHED_HOME/python" -mkdir -p "$NO_TEMP_PYTHON" -cat >"$NO_TEMP_PYTHON/sitecustomize.py" <<'EOF' -import tempfile def no_space(*args, **kwargs): raise OSError(28, "No space left on device") -tempfile.TemporaryFile = no_space -EOF - -result=$(HOME="$BATCHED_HOME" CODEX_HOME="$BATCHED_HOME/.codex" XDG_CACHE_HOME="$BATCHED_HOME/.cache" XDG_DATA_HOME="$BATCHED_HOME/.local/share" \ - PYTHONPATH="$NO_TEMP_PYTHON" PATH="$BATCHED_HOME/bin:$PATH" "$collector" --force) - -[[ $(jq -c '[.limits[] | .percent]' <<<"$result") == '[0.05]' ]] || - fail "Codex collector probes without temporary space" "$result" -pass "Codex collector probes without temporary space" - -# The lazy launcher at ~/.local/bin/codex runs `mise use -g` when executed, so -# a read-only usage probe on a machine without Codex must never spawn it. A -# private tools dir keeps a real codex or mise on the host out of the probe: -# PATH holds only the interpreter and file tools the collector may exec. -LAUNCH_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -SAFE_PATH="$LAUNCH_HOME/tools" -mkdir -p "$SAFE_PATH" -for tool in python3 rg jq; do - if command -v "$tool" >/dev/null; then - ln -s "$(command -v "$tool")" "$SAFE_PATH/$tool" - fi -done -mkdir -p "$LAUNCH_HOME/bin" "$LAUNCH_HOME/.local/bin" "$LAUNCH_HOME/.codex" -touch "$LAUNCH_HOME/.codex/auth.json" -cat >"$LAUNCH_HOME/.local/bin/codex" <<'LAUNCHER' -#!/bin/bash -export MISE_MINIMUM_RELEASE_AGE=0 -mise use -g --quiet "npm:@openai/codex" || exit 1 -exec mise x "npm:@openai/codex" -- codex "$@" -LAUNCHER -chmod +x "$LAUNCH_HOME/.local/bin/codex" -cat >"$LAUNCH_HOME/bin/mise" <<'STUB' -#!/bin/bash -printf '%s\n' "$*" >>"$MISE_CALLS_FILE" -exit 1 -STUB -chmod +x "$LAUNCH_HOME/bin/mise" - -result=$(HOME="$LAUNCH_HOME" CODEX_HOME="$LAUNCH_HOME/.codex" MISE_CALLS_FILE="$LAUNCH_HOME/mise-calls" XDG_DATA_HOME="$LAUNCH_HOME/.local/share" \ - PATH="$LAUNCH_HOME/bin:$SAFE_PATH" "$collector" --limits-only) - -[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex unavailable" ]] || - fail "Codex collector reports Codex unavailable when only the launcher exists" "$result" -# The launcher would log `use -g ...` through the mise stub if it ever ran; -# `which codex` is the only permitted call. -[[ ! -s $LAUNCH_HOME/mise-calls || $(cat "$LAUNCH_HOME/mise-calls") == "which codex" ]] || - fail "Codex collector must not execute the lazy codex launcher" "$(cat "$LAUNCH_HOME/mise-calls" 2>/dev/null)" -pass "Codex collector resolves through mise which instead of running the launcher" - -# When mise reports an installed binary, that binary is probed, not the -# launcher that shadows it on PATH. -cp "$TEST_HOME/bin/codex" "$LAUNCH_HOME/real-codex" -cat >"$LAUNCH_HOME/bin/mise" <>"$LAUNCH_HOME/mise-calls" -echo "$LAUNCH_HOME/real-codex" -STUB - -rm -f "$LAUNCH_HOME/mise-calls" -result=$(HOME="$LAUNCH_HOME" CODEX_HOME="$LAUNCH_HOME/.codex" CODEX_ARGS_FILE="$LAUNCH_HOME/codex-args" XDG_DATA_HOME="$LAUNCH_HOME/.local/share" \ - PATH="$LAUNCH_HOME/bin:$SAFE_PATH" "$collector" --limits-only) - -[[ -f $LAUNCH_HOME/codex-args && $(cat "$LAUNCH_HOME/mise-calls") == "which codex" ]] || - fail "Codex collector probes the binary mise which reports" "$result" -pass "Codex collector probes the mise-resolved binary" - -# A symlink at the launcher path is the user's own binary, so it is probed -# directly without asking mise at all. -LINK_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -mkdir -p "$LINK_HOME/bin" "$LINK_HOME/.local/bin" "$LINK_HOME/.codex" -touch "$LINK_HOME/.codex/auth.json" -ln -s "$TEST_HOME/bin/codex" "$LINK_HOME/.local/bin/codex" -cat >"$LINK_HOME/bin/mise" <>"$LINK_HOME/mise-calls" -echo "$LINK_HOME/real-codex" -STUB -chmod +x "$LINK_HOME/bin/mise" - -result=$(HOME="$LINK_HOME" CODEX_HOME="$LINK_HOME/.codex" CODEX_ARGS_FILE="$LINK_HOME/codex-args" XDG_DATA_HOME="$LINK_HOME/.local/share" \ - PATH="$LINK_HOME/bin:$SAFE_PATH" "$collector" --limits-only) - -[[ -f $LINK_HOME/codex-args && ! -s $LINK_HOME/mise-calls ]] || - fail "Codex collector probes a symlinked codex without invoking mise" "$result" -pass "Codex collector probes a user-owned symlink at the launcher path" - -# A mise shim is a symlink to the mise binary itself, so it resolves to mise, -# not to an installed codex — running it would exec `mise x` and install the -# tool. It must be treated as lazy despite being a symlink. -SHIM_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -mkdir -p "$SHIM_HOME/bin" "$SHIM_HOME/.local/share/mise/shims" -cat >"$SHIM_HOME/bin/mise" <>"$SHIM_HOME/mise-calls" -exit 1 -STUB -chmod +x "$SHIM_HOME/bin/mise" -ln -s "$SHIM_HOME/bin/mise" "$SHIM_HOME/.local/share/mise/shims/codex" - -result=$(HOME="$SHIM_HOME" CODEX_HOME="$SHIM_HOME/.codex" MISE_CALLS_FILE="$SHIM_HOME/mise-calls" XDG_DATA_HOME="$SHIM_HOME/.local/share" \ - PATH="$SHIM_HOME/bin:$SAFE_PATH" "$collector" --limits-only) - -[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex unavailable" ]] || - fail "Codex collector reports Codex unavailable when only a mise shim exists" "$result" -[[ ! -s $SHIM_HOME/mise-calls || $(cat "$SHIM_HOME/mise-calls") == "which codex" ]] || - fail "Codex collector must not execute a mise shim" "$(cat "$SHIM_HOME/mise-calls" 2>/dev/null)" -pass "Codex collector treats a shim symlink to mise as lazy" - -# Without Codex credentials, account/read can only fail — and starting the -# app-server is not free: it syncs the plugin list, a git fetch per refresh -# that leaves ~/.codex/.tmp/git-* folders behind. The collector must not -# spawn codex at all. -NOAUTH_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -mkdir -p "$NOAUTH_HOME/bin" -cp "$TEST_HOME/bin/codex" "$NOAUTH_HOME/bin/codex" - -result=$(HOME="$NOAUTH_HOME" CODEX_HOME="$NOAUTH_HOME/.codex" CODEX_ARGS_FILE="$NOAUTH_HOME/codex-args" XDG_DATA_HOME="$NOAUTH_HOME/.local/share" \ - PATH="$NOAUTH_HOME/bin:$PATH" env -u OPENAI_API_KEY -u CODEX_API_KEY -u CODEX_ACCESS_TOKEN "$collector" --limits-only) - -[[ ! -e $NOAUTH_HOME/codex-args ]] || - fail "Codex collector does not spawn app-server without credentials" "$result" -[[ $(jq -r '.usageStatusText' <<<"$result") == "Waiting for auth" ]] || - fail "Codex collector waits for auth without credentials" "$result" -pass "Codex collector does not spawn app-server without credentials" - -# A non-file credentials store keeps credentials outside auth.json, so the -# probe must still run. -KEYRING_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -mkdir -p "$KEYRING_HOME/bin" "$KEYRING_HOME/.codex" -cp "$TEST_HOME/bin/codex" "$KEYRING_HOME/bin/codex" -printf 'cli_auth_credentials_store = "keyring"\n' >"$KEYRING_HOME/.codex/config.toml" - -result=$(HOME="$KEYRING_HOME" CODEX_HOME="$KEYRING_HOME/.codex" CODEX_ARGS_FILE="$KEYRING_HOME/codex-args" XDG_DATA_HOME="$KEYRING_HOME/.local/share" \ - PATH="$KEYRING_HOME/bin:$PATH" env -u OPENAI_API_KEY -u CODEX_API_KEY -u CODEX_ACCESS_TOKEN "$collector" --limits-only) - -[[ -e $KEYRING_HOME/codex-args ]] || - fail "Codex collector probes the app-server when the keyring store is configured" "$result" -pass "Codex collector probes the app-server when the keyring store is configured" - -# The setting is TOML, so a single-quoted value names the keyring just as well, -# while an explicit file store with no auth.json is still signed out. -for store in "'keyring'" '"file"'; do - rm -f "$KEYRING_HOME/codex-args" - printf 'model = "gpt-5"\ncli_auth_credentials_store = %s\n' "$store" >"$KEYRING_HOME/.codex/config.toml" - HOME="$KEYRING_HOME" CODEX_HOME="$KEYRING_HOME/.codex" CODEX_ARGS_FILE="$KEYRING_HOME/codex-args" XDG_DATA_HOME="$KEYRING_HOME/.local/share" \ - PATH="$KEYRING_HOME/bin:$PATH" env -u OPENAI_API_KEY -u CODEX_API_KEY -u CODEX_ACCESS_TOKEN "$collector" --limits-only >/dev/null - if [[ $store == "'keyring'" ]]; then - [[ -e $KEYRING_HOME/codex-args ]] || fail "Codex collector reads a single-quoted keyring store as TOML" - else - [[ ! -e $KEYRING_HOME/codex-args ]] || fail "Codex collector does not probe an empty file store" - fi -done -pass "Codex collector reads the credentials store as TOML" - -# A CODEX_ACCESS_TOKEN is credentials even without auth.json. -TOKEN_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -mkdir -p "$TOKEN_HOME/bin" "$TOKEN_HOME/.codex" -cp "$TEST_HOME/bin/codex" "$TOKEN_HOME/bin/codex" - -result=$(HOME="$TOKEN_HOME" CODEX_HOME="$TOKEN_HOME/.codex" CODEX_ARGS_FILE="$TOKEN_HOME/codex-args" XDG_DATA_HOME="$TOKEN_HOME/.local/share" \ - PATH="$TOKEN_HOME/bin:$PATH" env -u OPENAI_API_KEY -u CODEX_API_KEY CODEX_ACCESS_TOKEN=x "$collector" --limits-only) - -[[ -e $TOKEN_HOME/codex-args ]] || - fail "Codex collector probes the app-server when CODEX_ACCESS_TOKEN is set" "$result" -pass "Codex collector probes the app-server when CODEX_ACCESS_TOKEN is set" - -# An API key alone does not log the app-server in. -APIKEY_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -mkdir -p "$APIKEY_HOME/bin" "$APIKEY_HOME/.codex" -cp "$TEST_HOME/bin/codex" "$APIKEY_HOME/bin/codex" - -result=$(HOME="$APIKEY_HOME" CODEX_HOME="$APIKEY_HOME/.codex" CODEX_ARGS_FILE="$APIKEY_HOME/codex-args" XDG_DATA_HOME="$APIKEY_HOME/.local/share" \ - PATH="$APIKEY_HOME/bin:$PATH" env -u CODEX_API_KEY -u CODEX_ACCESS_TOKEN OPENAI_API_KEY=x "$collector" --limits-only) - -[[ ! -e $APIKEY_HOME/codex-args ]] || - fail "Codex collector does not spawn app-server for an API key alone" "$result" -pass "Codex collector does not spawn app-server for an API key alone" - -# A codex that exits before speaking the protocol (rejected flag, crash, etc.) -# must not leave the panel showing the bare RPC method name "initialize". -EXIT_HOME=$(signed_in_home) -cat >"$EXIT_HOME/bin/codex" <<'EOF' -#!/bin/bash -echo "error: invalid value 'untrusted' for '--ask-for-approval '" >&2 -echo " [possible values: on-request, never]" >&2 -exit 2 -EOF -chmod +x "$EXIT_HOME/bin/codex" - -result=$(HOME="$EXIT_HOME" CODEX_HOME="$EXIT_HOME/.codex" XDG_CACHE_HOME="$EXIT_HOME/.cache" XDG_DATA_HOME="$EXIT_HOME/.local/share" \ - PATH="$EXIT_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.usageStatusText' <<<"$result") == "Codex limits unavailable" ]] || - fail "Codex collector reports limits unavailable when app-server rejects argv" "$result" -help=$(jq -r '.authHelpText' <<<"$result") -[[ $help == *"invalid value 'untrusted'"* ]] || - fail "Codex collector surfaces the CLI's own error" "$result" -[[ $help != "initialize" ]] || - fail "Codex collector must not leak the raw RPC method name" "$result" -pass "Codex collector surfaces a rejected app-server call instead of the RPC method name" - -# EOF on stdout during initialize (process died) is reported as an exit, not a bare method. -DEAD_HOME=$(signed_in_home) -cat >"$DEAD_HOME/bin/codex" <<'EOF' -#!/bin/bash -exec 1>&- -exec sleep 30 -EOF -chmod +x "$DEAD_HOME/bin/codex" - -result=$(HOME="$DEAD_HOME" CODEX_HOME="$DEAD_HOME/.codex" XDG_CACHE_HOME="$DEAD_HOME/.cache" XDG_DATA_HOME="$DEAD_HOME/.local/share" \ - PATH="$DEAD_HOME/bin:$PATH" "$collector" --limits-only) - -help=$(jq -r '.authHelpText' <<<"$result") -[[ $help == "Codex app-server exited before initialize" ]] || - fail "Codex collector identifies an app-server that exits during startup" "$result" -pass "Codex collector identifies an app-server that exits during startup" - -# Failure while sending the initialized notification after initialize answered. -HALF_HOME=$(signed_in_home) -cat >"$HALF_HOME/bin/codex" <<'EOF' -#!/bin/bash -read -r request -exec 0<&- -jq -cn --argjson id "$(jq -r '.id' <<<"$request")" '{id: $id, result: {}}' -exec sleep 30 -EOF -chmod +x "$HALF_HOME/bin/codex" - -result=$(HOME="$HALF_HOME" CODEX_HOME="$HALF_HOME/.codex" XDG_CACHE_HOME="$HALF_HOME/.cache" XDG_DATA_HOME="$HALF_HOME/.local/share" \ - PATH="$HALF_HOME/bin:$PATH" "$collector" --limits-only) - -help=$(jq -r '.authHelpText' <<<"$result") -[[ $help == "Codex app-server exited before initialized" ]] || - fail "Codex collector translates failure to send the initialized notification" "$result" -pass "Codex collector translates failure to send the initialized notification" - -# Silent clean exit with no stderr: fall back to the login hint. -SILENT_HOME=$(signed_in_home) -cat >"$SILENT_HOME/bin/codex" <<'EOF' -#!/bin/bash -exit 0 -EOF -chmod +x "$SILENT_HOME/bin/codex" - -result=$(HOME="$SILENT_HOME" CODEX_HOME="$SILENT_HOME/.codex" XDG_CACHE_HOME="$SILENT_HOME/.cache" XDG_DATA_HOME="$SILENT_HOME/.local/share" \ - PATH="$SILENT_HOME/bin:$PATH" "$collector" --limits-only) - -help=$(jq -r '.authHelpText' <<<"$result") -[[ $help == "Run \`codex login\` to authenticate." ]] || - fail "Codex collector falls back to the login hint when the CLI is silent" "$result" -pass "Codex collector falls back to the login hint when the app-server says nothing" - -# Live app-server that stalls on account/rateLimits/read: keep a clear stall message, not login. -STALL_HOME=$(signed_in_home) -cat >"$STALL_HOME/bin/codex" <<'EOF' -#!/bin/bash -while read -r request; do - id=$(jq -r '.id // empty' <<<"$request") - method=$(jq -r '.method // empty' <<<"$request") - case "$method" in - initialize) jq -cn --argjson id "$id" '{id: $id, result: {}}' ;; - account/rateLimits/read) : ;; - esac -done -EOF -chmod +x "$STALL_HOME/bin/codex" - -result=$(HOME="$STALL_HOME" CODEX_HOME="$STALL_HOME/.codex" XDG_CACHE_HOME="$STALL_HOME/.cache" XDG_DATA_HOME="$STALL_HOME/.local/share" \ - PATH="$STALL_HOME/bin:$PATH" "$collector" --limits-only) - -help=$(jq -r '.authHelpText' <<<"$result") -[[ $help != "Run \`codex login\` to authenticate." ]] || - fail "Codex collector must not blame auth when the app-server is merely stalled" "$result" -[[ $help == "Codex app-server did not answer account/rateLimits/read" ]] || - fail "Codex collector names the stalled RPC method clearly" "$result" -[[ $help != "account/rateLimits/read" && $help != "initialize" ]] || - fail "Codex collector must not leak a bare method name" "$result" -pass "Codex collector names a stalled RPC instead of leaking the method name" - -# A stalled app-server that has logged to stderr is still running: its logging -# is not why it stopped, and it has not exited. -NOISY_HOME=$(signed_in_home) -sed 's/^while read/echo "WARN codex_core: startup notice" >\&2\nwhile read/' "$STALL_HOME/bin/codex" >"$NOISY_HOME/bin/codex" -chmod +x "$NOISY_HOME/bin/codex" - -result=$(HOME="$NOISY_HOME" CODEX_HOME="$NOISY_HOME/.codex" XDG_CACHE_HOME="$NOISY_HOME/.cache" XDG_DATA_HOME="$NOISY_HOME/.local/share" \ - PATH="$NOISY_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.authHelpText' <<<"$result") == "Codex app-server did not answer account/rateLimits/read" ]] || - fail "Codex collector reports a stall, not an exit, when a live app-server has logged" "$result" -pass "Codex collector reports a stall, not an exit, when a live app-server has logged" - -# A CLI that logs plenty before failing must still show the failure, not the logging. -CHATTY_HOME=$(signed_in_home) -cat >"$CHATTY_HOME/bin/codex" <<'EOF' -#!/bin/bash -for i in {1..20}; do echo "WARN codex_core::config: ignoring unknown key number $i" >&2; done -echo "error: failed to start app-server" >&2 -exit 1 -EOF -chmod +x "$CHATTY_HOME/bin/codex" - -result=$(HOME="$CHATTY_HOME" CODEX_HOME="$CHATTY_HOME/.codex" XDG_CACHE_HOME="$CHATTY_HOME/.cache" XDG_DATA_HOME="$CHATTY_HOME/.local/share" \ - PATH="$CHATTY_HOME/bin:$PATH" "$collector" --limits-only) - -[[ $(jq -r '.authHelpText' <<<"$result") == "codex app-server exited: "*"error: failed to start app-server" ]] || - fail "Codex collector keeps the CLI's final error past its startup logging" "$result" -pass "Codex collector keeps the CLI's final error past its startup logging" - -# Codex CLI can front any OpenAI-compatible backend (`--oss`, or a custom -# model_provider in config.toml). Those rollouts land in the same sessions -# directory but spend a local box or a third party, never this subscription. -PROVIDER_HOME=$(mktemp -d "$SCRATCH/home.XXXXXX") -mkdir -p "$PROVIDER_HOME/bin" "$PROVIDER_HOME/.codex/sessions/$(date +%Y/%m/%d)" -cp "$TEST_HOME/bin/codex" "$PROVIDER_HOME/bin/codex" - -provider_session() { - local name=$1 meta=$2 model=$3 input=$4 output=$5 - local file="$PROVIDER_HOME/.codex/sessions/$(date +%Y/%m/%d)/rollout-$name.jsonl" - [[ -n $meta ]] && echo "{\"timestamp\":\"$timestamp\",\"type\":\"session_meta\",\"payload\":{\"model_provider\":\"$meta\"}}" >"$file" - cat >>"$file" <"$session" <"$incremental_session" <"$incremental_pi" <>"{home()}/mise-calls"\n' + echo) + + +def mise_calls(): + path = home() / "mise-calls" + return path.read_text().splitlines() if path.exists() else [] + + +new_home(None, auth="empty", path=TOOLS) +stub(home() / ".local" / "bin" / "codex", LAUNCHER) +mise() +record = codex("--limits-only") +expect("Codex never runs the lazy launcher, asking only `mise which`, and says its limits are unavailable", + [record["usageStatusText"], mise_calls() in ([], ["which codex"])], ["Codex limits unavailable", True], record) + +shutil.copy(FAKE_CODEX, home() / "real-codex") +mise(home() / "real-codex") +(home() / "mise-calls").unlink(missing_ok=True) +codex("--limits-only", env={"CODEX_ARGS_FILE": home() / "codex-args"}) +expect("Codex probes the binary `mise which` names", [(home() / "codex-args").exists(), mise_calls()], [True, ["which codex"]]) + +new_home(None, auth="empty", path=TOOLS) +(home() / ".local" / "bin").mkdir(parents=True) +(home() / ".local" / "bin" / "codex").symlink_to(FAKE_CODEX) +mise(home() / "real-codex") +codex("--limits-only", env={"CODEX_ARGS_FILE": home() / "codex-args"}) +expect("Codex probes a symlink at the launcher's path without asking mise", [(home() / "codex-args").exists(), mise_calls()], [True, []]) + +new_home(None, auth=False, path=TOOLS) +shim = mise() +(home() / ".local" / "share" / "mise" / "shims").mkdir(parents=True) +(home() / ".local" / "share" / "mise" / "shims" / "codex").symlink_to(shim) +record = codex("--limits-only") +expect("Codex never runs a mise shim, and says its limits are unavailable", + [record["usageStatusText"], mise_calls() in ([], ["which codex"])], ["Codex limits unavailable", True], record) + +# Without credentials account/read can only fail, and starting the app-server +# isn't free (it syncs the plugin list, a git fetch each time), so it isn't +# started. An API key alone doesn't sign it in, but a CODEX_ACCESS_TOKEN does, +# and so does a store outside auth.json, read as TOML; an explicit file store +# with no auth.json is still signed out. +new_home(auth=False) + + +def probed(config=None, **env): + config_file = home() / ".codex" / "config.toml" + if config is None: + config_file.unlink(missing_ok=True) + else: + config_file.write_text(config) + (home() / "codex-args").unlink(missing_ok=True) + record = codex("--limits-only", env={"CODEX_ARGS_FILE": home() / "codex-args", **env}) + return (home() / "codex-args").exists(), record + + +started_bare, record = probed() +expect("Codex waits for a sign-in without starting the app-server when there are no credentials", + [started_bare, record["usageStatusText"]], [False, "Waiting for auth"], record) +expect("Codex starts the app-server for a token or a store outside auth.json, and not for an API key or an empty file store", { + "apiKey": probed(OPENAI_API_KEY="x")[0], + "token": probed(CODEX_ACCESS_TOKEN="x")[0], + "keyring": probed('cli_auth_credentials_store = "keyring"\n')[0], + "singleQuoted": probed("model = \"gpt-5\"\ncli_auth_credentials_store = 'keyring'\n")[0], + "fileStore": probed('model = "gpt-5"\ncli_auth_credentials_store = "file"\n')[0], +}, {"apiKey": False, "token": True, "keyring": True, "singleQuoted": True, "fileStore": False}) + +# An app-server that fails says why in the CLI's own words, never as a bare +# RPC method name: one rejecting its arguments, one closing its output during +# startup, one closing its input after initialize, one saying nothing (the +# login hint), one stalled (named as a stall, its logging no reason), and one +# logging plenty before its final error. +shebang, fake_body = FAKE_CODEX.read_text().split("\n", 1) +noisy = stub(Path(os.environ["SANDBOX"]) / "noisy-codex", f"{shebang}\nimport sys\nprint('WARN codex_core: startup notice', file=sys.stderr, flush=True)\n{fake_body}") +failures = { + "rejected": ("echo \"error: invalid value 'untrusted' for '--ask-for-approval '\" >&2\necho ' [possible values: on-request, never]' >&2\nexit 2\n", {}), + "dead": ("exec 1>&-\nexec sleep 30\n", {}), + "half": ("read -r request\nexec 0<&-\njq -cn --argjson id \"$(jq -r '.id' <<<\"$request\")\" '{id: $id, result: {}}'\nexec sleep 30\n", {}), + "silent": (FAKE_CODEX, {"CODEX_FAKE": "silent"}), + "stalled": (FAKE_CODEX, {"CODEX_FAKE": "stall"}), + "noisy": (noisy, {"CODEX_FAKE": "stall"}), + "chatty": ("for i in {1..20}; do echo \"WARN codex_core::config: ignoring unknown key number $i\" >&2; done\necho 'error: failed to start app-server' >&2\nexit 1\n", {}), } - -result=$(run_incremental) -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "20" ]] || - fail "Codex collector counts both session sources on a first scan" "$result" - -file_cache=$(ls "$INCREMENTAL_HOME/.cache/omarchy/agent-usage/"codex-files-*.json 2>/dev/null | head -n 1) -[[ -n $file_cache && -s $file_cache ]] || - fail "Codex collector writes a per-file cache on first scan" -[[ $(jq --arg path "$incremental_session" -r '.files[$path].days | length' "$file_cache") == "1" ]] || - fail "Codex collector records the native session file it read" "$(cat "$file_cache")" -# A forked pi session repeats its parent's messages, so pi sessions are only -# deduplicated across the whole tree at once, never replayed per file. -[[ $(jq --arg path "$incremental_pi" -r '.files | has($path)' "$file_cache") == "false" ]] || - fail "Codex collector keeps pi sessions out of the per-file cache" "$(cat "$file_cache")" -pass "Codex collector records per-file totals as it scans" - -# Unreadable, but unchanged in mtime and size: a run that still reports the -# same totals can only have replayed them. -chmod 000 "$incremental_session" -expire_scan_cache -result=$(run_incremental) -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "20" ]] || - fail "Codex collector rereads session files it has already counted" "$result" -pass "Codex collector replays unchanged session files instead of rereading them" - -# --force means the history itself is re-read, per-file records included. -result=$(run_incremental --force) -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "5" ]] || - fail "Codex collector --force reuses per-file records" "$result" -pass "Codex collector --force rereads the history" - -chmod 644 "$incremental_session" -# The --force run above could not read the session, so cache it again first, -# or the append below is read from scratch whether or not the cache noticed it. -expire_scan_cache -run_incremental >/dev/null -cat >>"$incremental_session" <"$zone_session" </dev/null -expire_scan_cache -TZ=America/Los_Angeles run_incremental >/dev/null -file_cache=$(ls "$INCREMENTAL_HOME/.cache/omarchy/agent-usage/"codex-files-*.json | head -n 1) -[[ $(jq --arg path "$zone_session" -r '.files[$path].days | keys[0]' "$file_cache") == "$(TZ=America/Los_Angeles date -d "$zone_timestamp" +%Y-%m-%d)" ]] || - fail "Codex collector keeps the old timezone's days after a timezone change" "$(cat "$file_cache")" -pass "Codex collector re-reads per-file records after a timezone change" - -# A line that stops the read keeps the usage read before it, as it always has. -broken_session="$INCREMENTAL_HOME/.codex/sessions/broken.jsonl" -cat >"$broken_session" <>{gh_log}\n[[ $* == "auth token --hostname github.com" ]] && echo gh-token\n'), ("broken-gh", "exit 1\n")): + stub = home() / name / "gh" + stub.parent.mkdir() + stub.write_text("#!/bin/bash\n" + body) + stub.chmod(0o755) +hosts = config_home() / "gh" / "hosts.yml" +hosts.parent.mkdir(parents=True) +hosts.write_text("github.com:\n users:\n tester:\n git_protocol: ssh\n user: tester\n") +GH_ROUTES = {QUOTA_URL: quota("individual", 300, 30)} +os.environ["PATH"] = f"{home() / 'gh'}:{os.environ['PATH']}" + + +def gh_runs(): + return len(gh_log.read_text().splitlines()) if gh_log.exists() else 0 + + +os.environ["GH_HOST"] = "github.example.com" +record = copilot(routes=GH_ROUTES) +del os.environ["GH_HOST"] +expect("the quota is asked as gh's github.com account", [meter(record), gh_runs(), copilot.asked], + [{"limits": [["Monthly allowance", 0.1, True]], "stale": False, "status": "", "retry": False, "tier": "Pro"}, 1, ["Bearer gh-token"]], record) +copilot(routes=GH_ROUTES, force=False) +copilot("--limits-only", routes=GH_ROUTES) +expect("a reused answer runs no gh", [copilot.asked, gh_runs()], [[], 1]) + +# Another gh account has nothing kept; how it failed decides what it says. +hosts.write_text("github.com:\n user: someone-else\n") +failed = {mode: copilot(routes=GH_ROUTES, mode=mode) for mode in ("down", 500, 401)} +expect("with nothing kept, a failure says what failed", + [[record["usageStatusText"], record.get("retryAdvised", False), record["authHelpText"] != ""] for record in failed.values()], + [["Copilot limits unavailable", True, True], ["Copilot limits unavailable", False, True], ["Sign-in expired", False, True]], failed) +expect("an error status is the detail, not a sign-in hint", + ["status 500" in failed[500]["authHelpText"], "gh auth login" in failed[500]["authHelpText"]], [True, False], failed[500]) + +# gh still signed in to the first account but failing to hand over its token +# keeps that account's last allowance, stale. +hosts.write_text("github.com:\n user: tester\n") +real_path = os.environ["PATH"] +os.environ["PATH"] = f"{home() / 'broken-gh'}:{real_path}" +record = copilot(routes=GH_ROUTES) +expect("gh failing to hand over its token keeps the last allowance, stale", + [meter(record), "gh auth status" in record["authHelpText"], copilot.asked], + [{"limits": [["Monthly allowance", 0.1, True]], "stale": True, "status": "", "retry": False, "tier": "Pro"}, True, []], record) + +# No token anywhere, with the CLI in use: waiting for a sign-in. +os.environ["PATH"] = str(home() / "no-gh") +record = copilot(routes=GH_ROUTES) +expect("no GitHub token anywhere waits for a sign-in", [record["usageStatusText"], "gh auth login" in record["authHelpText"]], ["Waiting for auth", True], record) +os.environ["PATH"] = real_path + +# An answer under the old copilot-quota-* name is adopted once. +os.environ["COPILOT_QUOTA_TOKEN"] = "legacy-token" +digest = hashlib.sha256(b"legacy-token").hexdigest()[:16] +legacy = cache_root() / f"copilot-quota-{digest}.json" +write_json(legacy, {"available": True, "entitlement": 100, "used": 40, "resetsAt": RESET.isoformat(), "planLabel": "Pro", "fetchedAtMs": 1000}) +record = copilot(routes=GH_ROUTES, mode="down") +expect("an answer cached under the old name is adopted", + [meter(record)["limits"], record["limitsStale"], legacy.exists(), (cache_root() / f"copilot-limits-{digest}.json").exists()], + [[["Monthly allowance", 0.4, True]], True, False, True], record) +del os.environ["COPILOT_QUOTA_TOKEN"] + +# ------------------------------------------------------------ the token + +# An exported token comes ahead of gh: COPILOT_QUOTA_TOKEN, then GH_TOKEN, +# then GITHUB_TOKEN, each asked with as it is. +sandbox() +copilot_home().mkdir() +configured(monthlyCredits=3000) +asked, gh_before = [], gh_runs() +for name, token in (("GITHUB_TOKEN", "github-token"), ("GH_TOKEN", "gh-env-token"), ("COPILOT_QUOTA_TOKEN", "quota-token")): + os.environ[name] = token + copilot(routes=GH_ROUTES) + asked.append(copilot.asked) +expect("the quota token comes from COPILOT_QUOTA_TOKEN, GH_TOKEN, then GITHUB_TOKEN, ahead of gh", + [asked, gh_runs() - gh_before], [[["Bearer github-token"], ["Bearer gh-env-token"], ["Bearer quota-token"]], 0]) + +# "remote": true asks GitHub even where the Copilot CLI never ran. +configured(remote=True) +record = copilot("--force", "--copilot-home", str(home() / "never-ran"), routes=GH_ROUTES) +expect("remote: true asks GitHub without a Copilot CLI home", [copilot.asked, meter(record)["limits"]], [["Bearer quota-token"], [["Monthly allowance", 0.1, True]]], record) + +# An allowance missing an amount, or with one that can't be read, is a failed +# check that keeps the last allowance, stale; an explicitly unlimited one +# needs no amounts. +def snapshot_without(*keys, **changes): + answer = quota("individual", 300, 30) + snapshot = answer["quota_snapshots"]["premium_interactions"] + for key in keys: + del snapshot[key] + snapshot.update(changes) + return {QUOTA_URL: answer} + + +incomplete = { + "no entitlement": snapshot_without("entitlement"), + "no credits_used": snapshot_without("credits_used"), + "unreadable entitlement": snapshot_without(entitlement="lots"), +} +kept = {name: meter(copilot(routes=routes)) for name, routes in incomplete.items()} +expect("an incomplete allowance keeps the last one, stale", + kept, {name: {"limits": [["Monthly allowance", 0.1, True]], "stale": True, "status": "", "retry": False, "tier": "Pro"} for name in incomplete}, kept) +record = copilot(routes=snapshot_without("entitlement", "credits_used", unlimited=True)) +expect("an unlimited allowance needs no amounts", meter(record), + {"limits": [], "stale": False, "status": "", "retry": False, "tier": "Pro"}, record) + +# GitHub's own snapshot, as the CLI's CopilotUserResponseQuotaSnapshots types +# it, carries no credits_used: what is used is the entitlement less what +# remains, read from remaining, quota_remaining or percent_remaining. +def upstream(**remaining): + answer = quota("individual", 300, 0, overage_count=0, overage_permitted=False, quota_id="premium_interactions", + timestamp_utc=NOW, has_quota=True, token_based_billing=True, **remaining) + del answer["quota_snapshots"]["premium_interactions"]["credits_used"] + return {QUOTA_URL: answer} + + +shapes = { + "remaining": upstream(remaining=225, quota_remaining=225, percent_remaining=75), + "quota_remaining": upstream(quota_remaining=225), + "percent_remaining": upstream(percent_remaining=75), +} +read = {name: meter(copilot(routes=routes)) for name, routes in shapes.items()} +expect("an upstream allowance is used by what it has left", + read, {name: {"limits": [["Monthly allowance", 0.25, True]], "stale": False, "status": "", "retry": False, "tier": "Pro"} for name in shapes}, read) + +# A reset date that can't be read is no reset date, both live and once kept. +damaged_reset = quota("individual", 300, 30) +damaged_reset["quota_reset_date_utc"] = "\u00b2" +live = meter(copilot(routes={QUOTA_URL: damaged_reset})) +kept = meter(copilot(routes={QUOTA_URL: damaged_reset}, mode=500)) +expect("an unreadable reset date reads as none, live and kept", + [live["limits"], kept["limits"], kept["stale"]], [[["Monthly allowance", 0.1, False]], [["Monthly allowance", 0.1, False]], True], [live, kept]) + +for name in ("COPILOT_QUOTA_TOKEN", "GH_TOKEN", "GITHUB_TOKEN"): + del os.environ[name] +# ---------------------------------------------------- the current store + +# A current CLI's store, its schema as Copilot CLI 1.0.102 writes it: who +# started each request (initiator "user" for a typed prompt, "agent" for the +# agent's own follow-up, as a tool loop's next step), and GitHub's usage +# report on each request it served. +def current_store(rows): + path = copilot_home() / "session-store.db" + path.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(path) + conn.execute("CREATE TABLE IF NOT EXISTS assistant_usage_events (id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, turn_index INTEGER, agent_id TEXT," + " parent_tool_call_id TEXT, model TEXT NOT NULL, copilot_usage_model TEXT, input_tokens INTEGER, output_tokens INTEGER, cache_read_tokens INTEGER," + " cache_write_tokens INTEGER, reasoning_tokens INTEGER, total_nano_aiu INTEGER, request_multiplier REAL, duration_ms INTEGER," + " time_to_first_token_ms INTEGER, output_ttft_ms REAL, inter_token_latency_ms INTEGER, initiator TEXT, api_endpoint TEXT, reasoning_effort TEXT," + " finish_reason TEXT, content_filter_triggered INTEGER, token_details_json TEXT, created_at TEXT DEFAULT (datetime('now')))") + conn.executemany("INSERT INTO assistant_usage_events (session_id, turn_index, model, copilot_usage_model, input_tokens, output_tokens, cache_read_tokens," + " cache_write_tokens, reasoning_tokens, total_nano_aiu, request_multiplier, initiator, api_endpoint, finish_reason, created_at)" + " VALUES (?, ?, ?, ?, ?, ?, ?, 0, 5, ?, ?, ?, '/chat/completions', ?, ?)", rows) + conn.commit() + conn.close() + + +sandbox() +configured(remote=False) +current_store([ + ("sess-a", 0, "gpt-5.2", "gpt-5.2", 1000, 100, 100, 1_000_000_000, 1.0, "user", "stop", NOW), + ("sess-b", 0, "gpt-5.2", "gpt-5.2", 1000, 100, 100, 1_000_000_000, 1.0, "user", "tool_calls", NOW), + ("sess-b", 0, "gpt-5.2", "gpt-5.2", 1000, 100, 100, 0, 1.0, "agent", "stop", NOW), + ("sess-a", 1, "gpt-5.2", "gpt-5.2", 1000, 100, 100, 1_000_000_000, 1.0, "user", "stop", NOW), +]) +record = copilot() +expect("a current store counts the prompts typed, not the agent's follow-up calls, and shows them", + {key: record[key] for key in ("ready", "hasPromptStats", "todayPrompts", "totalPrompts", "todaySessions", "todayTotalTokens")}, + {"ready": True, "hasPromptStats": True, "todayPrompts": 3, "totalPrompts": 3, "todaySessions": 2, "todayTotalTokens": 4400}, record) + +# A request sent to a custom provider (BYOK, COPILOT_PROVIDER_BASE_URL) has +# no GitHub usage report: GitHub billed nothing, so it isn't Copilot's. +current_store([ + ("sess-byok", 0, "mock-model", None, 21372, 251, 271, None, 0.0, "user", "stop", NOW), + ("sess-byok", 1, "mock-model", None, 21675, 254, 274, None, 0.0, "user", "stop", NOW), +]) +record = copilot() +expect("a request to a custom provider isn't counted as Copilot's", + {key: record[key] for key in ("todayPrompts", "todaySessions", "todayTotalTokens")} | {"models": sorted(record["modelUsage"])}, + {"todayPrompts": 3, "todaySessions": 2, "todayTotalTokens": 4400, "models": ["gpt-5.2"]}, record) + +# A machine where Copilot only ever ran against a custom provider has nothing +# of Copilot's to show, and its transcripts aren't read in the store's place. +sandbox() +configured(remote=False) +current_store([("sess-byok", 0, "mock-model", None, 21372, 251, 271, None, 0.0, "user", "stop", NOW)]) +transcript("sess-byok", [("user.message", {"content": "Say hi"}, NOW), + ("session.shutdown", {"modelMetrics": {"mock-model": {"usage": {"inputTokens": 21372, "outputTokens": 251}}}}, NOW)]) +record = copilot() +expect("a store of custom-provider requests alone leaves Copilot empty", + [record["ready"], record["hasLocalStats"], record["totalPrompts"], record["todayTotalTokens"]], [False, False, 0, 0], record) +PY diff --git a/test/shell.d/agent-usage-cursor-test.sh b/test/shell.d/agent-usage-cursor-test.sh new file mode 100755 index 00000000000..d4a5118e012 --- /dev/null +++ b/test/shell.d/agent-usage-cursor-test.sh @@ -0,0 +1,299 @@ +#!/bin/bash + +# The Cursor collector's own sources: the dashboard RPCs the Cursor CLI calls +# (the period meters, the plan, and a week of daily token aggregations), the +# per-account token scan carried past midnight and through failed checks, and +# the editor or CLI being installed at all. The unused-machine record is held +# centrally by the partition test, the Cursor sign-in order and the editor's +# untouched database by the shared test, and the limits and scan caches by +# the contract checks. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +# Day windows resolve in local time, so pin the zone or the fixtures below +# would shift by a day depending on where the test runs. +export TZ=UTC + +python3 - <<'PY' +from agent_usage_fixtures import * +import collectors +from datetime import date + +RPC = "api2.cursor.sh/aiserver.v1.DashboardService/" +TODAY = datetime.now().astimezone().date() +YESTERDAY = TODAY - timedelta(days=1) + +# Spend well past the included limit, the way a plan with provider bonus usage +# reports it: 609.86 spent against a 400.00 allowance, while Cursor's own +# meters say 18.76% / 11.81% / 100%. +PERIOD = { + "billingCycleEnd": "4092099926000", + "planUsage": {"totalSpend": 60986, "includedSpend": 40000, "limit": 40000, "autoPercentUsed": 11.81, "apiPercentUsed": 100, "totalPercentUsed": 18.76}, + "spendLimitUsage": {"individualLimit": "1000", "individualUsed": "951", "limitType": "user"}, +} +AGGREGATIONS = { + TODAY.isoformat(): [{"modelIntent": "default", "inputTokens": "10", "outputTokens": "5", "cacheReadTokens": "100", "cacheWriteTokens": "1", "totalCents": None}], + YESTERDAY.isoformat(): [ + {"modelIntent": "default", "inputTokens": "2", "outputTokens": "3", "cacheReadTokens": "4", "cacheWriteTokens": "0"}, + {"modelIntent": "kimi-k3-max", "inputTokens": "1", "outputTokens": "1", "cacheReadTokens": "0", "cacheWriteTokens": "0"}, + ], +} +grown = {"tokens": 0} + + +def window_start(data): + return datetime.fromtimestamp(int(json.loads(data)["startDate"]) / 1000).astimezone() + + +# A day's aggregations, plus whatever the scan contract has grown today by. +def aggregations(request): + day = window_start(request.data).date() + rows = list(AGGREGATIONS.get(day.isoformat(), [])) + if day == TODAY and grown["tokens"]: + rows.append({"modelIntent": "grown", "inputTokens": str(grown["tokens"])}) + return {"aggregations": rows} + + +ROUTES = { + RPC + "GetCurrentPeriodUsage": PERIOD, + RPC + "GetPlanInfo": {"planInfo": {"planName": "Ultra"}}, + RPC + "GetAggregatedUsageEvents": aggregations, +} + + +def cli_signed_in(user="u-main", expires_in=3600): + write_json(config_home() / "cursor" / "auth.json", {"accessToken": jwt({"sub": user, "exp": time.time() + expires_in})}) + + +def cursor(*args, routes=None, mode="up", **env): + saved = dict(os.environ) + os.environ.update(env) + try: + with FakeHttp(ROUTES if routes is None else routes, mode=mode) as http: + record = collect("cursor", *(args or ("--force",))) + finally: + os.environ.clear() + os.environ.update(saved) + cursor.requests = http.requests + return record + + +def meters(record): + return [{"label": limit["label"], "percent": round(limit["percent"] * 10000)} for limit in record["limits"]] + + +# --------------------------------------------------------------- the answer + +sandbox() +cli_signed_in() +record = cursor() +expect("the record is account-wide, with no prompt counts, and the plan from GetPlanInfo", + {key: record.get(key) for key in ("ready", "hasPromptStats", "scope", "tierLabel", "usageStatusText")}, + {"ready": True, "hasPromptStats": False, "scope": "account", "tierLabel": "Ultra", "usageStatusText": ""}, record) +expect("the meters are Cursor's own, not spend over the allowance", meters(record), [ + {"label": "Included total", "percent": 1876}, {"label": "Auto models", "percent": 1181}, + {"label": "Named models", "percent": 10000}, {"label": "On-demand $9.51 / $10.00", "percent": 9510}, +], record) +expect("the reset is the billing cycle's end", record["limits"][0]["resetsAt"], "2099-09-03T06:25:26+00:00", record) +starts = [window_start(request["data"]) for request in cursor.requests if request["url"].endswith("GetAggregatedUsageEvents")] +expect("tokens are asked for in seven daily windows, each from local midnight", + [len(starts), {(start.hour, start.minute, start.second) for start in starts}], [7, {(0, 0, 0)}]) +expect("the week's token series is built from the windows", [day["messageCount"] for day in record["recentDays"]], [0, 0, 0, 0, 0, 11, 116], record) +expect("Cursor's token names map onto the record's", record["modelUsage"], { + "default": {"inputTokens": 12, "outputTokens": 8, "cacheReadInputTokens": 104, "cacheCreationInputTokens": 1}, + "kimi-k3-max": {"inputTokens": 1, "outputTokens": 1, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, +}, record) +expect("today stands apart from the rest of the week", + [record["todayTotalTokens"], record["todayTokensByModel"], record["activeDays"]], [116, {"default": 116}, 2], record) + +# A failed check keeps the week of tokens on disk along with the limits. +down = cursor(mode="down") +expect("a failed check keeps the last week of tokens and the plan", + [down["recentDays"] == record["recentDays"], down["tierLabel"], down["limitsStale"]], [True, "Ultra", True], down) + +# Another account's scan never stands in: signed in as someone else with the +# dashboard down, there are no tokens to show. +other = cursor(mode="down", CURSOR_API_KEY="someone_else") +expect("another account's numbers are never carried", + {"limits": other["limits"], "tierLabel": other["tierLabel"], "today": other["todayTotalTokens"], + "counts": [day["messageCount"] for day in other["recentDays"]], "last": other["recentDays"][6]["date"], "modelUsage": other["modelUsage"]}, + {"limits": [], "tierLabel": "", "today": 0, "counts": [0] * 7, "last": TODAY.isoformat(), "modelUsage": {}}, other) + +# A kept limit whose billing cycle has ended comes back at zero, and an +# on-demand meter among them is renamed for its new spend. +module = load_collector("cursor") +key = collectors.account_key(collectors.cursor_token()) +limits_path = module.limits_cache(key).path +good = json.loads(limits_path.read_text()) +past = iso_in(-3600) +write_json(limits_path, dict(good, limits=[dict(good["limits"][0], resetsAt=past)] + good["limits"][1:3] + [dict(good["limits"][3], resetsAt=past)])) +lapsed = cursor(mode="down") +write_json(limits_path, good) +expect("a kept limit whose cycle ended resets, the rest stand", + [[limit["label"], round(limit["percent"] * 100), limit["resetsAt"]] for limit in lapsed["limits"][:2]], + [["Included total", 0, ""], ["Auto models", 12, "2099-09-03T06:25:26+00:00"]], lapsed) +expect("an on-demand meter that reset is renamed with its new spend", + {name: lapsed["limits"][3][name] for name in ("label", "title", "percent", "resetsAt")}, + {"label": "On-demand $0.00 / $10.00", "title": "On-demand $0.00 / $10.00 Monthly", "percent": 0.0, "resetsAt": ""}, lapsed) + +# A scan from yesterday moves onto this week's dates, and its today figures +# are cleared, since none of it happened today. Each day's usage rides along, +# so the day that falls out of the week leaves the model totals and active +# days as well as the chart. +stats_path = module.stats_cache(key) +scan = json.loads(stats_path.read_text()) +yesterday_scan = dict(scan, scanDate=YESTERDAY.isoformat()) +yesterday_scan["stats"] = dict(scan["stats"], recentDays=[{"date": (TODAY - timedelta(days=offset)).isoformat(), "messageCount": offset} for offset in range(7, 0, -1)]) +yesterday_scan["stats"]["days"] = [[(TODAY - timedelta(days=offset)).isoformat(), {"carried-model": {"inputTokens": offset}}] for offset in range(7, 0, -1)] +write_json(stats_path, yesterday_scan) +overnight = module.carried_stats(key, TODAY) +expect("a scan from yesterday moves onto today's week", + [overnight["todayTotalTokens"], overnight["todayTokensByModel"], [day["messageCount"] for day in overnight["recentDays"]], overnight["recentDays"][6]["date"]], + [0, {}, [6, 5, 4, 3, 2, 1, 0], TODAY.isoformat()], overnight) +expect("the day that left the week leaves the model totals", + [overnight["modelUsage"]["carried-model"]["inputTokens"], overnight["activeDays"], len(overnight["activeDates"])], [21, 6, 6], overnight) +# A scan cached before days were kept still moves onto the new week. +write_json(stats_path, dict(yesterday_scan, stats={k: v for k, v in yesterday_scan["stats"].items() if k != "days"})) +overnight = module.carried_stats(key, TODAY) +expect("an older cached scan moves onto today's week too", + [overnight["todayTotalTokens"], [day["messageCount"] for day in overnight["recentDays"]]], [0, [6, 5, 4, 3, 2, 1, 0]], overnight) + +# Days that spring forward or fall back are 23 or 25 hours long. +saved_tz = os.environ.get("TZ") +os.environ["TZ"] = "America/New_York" +time.tzset() +try: + spring = module.day_windows(date(2026, 3, 9), 2) + fall = module.day_windows(date(2026, 11, 1), 1) +finally: + if saved_tz is None: + os.environ.pop("TZ", None) + else: + os.environ["TZ"] = saved_tz + time.tzset() +expect("daily windows are sized across DST", + [[(end - start) // 3_600_000 for _, start, end in spring + fall], spring[0][2] == spring[1][1]], [[23, 24, 25], True]) + +# ------------------------------------------------------------- the sign-in + +# Cursor installed but signed in nowhere: the agent is in use, so it says it +# is waiting for a sign-in. +sandbox() +(config_home() / "cursor").mkdir(parents=True) +record = cursor() +expect("Cursor installed but signed in nowhere waits for a sign-in", + [record["ready"], record["usageStatusText"], record["limits"], record["authHelpText"] != ""], [True, "Waiting for auth", [], True], record) + +# A lapsed sign-in isn't sent: the CLI or editor renews it when it runs. +sandbox() +cli_signed_in() +cursor() +cli_signed_in(expires_in=-60) +record = cursor() +expect("a lapsed sign-in pauses the limits without asking", + [record["usageStatusText"], record["limitsStale"], len(record["limits"]), cursor.requests], ["Limits paused", True, 4, []], record) + +# ------------------------------------------------------------ the endpoint + +# CURSOR_API_ENDPOINT moves the endpoint, but the sign-in only ever goes over +# HTTPS, or plain HTTP to loopback. +sandbox() +cli_signed_in() +for endpoint, sent in (("http://cursor.example.com", False), ("ftp://cursor.example.com", False), ("https://user:pw@cursor.example.com", False), + ("https://cursor.example.com", True), ("http://127.0.0.1:9", True)): + with FakeHttp({"/aiserver.v1.DashboardService/": PERIOD}) as http: + os.environ["CURSOR_API_ENDPOINT"] = endpoint + try: + record = collect("cursor", "--force") + finally: + del os.environ["CURSOR_API_ENDPOINT"] + expect(f"the Cursor sign-in {'goes' if sent else 'never goes'} to {endpoint}", bool(http.requests), sent, record) + if not sent: + expect(f"{endpoint} reads as unavailable and says why", [record["usageStatusText"], "HTTPS" in record["authHelpText"]], ["Cursor limits unavailable", True], record) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def grow(): + grown["tokens"] += 1 + + +# A probe asks for the period and the plan, and a forced one scans the +# week's seven days too. +adapter = Adapter("cursor", plant=cli_signed_in, answer=ROUTES, requests_per_probe=9, grow=grow) +check_limits_contract(adapter) +check_scan_cache(adapter) +grown["tokens"] = 0 + +# Without a billing cycle end there is no reset time, as "" rather than null, +# and every meter is titled after the billing period. +sandbox() +cli_signed_in() +routes = dict(ROUTES) +routes[RPC + "GetCurrentPeriodUsage"] = {"planUsage": {"totalPercentUsed": 40, "autoPercentUsed": 20, "apiPercentUsed": 60}, + "spendLimitUsage": {"individualLimit": "2000", "individualUsed": "500"}} +record = cursor(routes=routes) +expect("the meters are titled after the billing period", + [limit["title"] for limit in record["limits"]], ["Monthly", "Auto models Monthly", "Named models Monthly", "On-demand $5.00 / $20.00 Monthly"], record) +expect("an unknown reset is an empty string", sorted({limit["resetsAt"] for limit in record["limits"]}), [""], record) + +# A plan whose included allowance is weekly, as GetPlanInfoResponse.PlanInfo +# says with included_usage_period and included_usage_resets_at, titles and +# dates its included meters by the week; the on-demand budget stays on the +# billing cycle. +weekly_routes = dict(routes, **{ + RPC + "GetCurrentPeriodUsage": dict(PERIOD, planUsage={"totalPercentUsed": 40, "autoPercentUsed": 20, "apiPercentUsed": 60}), +}) +read = {} +for period, reset in (("INCLUDED_USAGE_PERIOD_WEEKLY", "4091840726000"), ("WEEKLY", "4091840726"), (2, 4091840726000)): + weekly_routes[RPC + "GetPlanInfo"] = {"planInfo": {"planName": "Ultra", "includedUsagePeriod": period, "includedUsageResetsAt": reset}} + record = cursor(routes=weekly_routes) + read[str(period)] = [[limit["title"], limit["resetsAt"]] for limit in record["limits"]] +WEEK_END, CYCLE_END = "2099-08-31T06:25:26+00:00", "2099-09-03T06:25:26+00:00" +expect("a weekly included allowance reads Weekly and resets with the week", read, { + period: [["Weekly", WEEK_END], ["Auto models Weekly", WEEK_END], ["Named models Weekly", WEEK_END], ["On-demand $9.51 / $10.00 Monthly", CYCLE_END]] + for period in ("INCLUDED_USAGE_PERIOD_WEEKLY", "WEEKLY", "2") +}, read) + +# A weekly plan whose GetPlanInfo fails can't say its week's reset, so the +# known weekly limits are kept, stale, rather than read as Monthly. +record = cursor(routes=dict(weekly_routes, **{RPC + "GetPlanInfo": 500})) +expect("a weekly plan that can't be checked keeps its weekly limits, stale", + [[[limit["title"], limit["resetsAt"]] for limit in record["limits"]], record["limitsStale"], record["tierLabel"]], + [read["2"], True, "Ultra"], record) + +# GetPlanInfo failing leaves the plan the account last had. +cursor(routes=routes) +record = cursor(routes=dict(routes, **{RPC + "GetPlanInfo": 500})) +expect("the plan is kept when GetPlanInfo fails", [record["tierLabel"], record["limitsStale"]], ["Ultra", False], record) + +# A failure that isn't Cursor's error keeps the limits and the tokens too. +with contextlib.redirect_stderr(io.StringIO()): + record = cursor(routes=dict(routes, **{RPC + "GetCurrentPeriodUsage": RuntimeError("boom")})) +expect("an unexpected failure keeps the limits and the tokens", + [len(record["limits"]), record["limitsStale"], record["usageStatusText"], record["todayTotalTokens"], record.get("retryAdvised", False)], + [4, True, "", 116, False], record) + +# A usage section of the wrong type is a reply this doesn't read: a failed +# check that keeps the known limits, stale, rather than no meters. A plan +# with nothing in it, its sections left out or null, is a real answer. +wrong = {name: cursor(routes=dict(routes, **{RPC + "GetCurrentPeriodUsage": dict(routes[RPC + "GetCurrentPeriodUsage"], **{name: []})})) + for name in ("planUsage", "spendLimitUsage")} +expect("a wrongly typed usage section keeps the known limits, stale", + {name: [len(record["limits"]), record["limitsStale"], record["authHelpText"] != ""] for name, record in wrong.items()}, + {"planUsage": [4, True, True], "spendLimitUsage": [4, True, True]}, wrong) +empty = [cursor(routes=dict(routes, **{RPC + "GetCurrentPeriodUsage": period})) for period in ({}, {"planUsage": None, "spendLimitUsage": None})] +expect("an empty plan is an answer with no meters", + [[record["limits"], record["limitsStale"], record["usageStatusText"]] for record in empty], [[[], False, ""]] * 2, empty) + +# Untrusted numbers never break the record. +expect("untrusted numbers read as nothing", + [module.fraction("nan"), module.fraction("inf"), module.fraction(True), module.number("1e400"), module.iso("9" * 30)], [0.0, 0.0, 0.0, 0, ""]) +PY diff --git a/test/shell.d/agent-usage-deepseek-balance-test.sh b/test/shell.d/agent-usage-deepseek-balance-test.sh new file mode 100755 index 00000000000..0ec0a434f54 --- /dev/null +++ b/test/shell.d/agent-usage-deepseek-balance-test.sh @@ -0,0 +1,207 @@ +#!/bin/bash + +# The DeepSeek collector's prepaid balance: DeepSeek's balance answer, the +# estimated ledger a funded amount adds, and the balance kept past a failed +# check. A balance has no windows, so check_limits_contract (which holds +# limits) doesn't apply; the same contract is held for the balance here. The +# unused-machine record is held centrally by the partition test, the limits +# cache's mechanics by the shared test, and the scan cache by the scanner +# test. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +BALANCE = "api.deepseek.com/user/balance" +answer = {"total": "42.50", "currency": "CNY"} + + +def balance(request): + return {"is_available": True, "balance_infos": [ + {"currency": answer["currency"], "total_balance": answer["total"], "granted_balance": "2.50", "topped_up_balance": "40.00"}, + ]} + + +def deepseek(mode="up", args=("--force",)): + with FakeHttp({BALANCE: balance}, mode=mode) as http: + record = collect("deepseek", *args) + return record, http + + +sandbox() + +# No key, no request: local usage alone never reaches DeepSeek. +record, http = deepseek() +expect("DeepSeek collector asks nothing without a key", [http.count, "balance" in record], [0, False], record) + +os.environ["DEEPSEEK_API_KEY"] = "key-123" +record, http = deepseek() +check("DeepSeek collector reads the live balance, asking with the key", + record["ready"] is True and record["tierLabel"] == "Prepaid" and record["balance"]["remaining"] == 42.5 + and record["balance"]["currency"] == "CNY" and record["balance"]["estimated"] is False + and record["limitsStale"] is False and record["limitsFetchedAt"] > 0 and record["usageStatusText"] == "" + and [(r["url"], r["headers"].get("authorization")) for r in http.requests] == [("https://api.deepseek.com/user/balance", "Bearer key-123")], + record) +fetched_at = record["limitsFetchedAt"] + +# A panel that opens and shuts repeatedly must not ask every time, a +# --limits-only run included. +for args in ((), ("--limits-only",)): + record, http = deepseek(500, args) + check(f"DeepSeek collector reuses a balance it just read on a {'--limits-only' if args else 'normal'} run", + http.count == 0 and record["balance"]["remaining"] == 42.5 and record["limitsStale"] is False and record["limitsFetchedAt"] == fetched_at, + record) + +answer["total"] = "40.00" +record, http = deepseek() +expect("DeepSeek collector --force asks past the reuse window", [http.count, record["balance"]["remaining"]], [1, 40.0], record) + +# The estimated ledger comes from fundedAmount, applied to a kept balance too. +write_json(config_home() / "omarchy" / "agents" / "deepseek.json", {"fundedAmount": 50}) +answer.update(total="8.15", currency="USD") +record, _ = deepseek() +expect("DeepSeek collector derives the estimated ledger from fundedAmount", + record["balance"], {"remaining": 8.15, "funded": 50.0, "spent": 41.85, "currency": "USD", "estimated": True}, record) +fetched_at = record["limitsFetchedAt"] + +# A failed check keeps the last good balance, marked stale with when it was +# read; only an unreachable service asks for a quick retry. +cases = { + "can't be reached": ("down", True, "Couldn't reach", ""), + "times out": ("timeout", True, "Couldn't reach", ""), + "answers with a server error": (500, False, "HTTP 500", ""), + "sends a garbled answer": ("garbage", False, "invalid balance response", ""), + # A refused key asks for a new sign-in. + "refuses the key": (401, False, "", "Sign-in expired"), +} +for name, (mode, retry, help_text, status) in cases.items(): + record, _ = deepseek(mode) + check(f"DeepSeek collector keeps the last balance, stale, when DeepSeek {name}", + record["balance"]["remaining"] == 8.15 and record["balance"]["estimated"] is True and record["limitsStale"] is True + and record["limitsFetchedAt"] == fetched_at and record["usageStatusText"] == status + and bool(record.get("retryAdvised")) is retry and help_text in record["authHelpText"], record) + +# With nothing kept, a failed check says the balance is unavailable. +for path in limits_cache_files("deepseek"): + path.unlink() +record, _ = deepseek(500) +check("DeepSeek collector says the balance is unavailable with nothing kept", + record["usageStatusText"] == "Balance unavailable" and "balance" not in record and record["limitsStale"] is False + and record["limitsFetchedAt"] == 0 and record["ready"] is True and record["tierLabel"] == "Prepaid", record) + +# Another key never shows this one's balance. +deepseek() +os.environ["DEEPSEEK_API_KEY"] = "key-other" +record, _ = deepseek("down") +check("DeepSeek collector keeps a balance per key", + "balance" not in record and record["usageStatusText"] == "Balance unavailable" and record.get("retryAdvised") is True, record) + +# An answer without a readable total_balance is a failed check, not a balance +# of zero: it keeps the last balance, stale, and never replaces it. +deepseek() +# A total too large for the record's number is as unreadable as no total. +for incomplete in ([{}], [{"currency": "USD", "total_balance": None}], [{"total_balance": "NaN"}], [{"total_balance": True}], [{"total_balance": "1e400"}]): + with FakeHttp({BALANCE: {"is_available": True, "balance_infos": incomplete}}): + record = collect("deepseek", "--force") + check(f"DeepSeek collector keeps the last balance, stale, when DeepSeek answers {incomplete}", + record["balance"]["remaining"] == 8.15 and record["limitsStale"] is True and record["usageStatusText"] == "" + and "total_balance" in record["authHelpText"], record) + record, _ = deepseek(500) + expect(f"DeepSeek collector never caches {incomplete} over the last balance", [record["balance"]["remaining"], record["limitsStale"]], [8.15, True], record) + +# DEEPSEEK_API_BASE_URL moves the endpoint, never the key onto plain HTTP to +# another machine, nor into a URL that isn't one. +for url in ("http://evil.example", "not a url"): + os.environ["DEEPSEEK_API_BASE_URL"] = url + record, http = deepseek() + expect(f"DeepSeek collector sends its key to no untrusted endpoint ({url})", + [http.count, record["usageStatusText"], "HTTPS" in record["authHelpText"]], [0, "Balance unavailable", True], record) +del os.environ["DEEPSEEK_API_BASE_URL"] + +# A balance never rolls back one another run stamped while this one waited on +# the network. +def later_run(request): + edit_limits_cache("deepseek", fetched_ago=-60) + return balance(request) + + +for path in limits_cache_files("deepseek"): + path.unlink() +deepseek() +with FakeHttp({BALANCE: later_run}): + collect("deepseek", "--force") +expect("DeepSeek collector never overwrites a newer cached balance with its own", + [json.loads(path.read_text())["fetchedAtMs"] > ms_in(30) for path in limits_cache_files("deepseek")], [True]) + +# Untrusted amounts read as nothing rather than as a number. +expect("DeepSeek collector reads untrusted amounts as zero", + [float(load_collector("deepseek").money(value)) for value in ("NaN", "Infinity", "1e400", True, "abc", None, "3.5")], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.5]) + +# The key pi keeps, where PI_CODING_AGENT_DIR moves it, or omp's. Only a +# literal key: a `!command` or `$VAR` entry is pi's to resolve, never sent. +sandbox() + + +def saved_key(): + return load_collector("deepseek").api_key() + + +pi_auth(home() / ".omp" / "agent", "deepseek", "key-omp") +expect("DeepSeek collector reads the key omp saved", saved_key(), "key-omp") +pi_auth(None, "deepseek", "key-pi") +expect("DeepSeek collector prefers the key pi saved", saved_key(), "key-pi") +os.environ["PI_CODING_AGENT_DIR"] = str(home() / "pi-moved") +pi_auth(home() / "pi-moved", "deepseek", "key-moved") +expect("DeepSeek collector reads pi's key where PI_CODING_AGENT_DIR moves it", saved_key(), "key-moved") +pi_auth(home() / "pi-moved", "deepseek", "!pass show deepseek") +pi_auth(home() / ".omp" / "agent", "deepseek", "$DEEPSEEK_KEY") +expect("DeepSeek collector never takes a pi !command or $VAR entry as a key", saved_key(), "") +opencode_auth(None, "deepseek", "key-opencode") +expect("DeepSeek collector reads the key opencode saved", saved_key(), "key-opencode") +pi_auth(home() / ".omp" / "agent", "deepseek", "key-omp") +expect("DeepSeek collector prefers omp's key to opencode's", saved_key(), "key-omp") + +# The balance is kept per API host: another host, the same key or not, is +# another account, and never shows this one's balance. +sandbox() +os.environ["DEEPSEEK_API_KEY"] = "key-123" +deepseek() +with FakeHttp(mode="down") as http: + other = collect("deepseek", "--force", "--api-base-url", "https://deepseek.example") +record, _ = deepseek("down") +expect("DeepSeek collector keeps one key's balance per API host", + [http.requests[0]["url"], "balance" in other, other["usageStatusText"], record["balance"]["remaining"], record["limitsStale"]], + ["https://deepseek.example/user/balance", False, "Balance unavailable", float(answer["total"]), True], other) + +# balance_infos holds one entry per currency, CNY or USD (DeepSeek's +# get-user-balance reference): the funded one is the balance, so an empty CNY +# entry never reads as an empty account beside USD credit. +sandbox() +os.environ["DEEPSEEK_API_KEY"] = "key-123" +with FakeHttp({BALANCE: {"is_available": True, "balance_infos": [ + {"currency": "CNY", "total_balance": "0.00", "granted_balance": "0.00", "topped_up_balance": "0.00"}, + {"currency": "USD", "total_balance": "12.50", "granted_balance": "0.00", "topped_up_balance": "12.50"}, +]}}): + record = collect("deepseek", "--force") +expect("DeepSeek collector reads the currency that holds the credit", + [record["balance"]["remaining"], record["balance"]["currency"], record["usageStatusText"]], [12.5, "USD", ""], record) + +# is_available false is DeepSeek saying the balance no longer covers a call: +# with no funded figure the gauge raises no alarm, so the status does, and a +# reused answer keeps saying so. +with FakeHttp({BALANCE: {"is_available": False, "balance_infos": [ + {"currency": "CNY", "total_balance": "0.00", "granted_balance": "0.00", "topped_up_balance": "0.00"}, +]}}): + record = collect("deepseek", "--force") +expect("DeepSeek collector reads an unavailable balance as Out of credits", + [record["balance"]["remaining"], record["usageStatusText"], "Top up" in record["authHelpText"], record.get("retryAdvised")], + [0.0, "Out of credits", True, None], record) +record, http = deepseek(500, ()) +expect("DeepSeek collector keeps saying Out of credits from the balance it reuses", [http.count, record["usageStatusText"]], [0, "Out of credits"], record) +PY 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..21909e8c16c --- /dev/null +++ b/test/shell.d/agent-usage-deepseek-scanner-test.sh @@ -0,0 +1,307 @@ +#!/bin/bash + +# The DeepSeek collector's local stats: pi, omp and opencode usage on the +# deepseek provider, DeepSeek Harness's own session logs, and the key dsh +# keeps. How pi and opencode stores are read (roots, forks, both opencode +# stores, OPENCODE_DB, opencode-v2.db, untrusted values) is held centrally by +# the shared and partition tests, the unused-machine record by the partition +# test, and the scan cache by the contract check. The balance is +# agent-usage-deepseek-balance-test.sh's. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +import shutil + +from agent_usage_fixtures import * + + +def deepseek(*args): + with FakeHttp(): + return collect("deepseek", *(args or ("--force",))) + + +def sessions(): + return home() / ".pi" / "agent" / "sessions" + + +# A machine that never used DeepSeek has no local stats to show. +sandbox() +record = deepseek() +expect("DeepSeek with no usage says it has no local stats", [record["ready"], record["hasLocalStats"]], [False, False], record) + +# ------------------------------------------------- pi, omp and opencode + +# Only the deepseek provider counts, from pi, omp and opencode alike, with +# reasoning folded into output and cache kept apart. A day active in more +# than one of them is still one active day. Local usage alone puts DeepSeek +# in use; without a key there is no plan and no balance section. +sandbox() +write_pi_session(sessions() / "project" / "pi.jsonl", [ + pi_reply("deepseek", id="pi-1", model="deepseek-pi", usage={"input": 10, "output": 4, "cacheRead": 3, "cacheWrite": 2, "totalTokens": 19}), + pi_reply("anthropic", id="pi-2", model="claude-test", usage={"input": 999, "output": 999}), + pi_reply("openai-codex", id="pi-3", model="gpt-test", usage={"input": 999, "output": 999}), +]) +write_pi_session(home() / ".omp" / "agent" / "sessions" / "project" / "omp.jsonl", [ + pi_reply("deepseek", id="omp-1", model="deepseek-omp", usage={"input": 20, "output": 5, "cacheRead": 4, "cacheWrite": 1, "totalTokens": 30}), +]) +opencode_db(tables=("message",), message=[ + ("m-1", "ses_1", oc_v1("deepseek", {"input": 100, "output": 20, "reasoning": 5, "cache": {"read": 30, "write": 2}}, model="deepseek-model")), + ("m-2", "ses_1", oc_v1("openai", {"input": 999, "output": 999}, model="gpt-model")), + ("m-3", "ses_1", oc_v1("anthropic", {"input": 999, "output": 999}, model="claude-model")), + ("m-4", "ses_1", oc_v1("deepseek-proxy", {"input": 999, "output": 999}, model="deepseek-model")), +]) +record = deepseek() +expect("DeepSeek counts its own provider's usage from pi, omp and opencode, one active day across them", { + key: record[key] for key in ("todayTotalTokens", "todayPrompts", "todaySessions", "totalSessions", "activeDays", "modelUsage") +}, { + "todayTotalTokens": 206, "todayPrompts": 3, "todaySessions": 3, "totalSessions": 3, "activeDays": 1, "modelUsage": { + "deepseek-pi": {"inputTokens": 10, "outputTokens": 4, "cacheReadInputTokens": 3, "cacheCreationInputTokens": 2}, + "deepseek-omp": {"inputTokens": 20, "outputTokens": 5, "cacheReadInputTokens": 4, "cacheCreationInputTokens": 1}, + "deepseek-model": {"inputTokens": 100, "outputTokens": 25, "cacheReadInputTokens": 30, "cacheCreationInputTokens": 2}, + }, +}, record) +expect("DeepSeek in use without a key prints a valid record with no plan and no balance", + {key: record.get(key) for key in ("id", "ready", "tierLabel", "hasPromptStats", "usageStatusText", "authHelpText")} | {"balance": "balance" in record}, + {"id": "deepseek", "ready": True, "tierLabel": "", "hasPromptStats": True, "usageStatusText": "", "authHelpText": "", "balance": False}, record) + +# -------------------------------------------------------- DeepSeek Harness + +# dsh keeps one event log per session under $DSH_HOME/sessions, Zstandard- +# compressed by default. Only its DeepSeek routes count; a retry's settled +# answer replaces the streamed sample of the same step; a fork's inherited +# history is its parent's; a torn final record is dropped; and only the +# newest generation of a session is read. +sandbox() +os.environ["DSH_HOME"] = str(home() / "dsh") +now = ms_in(0) + +try: + from compression import zstd +except ImportError: + zstd = None + + +def header(session, seeded=False): + return {"type": "session", "version": 4, "id": session, "createdAt": now, "isSeeded": seeded, "delegationDepth": 0} + + +# A message someone typed, as dsh logs it before the turn's first request; +# dsh's own turns (a schedule, a goal round) carry another source. +def typed(seq=1, kind="user"): + return {"type": "user/message", "seq": seq, "time": now, "data": {"content": [{"type": "text", "text": "go"}], "source": {"kind": kind}, "role": "user", "id": f"m{seq}"}} + + +def route(provider, model, seq=1): + return {"type": "request/header", "seq": seq, "time": now, "data": {"header": {"config": {"provider": provider, "model": model}}, "reason": "initial"}} + + +def message(turn, step, usage, seq=2, kind="assistant/message"): + data = {"turn": turn, "step": step, "stream": []} + if kind == "assistant/message": + data["usage"] = usage + else: + data["stream"] = [{"type": "chunk", "chunk": {"type": "usage", "usage": usage}}] + return {"type": kind, "seq": seq, "time": now, "data": data} + + +def usage(input, output=0, read=0, write=0): + return {"inputTokens": input, "outputTokens": output, "cacheReadTokens": read, "cacheWriteTokens": write} + + +def write_log(path, records, compressed, torn=b""): + path.parent.mkdir(parents=True, exist_ok=True) + lines = [(json.dumps(record) + "\n").encode() for record in records] + if compressed: + path.write_bytes(b"".join(zstd.compress(line) for line in lines) + torn) + else: + path.write_bytes(b"".join(lines) + torn) + + +compressed = zstd is not None +suffix = ".jsonl.zstd" if compressed else ".jsonl" +project = home() / "dsh" / "sessions" / "--home-me-code--" + +# A DeepSeek session: one step streamed a sample, then settled with the real +# count; a retry of step 2 adds a second billed attempt. +write_log(project / "s1" / ("session.v4" + suffix), [ + header("s1"), + typed(), + route("deepseek-official", "deepseek-v4-pro"), + message(1, 1, usage(1, 1), kind="assistant/attempt"), + message(1, 1, usage(100, 10, 20, 0)), + message(1, 2, usage(5, 1), kind="assistant/attempt"), + {"type": "llm/retry-started", "seq": 9, "time": now, "data": {"turn": 1, "step": 2}}, + message(1, 2, usage(7, 3)), +], compressed, torn=b'{"type":"assistant/mess' if not compressed else b"\x28\xb5\x2f\xfd") +# An older generation of the same session is migration history only. +write_log(project / "s1" / "session.jsonl", [header("s1"), route("deepseek-official", "deepseek-v4-pro"), message(1, 1, usage(999999))], False) +# A route billed by someone else. +write_log(project / "s2" / ("session.v4" + suffix), [header("s2"), route("pi-ai:anthropic", "claude-opus-5"), message(1, 1, usage(888888))], compressed) +# A fork of s1: its copy of the parent's step is inherited, and only what +# follows its seed marker was spent in it. +write_log(project / "s3" / ("session.v4" + suffix), [ + header("s3", seeded=True), + typed(), + route("deepseek-official", "deepseek-v4-pro"), + message(1, 1, usage(100, 10, 20, 0)), + {"type": "session/end-seed", "seq": 4, "time": now, "data": {"inherited": True}}, + typed(seq=5), + route("deepseek-account", "deepseek-v4-flash", seq=5), + message(2, 1, usage(4000, 6, 0, 1)), +], compressed) +# A root written with compression off reads the same, torn tail and all. +write_log(project / "s4" / "session.v4.jsonl", [header("s4"), typed(), route("deepseek-official", "deepseek-chat"), message(1, 1, usage(30000))], False, + torn=b'{"type":"assistant/mess') + +record = deepseek() +expect("DeepSeek counts DeepSeek Harness sessions on its DeepSeek routes", + {key: record[key] for key in ("todayTotalTokens", "totalPrompts", "totalSessions", "modelUsage")}, { + "todayTotalTokens": 34153, "totalPrompts": 3, "totalSessions": 3, "modelUsage": { + "deepseek-v4-pro": {"inputTokens": 112, "outputTokens": 14, "cacheReadInputTokens": 20, "cacheCreationInputTokens": 0}, + "deepseek-v4-flash": {"inputTokens": 4000, "outputTokens": 6, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 1}, + "deepseek-chat": {"inputTokens": 30000, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, + }, +}, record) + +# A line nested too deep to parse costs only that line (or, as the header, +# that log), never the rest of the scan. +DEEP = b'{"type":"assistant/message","data":' + b"[" * 100000 + b"]" * 100000 + b"}\n" +deep = project / "s5" / "session.v4.jsonl" +write_log(deep, [header("s5"), route("deepseek-official", "deepseek-deep")], False) +with deep.open("ab") as handle: + handle.write(DEEP + (json.dumps(message(1, 1, usage(50))) + "\n").encode()) +(project / "s6").mkdir() +(project / "s6" / "session.v4.jsonl").write_bytes(DEEP + (json.dumps(message(1, 1, usage(70))) + "\n").encode()) +record = deepseek() +expect("DeepSeek skips a dsh line nested too deep to parse, keeping the rest of the scan", + [record["todayTotalTokens"], record["modelUsage"].get("deepseek-deep", {}).get("inputTokens")], [34203, 50], record) +shutil.rmtree(project / "s5") +shutil.rmtree(project / "s6") + +# dsh's pi-ai adapter names a route by its pi-ai provider id, as dsh 0.2 logs +# it: pi-ai's `deepseek` provider is DeepSeek's API, and `openrouter` (or a +# hand-declared gateway) is billed by whoever runs it. +write_log(project / "s9" / ("session.v4" + suffix), [ + header("s9"), + route("deepseek", "deepseek-chat"), + message(1, 1, {"inputTokens": 5234, "outputTokens": 55, "cacheReadTokens": 75, "totalTokens": 5364}), + route("openrouter", "z-ai/glm-4.6", seq=3), + message(2, 1, {"inputTokens": 7000, "outputTokens": 68, "cacheReadTokens": 300, "totalTokens": 7368}, seq=4), +], compressed) +record = deepseek() +expect("DeepSeek counts dsh's pi-ai deepseek route, not its other pi-ai routes", + [record["todayTotalTokens"], record["modelUsage"].get("deepseek-chat", {}).get("inputTokens"), "z-ai/glm-4.6" in record["modelUsage"]], + [34153 + 5364, 30000 + 5234, False], record) +shutil.rmtree(project / "s9") + +# A prompt is a typed message, not a billed step: dsh 0.2.1's headless run +# of one typed prompt that read a file took three steps, the same session +# resumed with --session-id took one more typed turn, and a turn dsh started +# on its own schedule answers no one. +write_log(project / "s10" / ("session.v4" + suffix), [ + header("s10"), + {"type": "agent/inbox/spliced", "seq": 1, "time": now, "data": {"target": "next-turn", "start": 0, "inserted": [{"source": {"kind": "user"}, "role": "user", "id": "m2"}]}}, + {"type": "turn/start", "seq": 2, "time": now, "data": {"turn": 1}}, + typed(seq=3), + {"type": "user/message", "seq": 4, "time": now, "data": {"source": {"kind": "runtime-context", "form": "snapshot"}, "role": "user", "id": "m4"}}, + route("deepseek-official", "deepseek-flash", seq=5), + message(1, 1, usage(12417), seq=6), + message(1, 2, usage(14423), seq=7), + message(1, 3, usage(15426), seq=8), + {"type": "session/end-seed", "seq": 9, "time": now, "data": {}}, + {"type": "turn/start", "seq": 10, "time": now, "data": {"turn": 2}}, + typed(seq=11), + route("deepseek-official", "deepseek-flash", seq=12), + message(2, 1, usage(16429), seq=13), + {"type": "turn/start", "seq": 14, "time": now, "data": {"turn": 3}}, + typed(seq=15, kind="schedule"), + route("deepseek-official", "deepseek-flash", seq=16), + message(3, 1, usage(1000), seq=17), +], compressed) +record = deepseek() +expect("DeepSeek counts a dsh prompt per typed message, not per step or scheduled turn", + [record["totalPrompts"], record["modelUsage"].get("deepseek-flash", {}).get("inputTokens")], [3 + 2, 12417 + 14423 + 15426 + 16429 + 1000], record) +shutil.rmtree(project / "s10") + +# A frame that inflates past the cap (64 MiB, far beyond any real log) is +# never read in full: the log counts nothing, and the scan, incomplete, isn't +# kept for the next run. +if compressed: + bomb = project / "s7" / "session.v4.jsonl.zstd" + bomb.parent.mkdir() + opening = "".join(json.dumps(line) + "\n" for line in (header("s7"), route("deepseek-official", "deepseek-bomb"), message(1, 1, usage(60)))) + bomb.write_bytes(zstd.compress(opening.encode() + b" " * (65 << 20))) + for path in scan_cache_files("deepseek"): + path.unlink() + with contextlib.redirect_stderr(io.StringIO()): + record = deepseek() + expect("DeepSeek stops a dsh log that inflates past its cap, and doesn't cache the scan", + [record["todayTotalTokens"], "deepseek-bomb" in record["modelUsage"], scan_cache_files("deepseek")], [34153, False, []], record) + shutil.rmtree(project / "s7") +else: + print("ok - DeepSeek stops a dsh log that inflates past its cap # SKIP no Zstandard support") + +# A session directory the walk can't list has logs the scan missed, so the +# scan isn't kept: once it's readable again, the next plain run counts it. +if os.geteuid() != 0: + locked = project / "s8" + write_log(locked / "session.v4.jsonl", [header("s8"), route("deepseek-official", "deepseek-locked"), message(1, 1, usage(80))], False) + for path in scan_cache_files("deepseek"): + path.unlink() + locked.chmod(0) + try: + record = deepseek() + finally: + locked.chmod(0o755) + expect("DeepSeek doesn't cache a dsh scan that couldn't list every session directory", + ["deepseek-locked" in record["modelUsage"], scan_cache_files("deepseek")], [False, []], record) + record = deepseek("--limits-only") + expect("DeepSeek counts the dsh session once its directory is readable again", record["modelUsage"].get("deepseek-locked", {}).get("inputTokens"), 80, record) + shutil.rmtree(locked) +else: + print("ok - DeepSeek doesn't cache a dsh scan that couldn't list every session directory # SKIP root reads anything") + +# dsh keeps its key in its managed .credentials.yaml, or a .env beside it. +(home() / "dsh" / ".credentials.yaml").write_text("""# managed by dsh +version: 1 +refs: + OTHER_KEY: nope + DEEPSEEK_API_KEY: "sk-from-dsh" # stored by the Models page +""") +collector = load_collector("deepseek") +first = collector.api_key() +(home() / "dsh" / ".credentials.yaml").unlink() +(home() / "dsh" / ".env").write_text("export DEEPSEEK_API_KEY='sk-from-env-file'\n") +expect("DeepSeek reads the key dsh stores, from its credentials or its .env", [first, collector.api_key()], ["sk-from-dsh", "sk-from-env-file"]) +del os.environ["DSH_HOME"] + +# Claude Code pointed at DeepSeek's Anthropic API is DeepSeek's; with no +# endpoint, only a DeepSeek model's reply is. +sandbox() +claude_code_session("ds", [("deepseek-chat", 300), ("claude-sonnet-4-5", 40)]) +record = deepseek() +expect("DeepSeek counts Claude Code's DeepSeek replies, and no Claude one", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"deepseek-chat": 300}, 1, 1], record) +claude_code_session("ds", [("deepseek-chat", 300), ("claude-sonnet-4-5", 40)], base_url="https://api.deepseek.com/anthropic") +record = deepseek() +expect("DeepSeek counts every Claude Code reply while Claude Code is pointed at DeepSeek", + [record["todayTokensByModel"], record["todayPrompts"]], [{"deepseek-chat": 300, "claude-sonnet-4-5": 40}, 2], record) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_reply(): + add_reply.count = getattr(add_reply, "count", 0) + 1 + write_pi_session(sessions() / "grown" / f"{add_reply.count}.jsonl", [pi_reply("deepseek", 5, id=f"grown-{add_reply.count}")]) + + +check_scan_cache(Adapter("deepseek", grow=add_reply, pi_root=sessions)) +PY diff --git a/test/shell.d/agent-usage-devin-scanner-test.sh b/test/shell.d/agent-usage-devin-scanner-test.sh new file mode 100644 index 00000000000..33bbd4ec9f9 --- /dev/null +++ b/test/shell.d/agent-usage-devin-scanner-test.sh @@ -0,0 +1,324 @@ +#!/bin/bash + +# The Devin collector's own sources: the CLI's session database (and the +# older per-session layout), and the user-status caches the CLI writes, which +# are its quota. pi, omp and opencode usage on the devin provider is read the +# shared way and held by the shared and partition tests, as is the +# unused-machine record; the scan cache and the limits go through the +# contract checks. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +import base64 + +from agent_usage_fixtures import * + +now = int(time.time()) + + +def devin(*args): + with FakeHttp(): + return collect("devin", *(args or ("--force",))) + + +def cli_dir(): + return data_home() / "devin" / "cli" + + +def status_dir(): + return Path(os.environ["XDG_CACHE_HOME"]) / "devin" / "cli" + + +def credentials(): + return data_home() / "devin" / "credentials.toml" + + +# The CLI's sessions.db: one message_nodes row per reply, its metrics in the +# chat message's metadata, and one prompt_history row per typed prompt. +def node(session, created, model, inp, out, cache_read=0, cache_write=0, chat=None): + metrics = {"input_tokens": inp, "output_tokens": out, "cache_read_tokens": cache_read, "cache_creation_tokens": cache_write} + return (session, chat if chat is not None else json.dumps({"metadata": {"metrics": metrics, "generation_model": model}}), created) + + +def sessions_db(nodes=(), prompts=()): + path = cli_dir() / "sessions.db" + path.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(path) + conn.execute("CREATE TABLE IF NOT EXISTS message_nodes (row_id INTEGER PRIMARY KEY, session_id TEXT, node_id INTEGER, parent_node_id INTEGER, chat_message TEXT, created_at INTEGER, metadata TEXT)") + conn.execute("CREATE TABLE IF NOT EXISTS prompt_history (id INTEGER PRIMARY KEY, content TEXT, timestamp INTEGER, session_id TEXT, is_shell INTEGER)") + conn.executemany("INSERT INTO message_nodes (session_id, chat_message, created_at) VALUES (?, ?, ?)", nodes) + conn.executemany("INSERT INTO prompt_history (content, timestamp, session_id, is_shell) VALUES (?, ?, ?, 0)", prompts) + conn.commit() + conn.close() + return path + + +# A user-status cache as the CLI writes it: GetUserStatus's protobuf fields, +# base64 in a small JSON file. +def varint(value): + value &= (1 << 64) - 1 + out = b"" + while True: + byte = value & 0x7F + value >>= 7 + out += bytes([byte | 0x80]) if value else bytes([byte]) + if not value: + return out + + +def field(number, data): + if isinstance(data, int): + return varint(number << 3) + varint(data) + return varint((number << 3) | 2) + varint(len(data)) + data + + +def write_status(name, fetched_at, **options): + plan = field(2, options["plan"].encode()) if options.get("plan") else b"" + for key, number in (("tier", 1), ("billing", 35), ("hideDaily", 36), ("hideWeekly", 37), ("promptLimit", 12)): + if key in options: + plan += field(number, int(options[key])) + status = field(1, plan) + for key, number in (("daily", 14), ("weekly", 15), ("dailyReset", 17), ("weeklyReset", 18), ("promptUsed", 6)): + if key in options: + status += field(number, int(options[key])) + if "planEnd" in options: + status += field(3, field(1, int(options["planEnd"]))) + payload = field(13, status) + return write_json(status_dir() / f"user_status.{name}.bin", + {"version": 1, "identity_digest": name, "fetched_at_secs": fetched_at, "payload": base64.b64encode(payload).decode()}) + + +def signed_in_at(moment): + credentials().parent.mkdir(parents=True, exist_ok=True) + if not credentials().exists(): + credentials().write_text('api_server_url = "https://example.invalid"\n') + os.utime(credentials(), (moment, moment)) + + +def windows(record): + return [[entry["label"], entry["percent"]] for entry in record["limits"]] + + +def day(offset): + return (datetime.now() + timedelta(days=offset)).strftime("%Y-%m-%d") + + +# ------------------------------------------------------- the CLI's database + +# Untrusted fields skip themselves, or the row, never the scan: a boolean, a +# negative count and a string count read as nothing; a row with a broken +# timestamp, no metrics, or no JSON at all adds nothing. +sandbox() +sessions_db(nodes=[ + node("s1", now, "glm-5-2", 100, 20, 60, 0), + node("s1", now, "glm-5-2", 80, 10, 50, 0), + node("s1", now, "devin-pro", 200, 40, 0, 10), + node("s2", now - 86400, "glm-5-2", 50, 5), + node("s1", now, "glm-5-2", True, -5, "7", 0), + node("s1", "yesterday-ish", "glm-5-2", 999, 999), + node("s1", 10**15, "glm-5-2", 999, 999), + node("s1", now, None, 0, 0, chat=json.dumps({"metadata": {}})), + node("s1", now, None, 0, 0, chat='{not json "metrics"'), +], prompts=[("a", now, "s1"), ("b", now, "s1"), ("c", now, "s3"), ("d", now - 86400, "s2"), ("e", now - 86400, "s1")]) + +# (100+20+60) + (80+10+50) + (200+40+10) = 570. Prompts come from +# prompt_history, and a session spanning midnight counts once. +record = devin() +expect("Devin reads the CLI's session database, skipping untrusted fields", { + "name": record["name"], + "today": record["todayTotalTokens"], + "counts": [record[key] for key in ("todayPrompts", "todaySessions", "totalPrompts", "totalSessions", "activeDays")], + "models": record["modelUsage"], + "yesterday": next(entry["messageCount"] for entry in record["recentDays"] if entry["date"] == day(-1)), +}, { + "name": "Devin", + "today": 570, + "counts": [3, 2, 5, 3, 2], + "models": { + "glm-5-2": {"inputTokens": 230, "outputTokens": 35, "cacheReadInputTokens": 110, "cacheCreationInputTokens": 0}, + "devin-pro": {"inputTokens": 200, "outputTokens": 40, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 10}, + }, + "yesterday": 55, +}, record) + +# In use without any sign-in: the quota has nowhere to come from. +expect("Devin in use without a sign-in waits for auth", [record["ready"], record["usageStatusText"], record["limits"]], [True, "Waiting for auth", []], record) +expect("Devin's sessions here are local stats", record["hasLocalStats"], True, record) + +# A database it cannot read costs the stats, never the record, and is not +# cached as the whole story. +good = (cli_dir() / "sessions.db").read_bytes() +(cli_dir() / "sessions.db").write_text("not a database") +for path in scan_cache_files("devin"): + path.unlink() +with contextlib.redirect_stderr(io.StringIO()): + record = devin() +expect("Devin survives an unreadable session database and doesn't cache that scan", + [record["id"], record["todayTotalTokens"], scan_cache_files("devin")], ["devin", 0, []], record) +expect("an unreadable session database says Devin is unavailable rather than hiding it", + [record["ready"], record["usageStatusText"], str(cli_dir() / "sessions.db") in record["authHelpText"]], [True, "Devin unavailable", True], record) +(cli_dir() / "sessions.db").write_bytes(good) +(cli_dir() / "sessions.db").chmod(0) +with contextlib.redirect_stderr(io.StringIO()): + record = devin() +(cli_dir() / "sessions.db").chmod(0o644) +expect("a session database it may not open says Devin is unavailable", [record["ready"], record["usageStatusText"]], [True, "Devin unavailable"], record) + +# ------------------------------------------------------- the older layout + +# Builds without the cli/ directory kept one lifetime total per session. +sandbox() +old = data_home() / "devin" / "cli_sessions.db" +old.parent.mkdir(parents=True) +conn = sqlite3.connect(old) +conn.execute("CREATE TABLE sessions (id TEXT, working_directory TEXT, backend_type TEXT, model TEXT, permission_mode TEXT, created_at INTEGER, last_activity_at INTEGER, total_tokens INTEGER, current_epoch INTEGER)") +conn.execute("CREATE TABLE messages (id TEXT, session_id TEXT, sequence_number INTEGER, role TEXT, content TEXT, tool_call_id TEXT, tool_calls_json TEXT, timestamp INTEGER, compaction_epoch INTEGER)") +conn.execute("INSERT INTO sessions (id, model, created_at, last_activity_at, total_tokens) VALUES ('o1', 'swe-1.5', ?, ?, 4000)", (now - 86400, now)) +conn.execute("INSERT INTO messages (id, session_id, role, timestamp) VALUES ('1', 'o1', 'user', ?), ('2', 'o1', 'assistant', ?), ('3', 'o1', 'assistant', ?)", (now, now, now)) +conn.commit() +conn.close() +record = devin() +expect("Devin falls back to the older per-session layout", + [record["todayTotalTokens"], record["todayPrompts"], record["totalSessions"], record["modelUsage"]["swe-1.5"]["inputTokens"]], [4000, 2, 1, 4000], record) + +# Once the CLI's own database exists, the older files are not read again. +sessions_db() +expect("Devin reads only the CLI's own database once it exists", devin()["todayTotalTokens"], 0) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def plant_status(): + write_status("a", int(time.time()), plan="Pro", daily=69, weekly=53, dailyReset=now + 3600, weeklyReset=now + 4 * 86400) + + +def add_reply(): + sessions_db(nodes=[node("grown", int(time.time()), "glm-5-2", 5, 0)]) + + +# Devin's quota is the CLI's own cache, read on every run: there is no +# request to count, retry or refuse, no check of its own to fail stale or cut +# short, and no limits cache of the collector's to elapse (its own +# elapsed-window case is below). +adapter = Adapter("devin", plant=plant_status, grow=add_reply, requests_per_probe=0, retry_modes=(), refused_status=None, + skip=("stale", "cut", "elapsed"), pi_root=lambda: home() / ".pi" / "agent" / "sessions") +fresh = check_limits_contract(adapter) +expect("Devin's plan and quota come from the CLI's user-status cache, with the service's own reset times", { + "tierLabel": fresh["tierLabel"], "windows": windows(fresh), "fetchedAt": fresh["limitsFetchedAt"] // 1000 >= now, + "reset": fresh["limits"][0]["resetsAt"], "retry": "retryAdvised" in fresh, +}, { + "tierLabel": "Pro", "windows": [["Daily", 0.31], ["Weekly", 0.47]], "fetchedAt": True, + "reset": datetime.fromtimestamp(now + 3600, timezone.utc).isoformat(), "retry": False, +}, fresh) +check_scan_cache(adapter) + +# ------------------------------------------------------------- the quota + +sandbox() + +# The quota is read on every run, past the scan reuse, so a fetch the CLI +# just made shows without --force. +write_status("a", now, plan="Pro", daily=69, weekly=53, dailyReset=now + 3600, weeklyReset=now + 4 * 86400) +devin() +write_status("a", now, plan="Pro", daily=50, weekly=53, dailyReset=now + 3600, weeklyReset=now + 4 * 86400) +expect("Devin reads the quota cache on every run", devin("--limits-only")["limits"][0]["percent"], 0.5) + +# Spent windows: proto3 leaves a zero off the wire, so a dated window with no +# percent has nothing left. +write_status("a", now, plan="Pro", dailyReset=now + 3600) +expect("Devin reads a missing percent on a dated window as spent", windows(devin()), [["Daily", 1.0]]) + +# A cache from more than fifteen minutes ago still describes the account, but +# is reported stale with its age; a window that reset since is back at zero. +write_status("a", now - 7200, plan="Pro", daily=10, weekly=40, dailyReset=now - 60, weeklyReset=now + 86400) +record = devin() +expect("Devin reports an old quota stale, with elapsed windows at zero", { + "stale": record["limitsStale"], "fetchedAt": record["limitsFetchedAt"], "status": record["usageStatusText"], + "windows": [[entry["label"], entry["percent"], entry["resetsAt"] != ""] for entry in record["limits"]], +}, {"stale": True, "fetchedAt": (now - 7200) * 1000, "status": "", "windows": [["Daily", 0.0, False], ["Weekly", 0.6, True]]}, record) + +# The fetch stamp is floored to the millisecond, never rounded up. +fraction = now - 7200 + 0.9996 +write_status("a", fraction, plan="Pro", daily=10, weekly=40, weeklyReset=now + 86400) +expect("Devin floors the fetch stamp", devin()["limitsFetchedAt"], int(fraction * 1000)) + +# Fetches from before the last sign-in describe another account. +signed_in_at(now + 600) +write_status("a", now, plan="Pro", daily=69, weekly=53, dailyReset=now + 3600, weeklyReset=now + 86400) +record = devin() +expect("Devin drops quota fetched before the last sign-in and says how to refresh it", + [record["ready"], record["tierLabel"], record["limits"], record["usageStatusText"], "Start Devin" in record["authHelpText"]], + [True, "", [], "Devin limits unavailable", True], record) + +# Of the caches fetched since, the newest wins, whichever name sorts first. +signed_in_at(now - 600) +for name in ("0", "b"): + write_status(name, now - 300, plan="Max", daily=0, weekly=0, dailyReset=now + 3600) +expect("Devin serves the newest fetch under the current sign-in, sorted between older ones", devin()["tierLabel"], "Pro") +for name in ("0", "b"): + (status_dir() / f"user_status.{name}.bin").unlink() + +# The old account's last fetch can land in the seconds before a switch, inside +# the allowance for the new account's first fetch. Having been served under +# the old sign-in, it is known to be the old account's and stays out. +signed_in_at(now - 3600) +write_status("a", now - 100, plan="Pro", daily=69, weekly=53, dailyReset=now + 3600, weeklyReset=now + 86400) +expect("Devin serves the account signed in", devin()["tierLabel"], "Pro") +signed_in_at(now - 90) +record = devin() +expect("Devin drops the old account's fetch from just before a switch", + [record["tierLabel"], record["limits"], record["usageStatusText"]], ["", [], "Devin limits unavailable"], record) +write_status("b", now - 95, plan="Max", daily=0, weekly=0, dailyReset=now + 3600) +expect("Devin serves the new account's first fetch from just before the switch", devin()["tierLabel"], "Max") +(status_dir() / "user_status.b.bin").unlink() +write_status("a", now, plan="Pro", daily=69, weekly=53, dailyReset=now + 3600, weeklyReset=now + 86400) +expect("Devin serves a fetch made after the sign-in", devin()["tierLabel"], "Pro") + +# A plan can hide a window; credit buckets run to the plan's end; a plan +# without a name of its own is named by its tier. +write_status("a", now, tier=16, hideDaily=1, daily=10, weekly=80, dailyReset=now + 3600, weeklyReset=now + 86400, + promptLimit=500, promptUsed=125, planEnd=now + 20 * 86400) +record = devin() +expect("Devin honors hidden windows, credit buckets and tier names", + [record["tierLabel"], windows(record)], ["Devin Pro", [["Weekly", 0.2], ["Prompt credits", 0.25]]], record) + +write_status("a", now, plan="Pro", hideWeekly=1, daily=10, weekly=80, dailyReset=now + 3600, weeklyReset=now + 86400) +expect("Devin honors a hidden weekly window", windows(devin()), [["Daily", 0.9]]) + +# TeamsTier, as the Devin protocol numbers it, has 8 and 11 too. +names = [] +for tier in (8, 11): + write_status("a", now, tier=tier, daily=10, weekly=80, dailyReset=now + 3600, weeklyReset=now + 86400) + names.append(devin()["tierLabel"]) +expect("Devin names every tier the protocol numbers", names, ["Pro Ultimate", "Enterprise Saas Pooled"]) + +# A fetch stamped ahead of the clock (one that went back) isn't current. +write_status("a", now + 3600, plan="Pro", daily=10, weekly=80, dailyReset=now + 7200, weeklyReset=now + 86400) +record = devin() +expect("Devin reports a fetch stamped in the future stale", [record["limitsStale"], windows(record)], [True, [["Daily", 0.9], ["Weekly", 0.2]]], record) + +# A credit plan leaves the quota percents at their defaults with no reset: +# those windows are not meters. +write_status("a", now, plan="Pro", billing=1) +expect("Devin shows no quota windows on a credit plan", devin()["limits"], []) + +# A broken cache file is no quota at all. +write_json(status_dir() / "user_status.a.bin", {"fetched_at_secs": True, "payload": "%%%"}) +expect("Devin ignores a broken quota cache", devin()["usageStatusText"], "Devin limits unavailable") + +# A sign-in with no sessions on this machine shows, with no local stats. +sandbox() +signed_in_at(now - 600) +write_status("a", now, plan="Pro", daily=10, weekly=40, weeklyReset=now + 86400) +record = devin() +expect("a Devin sign-in with no sessions here has no local stats", + [record["ready"], record["totalSessions"], record["hasLocalStats"]], [True, 0, False], record) +PY diff --git a/test/shell.d/agent-usage-fireworks-scanner-test.sh b/test/shell.d/agent-usage-fireworks-scanner-test.sh index 6440ce1164f..8aeadbaedc4 100755 --- a/test/shell.d/agent-usage-fireworks-scanner-test.sh +++ b/test/shell.d/agent-usage-fireworks-scanner-test.sh @@ -1,280 +1,342 @@ #!/bin/bash -source "$(dirname "$0")/base-test.sh" - -require_command jq -require_command python3 - -TEST_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME"' EXIT - -auth_file="$TEST_HOME/auth.ini" -cat >"$auth_file" <<'EOF' -[default] -api_key = fw_test -account_id = example -EOF - -mkdir -p "$TEST_HOME/.config/omarchy/agents" -cat >"$TEST_HOME/.config/omarchy/agents/fireworks.json" <<'EOF' -{ - "accountId": "example", - "fundedAmount": 20, - "fundedAt": "2026-07-01" -} -EOF - -# 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_DATA_HOME="$TEST_HOME/.local/share" \ - PI_CODING_AGENT_DIR="$TEST_HOME/.pi/agent" FIREWORKS_API_KEY="" FIREWORKS_AUTH_PATH="$TEST_HOME/missing.ini" "$ROOT/bin/omarchy-agent-usage-fireworks") - -[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + (.hasPromptStats | tostring)' <<<"$no_key") == "fireworks:false:false" ]] || - fail "Fireworks collector prints a valid record without credentials" "$no_key" -pass "Fireworks collector prints a valid record without credentials" - -result=$(python3 - "$ROOT/bin/omarchy-agent-usage-fireworks" "$auth_file" "$TEST_HOME/.config" "$TEST_HOME/.local/share" "$TEST_HOME/.pi/agent" <<'PY' -import importlib.machinery -import importlib.util -import json -import os -import sys -import time -from datetime import date -from decimal import Decimal -from pathlib import Path +# The Fireworks collector's own sources: the key from FIREWORKS_API_KEY, +# firectl's auth.ini, or a pi or opencode login; the billing API's usage +# buckets, paged; and the balance, live where the key can read the ledger and +# estimated from configured funding where it can't. The unused-machine record +# is held centrally by the partition test. +# +# Fireworks keeps a balance and no limit windows, and reuses its usage rows +# through a limits cache rather than a scan cache, so neither contract check +# fits it as they stand: check_balance_contract below holds the balance to the +# same probe contract, and the usage reuse windows have a check of their own. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" -collector_path = str(Path(sys.argv[1])) -auth_path = Path(sys.argv[2]) -os.environ["XDG_CONFIG_HOME"] = sys.argv[3] -os.environ["XDG_DATA_HOME"] = sys.argv[4] +agent_sandbox # 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. The env account id -# would override the config file, so it must not leak in from the runner. -os.environ["TZ"] = "UTC" -time.tzset() -os.environ.pop("FIREWORKS_ACCOUNT_ID", None) - -loader = importlib.machinery.SourceFileLoader("fireworks_collector", collector_path) -spec = importlib.util.spec_from_loader(loader.name, loader) -scanner = importlib.util.module_from_spec(spec) -loader.exec_module(scanner) -RealFireworksClient = scanner.FireworksClient - -payload = { - "serverlessCosts": [ - { - "startTime": "2026-07-31T00:00:00Z", - "promptTokens": "100", - "cachedPromptTokens": "40", - "uncachedPromptTokens": "60", - "completionTokens": "20", - "group": {"model_name": "accounts/fireworks/models/kimi-k2p5"}, - }, - { - "startTime": "2026-07-30T00:00:00Z", - "promptTokens": "300", - "cachedPromptTokens": "0", - "completionTokens": "50", - "group": {"model_name": "accounts/fireworks/models/deepseek-v3p2"}, - }, - { - "startTime": "2026-07-20T00:00:00Z", - "promptTokens": "500", - "cachedPromptTokens": "0", - "completionTokens": "100", - "group": {"model_name": "accounts/fireworks/models/kimi-k2p5"}, - }, - ] -} +# would shift by a day depending on where the test runs. +export TZ=UTC -summary = scanner.summarize_usage(payload, date(2026, 7, 31)) -api_key, account_id = scanner.read_auth_file(auth_path) -summary["apiKey"] = api_key -summary["accountId"] = account_id -summary["money"] = float(scanner.money_value({"units": "12", "nanos": 430000000})) - -# The opencode key only wins when no explicit key, firectl login, or pi -# login exists. Pin the pi dir first so a developer's real ~/.pi/agent -# credentials cannot leak into the fixture run. -pi_dir = Path(sys.argv[5]) -os.environ["PI_CODING_AGENT_DIR"] = str(pi_dir) -data_home = Path(os.environ["XDG_DATA_HOME"]) -opencode_auth = data_home / "opencode" / "auth.json" -opencode_auth.parent.mkdir(parents=True, exist_ok=True) -opencode_auth.write_text(json.dumps({"fireworks-ai": {"type": "api", "key": "fw_opencode"}})) -os.environ.pop("FIREWORKS_API_KEY", None) -opencode_key, _ = scanner.credentials(Path("/nonexistent/auth.ini"), {}) -firectl_key, _ = scanner.credentials(auth_path, {}) -summary["opencodeFallback"] = opencode_key == "fw_opencode" and firectl_key == "fw_test" - -# A pi login is the next rung below firectl, above opencode. -pi_dir.mkdir(parents=True, exist_ok=True) -(pi_dir / "auth.json").write_text(json.dumps({"fireworks": {"type": "api_key", "key": "fw_pi"}})) -pi_key, _ = scanner.credentials(Path("/nonexistent/auth.ini"), {}) -firectl_still_wins, _ = scanner.credentials(auth_path, {}) -summary["piFallback"] = pi_key == "fw_pi" and firectl_still_wins == "fw_test" - -# pi key syntax: $ENV references resolve, missing ones stay unresolved, and -# !command lookups are skipped and fall through to opencode. -os.environ["FW_TEST_ENV_KEY"] = "fw_from_env" -summary["piEnvKey"] = scanner.resolve_pi_key("$FW_TEST_ENV_KEY") == "fw_from_env" -summary["piMissingEnvKey"] = scanner.resolve_pi_key("$FW_DEFINITELY_MISSING") == "" -(pi_dir / "auth.json").write_text(json.dumps({"fireworks": {"type": "api_key", "key": "!secret-tool lookup fw"}})) -command_key, _ = scanner.credentials(Path("/nonexistent/auth.ini"), {}) -summary["piCommandKeySkipped"] = command_key == "fw_opencode" - -class WorkingClient: - def __init__(self, api_key, base_url): - pass - - def request(self, path, query=None, body=None): - raise scanner.FireworksError("The Fireworks API key cannot read billing data") - - def usage(self, account_id, start_day, end_day): - return payload - - def account(self, account_id): - return {} - - def spent(self, account_id, start_at, end_at): - return Decimal("8.60") - -class BalanceFailureClient(WorkingClient): - def spent(self, account_id, start_at, end_at): - raise scanner.FireworksError("Billing scope denied") - -class LiveBalanceClient(WorkingClient): - def request(self, path, query=None, body=None): - assert path.endswith(":getBalance") - return {"balance": {"units": "12", "nanos": 500000000}} - -os.environ["FIREWORKS_API_KEY"] = "fw_test" - -scanner.FireworksClient = WorkingClient -record = scanner.scan("https://example.invalid", auth_path) -summary["record"] = { - "schemaVersion": record["schemaVersion"], - "id": record["id"], - "ready": record["ready"], - "hasPromptStats": record["hasPromptStats"], - "scope": record["scope"], - "tierLabel": record["tierLabel"], - "limits": record["limits"], - "balance": record["balance"], -} +python3 - <<'PY' +from agent_usage_fixtures import * +from datetime import date -# billingUsage pages long ranges; usage() must follow continuation tokens. -pages = { - "": {"serverlessCosts": [{"startTime": "2026-07-30T00:00:00Z"}], "nextPageToken": "p2"}, - "p2": {"serverlessCosts": [{"startTime": "2026-07-31T00:00:00Z"}]}, +API = "https://api.fireworks.ai" +TODAY = time.strftime("%Y-%m-%d") +ROWS = [{"startTime": TODAY + "T00:00:00Z", "promptTokens": "100", "completionTokens": "20", "group": {"model_name": "accounts/fireworks/models/kimi-k2p5"}}] +LIVE = {"balance": {"units": "12", "nanos": 500000000}} +ROUTES = { + ":getBalance": LIVE, + "/billingUsage": {"serverlessCosts": ROWS}, + "/usageCosts:query": {"subtotal": {"units": "8", "nanos": 600000000}}, } -paging_client = RealFireworksClient("fw_test", "https://example.invalid") -paging_queries = [] -def paged_request(path, query=None, body=None): - paging_queries.append(dict(query or {})) - return pages[str((query or {}).get("pageToken") or "")] -paging_client.request = paged_request -merged = paging_client.usage("example", date(2026, 7, 1), date(2026, 8, 1)) -summary["paginationMerges"] = ( - len(merged["serverlessCosts"]) == 2 - and len(paging_queries) == 2 - and paging_queries[0]["startTime"] == "2026-07-01T00:00:00Z" -) - -scanner.FireworksClient = LiveBalanceClient -live = scanner.scan("https://example.invalid", auth_path) -summary["liveBalance"] = live["balance"] - -scanner.FireworksClient = BalanceFailureClient -scanned = scanner.scan("https://example.invalid", auth_path) -summary["balanceFailurePreservesTokens"] = ( - scanned["ready"] is True - and "balance" not in scanned - and scanned["modelUsage"]["kimi-k2.5"]["outputTokens"] == 120 - and scanned["usageStatusText"] == "Balance unavailable" -) +# A key without the billing role: the ledger is gated, so the balance is +# estimated from the configured funding and the rated costs since. +ESTIMATED = dict(ROUTES, **{":getBalance": 403}) +# Nor can it read the costs, so there is no balance at all. +NO_BALANCE = dict(ROUTES, **{":getBalance": 403, "/usageCosts:query": 403, "/billing/summary": 403}) + + +def signed_in(key="fw_test"): + auth = home() / ".fireworks" / "auth.ini" + auth.parent.mkdir(parents=True, exist_ok=True) + auth.write_text(f"[default]\napi_key = {key}\naccount_id = example\n") + write_json(config_home() / "omarchy" / "agents" / "fireworks.json", {"accountId": "example", "fundedAmount": 20, "fundedAt": "2026-07-01"}) + + +def fireworks(*args, routes=None, mode="up", force=True, **env): + saved = dict(os.environ) + os.environ.update(env) + try: + with FakeHttp(ROUTES if routes is None else routes, mode=mode) as http: + record = collect("fireworks", *(args or (("--force",) if force else ()))) + finally: + os.environ.clear() + os.environ.update(saved) + fireworks.asked = ["balance" if ":getBalance" in request["url"] else "usage" for request in http.requests] + return record + + +def remaining(record): + return (record.get("balance") or {}).get("remaining") + + +# check_limits_contract held to the balance instead of limit windows: a fresh +# probe, the 15-second reuse for a normal and a --limits-only run, --force, a +# failed check, an answer cut short or garbled and an unreachable service +# keeping the last balance stale, and a refused key. A candidate for the +# shared fixtures. +def check_balance_contract(plant, per_probe): + for path in limits_cache_files("fireworks"): + path.unlink() + plant() + fresh = fireworks() + check("fireworks: a fresh balance is live, stamped, healthy, and asked for once per probe", + remaining(fresh) is not None and fresh["limitsStale"] is False and fresh["limitsFetchedAt"] > 0 and fresh["usageStatusText"] == "" + and fresh["authHelpText"] == "" and not fresh.get("retryAdvised") and len(fireworks.asked) == per_probe, json.dumps(fresh)) + for args in ((), ("--limits-only",)): + reused = fireworks(*args, mode=500, force=False) + check(f"fireworks: a {'--limits-only' if args else 'normal'} run inside the probe interval reuses the answer, live", + fireworks.asked == [] and reused.get("balance") == fresh["balance"] and reused["limitsFetchedAt"] == fresh["limitsFetchedAt"] + and reused["limitsStale"] is False and reused["usageStatusText"] == "", json.dumps(reused)) + forced = fireworks() + check("fireworks: --force asks again inside the probe interval", len(fireworks.asked) == per_probe and forced["limitsStale"] is False, json.dumps(forced)) + kept = fireworks(mode=500) + check("fireworks: a failed check keeps the last balance, stale, with when it was measured", + kept.get("balance") == forced["balance"] and kept["limitsStale"] is True and kept["limitsFetchedAt"] == forced["limitsFetchedAt"] + and kept["usageStatusText"] == "" and not kept.get("retryAdvised"), json.dumps(kept)) + for mode in ("cut", "garbage"): + answered = fireworks(mode=mode) + check(f"fireworks: an answer {'cut short' if mode == 'cut' else 'that is not JSON'} keeps the last balance, stale, and advises no retry", + answered.get("balance") == forced["balance"] and answered["limitsStale"] is True and answered["limitsFetchedAt"] == forced["limitsFetchedAt"] + and answered["usageStatusText"] == "" and not answered.get("retryAdvised"), json.dumps(answered)) + for mode in ("down", "timeout"): + unreachable = fireworks(mode=mode) + check(f"fireworks: an unreachable service ({mode}) advises a retry and keeps the last balance, stale", + unreachable.get("retryAdvised") is True and unreachable.get("balance") == forced["balance"] and unreachable["limitsStale"] is True + and unreachable["usageStatusText"] == "", json.dumps(unreachable)) + refused = fireworks(mode=401) + check("fireworks: a refused key says 'Sign-in expired', advises no retry, and shows nothing as live", + refused["usageStatusText"] == "Sign-in expired" and not refused.get("retryAdvised") and refused["limitsStale"] is True, json.dumps(refused)) + return fresh + + +# ---------------------------------------------------------------- contract + +sandbox() +fresh = check_balance_contract(signed_in, 2) +expect("the record is account-wide and prepaid, with tokens and no prompt counts or limits", + {key: fresh.get(key) for key in ("ready", "hasPromptStats", "scope", "tierLabel", "limits", "todayTotalTokens", "hasLocalStats")}, + {"ready": True, "hasPromptStats": False, "scope": "account", "tierLabel": "Prepaid", "limits": [], "todayTotalTokens": 120, "hasLocalStats": True}, fresh) +expect("the live getBalance ledger is preferred when the key can read it", fresh["balance"], + {"remaining": 12.5, "funded": 20.0, "spent": 7.5, "currency": "USD", "estimated": False}, fresh) + +# The usage rows are this record's stats, reused like a local scan: briefly +# on a normal run and for 15 minutes on the near-limit refresh, which probes +# the balance like a normal run. +fireworks() +for path in cache_root().glob("fireworks-*.json"): + data = json.loads(path.read_text()) + write_json(path, dict(data, fetchedAtMs=data["fetchedAtMs"] - 60_000)) +fireworks("--limits-only") +limits_only = fireworks.asked +fireworks(force=False) +expect("a minute-old answer serves --limits-only its usage but not its balance, and a normal run neither", + [limits_only, fireworks.asked], [["balance"], ["usage"]]) + +# ------------------------------------------------------------- the balance + +# A key that can't read the ledger gets an estimate from the configured +# funding, less the costs rated since. +sandbox() +signed_in() +record = fireworks(routes=ESTIMATED) +expect("the balance is estimated from configured funding", record["balance"], + {"remaining": 11.4, "funded": 20.0, "spent": 8.6, "currency": "USD", "estimated": True}, record) + +# A failed balance check keeps the last good balance, stale, with the status +# left to limitsStale; with nothing kept it says the balance is unavailable, +# and the tokens still show. +record = fireworks(routes=NO_BALANCE) +expect("a failed balance check keeps the last balance, stale", + [remaining(record), record["limitsStale"], record["usageStatusText"], record["authHelpText"], record.get("retryAdvised")], + [11.4, True, "", "The Fireworks API key cannot read billing data", None], record) +for path in limits_cache_files("fireworks"): + path.unlink() +record = fireworks(routes=NO_BALANCE) +expect("a failed balance check with nothing kept says so and keeps the tokens", + [record["ready"], "balance" in record, record["modelUsage"]["kimi-k2.5"]["outputTokens"], record["usageStatusText"], record["limitsStale"]], + [True, False, 20, "Balance unavailable", True], record) + +# Only the billing-role gate (403) means the key has no live balance; any +# other failed or unreadable getBalance keeps the last live one, stale, never +# replacing it with an estimate. An overdrawn ledger has nothing left. +sandbox() +signed_in() +fireworks() +for name, answer in (("answers 500", 500), ("is cut short", lambda request: http.client.IncompleteRead(b"{", 9)), + ("answers an amount past a float", {"balance": {"units": "1e400"}})): + record = fireworks(routes=dict(ROUTES, **{":getBalance": answer})) + expect(f"a getBalance that {name} keeps the last live balance, stale", + [record["balance"], record["limitsStale"], record["usageStatusText"]], + [{"remaining": 12.5, "funded": 20.0, "spent": 7.5, "currency": "USD", "estimated": False}, True, ""], record) +record = fireworks(routes=dict(ROUTES, **{":getBalance": {"balance": {"units": "-5"}}})) +expect("an overdrawn live balance reads as nothing left", record["balance"], + {"remaining": 0.0, "funded": 20.0, "spent": 20.0, "currency": "USD", "estimated": False}, record) + +# FIREWORKS_API_BASE_URL moves the endpoint, never the key onto plain HTTP to +# another machine, nor into a URL that isn't one. +for url in ("http://evil.example", "not a url"): + record = fireworks(FIREWORKS_API_BASE_URL=url) + expect(f"Fireworks sends its key to no untrusted endpoint ({url})", + [fireworks.asked, record["limitsStale"], "HTTPS" in record["authHelpText"]], [[], True, True], record) + +# Usage and balance both unreachable with nothing kept: the stats source is +# unavailable. +sandbox() +signed_in() +record = fireworks(mode="down") +expect("an unreachable API with nothing kept says Fireworks is unavailable", + [record["ready"], record["usageStatusText"], record["limitsStale"], record.get("retryAdvised"), record.get("balance")], + [True, "Fireworks unavailable", True, True, None], record) + +# Another key, or the same key now pointed at another account, never shows +# these numbers. +fireworks() +record = fireworks(mode="down", FIREWORKS_API_KEY="fw_other") +expect("each key's answers are kept apart", [record.get("balance"), record["todayTotalTokens"], record["usageStatusText"]], [None, 0, "Fireworks unavailable"], record) +record = fireworks(mode="down", FIREWORKS_ACCOUNT_ID="elsewhere") +expect("each account's answers are kept apart", [record.get("balance"), record["todayTotalTokens"]], [None, 0], record) + +# ------------------------------------------------------------ the account + +# Without an account id anywhere, the key's one account is found and asked +# about; a key reaching several can't be guessed at. +sandbox() +(home() / ".fireworks").mkdir() +(home() / ".fireworks" / "auth.ini").write_text("[default]\napi_key = fw_discover\n") + + +def accounts(*names): + return {"/v1/accounts?": {"accounts": [{"name": f"accounts/{name}"} for name in names]}} + + +def asked_accounts(http): + return sorted({request["url"].split("/v1/accounts/")[1].split("/")[0].split(":")[0] for request in http.requests if "/v1/accounts/" in request["url"]}) + + +with FakeHttp(dict(accounts("acct-one"), **ROUTES)) as http: + record = collect("fireworks", "--force") +expect("Fireworks finds the key's one account and asks about it", + [asked_accounts(http), record["todayTotalTokens"], remaining(record), record["usageStatusText"]], [["acct-one"], 120, 12.5, ""], record) +with FakeHttp(dict(accounts("acct-one", "acct-two"), **ROUTES)) as http: + record = collect("fireworks", "--force") +expect("Fireworks asks for an accountId when the key reaches several accounts", + [asked_accounts(http), record["limitsStale"], "Set accountId" in record["authHelpText"]], [[], True, True], record) + +# An account list that isn't a list is a failed check: the kept answers stand. +with FakeHttp(dict(accounts("acct-one"), **ROUTES)): + collect("fireworks", "--force") +with FakeHttp(dict({"/v1/accounts?": {"accounts": None}}, **ROUTES)): + record = collect("fireworks", "--force") +expect("Fireworks keeps its last answers past an account list that isn't one", + [record["ready"], record["todayTotalTokens"], remaining(record), record["limitsStale"]], [True, 120, 12.5, True], record) + +# The key now reaching another account starts over: the first account's +# usage and balance never stand in for it when its own checks fail. +with FakeHttp(dict(accounts("acct-two"), **{":getBalance": 500, "/billingUsage": 500, "/usageCosts:query": 500})): + record = collect("fireworks", "--force") +expect("Fireworks never keeps another account's answers when the key moves", + [record["todayTotalTokens"], "balance" in record, record["usageStatusText"]], [0, False, "Fireworks unavailable"], record) + +# A key that can rate costs only through the billing summary still gets an +# estimate: the summary's line items are the spend since funding. +sandbox() +signed_in() +SUMMARY = dict(ROUTES, **{":getBalance": 403, "/usageCosts:query": 403, + "/billing/summary": {"lineItems": [{"totalCost": {"units": "3", "nanos": 250000000}}, {"totalCost": {"units": "1"}}, "x"]}}) +record = fireworks(routes=SUMMARY) +expect("Fireworks falls back to the billing summary to estimate the balance", record.get("balance"), + {"remaining": 15.75, "funded": 20.0, "spent": 4.25, "currency": "USD", "estimated": True}, record) + +# ------------------------------------------------------------ the answer + +module = load_collector("fireworks") + +# Today's uncached, cached and output tokens count once, cache apart in the +# model totals, over the 30-day window. +summary = module.summarize_usage({"serverlessCosts": [ + {"startTime": "2026-07-31T00:00:00Z", "promptTokens": "100", "cachedPromptTokens": "40", "uncachedPromptTokens": "60", "completionTokens": "20", + "group": {"model_name": "accounts/fireworks/models/kimi-k2p5"}}, + {"startTime": "2026-07-30T00:00:00Z", "promptTokens": "300", "cachedPromptTokens": "0", "completionTokens": "50", + "group": {"model_name": "accounts/fireworks/models/deepseek-v3p2"}}, + {"startTime": "2026-07-20T00:00:00Z", "promptTokens": "500", "cachedPromptTokens": "0", "completionTokens": "100", + "group": {"model_name": "accounts/fireworks/models/kimi-k2p5"}}, +]}, date(2026, 7, 31)) +expect("today's token categories count once, cache apart, over the 30-day window", + [summary["todayTotalTokens"], summary["recentDays"][-1]["messageCount"], summary["activeDays"], summary["modelUsage"]["kimi-k2.5"]], + [120, 120, 3, {"inputTokens": 560, "outputTokens": 120, "cacheReadInputTokens": 40, "cacheCreationInputTokens": 0}], summary) +expect("untrusted token fields are skipped", module.summarize_usage({"serverlessCosts": [ + {"startTime": TODAY + "T00:00:00Z", "promptTokens": True, "completionTokens": "NaN", "group": {"model_name": "a"}}, + {"startTime": TODAY + "T00:00:00Z", "promptTokens": 1e300, "completionTokens": "inf", "group": {"model_name": "b"}}, + {"startTime": TODAY + "T00:00:00Z", "promptTokens": -5, "completionTokens": 7, "group": {"model_name": "c"}}, +]})["modelUsage"], {"c": {"inputTokens": 0, "outputTokens": 7, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}}) +expect("a bucket time at the calendar's edge has no day, and its tokens still count", + [module.row_date({"startTime": "9999-12-31T23:00:00-05:00"}), module.row_date({"startTime": "0001-01-01T00:00:00+05:00"}), + module.summarize_usage({"serverlessCosts": [{"startTime": "9999-12-31T23:00:00-05:00", "completionTokens": "7"}]})["modelUsage"]["unknown"]["outputTokens"]], + ["", "", 7]) +expect("Money reads units and nanos", float(module.money_value({"units": "12", "nanos": 430000000})), 12.43) + +# billingUsage pages long ranges, so every continuation token is followed, +# from a window that opens at local midnight. +client = module.FireworksClient("fw_test", API) +with FakeHttp({"pageToken=p2": {"serverlessCosts": [{"startTime": "2026-07-31T00:00:00Z"}]}, + "/billingUsage": {"serverlessCosts": [{"startTime": "2026-07-30T00:00:00Z"}], "nextPageToken": "p2"}}) as http: + merged = client.usage("example", date(2026, 7, 1), date(2026, 8, 1)) +first = urllib.parse.parse_qs(urllib.parse.urlsplit(http.requests[0]["url"]).query) +expect("billingUsage continuation tokens are followed", [len(merged["serverlessCosts"]), len(http.requests), first["startTime"]], [2, 2, ["2026-07-01T00:00:00Z"]]) + +# A reply that isn't an object, or usage rows that aren't a list, fail the +# check rather than passing for an idle account; a missing list is an idle +# month, since the API leaves empty lists out. +def fails(call): + try: + call() + except module.FireworksError: + return True + return False + + +with FakeHttp({"/v1/accounts": []}): + odd_object = fails(lambda: client.request("/v1/accounts")) +with FakeHttp({"/billingUsage": {"serverlessCosts": "rows"}}): + odd_rows = fails(lambda: client.usage("example", date(2026, 7, 1), date(2026, 8, 1))) +with FakeHttp({"/usageCosts:query": 403, "/billing/summary": {"lineItems": None}}): + odd_items = fails(lambda: client.spent("example", "2026-07-01T00:00:00Z", "2026-07-31T00:00:00Z")) +with FakeHttp({"/billingUsage": {}}): + idle = client.usage("example", date(2026, 7, 1), date(2026, 8, 1)) +expect("a reply of the wrong shape fails, and a missing usage list reads as idle", [odd_object, odd_rows, odd_items, idle], [True, True, True, {"serverlessCosts": []}]) # East of Greenwich, a local-midnight bucket starts on the previous UTC date; -# the row must still land on the local day it names, and the query window -# must ask for local midnights expressed in UTC. +# the row still lands on the local day it names, and the query window asks +# for local midnights expressed in UTC. +saved_tz = os.environ.get("TZ") os.environ["TZ"] = "Etc/GMT-2" time.tzset() -summary["bucketDayIsLocal"] = scanner.row_date({"startTime": "2026-07-30T22:00:00Z"}) == "2026-07-31" -summary["windowIsLocalMidnight"] = scanner.local_midnight_utc(date(2026, 7, 31)) == "2026-07-30T22:00:00Z" -print(json.dumps(summary, separators=(",", ":"))) +try: + east = [module.row_date({"startTime": "2026-07-30T22:00:00Z"}), module.local_midnight_utc(date(2026, 7, 31))] +finally: + if saved_tz is None: + os.environ.pop("TZ", None) + else: + os.environ["TZ"] = saved_tz + time.tzset() +expect("buckets are dated by local day, and windows open at local midnight in UTC", east, ["2026-07-31", "2026-07-30T22:00:00Z"]) + +# ---------------------------------------------------------- the credentials + +# firectl's auth.ini, then a pi login, then opencode's; pi's key resolves a +# $VAR reference and leaves a !command lookup to pi. +sandbox() +auth = home() / ".fireworks" / "auth.ini" +signed_in() +expect("firectl's auth.ini gives the key and account", list(module.read_auth_file(auth)), ["fw_test", "example"]) +unreadable = home() / "latin1.ini" +unreadable.write_bytes(b"[default]\napi_key = fw_caf\xe9\n") +expect("an auth.ini that isn't UTF-8 reads as no key", list(module.read_auth_file(unreadable)), ["", ""]) +missing = home() / "missing.ini" +opencode_auth(None, "fireworks-ai", "fw_opencode") +expect("opencode's key is the last resort", [module.credentials(missing, {})[0], module.credentials(auth, {})[0]], ["fw_opencode", "fw_test"]) +pi_auth(None, "fireworks", "fw_pi") +expect("a pi login comes below firectl and above opencode", [module.credentials(missing, {})[0], module.credentials(auth, {})[0]], ["fw_pi", "fw_test"]) +os.environ["FW_TEST_ENV_KEY"] = "fw_from_env" +expect("pi's $VAR keys resolve, and a missing one stays unresolved", + [module.resolve_pi_key("$FW_TEST_ENV_KEY"), module.resolve_pi_key("$FW_DEFINITELY_MISSING")], ["fw_from_env", ""]) +pi_auth(None, "fireworks", "!secret-tool lookup fw") +expect("pi's !command keys are skipped for the next rung", module.credentials(missing, {})[0], "fw_opencode") PY -) - -[[ $(jq -r '.todayTotalTokens' <<<"$result") == "120" ]] || - fail "Fireworks collector totals today's uncached, cached, and output tokens once" "$result" -pass "Fireworks collector totals today's token categories once" - -[[ $(jq -c '.modelUsage["kimi-k2.5"]' <<<"$result") == '{"inputTokens":560,"outputTokens":120,"cacheReadInputTokens":40,"cacheCreationInputTokens":0}' ]] || - fail "Fireworks collector keeps cache separate in model totals" "$result" -pass "Fireworks collector keeps cache separate in model totals" - -[[ $(jq -r '.recentDays[-1].messageCount' <<<"$result") == "120" ]] || - fail "Fireworks collector builds the seven-day token series" "$result" -pass "Fireworks collector builds the seven-day token series" - -[[ $(jq -r '.activeDays' <<<"$result") == "3" ]] || - fail "Fireworks collector retains the 30-day model window" "$result" -pass "Fireworks collector retains the 30-day model window" - -[[ $(jq -r '.apiKey + ":" + .accountId' <<<"$result") == "fw_test:example" ]] || - fail "Fireworks collector reads firectl credentials" "$result" -pass "Fireworks collector reads firectl credentials" - -[[ $(jq -r '.money' <<<"$result") == "12.43" ]] || - fail "Fireworks collector parses Money units and nanos" "$result" -pass "Fireworks collector parses Money units and nanos" - -[[ $(jq -c '.record | {schemaVersion, id, ready, hasPromptStats, scope, tierLabel, limits}' <<<"$result") == '{"schemaVersion":1,"id":"fireworks","ready":true,"hasPromptStats":false,"scope":"account","tierLabel":"Prepaid","limits":[]}' ]] || - fail "Fireworks collector prints the display-ready record contract" "$result" -pass "Fireworks collector prints the display-ready record contract" - -[[ $(jq -r '.paginationMerges' <<<"$result") == "true" ]] || - fail "Fireworks collector follows billingUsage continuation tokens" "$result" -pass "Fireworks collector follows billingUsage continuation tokens" - -[[ $(jq -r '.windowIsLocalMidnight' <<<"$result") == "true" ]] || - fail "Fireworks collector requests local-midnight windows in UTC" "$result" -pass "Fireworks collector requests local-midnight windows in UTC" - -[[ $(jq -c '.record.balance' <<<"$result") == '{"remaining":11.4,"funded":20.0,"spent":8.6,"currency":"USD","estimated":true}' ]] || - fail "Fireworks collector estimates the balance from configured funding" "$result" -pass "Fireworks collector estimates the balance from configured funding" - -[[ $(jq -c '.liveBalance' <<<"$result") == '{"remaining":12.5,"funded":20.0,"spent":7.5,"currency":"USD","estimated":false}' ]] || - fail "Fireworks collector prefers the live getBalance ledger when the key can read it" "$result" -pass "Fireworks collector prefers the live getBalance ledger when the key can read it" - -[[ $(jq -r '.balanceFailurePreservesTokens' <<<"$result") == "true" ]] || - fail "Fireworks collector preserves tokens when balance lookup fails" "$result" -pass "Fireworks collector preserves tokens when balance lookup fails" - -[[ $(jq -r '.bucketDayIsLocal' <<<"$result") == "true" ]] || - fail "Fireworks collector dates buckets by local day east of Greenwich" "$result" -pass "Fireworks collector dates buckets by local day east of Greenwich" - -[[ $(jq -r '.opencodeFallback' <<<"$result") == "true" ]] || - fail "Fireworks collector falls back to the opencode key last" "$result" -pass "Fireworks collector falls back to the opencode key last" - -[[ $(jq -r '.piFallback' <<<"$result") == "true" ]] || - fail "Fireworks collector falls back to the pi login before opencode" "$result" -pass "Fireworks collector falls back to the pi login before opencode" - -[[ $(jq -r '.piEnvKey' <<<"$result") == "true" ]] || - fail "Fireworks collector resolves env-var pi keys" "$result" -pass "Fireworks collector resolves env-var pi keys" - -[[ $(jq -r '.piMissingEnvKey' <<<"$result") == "true" ]] || - fail "Fireworks collector treats missing env references as unresolved" "$result" -pass "Fireworks collector treats missing env references as unresolved" - -[[ $(jq -r '.piCommandKeySkipped' <<<"$result") == "true" ]] || - fail "Fireworks collector skips command-lookup pi keys and falls through" "$result" -pass "Fireworks collector skips command-lookup pi keys and falls through" diff --git a/test/shell.d/agent-usage-grok-test.sh b/test/shell.d/agent-usage-grok-test.sh index eeaf985eaf3..00a9c3961e6 100755 --- a/test/shell.d/agent-usage-grok-test.sh +++ b/test/shell.d/agent-usage-grok-test.sh @@ -1,180 +1,306 @@ #!/bin/bash +# The Grok collector's own sources: the sign-in and plan the Grok CLI keeps in +# its home, xAI's credits answer, and the session summaries and usage ledgers +# beside each session. The unused-machine record is held centrally by the +# partition test, and the limits and scan caches by the contract checks. + set -euo pipefail source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" -require_command jq -require_command python3 - -test_tmp=$(mktemp -d) -trap 'rm -rf "$test_tmp"' EXIT +agent_sandbox -export HOME="$test_tmp/home" -export XDG_STATE_HOME="$test_tmp/state" -export XDG_CACHE_HOME="$test_tmp/cache" +python3 - <<'PY' +from agent_usage_fixtures import * -future=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=6)).isoformat())') -past=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=1)).isoformat())') -period_end=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) + dt.timedelta(days=9)).isoformat())') +PERIOD_END = iso_in(9 * 86400) +USED = {"token-main": 42.0, "token-side": 7.0, "token-fresh": 0} +answer = {"percent": None, "end": None} -# A home signed in the way the Grok CLI leaves it. -signed_in() { - mkdir -p "$1" - jq -nc --arg token "$2" --arg user "$3" --arg expires "$4" \ - '{"https://auth.x.ai::client": {key: $token, user_id: $user, email: ($user + "@example.com"), expires_at: $expires}}' >"$1/auth.json" - printf '{"payload":"{\\"settings\\":{\\"subscription_tier_display\\":\\"%s\\"}}"}\n' "$5" >"$1/settings_cache.json" -} - -# xAI answers per token, so each home's probe is told apart by its sign-in. -collect() { - COLLECTOR="$ROOT/bin/omarchy-agent-usage-grok" python3 - "$@" <<'PY' -import importlib.machinery, importlib.util, io, json, os, sys - -loader = importlib.machinery.SourceFileLoader("collector", os.environ["COLLECTOR"]) -spec = importlib.util.spec_from_loader(loader.name, loader) -collector = importlib.util.module_from_spec(spec) -loader.exec_module(collector) - -used = {"token-main": 42.0, "token-side": 7.0, "token-fresh": 0} -end = os.environ["PERIOD_END"] # Shaped like xAI's answer: protobuf JSON under config, which drops a -# percentage of zero. -def urlopen(request, timeout=None): - assert request.full_url == collector.CREDITS_URL +# percentage of zero. A unified-billing account has no percentage there at +# all; its monthly quota comes on the billing answer without the credits +# format, its amounts money fields of {"val": n}. +def credits(request): token = request.get_header("Authorization").split(" ", 1)[1] - config = {"currentPeriod": {"type": "USAGE_PERIOD_TYPE_WEEKLY", "end": end}, "billingPeriodEnd": end} - if used.get(token): - config["creditUsagePercent"] = used[token] - return io.BytesIO(json.dumps({"config": config}).encode()) - -collector.urllib.request.urlopen = urlopen -sys.argv = ["omarchy-agent-usage-grok"] + (os.environ.get("COLLECT_ARGS") or "--force").split() -collector.main() -PY -} -export PERIOD_END="$period_end" + config = {"currentPeriod": {"type": "USAGE_PERIOD_TYPE_WEEKLY", "end": answer["end"] or PERIOD_END}, "billingPeriodEnd": PERIOD_END} + if USED.get(token): + config["creditUsagePercent"] = USED[token] + if answer["percent"] is not None: + config["creditUsagePercent"] = answer["percent"] + if token == "token-unified": + config["isUnifiedBillingUser"] = True + return {"config": config} + + +def monthly(request): + return {"config": {"isUnifiedBillingUser": True, "billingPeriodStart": iso_in(-20 * 86400), "billingPeriodEnd": PERIOD_END, + "monthlyLimit": {"val": 1000}, "used": {"val": 900}}} + + +ROUTES = {"cli-chat-proxy.grok.com/v1/billing?format=credits": credits, "cli-chat-proxy.grok.com/v1/billing": monthly} + + +# A home signed in the way the Grok CLI leaves it. +def signed_in(grok_home, token="token-main", user="u-main", expires=None, tier="X Premium+", **extra): + login = {"key": token, "user_id": user, "email": user + "@example.com", "expires_at": expires or iso_in(6 * 3600), **extra} + write_json(Path(grok_home) / "auth.json", {"https://auth.x.ai::client": login}) + write_json(Path(grok_home) / "settings_cache.json", {"payload": json.dumps({"settings": {"subscription_tier_display": tier}})}) + -record=$(collect) -[[ $(jq -c '{ready, tierLabel, limits}' <<<"$record") == '{"ready":false,"tierLabel":"","limits":[]}' ]] || - fail "a machine nobody signed in to Grok on gives an empty record" "$record" -pass "a machine nobody signed in to Grok on gives an empty record" +def summary(grok_home, name, active_at, prompts): + write_json(Path(grok_home) / "sessions" / "%2Fhome%2Fme" / name / "summary.json", {"last_active_at": active_at, "num_messages": prompts}) + + +def grok(*args): + with FakeHttp(ROUTES): + return collect("grok", *(args or ("--force",))) + + +def grok_home(): + return home() / ".grok" + + +# ---------------------------------------------------------------- the home + +# The partition test holds the rest of the unused-machine record; whether a +# record's stats are this machine's is each collector's to say. +sandbox() +record = grok() +expect("a machine without Grok has no local Grok stats", [record["ready"], record["hasLocalStats"]], [False, False], record) + +# A Grok home with no sign-in in it has been set up, so the sign-in is what's +# missing. +(home() / "bare-grok").mkdir() +os.environ["GROK_HOME"] = str(home() / "bare-grok") +record = grok() +expect("a signed-out Grok home waits for auth", [record["ready"], record["usageStatusText"]], [False, "Waiting for auth"], record) + +# Sessions on disk put Grok in use here even with no sign-in, which is then +# what's missing. +os.environ["GROK_HOME"] = str(home() / "orphan") +summary(home() / "orphan", "one", "2026-01-02T10:00:00Z", 2) +record = grok() +expect("Grok sessions with no sign-in are in use, waiting for auth", + [record["ready"], record["totalSessions"], record["usageStatusText"]], [True, 1, "Waiting for auth"], record) # A home chosen with GROK_HOME is the one read, as the CLI itself does. -signed_in "$test_tmp/custom-grok" token-main u-main "$future" "SuperGrok" -record=$(GROK_HOME="$test_tmp/custom-grok" collect) -[[ $(jq -c '{ready, tierLabel, percent: .limits[0].percent}' <<<"$record") == '{"ready":true,"tierLabel":"SuperGrok","percent":0.42}' ]] || - fail "a single Grok account in GROK_HOME is read from there" "$record" -pass "a single Grok account in GROK_HOME is read from there" -rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json - -signed_in "$HOME/.grok" token-main u-main "$future" "X Premium+" -record=$(collect) -[[ $(jq -c '{ready, tierLabel, stale: .limitsStale, label: .limits[0].label, percent: .limits[0].percent}' <<<"$record") == '{"ready":true,"tierLabel":"X Premium+","stale":false,"label":"Weekly","percent":0.42}' ]] || - fail "Grok's plan and credits come from its own home" "$record" -[[ -n $(jq -r '.limits[0].resetsAt' <<<"$record") ]] || fail "the credits window says when the period ends" "$record" -[[ $(jq -r '.authHelpText' <<<"$record") == "" ]] || fail "a successful Grok probe drops the login hint" "$record" -record=$(collect) -[[ $(jq -c '{percent: .limits[0].percent, help: .authHelpText}' <<<"$record") == '{"percent":0.42,"help":""}' ]] || - fail "reused Grok limits drop the login hint" "$record" -pass "Grok's plan and credits come from its own home" - -signed_in "$HOME/.grok" token-main u-main "$past" "X Premium+" -record=$(collect) -[[ $(jq -c '{usageStatusText, first: .limits[0].percent, stale: .limitsStale}' <<<"$record") == '{"usageStatusText":"Sign-in expired","first":0.42,"stale":true}' ]] || - fail "a lapsed sign-in keeps the last credits and says so" "$record" -pass "a lapsed sign-in keeps the last credits and says so" +signed_in(home() / "custom-grok", tier="SuperGrok") +os.environ["GROK_HOME"] = str(home() / "custom-grok") +record = grok() +expect("a single Grok account in GROK_HOME is read from there", + [record["ready"], record["tierLabel"], record["limits"][0]["percent"]], [True, "SuperGrok", 0.42], record) +del os.environ["GROK_HOME"] + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_session(): + add_session.count = getattr(add_session, "count", 0) + 1 + summary(grok_home(), f"grown-{add_session.count}", iso_in(0), 1) + + +adapter = Adapter("grok", plant=lambda: signed_in(grok_home()), answer=ROUTES, grow=add_session) +fresh = check_limits_contract(adapter) +expect("Grok's plan and weekly credits come from its own home and xAI's answer", + [fresh["tierLabel"], fresh["limits"][0]["label"], fresh["limits"][0]["percent"], bool(fresh["limits"][0]["resetsAt"])], + ["X Premium+", "Weekly", 0.42, True], fresh) +check_scan_cache(adapter) + +# ------------------------------------------------------------ the sign-in + +sandbox() +signed_in(grok_home()) +grok() +signed_in(grok_home(), expires=iso_in(-3600)) +record = grok() +expect("a lapsed sign-in keeps the last credits and says so", + [record["usageStatusText"], record["limits"][0]["percent"], record["limitsStale"]], ["Sign-in expired", 0.42, True], record) # A lapsed access token with a refresh token is routine: Grok renews it when it -# starts, and while it hasn't run nothing here has spent its allowance, so the -# last numbers stand as current, with no sign-in asked for. -jq '.[].refresh_token = "r"' "$HOME/.grok/auth.json" >"$test_tmp/auth.json" -mv "$test_tmp/auth.json" "$HOME/.grok/auth.json" -record=$(collect) -[[ $(jq -c '{usageStatusText, first: .limits[0].percent, stale: .limitsStale}' <<<"$record") == '{"usageStatusText":"","first":0.42,"stale":false}' ]] || - fail "a lapsed access token with a refresh token keeps showing the last numbers" "$record" -pass "a lapsed access token with a refresh token keeps showing the last numbers" - -# A week that reset while Grok sat idle starts over at 0%, a week later. -cache=$(ls "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json | head -1) -reset_past=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=1)).isoformat())') -jq --arg at "$reset_past" '.limits[0].resetsAt = $at' "$cache" >"$test_tmp/cache.json" -mv "$test_tmp/cache.json" "$cache" -record=$(collect) -next_reset=$(jq -r '.limits[0].resetsAt' <<<"$record") -[[ $(jq -r '.limits[0].percent' <<<"$record") == 0.0 ]] && - python3 -c 'import datetime as dt, sys; n = dt.datetime.fromisoformat(sys.argv[1]); d = n - dt.datetime.now(dt.timezone.utc); sys.exit(0 if dt.timedelta(days=6) < d < dt.timedelta(days=7) else 1)' "$next_reset" || - fail "a week that reset while Grok was idle starts over a week later" "$record" -pass "a week that reset while Grok was idle starts over a week later" +# starts, so the limits are paused, the last numbers kept as stale, with no +# sign-in asked for. +signed_in(grok_home(), expires=iso_in(-3600), refresh_token="r") +record = grok() +expect("a lapsed access token with a refresh token pauses the limits, keeping the last numbers stale", + [record["usageStatusText"], record["limits"][0]["percent"], record["limitsStale"], "Start Grok" in record["authHelpText"]], + ["Limits paused", 0.42, True, True], record) + +# A week that reset while Grok sat idle starts over at 0%, its next reset +# unknown until xAI says. +edit_limits_cache("grok", resets_past=True) +record = grok() +expect("a week that reset while Grok was idle starts over at zero", + [record["limits"][0]["percent"], record["limits"][0]["resetsAt"], record["limitsStale"]], [0.0, "", True], record) # Without cached numbers there's nothing to show, so it says how to get them. -rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json -record=$(collect) -[[ $(jq -c '{usageStatusText, limits}' <<<"$record") == '{"usageStatusText":"Limits paused","limits":[]}' ]] || - fail "a lapsed token with nothing cached says to start Grok" "$record" -pass "a lapsed token with nothing cached says to start Grok" +for path in limits_cache_files("grok"): + path.unlink() +record = grok() +expect("a lapsed token with nothing cached says to start Grok, with no limits to call stale", [record["usageStatusText"], record["limits"], record["limitsStale"]], ["Limits paused", [], False], record) # A refresh token past Grok's 30-day sign-in can't renew anything: signed out. -long_ago=$(python3 -c 'import datetime as dt; print((dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=31)).isoformat())') -jq --arg at "$long_ago" '.[].expires_at = $at' "$HOME/.grok/auth.json" >"$test_tmp/auth.json" -mv "$test_tmp/auth.json" "$HOME/.grok/auth.json" -record=$(collect) -[[ $(jq -r '.usageStatusText' <<<"$record") == "Sign-in expired" ]] || - fail "a sign-in lapsed past Grok's 30 days asks for a sign-in" "$record" -pass "a sign-in lapsed past Grok's 30 days asks for a sign-in" +signed_in(grok_home(), expires=iso_in(-31 * 86400), refresh_token="r") +record = grok() +expect("a sign-in lapsed past Grok's 30 days asks for a sign-in", record["usageStatusText"], "Sign-in expired", record) + +# ------------------------------------------------------------- the answer + +# A percentage that isn't a number is no answer, not a full allowance. +sandbox() +signed_in(grok_home()) +answer["percent"] = float("nan") +record = grok() +answer["percent"] = None +expect("a Grok percentage that isn't a number reads as unavailable", + [record["usageStatusText"], record["limits"]], ["Grok limits unavailable", []], record) # A period with nothing used yet comes without a percentage. -signed_in "$HOME/.grok" token-fresh u-main "$future" "X Premium+" -rm -f "$XDG_CACHE_HOME"/omarchy/agent-usage/grok-limits-*.json -record=$(collect) -[[ $(jq -c '.limits[0] | {label, percent}' <<<"$record") == '{"label":"Weekly","percent":0.0}' ]] || - fail "an untouched period reads as 0% rather than nothing known" "$record" -pass "an untouched period reads as 0% rather than nothing known" - -signed_in "$HOME/.grok" token-main u-main "$future" "X Premium+" -side="$XDG_STATE_HOME/omarchy/agents/accounts/grok/side" -signed_in "$side" token-side u-side "$future" "SuperGrok" -jq -n --arg side "$side" '{active: "side", switch: "auto", threshold: 90, accounts: [ - {id: "main", label: "Main", home: "", primary: true}, - {id: "side", label: "Side", home: $side, primary: false} -]}' >"$XDG_STATE_HOME/omarchy/agents/accounts/grok.json" -record=$(collect) -[[ $(jq -c '{tierLabel, first: .limits[0].percent, accounts: [.accounts[] | {id, active, plan, percent: .limits[0].percent}], switch: .accountSwitch}' <<<"$record") == '{"tierLabel":"SuperGrok","first":0.07,"accounts":[{"id":"main","active":false,"plan":"X Premium+","percent":0.42},{"id":"side","active":true,"plan":"SuperGrok","percent":0.07}],"switch":{"mode":"auto","threshold":90}}' ]] || - fail "every Grok account reports its own plan and credits" "$record" -pass "every Grok account reports its own plan and credits" - -# Sessions come from each session's summary; tokens and today's prompts from -# its usage.json, one entry per finished turn, with cached input kept apart. -sessions="$HOME/.grok/sessions/%2Fhome%2Fme" -now=$(python3 -c 'import datetime as dt; print(dt.datetime.now(dt.timezone.utc).isoformat())') -mkdir -p "$sessions/today-1" "$sessions/today-2" "$sessions/old" -jq -n --arg at "$now" '{last_active_at: $at, num_messages: 3}' >"$sessions/today-1/summary.json" -jq -n --arg at "$now" '{last_active_at: $at, num_messages: 2}' >"$sessions/today-2/summary.json" -jq -n '{last_active_at: "2026-01-02T10:00:00Z", num_messages: 5}' >"$sessions/old/summary.json" -jq -n --arg at "$now" '{turns: [ - {endedAt: $at, modelUsage: {"grok-4.7-build": {inputTokens: 1000, cachedReadTokens: 800, outputTokens: 50, cacheCreationTokens: 0, totalTokens: 1050}}}, - {endedAt: $at, modelUsage: {"grok-4.7-build": {inputTokens: 200, cachedReadTokens: 100, outputTokens: 10, cacheCreationTokens: 5, totalTokens: 210}}} -]}' >"$sessions/today-1/usage.json" -jq -n '{turns: [{endedAt: "2026-01-02T10:00:00Z", modelUsage: {"grok-4.6": {inputTokens: 500, cachedReadTokens: 0, outputTokens: 20, cacheCreationTokens: 0, totalTokens: 520}}}]}' >"$sessions/old/usage.json" -record=$(collect) -[[ $(jq -c '{hasLocalStats, todayPrompts, todaySessions, totalPrompts, totalSessions, activeDays, todayTotalTokens, todayTokensByModel, today: .recentDays[6].messageCount}' <<<"$record") == '{"hasLocalStats":true,"todayPrompts":2,"todaySessions":2,"totalPrompts":10,"totalSessions":3,"activeDays":2,"todayTotalTokens":1265,"todayTokensByModel":{"grok-4.7-build":1265},"today":1265}' ]] || - fail "Grok counts sessions, prompts, and tokens from its session files" "$record" -[[ $(jq -cS '.modelUsage' <<<"$record") == '{"grok-4.6":{"cacheCreationInputTokens":0,"cacheReadInputTokens":0,"inputTokens":500,"outputTokens":20},"grok-4.7-build":{"cacheCreationInputTokens":5,"cacheReadInputTokens":900,"inputTokens":300,"outputTokens":60}}' ]] || - fail "Grok's tokens by model keep cached input apart" "$record" -pass "Grok counts sessions, prompts, and tokens from its session files" - -# A limits-only refresh reuses the last scan rather than reading every summary. -mkdir -p "$sessions/later" -jq -n --arg at "$now" '{last_active_at: $at, num_messages: 1}' >"$sessions/later/summary.json" -record=$(COLLECT_ARGS="--limits-only" collect) -[[ $(jq -r '.totalSessions' <<<"$record") == 3 ]] || fail "a limits-only refresh reuses the session scan" "$record" -pass "a limits-only refresh reuses the session scan" - -# A scan from another day is never reused, so yesterday's sessions don't -# count as today's after midnight. -jq '.day = "2000-01-01"' "$XDG_CACHE_HOME/omarchy/agent-usage/grok-stats.json" >"$test_tmp/stats.json" -mv "$test_tmp/stats.json" "$XDG_CACHE_HOME/omarchy/agent-usage/grok-stats.json" -record=$(COLLECT_ARGS="--limits-only" collect) -[[ $(jq -r '.totalSessions' <<<"$record") == 4 ]] || fail "a scan from another day is made again" "$record" -pass "a scan from another day is made again" +for path in limits_cache_files("grok"): + path.unlink() +signed_in(grok_home(), token="token-fresh") +record = grok() +expect("an untouched period reads as 0% rather than nothing known", + [record["limits"][0]["label"], record["limits"][0]["percent"]], ["Weekly", 0.0], record) + +# A period already over without a percentage says nothing about this one: no +# answer, the last numbers kept as stale. +answer["end"] = iso_in(-3600) +record = grok() +answer["end"] = None +expect("an ended period without a percentage is no answer, keeping the last numbers", + [record["limits"][0]["percent"], record["limitsStale"]], [0.0, True], record) + +# A unified-billing account leaves the percentage out of the credits answer, +# its usage counted against the monthly quota the plain billing answer gives. +sandbox() +signed_in(grok_home(), token="token-unified") +record = grok() +expect("a unified-billing account reads its monthly quota, not a confident 0%", + [record["limits"][0]["label"], record["limits"][0]["percent"], record["limitsStale"], bool(record["limits"][0]["resetsAt"])], + ["Monthly", 0.9, False, True], record) + +# A unified-billing account can also carry a weekly percentage on the credits +# answer; Grok's /usage then shows the monthly included quota beside it. +sandbox() +signed_in(grok_home(), token="token-unified") +answer["percent"] = 30 +record = grok() +answer["percent"] = None +expect("a unified-billing account with weekly credits shows the monthly quota beside them", + [[limit["label"], limit["percent"]] for limit in record["limits"]], [["Weekly", 0.3], ["Monthly", 0.9]], record) + +# ------------------------------------------------------------- accounts + +sandbox() +signed_in(grok_home()) +side = state_home() / "omarchy" / "agents" / "accounts" / "grok" / "side" +signed_in(side, token="token-side", user="u-side", tier="SuperGrok") +registry("grok", [ + {"id": "main", "label": "Main", "home": "", "primary": True}, + {"id": "side", "label": "Side", "home": str(side), "primary": False}, +], active="side", switch="auto", threshold=90) +record = grok() +expect("every Grok account reports its own plan and credits", { + "tierLabel": record["tierLabel"], + "first": record["limits"][0]["percent"], + "accounts": [{"id": a["id"], "active": a["active"], "plan": a["plan"], "percent": a["limits"][0]["percent"]} for a in record["accounts"]], + "switch": record["accountSwitch"], +}, { + "tierLabel": "SuperGrok", + "first": 0.07, + "accounts": [{"id": "main", "active": False, "plan": "X Premium+", "percent": 0.42}, {"id": "side", "active": True, "plan": "SuperGrok", "percent": 0.07}], + "switch": {"mode": "auto", "threshold": 90}, +}, record) + +# ------------------------------------------------------------ odd shapes + +# A ledger or settings cache of the wrong shape costs that field, never the +# record. +sandbox() +signed_in(grok_home()) +write_json(grok_home() / "settings_cache.json", {"payload": json.dumps({"settings": "x"})}) +summary(grok_home(), "odd", iso_in(0), 1) +write_json(grok_home() / "sessions" / "%2Fhome%2Fme" / "odd" / "usage.json", {"turns": [{"endedAt": iso_in(0), "modelUsage": ["grok-4.7-build"]}]}) +record = grok() +expect("a list-shaped modelUsage and a non-object settings read as empty", + [record["ready"], record["tierLabel"], record["todayPrompts"], record["modelUsage"]], [True, "", 1, {}], record) + +# A session or turn time out of range (year 1, before 2000, past 2100) is an +# untrusted stamp: it skips that session or turn, never the record. +sandbox() +signed_in(grok_home()) +summary(grok_home(), "good", iso_in(0), 2) +write_json(grok_home() / "sessions" / "%2Fhome%2Fme" / "good" / "usage.json", {"turns": [ + {"endedAt": iso_in(0), "modelUsage": {"grok-4.7-build": {"inputTokens": 10, "outputTokens": 1}}}, + {"endedAt": "0001-01-01T00:00:00+00:00", "modelUsage": {"grok-4.7-build": {"inputTokens": 500}}}, + {"endedAt": "9999-12-31T23:59:59+00:00", "modelUsage": {"grok-4.7-build": {"inputTokens": 500}}}, +]}) +for name, stamp in (("year1", "0001-01-01T00:00:00+00:00"), ("1970", "1970-01-01T00:00:00Z"), ("far", "2999-01-01T00:00:00Z")): + summary(grok_home(), name, stamp, 7) +record = grok() +expect("Grok skips sessions and turns with out-of-range times", + [record["totalSessions"], record["totalPrompts"], record["todayPrompts"], record["todayTotalTokens"], record["activeDates"]], + [1, 1, 1, 11, [datetime.now().date().isoformat()]], record) + +# ---------------------------------------------------------- local stats + +# Sessions come from each session's summary; tokens and prompts from its +# usage.json, one entry per finished turn, with cached input kept apart. The +# summary's num_messages counts the model's replies too, so it isn't read. +sandbox() +signed_in(grok_home()) +now = iso_in(0) +sessions = grok_home() / "sessions" / "%2Fhome%2Fme" +summary(grok_home(), "today-1", now, 3) +summary(grok_home(), "today-2", now, 2) +summary(grok_home(), "old", "2026-01-02T10:00:00Z", 5) +write_json(sessions / "today-1" / "usage.json", {"turns": [ + {"endedAt": now, "modelUsage": {"grok-4.7-build": {"inputTokens": 1000, "cachedReadTokens": 800, "outputTokens": 50, "cacheCreationTokens": 0, "totalTokens": 1050}}}, + {"endedAt": now, "modelUsage": {"grok-4.7-build": {"inputTokens": 200, "cachedReadTokens": 100, "outputTokens": 10, "cacheCreationTokens": 5, "totalTokens": 210}}}, +]}) +write_json(sessions / "old" / "usage.json", {"turns": [ + {"endedAt": "2026-01-02T10:00:00Z", "modelUsage": {"grok-4.6": {"inputTokens": 500, "cachedReadTokens": 0, "outputTokens": 20, "cacheCreationTokens": 0, "totalTokens": 520}}}, +]}) +record = grok() +expect("Grok counts sessions, prompts, and tokens from its session files", + {key: record[key] for key in ("hasLocalStats", "todayPrompts", "todaySessions", "totalPrompts", "totalSessions", "activeDays", "todayTotalTokens", "todayTokensByModel")} + | {"today": record["recentDays"][6]["messageCount"]}, + {"hasLocalStats": True, "todayPrompts": 2, "todaySessions": 2, "totalPrompts": 3, "totalSessions": 3, "activeDays": 2, + "todayTotalTokens": 1265, "todayTokensByModel": {"grok-4.7-build": 1265}, "today": 1265}, record) +expect("Grok's tokens by model keep cached input apart", record["modelUsage"], { + "grok-4.6": {"inputTokens": 500, "outputTokens": 20, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, + "grok-4.7-build": {"inputTokens": 300, "outputTokens": 60, "cacheReadInputTokens": 900, "cacheCreationInputTokens": 5}, +}, record) + +# A fork repeats its source's turns verbatim ahead of its own, as `grok -r +# --fork-session` leaves its usage.json; each turn counts once. A +# subagent's turns are tokens spent, but no prompts typed. +sandbox() +signed_in(grok_home()) +sessions = grok_home() / "sessions" / "%2Fhome%2Fme" + + +def turn(number, seconds_ago, tokens): + ended = (datetime.now(timezone.utc) - timedelta(seconds=seconds_ago)).strftime("%Y-%m-%dT%H:%M:%S.620236893+00:00") + usage = {"inputTokens": tokens, "outputTokens": 10, "cachedReadTokens": 0, "cacheCreationTokens": 0, "reasoningTokens": 5, "totalTokens": tokens + 10, "modelCalls": 1} + return {"turnNumber": number, "endedAt": ended, **usage, "turnCount": 1, "primaryModelId": "grok-build", "modelUsage": {"grok-build": usage}} + + +first, second = turn(1, 4, 100), turn(2, 3, 200) +write_json(sessions / "source" / "summary.json", {"last_active_at": iso_in(-3), "num_messages": 7, "session_kind": "headless"}) +write_json(sessions / "source" / "usage.json", {"turns": [first, second]}) +write_json(sessions / "fork" / "summary.json", {"last_active_at": iso_in(-1), "num_messages": 11, "session_kind": "headless", + "parent_session_id": "source", "forked_at": iso_in(-2)}) +write_json(sessions / "fork" / "usage.json", {"turns": [first, second, turn(3, 1, 400)]}) +write_json(sessions / "child" / "summary.json", {"last_active_at": iso_in(-2), "num_messages": 5, "session_kind": "subagent"}) +write_json(sessions / "child" / "usage.json", {"turns": [turn(1, 2, 800)]}) +record = grok() +expect("a fork's copied turns count once, and a subagent's turns as tokens but no prompts", + [record["totalSessions"], record["totalPrompts"], record["todayPrompts"], record["todayTotalTokens"]], [3, 3, 3, 1540], record) +PY diff --git a/test/shell.d/agent-usage-grokbot-test.sh b/test/shell.d/agent-usage-grokbot-test.sh new file mode 100755 index 00000000000..d5987d45c10 --- /dev/null +++ b/test/shell.d/agent-usage-grokbot-test.sh @@ -0,0 +1,127 @@ +#!/bin/bash + +# The Grok Bot collector's own sources: the Cursor sign-in it shares with the +# Cursor collector, and the dashboard's GetSandUsageStatus answer, protobuf +# JSON that drops zeros and falses. The unused-machine record is held +# centrally by the partition test, the Cursor sign-in order by the shared +# test, and the limits cache by the contract checks. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * +import collectors + +RESETS = iso_in(3 * 86400 - 600) +service = {"answer": None} + + +# The dashboard's answer: a week 42.5% used on a Pro plan unless a test +# swaps it. +def sand_status(request): + if service["answer"] is not None: + return service["answer"] + return {"usagePercent": 42.5, "nextResetTimestampUtc": RESETS, "hasNonZeroIncludedLimit": True, "grokPlanLabel": "Pro"} + + +ROUTES = {"api2.cursor.sh/aiserver.v1.DashboardService/GetSandUsageStatus": sand_status} + + +# The Cursor CLI's sign-in, a JWT naming the account, an hour from lapsing +# unless told otherwise. +def cli_signed_in(user="u-main", expires_in=3600): + write_json(config_home() / "cursor" / "auth.json", {"accessToken": jwt({"sub": user, "exp": time.time() + expires_in})}) + + +def grokbot(*args, mode="up"): + with FakeHttp(ROUTES, mode=mode) as http: + record = collect("grokbot", *(args or ("--force",))) + grokbot.requests = http.requests + return record + + +def limit(record): + return {"ready": record["ready"], "stale": record["limitsStale"], "percent": record["limits"][0]["percent"] if record["limits"] else None} + + +# ---------------------------------------------------------------- contract + +sandbox() +fresh = check_limits_contract(Adapter("grokbot", plant=cli_signed_in, answer=ROUTES)) +expect("Grok Bot's weekly allowance and plan come from the Cursor sign-in's dashboard answer", + [fresh["tierLabel"], fresh["limits"][0]["label"], fresh["limits"][0]["percent"], fresh["hasPromptStats"], fresh["hasLocalStats"]], + ["Pro", "Weekly", 0.425, False, False], fresh) +expect("the weekly window resets when the dashboard says", fresh["limits"][0]["resetsAt"], datetime.fromisoformat(RESETS).isoformat(), fresh) + +# ------------------------------------------------------------ failed checks + +# A reply cut short or reset mid-reply keeps the last limit, and only one cut +# short (an answer) skips the sooner retry. An error status says which. +grokbot() +server = grokbot(mode=500) +expect("a failed check says the dashboard's status", "500" in server["authHelpText"], True, server) +for broken in ("cut", "reset"): + record = grokbot(mode=broken) + expect(f"a connection that breaks mid-reply ({broken}) keeps the last limit, stale", limit(record), {"ready": True, "stale": True, "percent": 0.425}, record) +expect("a reply cut short is an answer, not a reason to retry", bool(grokbot(mode="cut").get("retryAdvised")), False) + +# A cache that can't be written never costs the live answer. +cache_root().chmod(0o555) +try: + record = grokbot() +finally: + cache_root().chmod(0o755) +expect("an unwritable cache still gives the live limit", limit(record), {"ready": True, "stale": False, "percent": 0.425}, record) + +# A lapsed sign-in isn't sent: the CLI or editor renews it the next time it +# runs, so the last limit is kept with the limits paused. +cli_signed_in(expires_in=-60) +record = grokbot() +expect("a lapsed sign-in pauses the limit without asking", + [limit(record), record["usageStatusText"], len(grokbot.requests)], [{"ready": True, "stale": True, "percent": 0.425}, "Limits paused", 0], record) + +# Another account's limit never stands in for this one's. +cli_signed_in("u-other") +record = grokbot(mode=500) +expect("another account's cached limit is never shown", [record["ready"], record["limits"], record["tierLabel"]], [False, [], ""], record) + +# --------------------------------------------------------------- the answer + +# Protobuf JSON leaves out zero and false: an untouched week reads 0%, and an +# account with no included limit has no Grok Bot, so the panel skips it. +service["answer"] = {"nextResetTimestampUtc": RESETS, "hasNonZeroIncludedLimit": True} +record = grokbot() +expect("an untouched week reads as 0%", [record["limits"][0]["label"], record["limits"][0]["percent"]], ["Weekly", 0.0], record) +service["answer"] = {"usagePercent": 10} +record = grokbot() +expect("a Cursor account without Grok Bot gives an unused record", [record["ready"], record["limits"]], [False, []], record) +service["answer"] = {"usagePercent": "NaN", "hasNonZeroIncludedLimit": True} +expect("an unreadable percentage reads as 0%", grokbot()["limits"][0]["percent"], 0.0) + +# A limit with no known end is kept as it is through a failed check. +cli_signed_in("u-no-reset") +service["answer"] = {"usagePercent": 5, "hasNonZeroIncludedLimit": True} +record = grokbot() +expect("a limit with no known reset still shows while live", [limit(record), record["limits"][0]["resetsAt"]], [{"ready": True, "stale": False, "percent": 0.05}, ""], record) +record = grokbot(mode=500) +expect("a limit with no known reset is kept as it is through a failed check", + [limit(record), record["limits"][0]["resetsAt"]], [{"ready": True, "stale": True, "percent": 0.05}, ""], record) +service["answer"] = None + +# --------------------------------------------------------------- the sign-in + +# The account asked is the one the Cursor collector reads, whose order the +# shared test holds: here the CLI's sign-in over the editor's. +sandbox() +cli_signed_in("p-cli") +cursor_state_db(None, jwt({"sub": "p-editor", "exp": time.time() + 3600})) +grokbot() +expect("Grok Bot asks as the sign-in the Cursor collector reads", + [request["headers"]["authorization"] for request in grokbot.requests], ["Bearer " + collectors.cursor_token()]) +expect("which is the CLI's over the editor's", collectors.account_key(collectors.cursor_token()), "p-cli") +PY diff --git a/test/shell.d/agent-usage-helpers.sh b/test/shell.d/agent-usage-helpers.sh new file mode 100644 index 00000000000..f36febdf952 --- /dev/null +++ b/test/shell.d/agent-usage-helpers.sh @@ -0,0 +1,302 @@ +#!/bin/bash + +# Fixtures the agent usage collector tests share. Source it after +# base-test.sh: +# +# source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +# source "$SHELL_TEST_DIR/agent-usage-helpers.sh" +# agent_sandbox +# +# The Python half, agent_usage_fixtures.py beside it, is importable from any +# heredoc once agent_sandbox has run. docs/testing.md describes both. +# +# Not a test of its own: the shell runner only picks up *-test.sh files. + +if [[ ${BASH_SOURCE[0]} == "$0" ]]; then + echo "source test/shell.d/agent-usage-helpers.sh from a shell test; do not run it directly" >&2 + exit 1 +fi + +# Variables a developer's own shell may carry that point a collector at a +# real credential or store, matched by prefix, plus any *_API_KEY. Kept in +# step with DEVELOPER_PREFIXES in agent_usage_fixtures.py. +AGENT_DEVELOPER_PREFIXES='^(AGY_|ANTHROPIC_|CLAUDE_|CLOUD_CODE_|CODEX_|COPILOT_|CURSOR_|DEEPSEEK_|DEVIN_|DSH_|FIREWORKS_|GH_|GITHUB_|GROK_|HERMES_|KIMCHI_|KIMI_|MAVIS_|MINIMAX_|MISE_|MMX_|MUSE_|OLLAMA_|OPENAI_|OPENCLAW_|OPENCODE_|OPENROUTER_|PI_CODING_AGENT_|QWEN_|SYNTHETIC_|VIBE_|ZAI_|ZHIPU)' +AGENT_DEAD_PROXY=http://127.0.0.1:9 + +agent_sandbox_pids=() + +# Sets up the scratch machine every agent usage test runs on: +# +# - SANDBOX, a temporary directory removed on exit, with one EXIT trap that +# first stops every process handed to track_pid. +# - HOME and every XDG directory under $SANDBOX/home, exported. +# - Every developer variable naming an agent credential or store unset. +# - Requests that escape a fake go to a dead proxy; loopback stays direct. +# - PYTHONPATH reaching agent_usage_fixtures, and no bytecode written. +agent_sandbox() { + local name + + require_command python3 + SANDBOX=$(mktemp -d) + # Exported so the Python fixtures' sandbox() makes its homes in here too. + export SANDBOX + trap agent_sandbox_cleanup EXIT + + export HOME="$SANDBOX/home" + export XDG_CONFIG_HOME="$HOME/.config" + export XDG_DATA_HOME="$HOME/.local/share" + export XDG_STATE_HOME="$HOME/.local/state" + export XDG_CACHE_HOME="$HOME/.cache" + export XDG_RUNTIME_DIR="$HOME/.run" + mkdir -p "$XDG_CONFIG_HOME" "$XDG_DATA_HOME" "$XDG_STATE_HOME" "$XDG_CACHE_HOME" "$XDG_RUNTIME_DIR" "$SANDBOX/homes" + + while IFS= read -r name; do + if [[ $name == *_API_KEY || $name =~ $AGENT_DEVELOPER_PREFIXES ]]; then + unset "$name" + fi + done < <(compgen -e) + + export http_proxy=$AGENT_DEAD_PROXY https_proxy=$AGENT_DEAD_PROXY + export HTTP_PROXY=$AGENT_DEAD_PROXY HTTPS_PROXY=$AGENT_DEAD_PROXY + export no_proxy=127.0.0.1,localhost NO_PROXY=127.0.0.1,localhost + + export PYTHONPATH="$SHELL_TEST_DIR${PYTHONPATH:+:$PYTHONPATH}" + export PYTHONDONTWRITEBYTECODE=1 + + # What run_collector's isolated PATH holds besides a home's own bin/. + mkdir -p "$SANDBOX/tools" + for name in python3 rg jq sqlite3; do + if command -v "$name" >/dev/null; then + ln -s "$(command -v "$name")" "$SANDBOX/tools/$name" + fi + done +} + +# Records background processes (a stub server, a lock holder) for the EXIT +# trap to stop before the sandbox goes. +track_pid() { + agent_sandbox_pids+=("$@") +} + +agent_sandbox_cleanup() { + local pid + for pid in "${agent_sandbox_pids[@]}"; do + kill "$pid" 2>/dev/null || true + wait "$pid" 2>/dev/null || true + done + if [[ -n ${SANDBOX:-} && -d $SANDBOX ]]; then + # A test that made a file unreadable must not leave it behind. + chmod -R u+rwX "$SANDBOX" 2>/dev/null || true + rm -rf "$SANDBOX" + fi +} + +# Prints a new, empty home under the sandbox with its XDG directories and a +# bin/ for stubs: $SANDBOX/homes/, or a fresh unique name. +new_home() { + local home + if [[ -n ${1:-} ]]; then + home="$SANDBOX/homes/$1" + mkdir -p "$home" + else + home=$(mktemp -d "$SANDBOX/homes/home.XXXXXX") + fi + mkdir -p "$home/.config" "$home/.local/share" "$home/.local/state" "$home/.cache" "$home/.run" "$home/bin" + printf '%s\n' "$home" +} + +# run_collector [VAR=value ...] [-- args...] +# +# Runs bin/omarchy-agent-usage- as a process, in an environment holding +# nothing but and its XDG directories, OMARCHY_PATH, the caller's TZ, +# the dead proxy, and the VAR=value pairs given (a later PATH= replaces the default). PATH is +# /bin, then python3, rg, jq and sqlite3. Runs with --force unless args +# follow `--` (then exactly those, none included). Prints the record; stderr +# lands in /.run/.err and is shown when the collector fails, whose +# status is returned. +run_collector() { + local home="$1" id="$2" status=0 timeout_command + local -a vars=() args=(--force) + shift 2 + while (( $# > 0 )); do + if [[ $1 == "--" ]]; then + shift + args=("$@") + break + fi + vars+=("$1") + shift + done + timeout_command=$(command -v timeout) + mkdir -p "$home/.cache" "$home/.run" "$home/bin" + env -i \ + HOME="$home" \ + PATH="$home/bin:$SANDBOX/tools" \ + LANG=C.UTF-8 \ + ${TZ+TZ="$TZ"} \ + OMARCHY_PATH="$ROOT" \ + OMARCHY_LANGUAGE_FILE="$OMARCHY_LANGUAGE_FILE" \ + XDG_CONFIG_HOME="$home/.config" \ + XDG_DATA_HOME="$home/.local/share" \ + XDG_STATE_HOME="$home/.local/state" \ + XDG_CACHE_HOME="$home/.cache" \ + XDG_RUNTIME_DIR="$home/.run" \ + http_proxy=$AGENT_DEAD_PROXY https_proxy=$AGENT_DEAD_PROXY \ + HTTP_PROXY=$AGENT_DEAD_PROXY HTTPS_PROXY=$AGENT_DEAD_PROXY ALL_PROXY=$AGENT_DEAD_PROXY \ + no_proxy=127.0.0.1,localhost NO_PROXY=127.0.0.1,localhost \ + "${vars[@]}" \ + "$timeout_command" --kill-after=5 "${AGENT_COLLECTOR_TIMEOUT:-60}" "$ROOT/bin/omarchy-agent-usage-$id" "${args[@]}" \ + 2>"$home/.run/$id.err" || status=$? + if (( status != 0 )); then + cat "$home/.run/$id.err" >&2 + fi + return "$status" +} + +# expect_jq "" '' '' "$json" +# +# Passes when `jq -c ` of the JSON prints exactly , and +# fails showing both and the JSON otherwise. +expect_jq() { + local description="$1" filter="$2" expected="$3" json="$4" actual + actual=$(jq -c "$filter" <<<"$json" 2>&1) || true + if [[ $actual == "$expected" ]]; then + pass "$description" + else + fail "$description" "expected: $expected"$'\n'"actual: $actual"$'\n'"json: $json" + fi +} + +# stub_cli , with the script's body on stdin: an executable +# /bin/, a bash script unless the body brings its own #! line. +# Prints its path. +stub_cli() { + local home="$1" name="$2" body + body=$(cat) + mkdir -p "$home/bin" + if [[ $body == '#!'* ]]; then + printf '%s\n' "$body" >"$home/bin/$name" + else + printf '#!/bin/bash\n%s\n' "$body" >"$home/bin/$name" + fi + chmod +x "$home/bin/$name" + printf '%s\n' "$home/bin/$name" +} + +# fake_codex [mode] +# +# A stand-in Codex app-server at /bin/codex speaking its JSON-RPC on +# stdin and stdout. Its rate-limit answer follows CODEX_FAKE when set at run +# time, else (default ok): +# +# ok limits: CODEX_RATE_LIMITS (the rateLimits object) or 40% of a +# 5-hour window and 70% of a week, both resetting ahead +# backend an RPC error from OpenAI's side (a failed check) +# auth "not authenticated" (a refused sign-in) +# empty an answer with no limits in it +# died exits saying it couldn't reach OpenAI (unreachable) +# crashed exits with a panic of its own +# silent exits without a word +# stall never answers (the collector's 8-second timeout) +# +# account/read answers {planType: CODEX_PLAN or "pro"}; CODEX_ARGS_FILE gets +# the arguments it was started with. Every start appends its mode to +# /codex-starts, so a reused probe can be told from a new one. +# +# Unless /.codex/auth.json exists, it is written signed in to workspace +# ws-test as user-test, which the collector keys its limits cache by. +fake_codex() { + local home="$1" mode="${2:-ok}" + mkdir -p "$home/bin" "$home/.codex" + if [[ ! -e $home/.codex/auth.json ]]; then + python3 -c 'import sys, agent_usage_fixtures as f; f.write_json(sys.argv[1], {"tokens": {"account_id": "ws-test", "id_token": f.jwt({"https://api.openai.com/auth": {"chatgpt_user_id": "user-test"}})}})' \ + "$home/.codex/auth.json" >/dev/null || fail "sign the fake Codex home in" + fi + { + printf '#!%s\n' "$(command -v python3)" + python3 -c 'import json, sys; print(f"DEFAULT_MODE = {json.dumps(sys.argv[1])}\nSTARTS = {json.dumps(sys.argv[2])}")' "$mode" "$home/codex-starts" + cat <<'PY' +import json, os, sys, time + +mode = os.environ.get("CODEX_FAKE") or DEFAULT_MODE +with open(STARTS, "a") as log: + log.write(mode + "\n") +if os.environ.get("CODEX_ARGS_FILE"): + with open(os.environ["CODEX_ARGS_FILE"], "w") as args: + args.write("\0".join(sys.argv[1:])) +now = int(time.time()) + + +def reply(message, result=None, error=None): + body = {"id": message["id"]} + body["error" if error else "result"] = error or result + print(json.dumps(body), flush=True) + + +for line in sys.stdin: + message = json.loads(line) + if "id" not in message: + continue + method = message.get("method") + if method == "account/read": + reply(message, {"account": {"planType": os.environ.get("CODEX_PLAN", "pro")}}) + elif method == "account/rateLimits/read": + if mode == "died": + print("error sending request for url (https://chatgpt.com/backend-api/wham/usage)", file=sys.stderr) + sys.exit(1) + if mode == "crashed": + print("thread 'main' panicked at config.rs: unexpected argument", file=sys.stderr) + sys.exit(1) + if mode == "silent": + sys.exit(0) + if mode == "stall": + continue + if mode == "backend": + reply(message, error={"code": -32603, "message": "failed to fetch codex rate limits"}) + elif mode == "auth": + reply(message, error={"code": -32603, "message": "not authenticated"}) + elif mode == "empty": + reply(message, {"rateLimits": None, "rateLimitsByLimitId": {}}) + else: + limits = json.loads(os.environ["CODEX_RATE_LIMITS"]) if os.environ.get("CODEX_RATE_LIMITS") else { + "planType": "pro", + "primary": {"usedPercent": 40, "windowDurationMins": 300, "resetsAt": now + 3600}, + "secondary": {"usedPercent": 70, "windowDurationMins": 10080, "resetsAt": now + 7 * 86400}, + } + reply(message, {"rateLimits": limits}) + else: + reply(message, {}) +PY + } >"$home/bin/codex" + chmod +x "$home/bin/codex" +} + +# age_scan_cache +# +# Moves a scan cache's stamp into the past (negative: into the +# future), as if it were made that long ago. Given an agent id, every scan +# cache of it under XDG_CACHE_HOME. +age_scan_cache() { + python3 -c 'import sys, agent_usage_fixtures as f; f.age_scan_cache(sys.argv[1], float(sys.argv[2]))' "$1" "$2" >/dev/null || + fail "age the scan cache of $1" +} + +# edit_limits_cache +# +# Rewrites every limits cache of (-limits*.json) with the filter, +# which can use $now_ms, $past (an ISO time an hour gone) and $future (one an +# hour ahead): '.limits[].resetsAt = $past' elapses every window, and +# '.fetchedAtMs = $now_ms - 60000' ages the answer a minute. +edit_limits_cache() { + local id="$1" filter="$2" file now_ms past future found=0 + now_ms=$(python3 -c 'import time; print(int(time.time() * 1000))') + past=$(python3 -c 'import datetime as d; print((d.datetime.now(d.timezone.utc) - d.timedelta(hours=1)).isoformat())') + future=$(python3 -c 'import datetime as d; print((d.datetime.now(d.timezone.utc) + d.timedelta(hours=1)).isoformat())') + for file in "$XDG_CACHE_HOME/omarchy/agent-usage/$id"-limits*.json; do + [[ -f $file ]] || continue + found=1 + jq --argjson now_ms "$now_ms" --arg past "$past" --arg future "$future" "$filter" "$file" >"$file.edit" && + mv "$file.edit" "$file" || fail "edit the $id limits cache with $filter" + done + (( found )) || fail "$id keeps a limits cache to edit" +} 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..cc3becdf393 --- /dev/null +++ b/test/shell.d/agent-usage-hermes-test.sh @@ -0,0 +1,315 @@ +#!/bin/bash + +# The Hermes collector's own source: the state.db in Hermes' home and in each +# named profile under ~/.hermes/profiles. Hermes has no subscription, so it has +# no limits. The unused-machine record is held centrally by the partition test, +# and the scan cache by the contract check. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +import datetime as dt +import sqlite3 + +from agent_usage_fixtures import * + +SCHEMA = """ +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) +); +""" + + +# A Hermes state.db with one session today, its prompt and reply, and usage +# for each model given as (name, base), every count being base plus a +# per-field offset. +def state_db(path, *models): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(path) + conn.executescript(SCHEMA) + now = time.time() + conn.execute("INSERT INTO sessions VALUES ('session-1', ?)", (now,)) + conn.executemany("INSERT INTO messages VALUES (?, 'session-1', ?, ?)", [(1, "user", now), (2, "assistant", now)]) + for name, base in models: + conn.execute("INSERT INTO session_model_usage VALUES ('session-1', ?, ?, ?, ?, ?)", (name, base, base + 10, base + 20, base + 30)) + conn.commit() + conn.close() + return path + + +def hermes(*args): + return collect("hermes", *(args or ("--force",))) + + +# ---------------------------------------------------------------- state.db + +sandbox() +state_db(home() / ".hermes" / "state.db", ("gpt-6-astra", 10), ("gpt-6-sol", 11)) +record = hermes() +expect("Hermes reads each model's usage, today's totals, and its prompts and sessions from state.db", { + "ready": record["ready"], + "status": record["usageStatusText"], + "astra": record["modelUsage"]["gpt-6-astra"], + "solToday": record["todayTokensByModel"]["gpt-6-sol"], + "today": record["todayTotalTokens"], + "week": [len(record["recentDays"]), record["recentDays"][-1]["messageCount"]], + "counts": [record["totalPrompts"], record["todayPrompts"], record["totalSessions"], record["todaySessions"]], +}, { + "ready": True, + "status": "", + "astra": {"inputTokens": 10, "outputTokens": 20, "cacheReadInputTokens": 30, "cacheCreationInputTokens": 40}, + "solToday": 104, + "today": 204, + "week": [7, 204], + "counts": [1, 1, 1, 1], +}, record) + +# A state.db with no usage table still counts its sessions and prompts. +sandbox() +bare = state_db(home() / ".hermes" / "state.db") +conn = sqlite3.connect(bare) +conn.execute("DROP TABLE session_model_usage") +conn.commit() +conn.close() +record = hermes() +expect("a state.db without a usage table still counts its sessions and prompts", + [record["usageStatusText"], record["totalSessions"], record["todayPrompts"], record["modelUsage"]], ["", 1, 1, {}], record) + +# A start or a prompt stamped as text, or in the future, is no today. +sandbox() +odd = state_db(home() / ".hermes" / "state.db") +conn = sqlite3.connect(odd) +conn.execute("INSERT INTO sessions VALUES ('text', 'soon')") +conn.execute("INSERT INTO sessions VALUES ('future', ?)", (time.time() + 3 * 86400,)) +conn.executemany("INSERT INTO messages VALUES (?, 'text', 'user', ?)", [(3, "soon"), (4, time.time() + 3 * 86400)]) +conn.commit() +conn.close() +record = hermes() +expect("text and future times count in the totals but not today", + [record["totalSessions"], record["todaySessions"], record["totalPrompts"], record["todayPrompts"]], [3, 1, 3, 1], record) + +# A negative count is no count; it never nets out another row's usage. +sandbox() +netted = state_db(home() / ".hermes" / "state.db", ("gpt-6-astra", 10)) +conn = sqlite3.connect(netted) +conn.execute("INSERT INTO sessions VALUES ('session-2', ?)", (time.time(),)) +conn.execute("INSERT INTO session_model_usage VALUES ('session-2', 'gpt-6-astra', -5, 'many', 1e300, 0)") +conn.commit() +conn.close() +record = hermes() +expect("negative, text and huge counts are left out of the sums", + record["modelUsage"]["gpt-6-astra"], {"inputTokens": 10, "outputTokens": 20, "cacheReadInputTokens": 30, "cacheCreationInputTokens": 40}, record) + +# Named profiles live under ~/.hermes/profiles even when HERMES_HOME moves the +# default home elsewhere. +sandbox() +state_db(home() / "elsewhere" / "state.db", ("gpt-6-astra", 100)) +state_db(home() / ".hermes" / "profiles" / "work" / "state.db", ("gpt-6-sol", 1000)) +os.environ["HERMES_HOME"] = str(home() / "elsewhere") +record = hermes() +del os.environ["HERMES_HOME"] +expect("Hermes reads HERMES_HOME and every named profile", + [record["modelUsage"]["gpt-6-astra"]["inputTokens"], record["modelUsage"]["gpt-6-sol"]["inputTokens"], record["totalSessions"], record["totalPrompts"]], + [100, 1000, 2, 2], record) + +# Hermes expands ~ and $VARs in HERMES_HOME, and a custom home at +# /profiles/ keeps its sibling profiles under /profiles +# (hermes_constants._hermes_root_of). +sandbox() +state_db(home() / "srv" / "profiles" / "work" / "state.db", ("gpt-6-astra", 100)) +state_db(home() / "srv" / "profiles" / "play" / "state.db", ("gpt-6-sol", 1000)) +os.environ["HERMES_ROOT_FOR_TEST"] = "srv" +os.environ["HERMES_HOME"] = "~/$HERMES_ROOT_FOR_TEST/profiles/work" +record = hermes() +del os.environ["HERMES_HOME"], os.environ["HERMES_ROOT_FOR_TEST"] +expect("Hermes expands HERMES_HOME and reads the profiles beside a custom profile home", + [record["modelUsage"].get("gpt-6-astra", {}).get("inputTokens"), record["modelUsage"].get("gpt-6-sol", {}).get("inputTokens"), record["totalSessions"]], + [100, 1000, 2], record) + +# A shell with a profile active exports HERMES_HOME=~/.hermes/profiles/; +# the root ~/.hermes is still the default profile, so it counts too +# (hermes_cli.profiles.list_profiles), as with HERMES_HOME unset. +sandbox() +state_db(home() / ".hermes" / "state.db", ("gpt-6-astra", 100)) +state_db(home() / ".hermes" / "profiles" / "side" / "state.db", ("gpt-6-sol", 1000)) +os.environ["HERMES_HOME"] = str(home() / ".hermes" / "profiles" / "side") +record = hermes() +del os.environ["HERMES_HOME"] +expect("Hermes reads the default profile beside an active named profile", + [record["modelUsage"].get("gpt-6-astra", {}).get("inputTokens"), record["modelUsage"].get("gpt-6-sol", {}).get("inputTokens"), record["totalSessions"]], + [100, 1000, 2], record) + +# The same holds for a custom root's profile home: /state.db is its +# default profile. +sandbox() +state_db(home() / "srv" / "state.db", ("gpt-6-astra", 100)) +state_db(home() / "srv" / "profiles" / "work" / "state.db", ("gpt-6-sol", 1000)) +os.environ["HERMES_HOME"] = str(home() / "srv" / "profiles" / "work") +record = hermes() +del os.environ["HERMES_HOME"] +expect("Hermes reads a custom root's default profile beside its active profile", + [record["modelUsage"].get("gpt-6-astra", {}).get("inputTokens"), record["modelUsage"].get("gpt-6-sol", {}).get("inputTokens"), record["totalSessions"]], + [100, 1000, 2], record) + +# A custom root beside a used ~/.hermes: ~/.hermes is read whole, its own +# state.db (its default profile) with its named profiles, not the profiles +# alone. +sandbox() +state_db(home() / "custom-hermes" / "profiles" / "work" / "state.db", ("gpt-6-astra", 100)) +state_db(home() / ".hermes" / "state.db", ("gpt-6-sol", 1000)) +state_db(home() / ".hermes" / "profiles" / "side" / "state.db", ("gpt-6-luna", 10000)) +os.environ["HERMES_HOME"] = str(home() / "custom-hermes") +record = hermes() +del os.environ["HERMES_HOME"] +expect("Hermes reads ~/.hermes whole beside a custom root", + [record["modelUsage"].get(model, {}).get("inputTokens") for model in ("gpt-6-astra", "gpt-6-sol", "gpt-6-luna")] + [record["totalSessions"]], + [100, 1000, 10000, 3], record) + +# HERMES_DATA_DIR_SUFFIX names the default home. +sandbox() +state_db(home() / ".hermes-beta" / "state.db", ("gpt-6-astra", 100)) +os.environ["HERMES_DATA_DIR_SUFFIX"] = "-beta" +record = hermes() +del os.environ["HERMES_DATA_DIR_SUFFIX"] +expect("Hermes reads the default home HERMES_DATA_DIR_SUFFIX names", + record["modelUsage"].get("gpt-6-astra", {}).get("inputTokens"), 100, record) + +# A profile that is a link to the home is the home: counted once. +sandbox() +state_db(home() / ".hermes" / "state.db", ("gpt-6-astra", 100)) +(home() / ".hermes" / "profiles").mkdir() +(home() / ".hermes" / "profiles" / "default").symlink_to(home() / ".hermes") +record = hermes() +expect("Hermes counts a profile linked to its home once", + [record["modelUsage"]["gpt-6-astra"]["inputTokens"], record["totalSessions"], record["totalPrompts"]], [100, 1, 1], record) + +# A session counts where Hermes billed it: Nous Portal is Hermes', a server +# you run the local collector's, and Anthropic's API (or a Claude model on a +# local proxy), OpenRouter or a gateway nobody's here. A session Hermes kept +# no route for goes where its first request did, and one that never reached a +# provider is nobody's, its prompt and all. The rows are shaped as Hermes Agent +# wrote them against a mock of each API. +sandbox() +db = home() / ".hermes" / "state.db" +db.parent.mkdir(parents=True) +conn = sqlite3.connect(db) +conn.executescript(""" +CREATE TABLE sessions (id TEXT PRIMARY KEY, source TEXT, model TEXT, started_at REAL NOT NULL, billing_provider TEXT, billing_base_url TEXT, billing_mode TEXT); +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, billing_provider TEXT NOT NULL DEFAULT '', billing_base_url TEXT NOT NULL DEFAULT '', + billing_mode TEXT NOT NULL DEFAULT '', task TEXT NOT NULL DEFAULT '', + 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, billing_provider, billing_base_url, billing_mode, task) +); +""") +now = time.time() +ROUTES = [ + ("nous", "Hermes-4-405B", "nous", "https://inference-api.nousresearch.com/v1", 100), + ("custom", "hermes-mock-model", "custom", "http://127.0.0.1:18473/v1", 1000), + ("anthropic", "claude-mock-hermes", "anthropic", "http://127.0.0.1:18473", 10000), + ("router", "z-ai/glm-4.6", "openrouter", "https://openrouter.ai/api/v1", 100000), + ("gateway", "acme-1", "custom", "https://llm.acme.example/v1", 1000000), +] +for n, (session, model, provider, url, tokens) in enumerate(ROUTES): + conn.execute("INSERT INTO sessions VALUES (?, 'oneshot', ?, ?, ?, ?, NULL)", (session, model, now, provider, url)) + conn.execute("INSERT INTO messages VALUES (?, ?, 'user', ?)", (n + 1, session, now)) + conn.execute("INSERT INTO session_model_usage (session_id, model, billing_provider, billing_base_url, input_tokens) VALUES (?, ?, ?, ?, ?)", (session, model, provider, url, tokens)) + # The session title is a request of its own, on the session's route. + conn.execute("INSERT INTO session_model_usage (session_id, model, billing_provider, billing_base_url, task, input_tokens) VALUES (?, ?, ?, ?, 'title_generation', ?)", (session, model, provider, url + "/", tokens)) +# A session whose route Hermes didn't keep, and two that failed before any +# request Hermes accounted, one with the prompt typed into it. +conn.execute("INSERT INTO sessions VALUES ('unrouted', 'oneshot', 'hermes-mock-model', ?, NULL, NULL, NULL)", (now,)) +conn.execute("INSERT INTO session_model_usage (session_id, model, billing_provider, billing_base_url, input_tokens) VALUES ('unrouted', 'hermes-mock-model', 'custom', 'http://127.0.0.1:18473/v1', 7)") +conn.execute("INSERT INTO sessions VALUES ('failed', 'oneshot', 'claude-mock-hermes', ?, NULL, NULL, NULL)", (now,)) +conn.execute("INSERT INTO messages VALUES (99, 'failed', 'user', ?)", (now,)) +conn.execute("INSERT INTO messages VALUES (100, 'failed', 'assistant', ?)", (now,)) +conn.execute("INSERT INTO sessions VALUES ('failed-silent', 'oneshot', 'claude-mock-hermes', ?, NULL, NULL, NULL)", (now,)) +conn.commit() +conn.close() +record = hermes() +expect("Hermes counts the sessions Nous Portal billed, not ones that reached no provider", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"Hermes-4-405B": 200}, 1, 1], record) +local = collect("local", "--force") +expect("Local counts the Hermes sessions a server you run answered", + [local["todayTokensByModel"], local["todayPrompts"], local["todaySessions"]], [{"hermes-mock-model": 2007}, 1, 2], local) + +# An unreadable state.db says Hermes is unavailable, naming it in the help +# text, while the readable profiles still count, and that scan isn't cached. +sandbox() +state_db(home() / ".hermes" / "profiles" / "work" / "state.db", ("gpt-6-sol", 1000)) +(home() / ".hermes" / "state.db").write_text("not a database at all, just text that is long enough to read") +record = hermes() +expect("an unreadable state.db reports Hermes unavailable, the readable profiles still counted", + [record["ready"], record["usageStatusText"], record["hasLocalStats"], record["totalSessions"]], + [True, "Hermes unavailable", True, 1], record) +check("the help text names the unreadable state.db", + record["authHelpText"].startswith(f"Couldn't read Hermes' state.db ({home() / '.hermes' / 'state.db'}: "), record["authHelpText"]) +check("a scan that missed a database isn't cached", not scan_cache_files("hermes"), scan_cache_files("hermes")) + +# Every active date travels, not only the latest 400, and broken token counts +# or start times skip their row instead of inflating the record. +sandbox() +db = state_db(home() / ".hermes" / "state.db") +conn = sqlite3.connect(db) +# Local noon on each earlier day, so a daylight-saving change can't put two +# sessions on one date the way fixed 86400-second steps can near midnight. +today = dt.date.today() +noons = [dt.datetime.combine(today - dt.timedelta(days=n + 1), dt.time(12)).timestamp() for n in range(450)] +conn.executemany("INSERT INTO sessions VALUES (?, ?)", [(f"old-{n}", noons[n]) for n in range(450)]) +conn.execute("INSERT INTO sessions VALUES ('broken', 'not a time')") +conn.executemany("INSERT INTO session_model_usage VALUES (?, ?, ?, 0, 0, 0)", [ + ("old-1", "negative", -500), + ("old-2", "huge", 10**15), + ("broken", "untimed", 5), +]) +conn.commit() +conn.close() +record = hermes() +expect("Hermes emits every active date", [record["activeDays"], len(record["activeDates"])], [451, 451], record) +expect("Hermes skips broken token counts, and counts an untimed session's in the totals", sorted(record["modelUsage"]), ["untimed"], record) + +# A start time outside 2000-2100 is a damaged stamp: the session counts, but +# names no day. +sandbox() +db = state_db(home() / ".hermes" / "state.db") +conn = sqlite3.connect(db) +conn.executemany("INSERT INTO sessions VALUES (?, ?)", [("stray", 1), ("far", 12307334400)]) +conn.commit() +conn.close() +record = hermes() +expect("a start time outside 2000-2100 names no active day", + [record["totalSessions"], record["activeDates"]], [3, [dt.date.today().isoformat()]], record) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def plant(): + if not (home() / ".hermes" / "state.db").exists(): + state_db(home() / ".hermes" / "state.db", ("gpt-6-astra", 10)) + + +def add_session(): + add_session.count = getattr(add_session, "count", 0) + 1 + conn = sqlite3.connect(home() / ".hermes" / "state.db") + conn.execute("INSERT INTO sessions VALUES (?, ?)", (f"grown-{add_session.count}", time.time())) + conn.commit() + conn.close() + + +check_scan_cache(Adapter("hermes", plant=plant, grow=add_session)) +PY diff --git a/test/shell.d/agent-usage-kimchi-test.sh b/test/shell.d/agent-usage-kimchi-test.sh new file mode 100755 index 00000000000..47300cbd18d --- /dev/null +++ b/test/shell.d/agent-usage-kimchi-test.sh @@ -0,0 +1,419 @@ +#!/bin/bash + +# The Kimchi collector's own sources: the key and region the Kimchi CLI keeps +# in its config, the gateway's credits and budget answers (each kept on its +# own), and the pi-format sessions under its harness. The unused-machine +# record is held centrally by the partition test, the pi reader by the shared +# test, and the limits and scan caches by the contract checks. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +PERIOD = {"startTime": iso_in(-11 * 86400), "endTime": iso_in(20 * 86400)} + +# The gateway answers per key: credits for every key, a budget only for some, +# and a refusal for a revoked one. +CREDITS = { + "key-coder": {"serverless": True, "tier": "coder", "is_paid_tier": True, "remaining": "24.95", "has_credits": True}, + "key-budget": {"serverless": True, "tier": "TEAMS", "remaining": 30, "has_credits": True}, + "key-empty": {"serverless": True, "tier": "free", "remaining": "0", "has_credits": False}, + "key-sprint": {"serverless": True, "tier": "teams", "remaining": 12, "has_credits": True}, + "key-pair": {"serverless": True, "tier": "teams", "remaining": 5, "has_credits": True}, +} +BUDGETS = { + "key-sprint": {"period": {"startTime": iso_in(-7 * 86400), "endTime": iso_in(7 * 86400)}, "budgets": [ + {"scope": "TEAM_POOLED", "scopeId": "t-1", "budgetLimitUsd": "100", "totalSpendUsd": "10", "providerBudgets": [ + {"provider": "openai", "limitType": "CAPPED", "budgetLimitUsd": "20", "usageUsd": "5"}, + ]}, + ]}, + "key-budget": {"period": PERIOD, "budgets": [ + {"scope": "USER", "scopeId": "u-1", "budgetLimitUsd": "50", "totalSpendUsd": "20", "providerBudgets": [ + {"provider": "anthropic", "limitType": "PROVIDER_LIMIT_TYPE_CAPPED", "budgetLimitUsd": "10", "usageUsd": "2.5"}, + {"provider": "openai", "limitType": "UNLIMITED", "usageUsd": "4"}, + ]}, + ]}, + "key-pair": {"period": PERIOD, "budgets": [ + {"scope": "USER", "scopeId": "u-1", "budgetLimitUsd": "50", "totalSpendUsd": "10", "providerBudgets": [ + {"provider": "anthropic", "limitType": "CAPPED", "budgetLimitUsd": "10", "usageUsd": "9"}, + ]}, + {"scope": "TEAM_POOLED", "scopeId": "t-1", "budgetLimitUsd": "200", "totalSpendUsd": "50", "providerBudgets": [ + {"provider": "openai", "limitType": "CAPPED", "budgetLimitUsd": "40", "usageUsd": "4"}, + ]}, + ]}, +} +# An endpoint a test breaks on its own: its answer, whatever the key. +broken = {} +DOWN = urllib.error.URLError(OSError("Name or service not known")) + + +def gateway(resource, answers): + def answer(request): + if resource in broken: + return broken[resource] + key = request.get_header("Authorization").split(" ", 1)[1] + if key == "key-revoked": + return 401 + if key == "key-offline": + return DOWN + return answers.get(key, {"period": {"startTime": "", "endTime": ""}, "budgets": []}) + return answer + + +ROUTES = {"/v1/credits": gateway("credits", CREDITS), "/v1/budget": gateway("budget", BUDGETS)} + + +# A Kimchi login the way the CLI leaves it. +def signed_in(key="key-budget", region="us"): + write_json(home() / ".config" / "kimchi" / "config.json", {"apiKey": key, "region": region}) + + +def kimchi(*args, key=None, **env): + saved = dict(os.environ) + if key: + env["KIMCHI_API_KEY"] = key + os.environ.update(env) + try: + with FakeHttp(ROUTES) as http: + record = collect("kimchi", *(args or ("--force",))) + finally: + os.environ.clear() + os.environ.update(saved) + kimchi.asked = sorted({request["url"] for request in http.requests}) + return record + + +def harness(): + return home() / ".config" / "kimchi" / "harness" + + +def reply(id, model, usage, timestamp="now"): + return pi_reply("kimchi-dev", model=model, usage=usage, timestamp=timestamp, id=id) + + +def kept(record): + return {"stale": record["limitsStale"], "tierLabel": record["tierLabel"], "limits": record["limits"], + "credits": (record.get("balance") or {}).get("remaining"), "retry": record.get("retryAdvised"), "status": record["usageStatusText"]} + + +# ------------------------------------------------------------ the gateway + +# Without a key the gateway is never asked; the partition test holds the rest +# of the unused-machine record. +sandbox() +record = kimchi() +expect("no key means the gateway is never asked", [record["ready"], kimchi.asked], [False, []], record) + +signed_in("key-coder") +record = kimchi() +expect("a Coder plan reports its plan and prepaid credits", + {key: record.get(key) for key in ("ready", "tierLabel", "limits", "balance", "usageStatusText", "authHelpText")}, + {"ready": True, "tierLabel": "Coder", "limits": [], "balance": {"remaining": 24.95, "funded": 0, "spent": 0, "currency": "USD", "estimated": False}, + "usageStatusText": "", "authHelpText": ""}, record) +expect("the US region asks the US gateway", kimchi.asked, ["https://llm.kimchi.dev/v1/budget", "https://llm.kimchi.dev/v1/credits"]) +signed_in("key-coder", "eu") +kimchi() +expect("the EU region asks the EU gateway", kimchi.asked, ["https://llm.eu.kimchi.dev/v1/budget", "https://llm.eu.kimchi.dev/v1/credits"]) +signed_in("key-coder", "EU") +kimchi() +expect("a region is matched whatever its case", kimchi.asked, ["https://llm.eu.kimchi.dev/v1/budget", "https://llm.eu.kimchi.dev/v1/credits"]) +kimchi(KIMCHI_BASE_URL="https://gateway.example/openai/v1/") +expect("KIMCHI_BASE_URL names the gateway, without its OpenAI path", kimchi.asked, ["https://gateway.example/v1/budget", "https://gateway.example/v1/credits"]) + +# A gateway override never sends the key in the clear to another machine; +# plain HTTP is for a gateway on this one. +record = kimchi(KIMCHI_BASE_URL="http://gateway.example") +expect("a plain HTTP KIMCHI_BASE_URL is never sent the key", [kimchi.asked, record["usageStatusText"], "HTTPS" in record["authHelpText"]], [[], "Balance unavailable", True], record) +write_json(home() / ".config" / "kimchi" / "config.json", {"apiKey": "key-coder", "llmEndpoint": "http://gateway.example/openai/v1"}) +kimchi() +expect("a plain HTTP llmEndpoint is never sent the key", kimchi.asked, []) +kimchi(KIMCHI_BASE_URL="http://127.0.0.1:9") +expect("plain HTTP to loopback still carries the key", kimchi.asked, ["http://127.0.0.1:9/v1/budget", "http://127.0.0.1:9/v1/credits"]) + +# Kimchi 1.18's self-hosted login saves the region and its host but no +# llmEndpoint, and its gateway is /llm; KIMCHI_SELF_HOSTED_URL moves it. +# Without a host the key goes nowhere, never to the public gateway. +write_json(home() / ".config" / "kimchi" / "config.json", {"apiKey": "key-coder", "region": "self-hosted", "selfHostedUrl": "https://kimchi.example.internal/"}) +kimchi() +expect("a self-hosted login asks its own host's gateway", kimchi.asked, + ["https://kimchi.example.internal/llm/v1/budget", "https://kimchi.example.internal/llm/v1/credits"]) +kimchi(KIMCHI_SELF_HOSTED_URL="https://other.example.internal") +expect("KIMCHI_SELF_HOSTED_URL moves a self-hosted gateway", kimchi.asked, + ["https://other.example.internal/llm/v1/budget", "https://other.example.internal/llm/v1/credits"]) +write_json(home() / ".config" / "kimchi" / "config.json", {"apiKey": "key-coder", "region": "self-hosted"}) +record = kimchi() +expect("a self-hosted login with no host sends its key nowhere", + [kimchi.asked, record["usageStatusText"], record.get("retryAdvised")], [[], "Balance unavailable", None], record) +signed_in("key-coder") +kimchi(KIMCHI_REGION="self-hosted", KIMCHI_SELF_HOSTED_URL="https://ci.example.internal") +expect("KIMCHI_REGION picks the self-hosted region over the saved one", kimchi.asked, + ["https://ci.example.internal/llm/v1/budget", "https://ci.example.internal/llm/v1/credits"]) + +# ------------------------------------------------------------- contract + +sandbox() + + +def add_session(): + add_session.count = getattr(add_session, "count", 0) + 1 + write_pi_session(harness() / "sessions" / "--grown--" / f"{add_session.count}.jsonl", [reply(f"g-{add_session.count}", "glm-5.3", {"input": 1})]) + + +adapter = Adapter("kimchi", plant=signed_in, answer=ROUTES, requests_per_probe=2, grow=add_session, pi_root=lambda: harness() / "sessions") +fresh = check_limits_contract(adapter) +expect("a monthly spend budget becomes a Monthly limit, and its cap never poses as funding", + [fresh["tierLabel"], {key: fresh["limits"][0].get(key) for key in ("label", "title", "percent")}, fresh["balance"]], + ["Teams", {"label": "Budget", "title": "Monthly", "percent": 0.4}, {"remaining": 30.0, "funded": 0, "spent": 0, "currency": "USD", "estimated": False}], fresh) +expect("a capped provider is a scoped allowance on the budget's clock, an uncapped one isn't", + [{"title": limit["title"], "percent": limit["percent"]} for limit in fresh["limits"][1:]], [{"title": "Anthropic Monthly", "percent": 0.25}], fresh) +expect("the budget resets when its period ends", fresh["limits"][0]["resetsAt"], PERIOD["endTime"], fresh) +check_scan_cache(adapter) + +# ------------------------------------------------------------ the answer + +sandbox() +signed_in() +record = kimchi(key="key-sprint") +expect("a period with no clock name still shows its capped providers", + [{"title": limit.get("title"), "percent": limit["percent"]} for limit in record["limits"]], + [{"title": None, "percent": 0.1}, {"title": "Openai budget", "percent": 0.25}], record) + +# Two budgets on one clock: the panel hangs a cap on the meter its title +# names, so each budget is titled by its scope and its caps name the budget +# they belong to, instead of both caps landing on the last Monthly meter. +record = kimchi(key="key-pair") +expect("several budgets on one clock each keep their own provider caps", + [{"title": limit["title"], "percent": limit["percent"]} for limit in record["limits"]], + [{"title": "Budget Monthly", "percent": 0.2}, {"title": "Anthropic budget (Budget)", "percent": 0.9}, + {"title": "Team pool Monthly", "percent": 0.25}, {"title": "Openai budget (Team pool)", "percent": 0.1}], record) + +record = kimchi(key="key-empty") +expect("an empty balance shows, out of credits, with the rate limiting in the help", + [record["tierLabel"], record["balance"]["remaining"], record["usageStatusText"], record["authHelpText"]], + ["Community", 0.0, "Out of credits", "Kimchi's gateway is rate limiting requests until the account has credits again."], record) + +record = kimchi(key="key-revoked") +expect("a rejected key says to log in again in Kimchi", "Log in again in Kimchi" in record["authHelpText"], True, record) +record = kimchi(key="key-offline") +expect("an unreachable gateway with nothing kept says the balance is unavailable and asks for a sooner retry", + [record["usageStatusText"], record.get("retryAdvised")], ["Balance unavailable", True], record) + +# Another key never shows the last key's numbers. +kimchi(key="key-budget") +broken.update(credits=DOWN, budget=DOWN) +record = kimchi(key="key-coder") +broken.clear() +expect("another key never shows the last key's numbers", [record["tierLabel"], record["limits"]], ["", []], record) + +# One endpoint failing keeps that endpoint's last answer and takes the other's +# fresh one, without letting the half that came back replace the cache. +BUDGET_LIMITS = kimchi(key="key-budget")["limits"] +check("key-budget answers two limits", len(BUDGET_LIMITS) == 2, BUDGET_LIMITS) +broken["budget"] = DOWN +record = kimchi(key="key-budget") +expect("a failed budget endpoint keeps its last limits beside fresh credits", kept(record), + {"stale": True, "tierLabel": "Teams", "limits": BUDGET_LIMITS, "credits": 30.0, "retry": True, "status": ""}, record) +broken.clear() +broken["credits"] = DOWN +record = kimchi(key="key-budget") +expect("a failed credits endpoint keeps its last plan and balance beside fresh limits", kept(record), + {"stale": True, "tierLabel": "Teams", "limits": BUDGET_LIMITS, "credits": 30.0, "retry": True, "status": ""}, record) +broken.clear() + +# A reply of an unexpected shape is a failed check, not an empty account: the +# last answer stays up, stale, and a server that answered isn't retried sooner. +for odd in ("budget", "credits"): + for body in ({"message": "ok"}, [], None): + kimchi(key="key-budget") + broken[odd] = b"null" if body is None else body + record = kimchi(key="key-budget") + broken.clear() + expect(f"an unexpected {odd} reply ({json.dumps(body)}) keeps the last answer, stale", kept(record), + {"stale": True, "tierLabel": "Teams", "limits": BUDGET_LIMITS, "credits": 30.0, "retry": None, "status": ""}, record) + +# A part that fails before anything was cached is missing, not current: the +# record is stale, but a server that answered with an error isn't retried +# sooner. +for path in limits_cache_files("kimchi"): + path.unlink() +broken["budget"] = 500 +record = kimchi(key="key-sprint") +broken.clear() +expect("a budget endpoint failing with nothing cached leaves the record stale", + [record["limitsStale"], record["limits"], record["balance"]["remaining"], record.get("retryAdvised"), record["usageStatusText"]], [True, [], 12.0, None, ""], record) +for path in cache_root().glob("kimchi-credits-*.json"): + path.unlink() +broken["credits"] = 500 +record = kimchi(key="key-sprint") +broken.clear() +expect("a credits endpoint failing with nothing cached says credits are unavailable", + [record["limitsStale"], record.get("balance"), record["usageStatusText"], record.get("retryAdvised")], [True, None, "Balance unavailable", None], record) + +# Refreshes can overlap; an answer fetched earlier never replaces one fetched +# later, so a slow run finishing last can't roll the cache back. +kimchi(key="key-budget") +credits_cache = max(cache_root().glob("kimchi-credits-*.json"), key=lambda path: path.stat().st_mtime) +write_json(credits_cache, dict(json.loads(credits_cache.read_text()), fetchedAtMs=ms_in(60), tierLabel="Enterprise")) +kimchi(key="key-budget") +expect("an earlier answer never replaces a later one", json.loads(credits_cache.read_text())["tierLabel"], "Enterprise") +broken.update(credits=DOWN, budget=DOWN) +expect("an outage falls back to the latest answer", kimchi(key="key-budget")["tierLabel"], "Enterprise") + +# The same key asked of another gateway is another answer. +record = kimchi(key="key-budget", KIMCHI_BASE_URL="https://gateway.example") +broken.clear() +expect("another gateway never shows the last gateway's numbers", + [record["tierLabel"], record["limits"], record["usageStatusText"]], ["", [], "Balance unavailable"], record) + +# A budget cached under its old name, kimchi-budget-*, is still read, its +# elapsed window zeroed and one with no reset time kept as it is. +for path in cache_root().glob("kimchi-credits-*.json"): + path.unlink() +kimchi(key="key-budget") +limits_cache = max(limits_cache_files("kimchi"), key=lambda path: path.stat().st_mtime) +cached = json.loads(limits_cache.read_text()) +cached["limits"][0].update(resetsAt=iso_in(-3600), percent=0.9) +cached["limits"][1]["resetsAt"] = "" +write_json(limits_cache.with_name(limits_cache.name.replace("kimchi-limits-", "kimchi-budget-")), cached) +limits_cache.unlink() +broken["budget"] = DOWN +record = kimchi(key="key-budget") +broken.clear() +expect("a budget cached under its old name is read, elapsed at zero and an unclocked one as it was", + [[round(limit["percent"] * 100), limit["resetsAt"]] for limit in record["limits"]] + [record["limitsStale"]], [[0, ""], [25, ""], True], record) + +# ---------------------------------------------------------------- sessions + +# Tokens come from assistant messages in the harness's pi session files: +# cached input kept apart, a repeated entry counted once, user turns and other +# entries skipped. The one typed prompt today books on its first reply, not +# on each step; a file without user messages keeps one per reply. +sandbox() +sessions = harness() / "sessions" / "--home-me--" +now = iso_in(0) +write_pi_session(sessions / "today.jsonl", [ + {"type": "session", "id": "s-1"}, + {"type": "message", "id": "u-1", "timestamp": now, "message": {"role": "user", "content": "hi", "usage": {"input": 999}}}, + reply("a-1", "minimax-m3", {"input": 100, "output": 20, "cacheRead": 800, "cacheWrite": 5, "totalTokens": 925}), + reply("a-1", "minimax-m3", {"input": 100, "output": 20, "cacheRead": 800, "cacheWrite": 5, "totalTokens": 925}), + reply("a-2", "glm-5.3", {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": 75}), + {"type": "custom", "id": "c-1", "timestamp": now, "data": {"usage": {"input": 999}}}, + {"type": "message", "id": "a-3", "message": {"role": "assistant", "model": "glm-5.3", "timestamp": ms_in(0), "usage": {"input": 10, "output": 5}}}, +]) +write_pi_session(sessions / "old.jsonl", [reply("b-1", "minimax-m3", {"input": 300, "output": 30, "cacheRead": 0, "cacheWrite": 0}, timestamp="2026-01-02T10:00:00Z")]) +record = kimchi(key="key-coder") +expect("Kimchi counts sessions, prompts, and tokens from its session files", + {key: record[key] for key in ("hasLocalStats", "todayPrompts", "todaySessions", "totalPrompts", "totalSessions", "activeDays", "todayTotalTokens", "todayTokensByModel")} + | {"today": record["recentDays"][6]["messageCount"]}, + {"hasLocalStats": True, "todayPrompts": 1, "todaySessions": 1, "totalPrompts": 2, "totalSessions": 2, "activeDays": 2, + "todayTotalTokens": 1015, "todayTokensByModel": {"minimax-m3": 925, "glm-5.3": 90}, "today": 1015}, record) +expect("Kimchi's tokens by model keep cached input apart", record["modelUsage"], { + "minimax-m3": {"inputTokens": 400, "outputTokens": 50, "cacheReadInputTokens": 800, "cacheCreationInputTokens": 5}, + "glm-5.3": {"inputTokens": 85, "outputTokens": 5, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, +}, record) + +# Kimchi 1.18 names a reply's route as kimchi-dev/ on its gateway, +# but `kimchi --provider openrouter` (or a custom provider in models.json) +# runs a session that provider bills, which isn't Kimchi's usage. +write_pi_session(sessions / "routed.jsonl", [ + {"type": "session", "id": "s-3"}, + pi_reply("kimchi-dev/moonshotai", id="r-1", model="kimi-k2.6", usage={"input": 1234, "output": 51, "cacheRead": 71, "totalTokens": 1356}), + pi_reply("kimchi-experimental", id="r-2", model="glm-5.3", usage={"input": 4, "output": 1}), + pi_reply("openrouter", id="r-3", model="z-ai/glm-4.6", usage={"input": 5234, "output": 55, "cacheRead": 75, "totalTokens": 5364}), + pi_reply("kimchi-devious", id="r-4", model="glm-5.3", usage={"input": 999}), +]) +record = kimchi(key="key-coder") +expect("Kimchi counts only replies on its own gateway routes", + [record["todayTotalTokens"], record["todayTokensByModel"]], [1015 + 1356 + 5, {"minimax-m3": 925, "glm-5.3": 95, "kimi-k2.6": 1356}], record) +(sessions / "routed.jsonl").unlink() + +# KIMCHI_CODING_AGENT_DIR moves the harness, as it does for the CLI. +write_pi_session(home() / "elsewhere" / "sessions" / "--elsewhere--" / "s.jsonl", [reply("x-1", "kimi-k2.6", {"input": 7, "output": 3})]) +record = kimchi(key="key-coder", KIMCHI_CODING_AGENT_DIR=str(home() / "elsewhere")) +expect("KIMCHI_CODING_AGENT_DIR names the harness the sessions are read from", + [record["totalSessions"], record["todayTokensByModel"]], [1, {"kimi-k2.6": 10}], record) + +# Sessions on disk without a key: Kimchi is in use here and waits for a sign-in. +record = kimchi() +expect("local sessions without a key wait for auth", + {key: record[key] for key in ("ready", "usageStatusText", "authHelpText", "hasLocalStats")}, + {"ready": True, "usageStatusText": "Waiting for auth", "authHelpText": "Start Kimchi and log in, or set KIMCHI_API_KEY.", "hasLocalStats": True}, record) + +# ------------------------------------------------- Kimchi's gateway elsewhere + +# Kimchi's gateway bills whatever reaches it, through whichever agent: pi and +# omp sessions on its provider ids wherever they live (a models.json +# provider named kimchi, as pi 1.4.0 and omp 18.12.6 write them), harness +# sessions on plain `kimchi`, and the setups `kimchi claudecode`, `kimchi +# codex` and its opencode integration write, which point Claude Code, Codex +# and opencode at llm.kimchi.dev or llm.eu.kimchi.dev. A Claude or kimi-* +# model through the gateway is Kimchi's, not Anthropic's or Kimi's. +sandbox() +STAMP = today_stamp() +NOW_MS = ms_in(0) + + +def pi_turns(path, provider, model, tokens, omp=False): + switch = {"type": "model_change", "id": "m0", "parentId": None, "timestamp": STAMP} + switch.update({"runtime": {"version": "18.12.6"}, "model": f"{provider}/{model}"} if omp else {"provider": provider, "modelId": model}) + user = {"role": "user", "content": [{"type": "text", "text": "Say hi"}], "timestamp": NOW_MS} + if omp: + user["attribution"] = "user" + write_pi_session(path, [ + {"type": "session", "version": 3, "id": path.stem, "timestamp": STAMP, "cwd": "/home/me"}, + switch, + {"type": "message", "id": f"u-{path.stem}", "parentId": "m0", "timestamp": STAMP, "message": user}, + {"type": "message", "id": f"a-{path.stem}", "parentId": f"u-{path.stem}", "timestamp": STAMP, "message": { + "role": "assistant", "content": [{"type": "text", "text": "hi"}], "api": "openai-completions", "provider": provider, "model": model, + "usage": {"input": tokens, "output": 82, "cacheRead": 102, "cacheWrite": 0, "totalTokens": tokens + 184, "cost": {"total": 0}}, + "stopReason": "stop", "timestamp": NOW_MS}}, + ]) + + +pi_turns(home() / ".pi/agent/sessions/--home-me--/pi-kimchi.jsonl", "kimchi", "minimax-m3", 4201) +pi_turns(home() / ".pi/agent/sessions/--home-me--/pi-kimchi-dev.jsonl", "kimchi-dev", "kimi-k2.6", 5201) +pi_turns(home() / ".omp/agent/sessions/-/omp-kimchi.jsonl", "kimchi", "kimi-k2.5", 6201, omp=True) +pi_turns(home() / ".omp/agent/sessions/-/omp-upstream.jsonl", "kimchi-dev/zai", "glm-5.2", 7201, omp=True) +pi_turns(harness() / "sessions/--home-me--/harness-kimchi.jsonl", "kimchi", "glm-5.2", 8201) +# Someone else's gateway that only shares the name is not Kimchi's. +pi_turns(home() / ".pi/agent/sessions/--home-me--/pi-lookalike.jsonl", "kimchi-proxy", "minimax-m3", 999) +# opencode rows as opencode 1.19.3 writes them for the provider Kimchi's +# opencode setup adds. +opencode_db(message=[ + ("msg_u1", "ses_k", json.dumps({"role": "user", "time": {"created": NOW_MS}, "agent": "build", "model": {"providerID": "kimchi", "modelID": "minimax-m3"}})), + ("msg_a1", "ses_k", json.dumps({"parentID": "msg_u1", "role": "assistant", "mode": "build", "agent": "build", "cost": 0, "modelID": "minimax-m3", "providerID": "kimchi", + "tokens": {"total": 9201, "input": 9201, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + "time": {"created": NOW_MS + 1, "completed": NOW_MS + 2}, "finish": "stop"})), +]) +claude_code_session("kc", [("claude-sonnet-4-5", 3000), ("kimi-k2.6", 2387), ("deepseek-v4-flash", 3390)], base_url="https://llm.eu.kimchi.dev/anthropic") +os.environ["CODEX_HOME"] = str(home() / ".codex") +(home() / ".codex").mkdir() +(home() / ".codex/config.toml").write_text('model_provider = "kimchi"\nmodel = "minimax-m3"\n\n[model_providers.kimchi]\nname = "Kimchi Gateway"\nbase_url = "https://llm.kimchi.dev/openai/v1"\nwire_api = "responses"\n') +rollouts = home() / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") +rollouts.mkdir(parents=True) +(rollouts / "rollout-kimchi.jsonl").write_text("".join(json.dumps(line) + "\n" for line in ( + {"timestamp": STAMP, "type": "session_meta", "payload": {"model_provider": "kimchi"}}, + {"timestamp": STAMP, "type": "turn_context", "payload": {"model": "minimax-m3"}}, + {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "user_message", "message": "go"}}, + {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": 1000, "cached_input_tokens": 0, "output_tokens": 10}, "total_token_usage": {"input_tokens": 1000}}}}, +))) +record = kimchi() +os.environ.pop("CODEX_HOME") +expect("Kimchi counts its gateway's usage from pi, omp, opencode, Claude Code and Codex", + {key: record[key] for key in ("todayTotalTokens", "todayPrompts", "todaySessions", "todayTokensByModel")}, + {"todayTotalTokens": 4385 + 5385 + 6385 + 7385 + 8385 + 9201 + 3000 + 2387 + 3390 + 1010, "todayPrompts": 10, "todaySessions": 8, + "todayTokensByModel": {"minimax-m3": 4385 + 9201 + 1010, "kimi-k2.6": 5385 + 2387, "kimi-k2.5": 6385, "glm-5.2": 7385 + 8385, + "claude-sonnet-4-5": 3000, "deepseek-v4-flash": 3390}}, record) +for other in ("claude", "kimi", "deepseek", "zai", "minimax", "local"): + os.environ["CODEX_HOME"] = str(home() / ".codex") + counted = collect(other, "--force") + os.environ.pop("CODEX_HOME") + expect(f"{other} counts none of Kimchi's gateway usage", [counted.get("todayTotalTokens") or 0, counted.get("todayPrompts") or 0], [0, 0], counted) +PY diff --git a/test/shell.d/agent-usage-kimi-regression-test.sh b/test/shell.d/agent-usage-kimi-regression-test.sh index 9b5f9f92f1f..effebbb886b 100644 --- a/test/shell.d/agent-usage-kimi-regression-test.sh +++ b/test/shell.d/agent-usage-kimi-regression-test.sh @@ -1,52 +1,53 @@ #!/bin/bash +# The Kimi collector's own sources: the key in KIMI_API_KEY or the sign-in the +# Kimi CLI keeps (refreshed under the CLI's own lock), the hosts a token may +# go to, and Kimi's quota answers. Kimi Code's local stats are held by +# agent-usage-kimi-stats-test.sh, the unused-machine record centrally by the +# partition test, and the limits cache by the contract checks. +# +# Each case runs in a fresh sandbox against a real stand-in for Kimi's +# services, which every https:// request is carried to. Set +# OMARCHY_TEST_KIMI_CASE to run one case. + set -euo pipefail + source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox -python3 - "$ROOT/bin/omarchy-agent-usage-kimi" <<'PY' -import contextlib -import importlib.machinery -import importlib.util -import io -import json -import os +python3 - <<'PY' import stat -import sys -import tempfile -import threading -import time -import urllib.parse -import urllib.request -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - -requests = [] +from http.server import BaseHTTPRequestHandler + +from agent_usage_fixtures import * + state = {} -payload = {"usages": {"limit_5h": {"used_ratio": 0.25, "reset_time": "2026-10-15T12:00:00Z"}, +# A day out, so the window is still open whenever the test runs. +RESET_AT = (datetime.now(timezone.utc) + timedelta(days=1)).strftime("%Y-%m-%dT%H:%M:%SZ") +PAYLOAD = {"usages": {"limit_5h": {"used_ratio": 0.25, "reset_time": RESET_AT}, "limit_7d": {"used_ratio": "0.8"}, "limit_month_total": {"used_ratio": 1}, "limit_month_code": {"used_ratio": 0}}} -class Handler(BaseHTTPRequestHandler): - def log_message(self, *args): - pass +# Kimi's usage service (GET) and its login service (POST), as state says. +class Handler(BaseHTTPRequestHandler): def reply(self, status, data): self.send_response(status) self.end_headers() self.wfile.write(json.dumps(data).encode()) def do_GET(self): - requests.append((self.path, self.headers.get("Authorization"), None)) if state.get("redirect"): self.send_response(302) self.send_header("Location", "/sink") self.end_headers() else: - self.reply(state.get("usage_status", 200), payload) + self.reply(state.get("usage_status", 200), state.get("payload", PAYLOAD)) def do_POST(self): - body = self.rfile.read(int(self.headers["Content-Length"])).decode() - requests.append((self.path, None, urllib.parse.parse_qs(body))) + self.rfile.read(int(self.headers["Content-Length"])) if state.get("token_hook"): state["token_hook"]() if state.get("entered"): @@ -57,117 +58,162 @@ class Handler(BaseHTTPRequestHandler): else: self.reply(200, state.get("grant", {"access_token": "AT_new", "refresh_token": "RT_new", "expires_in": 3600})) -server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) -threading.Thread(target=server.serve_forever, daemon=True).start() -local_url = f"http://127.0.0.1:{server.server_port}" -real_opener = urllib.request.build_opener - -# Keep real urllib request/redirect behavior; only TLS transport goes to the fixture. -class LocalTLS(urllib.request.HTTPSHandler): - def https_open(self, request): - mapped = urllib.request.Request(local_url + urllib.parse.urlsplit(request.full_url).path, - data=request.data, headers=dict(request.header_items()), method=request.get_method()) - return real_opener(scanner.NoRedirectHandler()).open(mapped, timeout=10) - -urllib.request.build_opener = lambda *handlers: real_opener(LocalTLS(), *handlers) - -def load(root): - for key in list(os.environ): - if key.startswith("KIMI_"): - del os.environ[key] - os.environ.update(HOME=str(root / "home with spaces"), KIMI_CODE_HOME=str(root / "code home"), - XDG_CACHE_HOME=str(root / "cache")) - loader = importlib.machinery.SourceFileLoader("kimi_regression", sys.argv[1]) - spec = importlib.util.spec_from_loader(loader.name, loader) - module = importlib.util.module_from_spec(spec) - loader.exec_module(module) - return module + +LOCAL_URL, reqs = stub_server(Handler, tls=True) + + +def auth_header(): + return reqs[-1]["headers"].get("Authorization") if reqs else None + def credentials(token="AT_old", expiry=None): path = scanner.CREDENTIALS_PATH path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps({"access_token": token, "refresh_token": "RT_old", - "expires_at": time.time() - 10 if expiry is None else expiry})) + "expires_at": time.time() - 10 if expiry is None else expiry})) + def record(): - saved = sys.argv - sys.argv = [saved[1], "--force", "--limits-only"] - output = io.StringIO() + return run_main(scanner, "--force", "--limits-only") + + +# ----------------------------------------------------------------- contract + +@contextlib.contextmanager +def serve(mode): + if mode in ("down", "timeout", "cut", "garbage"): + with FakeHttp(mode=mode): + yield + return + state["usage_status"] = 200 if mode == "up" else mode + yield + + +def contract(root): + check_limits_contract(Adapter("kimi", plant=lambda: os.environ.update(KIMI_API_KEY="AT_env"), serve=serve, count=lambda: len(reqs))) + + +# -------------------------------------------------------------- credentials + +def env_key(root): + credentials("AT_file", time.time() + 3600) + os.environ["KIMI_API_KEY"] = "kimi_env" + expect("KIMI_API_KEY wins over the stored sign-in", [scanner.api_key(), len(reqs)], ["kimi_env", 0]) + + +def stores(root): + del os.environ["KIMI_CODE_HOME"] + legacy = home() / ".kimi" / "credentials" + write_json(legacy / "kimi-code.json", {}) + expect("the legacy kimi-cli store is the fallback", load_collector("kimi").credentials_dir(), legacy) + write_json(home() / ".kimi-code" / "credentials" / "kimi-code.json", {}) + expect("the kimi-code store wins over the legacy one", load_collector("kimi").credentials_dir(), home() / ".kimi-code" / "credentials") + + +# kimi-cli keeps its sign-in under KIMI_SHARE_DIR when that is set +# (kimi_cli/share.py), refreshed under the file lock beside it. +def share_dir(root): + del os.environ["KIMI_CODE_HOME"] + share = root / "share dir" + os.environ["KIMI_SHARE_DIR"] = str(share) try: - with contextlib.redirect_stdout(output): - assert scanner.main() == 0 + write_json(share / "credentials" / "kimi-code.json", {"access_token": "AT_old", "refresh_token": "RT_old", "expires_at": time.time() - 10}) + expect("the legacy kimi-cli store follows KIMI_SHARE_DIR", load_collector("kimi").credentials_dir(), share / "credentials") + scanner = load_collector("kimi") + token = scanner.api_key() + expect("a KIMI_SHARE_DIR sign-in is refreshed under its own file lock", + [token, (share / "credentials" / "kimi-code.lock").is_file(), (share / "oauth").exists()], ["AT_new", True, False]) finally: - sys.argv = saved - return json.loads(output.getvalue()) + del os.environ["KIMI_SHARE_DIR"] -def destination(root): - for url in ("https://outside.example/api/oauth/token", local_url + "/api/oauth/token"): - scanner.TOKEN_URL = url - try: - scanner.request_refresh("RT_old") - except scanner.KimiError: - pass - else: - raise AssertionError("refresh accepted an untrusted destination") - assert not requests, "refresh token reached an untrusted destination" def home_isolation(root): - default = Path.home() / ".kimi-code" / "credentials" + default = home() / ".kimi-code" / "credentials" default.mkdir(parents=True) (default / "kimi-code.json").write_text("{}") - assert scanner.credentials_dir() == root / "code home" / "credentials", "explicit home fell back to another account" + expect("an explicit KIMI_CODE_HOME never falls back to another account", scanner.credentials_dir(), root / "code home" / "credentials") + def static_token(root): + credentials("AT_fresh", time.time() + 3600) + expect("a fresh token is used as it is", [scanner.api_key(), len(reqs)], ["AT_fresh", 0]) credentials("AT_static", 0) - assert scanner.api_key() == "AT_static" and not requests, "non-expiring token was refreshed" + expect("a token that never expires is never refreshed", [scanner.api_key(), len(reqs)], ["AT_static", 0]) + + +def refresh(root): + credentials() + data = record() + persisted = json.loads(scanner.CREDENTIALS_PATH.read_text()) + expect("a lapsed token is refreshed, and the usage asked with the new one", + [[request["method"] + " " + request["path"] for request in reqs], auth_header(), bool(data["limits"])], + [["POST /api/oauth/token", "GET /coding/v1/usages"], "Bearer AT_new", True], data) + expect("the rotated single-use refresh token is kept", [persisted["access_token"], persisted["refresh_token"]], ["AT_new", "RT_new"], persisted) + def cli_lock(root): credentials() lock = root / "code home" / "oauth" / "kimi-code.lock" lock.mkdir(parents=True) + def peer(): time.sleep(0.3) credentials("AT_peer", time.time() + 3600) lock.rmdir() + thread = threading.Thread(target=peer) thread.start() try: token = scanner.api_key() finally: thread.join() - assert token == "AT_peer" and not requests, "collector raced the current CLI's refresh lock" + expect("the collector waits out the Kimi CLI's refresh lock and takes its token", [token, len(reqs)], ["AT_peer", 0]) + + +def stale_lock(root): + credentials() + lock = root / "code home" / "oauth" / "kimi-code.lock" + lock.mkdir(parents=True) + os.utime(lock, (0, 0)) + started = time.monotonic() + token = scanner.api_key() + expect("a lease left by a CLI that died is taken over, not waited out", + [token, time.monotonic() - started < 5, lock.exists()], ["AT_new", True, False]) + def lock_failure(root): credentials() (root / "code home" / "oauth").write_text("blocked") try: scanner.api_key() + refused = False except (scanner.KimiError, OSError): - pass - else: - raise AssertionError("refresh continued without a lock") - assert not requests, "refresh posted while lock acquisition failed" + refused = True + expect("no lock, no refresh, and nothing posted", [refused, len(reqs)], [True, 0]) + def heartbeat(root): credentials() state.update(delay=6, entered=threading.Event()) outcome = [] + lock = root / "code home" / "oauth" / "kimi-code.lock" + def refresh(): try: outcome.append(scanner.api_key()) except Exception as error: outcome.append(error) + thread = threading.Thread(target=refresh) thread.start() try: - assert state["entered"].wait(2), "refresh request did not arrive" + check("the refresh request arrives", state["entered"].wait(2)) time.sleep(5.2) - lock = root / "code home" / "oauth" / "kimi-code.lock" - assert lock.is_dir() and time.time() - lock.stat().st_mtime < 3, "CLI lock expired during a slow refresh" + check("the CLI lock is renewed through a slow refresh", lock.is_dir() and time.time() - lock.stat().st_mtime < 3) finally: thread.join() - assert outcome == ["AT_new"], "slow refresh failed" - assert not lock.exists(), "refresh retained its directory lease" + expect("a slow refresh succeeds and lets go of its lease", [outcome, lock.exists()], [["AT_new"], False]) + def private_write(root): credentials() @@ -177,40 +223,169 @@ def private_write(root): scanner.CREDENTIALS_PATH.with_suffix(".json.tmp").symlink_to(victim) modes, synced = [], [] fdopen, fsync = os.fdopen, os.fsync + def private_open(fd, *args, **kwargs): modes.append(stat.S_IMODE(os.fstat(fd).st_mode)) return fdopen(fd, *args, **kwargs) + os.fdopen = private_open os.fsync = lambda fd: (synced.append(fd), fsync(fd))[1] try: scanner.write_credentials({"access_token": "AT_new", "refresh_token": "RT_new"}) finally: os.fdopen, os.fsync = fdopen, fsync - assert victim.read_text() == "untouched", "credential write followed a predictable temporary symlink" - assert modes == [0o600] and synced, "credentials were not private and flushed before publication" - assert stat.S_IMODE(scanner.CREDENTIALS_DIR.stat().st_mode) == 0o700 - assert stat.S_IMODE(scanner.CREDENTIALS_PATH.stat().st_mode) == 0o600 + expect("a credential write never follows a predictable temporary symlink", victim.read_text(), "untouched") + check("credentials are private and flushed before they replace the old ones", modes == [0o600] and synced, f"modes {modes}, synced {synced}") + expect("the credential directory and file are private", + [stat.S_IMODE(scanner.CREDENTIALS_DIR.stat().st_mode), stat.S_IMODE(scanner.CREDENTIALS_PATH.stat().st_mode)], [0o700, 0o600]) + + +def concurrent_signin(root): + credentials() + state["token_hook"] = lambda: credentials("AT_signed_in", time.time() + 3600) + data = record() + expect("a refresh never overwrites a sign-in made meanwhile, and the usage asks for the new account", + [data["ready"], json.loads(scanner.CREDENTIALS_PATH.read_text())["access_token"], auth_header()], [True, "AT_signed_in", "Bearer AT_signed_in"], data) + + +def concurrent_logout(root): + credentials() + state["token_hook"] = scanner.CREDENTIALS_PATH.unlink + data = record() + expect("a refresh never restores an account signed out meanwhile, which then waits for a sign-in", + [scanner.CREDENTIALS_PATH.exists(), data["usageStatusText"]], [False, "Waiting for auth"], data) + + +def invalid_grant(root): + for grant in ({"access_token": "AT_new", "expires_in": 3600}, + {"access_token": True, "refresh_token": "RT_new", "expires_in": 3600}, + {"access_token": "AT_new", "refresh_token": "RT_new", "expires_in": 0}): + credentials() + before = scanner.CREDENTIALS_PATH.read_text() + state["grant"] = grant + data = record() + expect(f"an invalid grant is refused and keeps the usable credentials ({json.dumps(grant)})", + [data["usageStatusText"], data["limits"], scanner.CREDENTIALS_PATH.read_text() == before], ["Kimi limits unavailable", [], True], data) + + +def refused_refresh(root): + credentials() + state["token_status"] = 401 + data = record() + expect("a lapsed refresh token asks for a sign-in", data["usageStatusText"], "Sign-in expired", data) + check("a refused refresh never publishes a token", "RT_old" not in json.dumps(data), data) + + +def key_help(root): + os.environ["KIMI_API_KEY"] = "AT_invalid" + state["usage_status"] = 401 + data = record() + # AUTH_HELP names KIMI_API_KEY too, so the check is for this remedy's own words. + expect("a refused KIMI_API_KEY says to change it, which a login can't", data["authHelpText"], "Kimi rejected the credentials. Update or unset KIMI_API_KEY.", data) + + +def refused(root): + credentials("AT_live", time.time() + 3600) + good = record() + state["usage_status"] = 401 + data = record() + expect("a refused sign-in says so beside the last good limits, kept stale", + [data["usageStatusText"], data["authHelpText"], data["limits"] == good["limits"], data["limitsStale"]], + ["Sign-in expired", "Kimi rejected the credentials. Sign in again with `kimi login`.", True, True], data) + + +def waiting_for_auth(root): + (root / "code home").mkdir() + data = run_main(scanner) + expect("a Kimi home without a sign-in waits for one, asking nothing", + [data["ready"], data["usageStatusText"], "kimi login" in data["authHelpText"], len(reqs)], [True, "Waiting for auth", True, 0], data) + + +# ------------------------------------------------------------ where tokens go + +def trust(root): + try: + scanner.trusted_url("http://evil.example.com/usages") + refused = False + except scanner.KimiError: + refused = True + expect("plain HTTP is refused off loopback", refused, True) + expect("loopback HTTP is allowed but sent no token", scanner.trusted_url("http://127.0.0.1:9000/usages")[1], False) + expect("kimi.com is trusted with a token", scanner.trusted_url("https://api.kimi.com/coding/v1/usages")[1], True) + expect("a domain merely ending in kimi.com is sent no token", scanner.trusted_url("https://evilkimi.com/x")[1], False) + + +def destination(root): + for where, url in (("another domain", "https://outside.example/api/oauth/token"), ("loopback", LOCAL_URL + "/api/oauth/token")): + scanner.TOKEN_URL = url + try: + scanner.request_refresh("RT_old") + accepted = True + except scanner.KimiError: + accepted = False + expect(f"a refresh token never goes to an untrusted destination ({where})", [accepted, len(reqs)], [False, 0]) + + +def redirects(root): + os.environ["KIMI_API_KEY"] = "AT_env" + state["redirect"] = True + data = record() + expect("the usage request never follows a redirect", + [data["usageStatusText"], data["limits"], [request["path"] for request in reqs]], ["Kimi limits unavailable", [], ["/coding/v1/usages"]], data) + + +# ------------------------------------------------------------------ the answer def modern_quota(root): os.environ["KIMI_API_KEY"] = "AT_env" data = record() - assert [row["percent"] for row in data["limits"]] == [0.25, 0.8, 1, 0], "current CLI quota windows were dropped" - assert data["limits"][0]["resetsAt"] == "2026-10-15T12:00:00Z" - assert data["limits"][-1].get("title") == "Code Monthly", "monthly code allowance lost its scope" - assert not data["hasLocalStats"] and not data["hasPromptStats"], "quota was advertised as local history" + expect("the current CLI's quota windows are read, the monthly code allowance titled", + [[row["percent"] for row in data["limits"]], data["limits"][0]["resetsAt"], data["limits"][-1].get("title")], + [[0.25, 0.8, 1, 0], RESET_AT, "Code Monthly"], data) + expect("Kimi's quota is this machine's record, with no local stats while no session has run", + [data.get("scope"), data["hasLocalStats"], data["hasPromptStats"]], [None, False, True], data) scanner.USAGE_URL = "https://api.kimi.ai/coding/v1/usages" record() - assert requests[-1][1] == "Bearer AT_env", "global Kimi service lost its authorization header" + expect("the global Kimi service gets the authorization header too", auth_header(), "Bearer AT_env") + -def stale_limits(root): +def legacy_quota(root): os.environ["KIMI_API_KEY"] = "AT_env" - good = record() - state["usage_status"] = 503 - stale = record() - assert good["limits"] and stale["limits"] == good["limits"], "failed probe hid the provider" - assert stale["limitsStale"] and stale["limitsFetchedAt"] == good["limitsFetchedAt"] + state["payload"] = { + "usage": {"limit": 100, "used": 33, "reset_at": "2026-08-23T11:52:13.042087123Z"}, + "limits": [ + {"window": {"duration": 300, "timeUnit": "MINUTES"}, "detail": {"limit": 50, "remaining": 17, "reset_at": "2026-08-17T16:52:13Z"}}, + {"window": {"duration": 7, "timeUnit": "DAYS"}, "detail": {"name": "Weekly quota", "limit": 200, "used": 50}}, + {"window": {"duration": 0}, "detail": {"limit": 0, "used": 0}}, + ], + } + data = record() + # Used comes from remaining when only that is given, nanosecond stamps are + # trimmed, and a window without a limit is no window. + expect("the earlier quota answer reads its weekly limit and windows", data["limits"], [ + {"label": "Weekly limit", "percent": 0.33, "resetsAt": "2026-08-23T11:52:13.042087Z"}, + {"label": "Session (5-hour)", "percent": 0.66, "resetsAt": "2026-08-17T16:52:13Z"}, + {"label": "Weekly quota", "percent": 0.25, "resetsAt": ""}], data) + + +def odd_resets(root): + os.environ["KIMI_API_KEY"] = "AT_env" + odd = ("garbage", "2030-13-45", "99999999999999999999999", "9" * 5000, "-5", "1e10", True, [1]) + for value in odd: + state["payload"] = {"usages": {"limit_5h": {"used_ratio": 0.5, "reset_time": value}}} + data = record() + expect(f"a reset time that won't parse names no reset ({str(value)[:30]})", [row["resetsAt"] for row in data["limits"]], [""], data) + + +# ------------------------------------------------------------------ the cache + +def cache_identity(root): + os.environ["KIMI_API_KEY"] = "AT_env" + check("a probe keeps limits", record()["limits"]) os.environ["KIMI_API_KEY"] = "AT_other_account" - assert record()["limits"] == [], "cached quota crossed credential identities" + state["usage_status"] = 503 + expect("another key never falls back on this one's limits", record()["limits"], []) + def rotated_cache(root): credentials("AT_old", time.time() + 3600) @@ -219,82 +394,65 @@ def rotated_cache(root): state["usage_status"] = 503 first = record() second = record() - assert first["limits"] == good["limits"] and second["limits"] == good["limits"], "rotation discarded last-good limits" - assert first["limitsStale"] and first["limitsFetchedAt"] == good["limitsFetchedAt"] + expect("a refresh that rotates the token keeps the last good limits, stale", + [first["limits"] == good["limits"], second["limits"] == good["limits"], first["limitsStale"], first["limitsFetchedAt"] == good["limitsFetchedAt"]], + [True, True, True, True], first) -def cache_identity(root): + +def empty_answer(root): + os.environ["KIMI_API_KEY"] = "AT_env" + state["payload"] = {"usages": {}} + first = record() + second = run_main(scanner, "--limits-only") + expect("an answer without windows stands in for the next probe", + [first["limits"], second["limits"], second["usageStatusText"], len(reqs)], [[], [], "", 1], second) + + +def odd_answer(root): + os.environ["KIMI_API_KEY"] = "AT_env" + good = record() + state["payload"] = {"something": "else"} + data = record() + expect("an answer of another shape keeps the last good limits, stale", + [data["limits"] == good["limits"], data["limitsStale"], data["limitsFetchedAt"] == good["limitsFetchedAt"], data.get("retryAdvised")], + [True, True, True, None], data) + expect("an answer of another shape never replaces the cached one", run_main(scanner, "--limits-only")["limits"], good["limits"]) + + +def switched_cache(root): credentials("AT_account_a", time.time() + 3600) capture = scanner.api_key + def switched(*args, **kwargs): key = capture(*args, **kwargs) credentials("AT_account_b", time.time() + 3600) return key + scanner.api_key = switched - assert record()["limits"] + check("a probe across a sign-in change still answers", record()["limits"]) scanner.api_key = capture state["usage_status"] = 503 - assert record()["limits"] == [], "request quota was cached under a different token" - -def concurrent_signin(root): - credentials() - state["token_hook"] = lambda: credentials("AT_signed_in", time.time() + 3600) - assert record()["ready"] - assert json.loads(scanner.CREDENTIALS_PATH.read_text())["access_token"] == "AT_signed_in", "refresh overwrote a completed sign-in" - assert requests[-1][1] == "Bearer AT_signed_in", "usage ignored the newly signed-in account" + expect("an answer is never kept under a token it wasn't asked with", record()["limits"], []) -def concurrent_logout(root): - credentials() - state["token_hook"] = scanner.CREDENTIALS_PATH.unlink - assert not record()["ready"] and not scanner.CREDENTIALS_PATH.exists(), "refresh restored a signed-out account" -def invalid_grant(root): - for grant in ({"access_token": "AT_new", "expires_in": 3600}, - {"access_token": True, "refresh_token": "RT_new", "expires_in": 3600}, - {"access_token": "AT_new", "refresh_token": "RT_new", "expires_in": 0}): - credentials() - before = scanner.CREDENTIALS_PATH.read_text() - state["grant"] = grant - assert not record()["ready"], "invalid rotated token was accepted" - assert scanner.CREDENTIALS_PATH.read_text() == before, "invalid grant replaced usable credentials" - -def key_help(root): - os.environ["KIMI_API_KEY"] = "AT_invalid" - state["usage_status"] = 401 - assert "KIMI_API_KEY" in record()["authHelpText"], "login cannot repair an overriding API key" - -def error_redaction(root): - credentials() - state["token_status"] = 401 - assert "RT_old" not in json.dumps(record()), "remote refresh error published a token" - -def redirects(root): - os.environ["KIMI_API_KEY"] = "AT_env" - state["redirect"] = True - assert not record()["ready"] - assert len(requests) == 1 and requests[0][0] == "/coding/v1/usages", "usage followed a redirect" - -cases = {"destination": destination, "home": home_isolation, "static": static_token, - "cli_lock": cli_lock, "lock_failure": lock_failure, "heartbeat": heartbeat, - "private_write": private_write, "quota": modern_quota, "stale": stale_limits, - "rotation_cache": rotated_cache, "cache_identity": cache_identity, - "concurrent_signin": concurrent_signin, "concurrent_logout": concurrent_logout, - "invalid_grant": invalid_grant, - "key_help": key_help, "redaction": error_redaction, "redirects": redirects} -failed = False -for name, check in cases.items(): - selected = os.environ.get("OMARCHY_TEST_KIMI_CASE") +cases = {"contract": contract, "env_key": env_key, "stores": stores, "share_dir": share_dir, "home": home_isolation, "static": static_token, + "refresh": refresh, "cli_lock": cli_lock, "stale_lock": stale_lock,"lock_failure": lock_failure, "heartbeat": heartbeat, + "private_write": private_write, "concurrent_signin": concurrent_signin, "concurrent_logout": concurrent_logout, + "invalid_grant": invalid_grant, "refused_refresh": refused_refresh, "key_help": key_help, "refused": refused, + "waiting_for_auth": waiting_for_auth, "trust": trust, "destination": destination, "redirects": redirects, + "quota": modern_quota, "legacy_quota": legacy_quota, "cache_identity": cache_identity, + "rotation_cache": rotated_cache, "switched_cache": switched_cache, "empty_answer": empty_answer, + "odd_answer": odd_answer, "odd_resets": odd_resets} +selected = os.environ.get("OMARCHY_TEST_KIMI_CASE") +for name, case in cases.items(): if selected and name != selected: continue - requests.clear() + reqs.clear() state.clear() - with tempfile.TemporaryDirectory(prefix="omarchy-kimi-regression-") as directory: - scanner = load(Path(directory)) - try: - check(Path(directory)) - print(f"ok - Kimi {name}", flush=True) - except Exception as error: - print(f"not ok - Kimi {name}: {error}", flush=True) - failed = True -server.shutdown() -raise SystemExit(1 if failed else 0) + # A home and a Kimi home with spaces in their names, as people have them. + root = Path(tempfile.mkdtemp(prefix="kimi with spaces.", dir=os.environ["SANDBOX"])) + sandbox(root) + os.environ["KIMI_CODE_HOME"] = str(root / "code home") + scanner = load_collector("kimi") + case(root) PY diff --git a/test/shell.d/agent-usage-kimi-stats-test.sh b/test/shell.d/agent-usage-kimi-stats-test.sh new file mode 100755 index 00000000000..f90fcc2669f --- /dev/null +++ b/test/shell.d/agent-usage-kimi-stats-test.sh @@ -0,0 +1,407 @@ +#!/bin/bash + +# The Kimi collector's local stats: the wire.jsonl Kimi Code 2.x keeps for +# each agent of a session under its home (KIMI_CODE_HOME, else ~/.kimi-code), +# with a usage.record for every model request and a turn.prompt for every +# typed prompt, read once each though a fork copies its source's records. +# The records below are shaped as a real Kimi Code 2.3.0 run wrote them. The +# limits are held by agent-usage-kimi-regression-test.sh, and the scan cache +# by the contract check. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +freeze_today() +NOW_MS = int(datetime.fromisoformat(today_stamp()).timestamp() * 1000) +TWO_DAYS_AGO_MS = NOW_MS - 2 * 86400 * 1000 +WORKSPACE = "wd_proj-kimi_82e7a3eee59c" + + +def kimi(*args): + return collect("kimi", *(args or ("--force",))) + + +# A machine that never ran Kimi Code has no local stats to show, and nor does +# a Kimi home with no session in it. +sandbox() +record = kimi() +(home() / ".kimi-code").mkdir() +in_use = kimi() +expect("Kimi with no sessions says it has no local stats", + [record["ready"], record["hasLocalStats"], in_use["ready"], in_use["hasLocalStats"]], [False, False, True, False], in_use) + + +def prompt(text, prompt_id, at, origin=None, kind="turn.prompt"): + line = {"type": kind, "agentId": "main", "input": [{"type": "text", "text": text}], "origin": origin or {"kind": "user"}, "turnId": 0, "time": at} + line["promptId" if kind == "turn.prompt" else "messageId"] = prompt_id + return line + + +def usage(at, input_other, output, cache_read, cache_creation=0, model="mock/kimi-k2", scope="turn", agent="main"): + return {"type": "usage.record", "agentId": agent, "model": model, + "usage": {"inputOther": input_other, "output": output, "inputCacheRead": cache_read, "inputCacheCreation": cache_creation}, + "usageScope": scope, "time": at} + + +# One answered prompt, as Kimi Code logs it: the prompt, the request's usage, +# and the loop's step.end carrying the same usage, which isn't counted again. +def turn(text, prompt_id, at, input_other, output, cache_read, message_id): + counts = {"inputOther": input_other, "output": output, "inputCacheRead": cache_read, "inputCacheCreation": 0} + return [ + prompt(text, prompt_id, at), + {"type": "context.append_message", "agentId": "main", "message": {"role": "user", "content": [{"type": "text", "text": text}], "id": prompt_id, "toolCalls": [], "origin": {"kind": "user"}}, "time": at}, + {"type": "context.append_message", "agentId": "main", "message": {"role": "user", "content": [{"type": "text", "text": "date"}], "toolCalls": [], "origin": {"kind": "injection", "variant": "date_change"}}, "time": at}, + {"type": "llm.request", "agentId": "main", "kind": "loop", "provider": "openai", "model": "kimi-k2", "modelAlias": "mock/kimi-k2", "time": at}, + usage(at + 35, input_other, output, cache_read), + {"type": "context.append_loop_event", "agentId": "main", "event": {"type": "step.end", "turnId": "0", "step": 1, "finishReason": "end_turn", "usage": counts, "messageId": message_id}, "time": at + 37}, + {"type": "agent.message.appended", "message": {"meta": {"source": "llm", "usage": counts, "messageId": message_id}}, "time": at + 37, "kind": "event"}, + {"type": "turn.ended", "agentId": "main", "turnId": 0, "reason": "completed", "time": at + 38}, + ] + + +# Where the runs below sent mock/kimi-k2, unless a test says otherwise: a +# provider of Kimi's type with no base_url, Moonshot's own API. +MOCK_CONFIG = """[providers.mock] +type = "kimi" +api_key = "sk-test-dummy" + +[models."mock/kimi-k2"] +provider = "mock" +model = "kimi-k2" +max_context_size = 131072 +""" + + +def wire(session, lines, agent="main", root=None): + config = (root or home() / ".kimi-code") / "config.toml" + if not config.exists(): + config.parent.mkdir(parents=True, exist_ok=True) + config.write_text(MOCK_CONFIG) + path = (root or home() / ".kimi-code") / "sessions" / WORKSPACE / session / "agents" / agent / "wire.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a") as handle: + if path.stat().st_size == 0: + handle.write(json.dumps({"type": "metadata", "protocol_version": "1.5", "created_at": NOW_MS}) + "\n") + for line in lines: + handle.write((line if isinstance(line, str) else json.dumps(line)) + "\n") + return path + + +# ------------------------------------------------------------ Kimi Code runs + +sandbox() +wire("session_745e630c", turn("Say hi kimi one", "msg_01M5HQ7K1Y0EDFNMS0NRAH3PRD", NOW_MS, 1163, 51, 71, "chatcmpl-mock1")) +wire("session_2a051e63", turn("Say hi kimi two", "msg_01M5HQ7KMCT3DNQ9MJVCBKR52Z", NOW_MS + 600, 2162, 52, 72, "chatcmpl-mock2")) +record = kimi() +expect("Kimi counts Kimi Code's prompts, sessions and request usage from its wire logs", { + key: record[key] for key in ("ready", "hasLocalStats", "hasPromptStats", "todayPrompts", "todaySessions", "todayTotalTokens", "totalPrompts", "totalSessions", "modelUsage") +}, { + "ready": True, "hasLocalStats": True, "hasPromptStats": True, "todayPrompts": 2, "todaySessions": 2, "todayTotalTokens": 3571, "totalPrompts": 2, "totalSessions": 2, + "modelUsage": {"mock/kimi-k2": {"inputTokens": 3325, "outputTokens": 103, "cacheReadInputTokens": 143, "cacheCreationInputTokens": 0}}, +}, record) +check("Kimi's local stats are this machine's, not the account's", "scope" not in record, record) +expect("with no sign-in the limits still wait for one", record["usageStatusText"], "Waiting for auth", record) + +# A compaction's request is billed too, and a subagent logs its own requests +# to its own wire; only what someone typed is a prompt, slash commands +# included, while injections, a subagent's task and shell commands are not. +sandbox() +wire("session_a", [ + *turn("refactor it", "msg_a1", NOW_MS, 1000, 100, 0, "m1"), + prompt("also tests", "msg_a2", NOW_MS + 100, kind="turn.steer"), + prompt("/review", "msg_a3", NOW_MS + 200, origin={"kind": "skill_activation", "trigger": "user-slash"}), + prompt("ls", "msg_a4", NOW_MS + 300, origin={"kind": "shell_command", "phase": "input"}), + prompt("compacted", "msg_a5", NOW_MS + 400, origin={"kind": "compaction_summary"}), + usage(NOW_MS + 500, 400, 40, 0, 10, scope="session"), + usage(TWO_DAYS_AGO_MS, 7, 3, 0, model="kimi-code/kimi-for-coding"), +]) +wire("session_a", [ + {**prompt("look around", "msg_sub1", NOW_MS + 50, origin={"kind": "agent"}), "agentId": "agent-0"}, + usage(NOW_MS + 60, 300, 30, 0, agent="agent-0"), +], agent="agent-0") +record = kimi() +two_days_ago = (fixture_now().astimezone() - timedelta(days=2)).strftime("%Y-%m-%d") +expect("Kimi counts every billed request, subagents' and compactions' too, and only typed prompts", { + "today": [record["todayPrompts"], record["todaySessions"], record["todayTotalTokens"]], + "two days ago": next(day["messageCount"] for day in record["recentDays"] if day["date"] == two_days_ago), + "models": sorted(record["modelUsage"]), +}, {"today": [3, 1, 1880], "two days ago": 10, "models": ["kimi-code/kimi-for-coding", "mock/kimi-k2"]}, record) + +# A request counts where it was billed. Each usage.record names its model +# alias, and config.toml where that alias goes: the Kimi Code plan /login +# provisions is Kimi's, as is an alias the config doesn't name (KIMI_MODEL_NAME +# makes one up for a run), a server you run is the local collector's, and an +# Anthropic-type provider or a gateway nobody's here, a Claude model on a +# local proxy included. A prompt goes with the request that answered it. +# Shaped as Kimi Code 2.3.0 wrote them, `kimi -m anth/claude-via-kimi -p` +# included. +sandbox() +(home() / ".kimi-code").mkdir() +(home() / ".kimi-code" / "config.toml").write_text("""default_model = "kimi-code/kimi-for-coding" + +[providers."managed:kimi-code"] +type = "kimi" +base_url = "https://api.kimi.com/coding/v1" +api_key = "" + +[models."kimi-code/kimi-for-coding"] +provider = "managed:kimi-code" +model = "kimi-for-coding" +max_context_size = 262144 + +[providers.mock] +type = "kimi" +base_url = "http://127.0.0.1:18472/v1" +api_key = "sk-test-dummy" + +[models."mock/kimi-k2"] +provider = "mock" +model = "kimi-k2" +max_context_size = 131072 + +[providers.anth] +type = "anthropic" +base_url = "http://127.0.0.1:18491" +api_key = "sk-ant-dummy-kimi" + +[models."anth/claude-via-kimi"] +provider = "anth" +model = "claude-via-kimi" +max_context_size = 131072 + +[providers.claude] +type = "anthropic" +api_key = "sk-ant-dummy" + +[models."claude/opus"] +provider = "claude" +model = "claude-opus-4-1" +max_context_size = 200000 + +[providers.gateway] +type = "openai" +base_url = "https://llm.acme.example/v1" +api_key = "sk-acme" + +[models."gateway/k2"] +provider = "gateway" +model = "kimi-k2" +max_context_size = 131072 +""") +wire("session_routes", [ + prompt("plan", "msg_r1", NOW_MS), + usage(NOW_MS + 10, 1000, 10, 0, model="kimi-code/kimi-for-coding"), + prompt("local", "msg_r2", NOW_MS + 20), + usage(NOW_MS + 30, 2162, 52, 72, model="mock/kimi-k2"), + prompt("anthropic", "msg_r3", NOW_MS + 40), + usage(NOW_MS + 50, 11234, 60, 80, 40, model="anth/claude-via-kimi"), + usage(NOW_MS + 60, 5000, 50, 0, model="claude/opus"), + usage(NOW_MS + 70, 4000, 40, 0, model="gateway/k2"), + prompt("env", "msg_r4", NOW_MS + 80), + usage(NOW_MS + 90, 100, 1, 0, model="env-model"), + prompt("unanswered", "msg_r5", NOW_MS + 100), +]) +record = kimi() +expect("Kimi counts only the requests Kimi billed, with their prompts", [record["todayTokensByModel"], record["todayPrompts"]], + [{"kimi-code/kimi-for-coding": 1010, "env-model": 101}, 3], record) +local = collect("local", "--force") +expect("Local counts the Kimi Code requests a server you run answered, and their prompts", + [local["todayTokensByModel"], local["todayPrompts"]], [{"mock/kimi-k2": 2286}, 1], local) + +# A prompt no request answered goes where its requests were sent, answered or +# not, else where its agent was bound: a server you run that refused every +# attempt is local's, and Kimi's own plan Kimi's. Shaped as Kimi Code 2.3.0 +# wrote `kimi -m dead/kimi-dead -p` against a dead port, retries and all. +sandbox() +(home() / ".kimi-code").mkdir() +(home() / ".kimi-code" / "config.toml").write_text("""[providers."managed:kimi-code"] +type = "kimi" +base_url = "https://api.kimi.com/coding/v1" + +[models."kimi-code/kimi-for-coding"] +provider = "managed:kimi-code" +model = "kimi-for-coding" + +[providers.dead] +type = "kimi" +base_url = "http://127.0.0.1:18499/v1" +api_key = "sk-test-dummy-dead" + +[models."dead/kimi-dead"] +provider = "dead" +model = "kimi-dead" +max_context_size = 131072 +""") + + +def failed_turn(alias, prompt_id, at, attempts=2): + lines = [{"type": "profile.bind", "agentId": "main", "modelAlias": alias, "profileName": "agent", "thinkingEffort": "off", "systemPrompt": "You are Kimi Code CLI"}, prompt("hi", prompt_id, at)] + for step in range(1, attempts + 1): + lines += [ + {"type": "llm.request", "agentId": "main", "kind": "loop", "provider": "openai", "model": alias.split("/")[-1], "modelAlias": alias, "turnStep": f"0.{step}", "time": at + step}, + {"type": "context.append_loop_event", "agentId": "main", "event": {"type": "step.end", "turnId": "0", "step": step, "finishReason": "error"}, "time": at + step}, + ] + return lines + [{"type": "turn.ended", "agentId": "main", "turnId": 0, "reason": "failed", "error": {"code": "provider.connection_error", "message": "Connection error.", "name": "APIConnectionError", "retryable": True}, "time": at + 9}] + + +wire("session_dead", failed_turn("dead/kimi-dead", "msg_01M5J2C56CQXYMM9KX3ABF4D0H", NOW_MS)) +wire("session_plan", failed_turn("kimi-code/kimi-for-coding", "msg_p", NOW_MS)) +# Bound to the dead server, and stopped before any request. +wire("session_bound", failed_turn("dead/kimi-dead", "msg_b", NOW_MS, attempts=0)) +records = {cid: collect(cid, "--force") for cid in ("kimi", "local")} +expect("a Kimi Code prompt no request answered counts where its route goes", + {cid: [record["todayPrompts"], record["todaySessions"], record["todayTotalTokens"]] for cid, record in records.items()}, + {"kimi": [1, 1, 0], "local": [2, 2, 0]}, records) + +# Kimi Code sends nothing on an alias config.toml doesn't name, so one +# missing from it was removed since: a provider's, which nothing can place +# any more and counts nowhere, prompts included, or one of the plan's +# kimi-code/ models, which stays Kimi's. KIMI_MODEL_NAME's alias, never +# saved, is Kimi's too. Shaped as Kimi Code 2.3.0 wrote `kimi -m +# dead/kimi-dead -p` against a dead port and a tool loop on mock2/kimi-tool, +# both providers then deleted from config.toml. +sandbox() +(home() / ".kimi-code").mkdir() +(home() / ".kimi-code" / "config.toml").write_text('default_model = "kimi-code/kimi-for-coding"\n') +wire("session_dead", failed_turn("dead/kimi-dead", "msg_01M5J2C56CQXYMM9KX3ABF4D0H", NOW_MS)) +wire("session_bound", failed_turn("dead/kimi-dead", "msg_b", NOW_MS, attempts=0)) +wire("session_tool", [prompt("TOOLLOOP read note.txt", "msg_t1", NOW_MS), usage(NOW_MS + 10, 24702, 120, 0, model="mock2/kimi-tool")]) +wire("session_plan", [prompt("plan", "msg_p1", NOW_MS), usage(NOW_MS + 10, 1000, 10, 0, model="kimi-code/kimi-k2.5")]) +wire("session_env", [prompt("env", "msg_e1", NOW_MS), usage(NOW_MS + 10, 100, 1, 0, model="__kimi_env_model__")]) +records = {cid: collect(cid, "--force") for cid in ("kimi", "local")} +expect("a removed Kimi Code provider counts nowhere, while the plan's dropped models and KIMI_MODEL_NAME's stay Kimi's", + {cid: [record["todayPrompts"], record["todaySessions"], record["todayTokensByModel"]] for cid, record in records.items()}, + {"kimi": [2, 2, {"kimi-code/kimi-k2.5": 1010, "__kimi_env_model__": 101}], "local": [0, 0, {}]}, records) + +# A fork copies its source's wire, then marks it forked; the copies count +# once, under the source's session, while what the fork did after counts too. +# A fork with nothing of its own is no session of its own. +sandbox() +source = turn("start", "msg_s1", NOW_MS, 1000, 100, 0, "s1") +wire("session_f", source) +wire("session_0", [*source, {"type": "forked", "agentId": "main", "time": NOW_MS + 1000}, *turn("branch", "msg_f1", NOW_MS + 2000, 50, 5, 0, "f1")]) +wire("session_1", [*source, {"type": "forked", "agentId": "main", "time": NOW_MS + 1000}]) +record = kimi() +expect("Kimi counts a fork's copies of its source once, and what the fork added", + [record["totalPrompts"], record["totalSessions"], record["todayTotalTokens"]], [2, 2, 1155], record) + +# KIMI_CODE_HOME moves Kimi Code's home, and the default one isn't read then. +sandbox() +wire("session_default", turn("hi", "msg_d", NOW_MS, 1, 1, 0, "d")) +moved = home() / "kimi moved" +wire("session_moved", turn("hi", "msg_m", NOW_MS, 10, 1, 0, "m"), root=moved) +os.environ["KIMI_CODE_HOME"] = str(moved) +record = kimi() +del os.environ["KIMI_CODE_HOME"] +expect("Kimi reads the wire logs under KIMI_CODE_HOME", record["todayTotalTokens"], 11, record) + +# --------------------------------------------------------- untrusted fields + +sandbox() +wire("session_bad", [ + usage(NOW_MS, "100", True, -5), + usage(NOW_MS + 1, 1e30, 0, 0), + usage(NOW_MS + 2, float("nan"), 0, 0), + usage(0, 1000, 0, 0), + usage("yesterday", 1000, 0, 0), + {**usage(NOW_MS + 3, 1000, 0, 0), "time": None}, + {**usage(NOW_MS + 4, 1000, 0, 0), "model": ["m"]}, + {**usage(NOW_MS + 5, 1000, 0, 0), "usage": [1000]}, + prompt("hi", "msg_bad", 0), + prompt("hi", "msg_bad2", NOW_MS, origin="user"), + usage(NOW_MS + 6, 10, 2, 0), + # A line nested deeper than the parser will go. + "[" * 100000 + '"usage.record"', +]) +with wire("session_bad", []).open("a") as torn: + torn.write('{"type":"usage.record","agentId":"main","model":"mock/kimi-k2","usage":{"inputOther":') +record = kimi() +expect("Kimi skips bad counts, bad times and models, a line nested too deep, and a torn last line", + [record["todayPrompts"], record["todayTotalTokens"]], [0, 12], record) + +# A wire that can't be read leaves the others counted, and that scan isn't +# cached. +sandbox() +wire("session_ok", turn("hi", "msg_ok", NOW_MS, 10, 2, 0, "ok")) +(home() / ".kimi-code" / "sessions" / WORKSPACE / "session_broken" / "agents" / "main" / "wire.jsonl").mkdir(parents=True) +with contextlib.redirect_stderr(io.StringIO()): + record = kimi() +expect("an unreadable wire leaves the rest counted", record["todayTotalTokens"], 12, record) +check("a scan that couldn't read every wire isn't cached", not scan_cache_files("kimi")) + +# Claude Code pointed at Kimi's or Moonshot's API is Kimi's; with no endpoint, +# only a Kimi model's reply is. That usage alone puts Kimi in use here, with +# no sign-in to ask for its limits. +sandbox() +claude_code_session("k", [("kimi-k2-0905", 300), ("moonshot-v1-8k", 20), ("claude-sonnet-4-5", 40)]) +record = kimi() +expect("Kimi counts Claude Code's Kimi and Moonshot replies, and waits for a sign-in on them alone", + [record["todayTokensByModel"], record["todayPrompts"], record["ready"], record["usageStatusText"]], + [{"kimi-k2-0905": 300, "moonshot-v1-8k": 20}, 2, True, "Waiting for auth"], record) +claude_code_session("k", [("kimi-for-coding", 300), ("claude-sonnet-4-5", 40)], base_url="https://api.kimi.com/coding/") +record = kimi() +expect("Kimi counts every Claude Code reply while Claude Code is pointed at Kimi Code", + [record["todayTokensByModel"], record["todayPrompts"]], [{"kimi-for-coding": 300, "claude-sonnet-4-5": 40}, 2], record) + +# pi, omp and opencode on the Kimi Code plan count here: pi's own +# `kimi-coding` provider and models.dev's `kimi-code-plan-global` and +# `kimi-code-plan-cn`, and so does Moonshot's pay-as-you-go API. A gateway +# named like one of them doesn't. +sandbox() +write_pi_session(home() / ".pi" / "agent" / "sessions" / "project" / "pi.jsonl", [ + {"type": "message", "id": "u1", "timestamp": today_stamp(), "message": {"role": "user", "content": "hi"}}, + pi_reply("kimi-coding", 300, id="k1", model="k3", api="openai-completions"), + pi_reply("moonshotai", 999, id="m1", model="kimi-k2.6"), + pi_reply("kimi-coding-proxy", 999, id="p1", model="k3"), +]) +opencode_db(message=[ + ("msg_1", "ses_1", oc_v1("kimi-code-plan-global", 200, model="kimi-for-coding")), + ("msg_2", "ses_1", oc_v1("kimi-code-plan-cn", 20, model="kimi-for-coding-highspeed")), + ("msg_3", "ses_1", oc_v1("moonshotai-cn", 999, model="kimi-k2.6")), +]) +record = kimi() +expect("Kimi counts pi and opencode on the Kimi Code plan ids and Moonshot's pay-as-you-go API", + [record["todayTokensByModel"], record["ready"], record["usageStatusText"]], + [{"k3": 300, "kimi-k2.6": 1998, "kimi-for-coding": 200, "kimi-for-coding-highspeed": 20}, True, "Waiting for auth"], record) + +# A route change reaches the cached scan at once, even on the panel's +# --limits-only refresh that otherwise reuses a scan for minutes: the request +# on Moonshot's API was Kimi's, and once config.toml sends its alias to a +# server you run, it is not. Shaped as Kimi Code 2.3.0 wrote a tool loop on +# that alias. +sandbox() +(home() / ".kimi-code").mkdir() +(home() / ".kimi-code" / "config.toml").write_text('[providers.mock2]\ntype = "kimi"\n\n[models."mock2/kimi-tool"]\nprovider = "mock2"\nmodel = "kimi-tool"\n') +wire("session_tool", [prompt("TOOLLOOP read note.txt", "msg_t1", NOW_MS), usage(NOW_MS + 10, 24702, 120, 0, model="mock2/kimi-tool")]) +before = kimi("--limits-only") +(home() / ".kimi-code" / "config.toml").write_text("""[providers.mock2] +type = "kimi" +base_url = "http://127.0.0.1:18491/v1" + +[models."mock2/kimi-tool"] +provider = "mock2" +model = "kimi-tool" +""") +after = kimi("--limits-only") +expect("a Kimi Code route change reaches Kimi's cached scan at once", + [before["todayTokensByModel"], after["todayTokensByModel"], after["todayPrompts"]], [{"mock2/kimi-tool": 24822}, {}, 0], after) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_turn(): + add_turn.count = getattr(add_turn, "count", 0) + 1 + wire(f"session_grown_{add_turn.count}", turn("again", f"msg_g{add_turn.count}", NOW_MS + add_turn.count, 10, 1, 0, f"g{add_turn.count}")) + + +check_scan_cache(Adapter("kimi", grow=add_turn)) +PY diff --git a/test/shell.d/agent-usage-kimi-test.sh b/test/shell.d/agent-usage-kimi-test.sh deleted file mode 100644 index d1627f5c909..00000000000 --- a/test/shell.d/agent-usage-kimi-test.sh +++ /dev/null @@ -1,151 +0,0 @@ -#!/bin/bash - -source "$(dirname "$0")/base-test.sh" - -require_command jq -require_command python3 - -TEST_HOME=$(mktemp -d) -trap 'rm -rf "$TEST_HOME"' EXIT - -# 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" KIMI_CODE_HOME="$TEST_HOME/kimi-code" KIMI_API_KEY="" \ - "$ROOT/bin/omarchy-agent-usage-kimi") - -[[ $(jq -r '.id + ":" + (.ready | tostring) + ":" + (.scope)' <<<"$no_key") == "kimi:false:account" ]] || - fail "Kimi collector prints a valid record without credentials" "$no_key" -pass "Kimi collector prints a valid record without credentials" - -result=$(TEST_HOME="$TEST_HOME" python3 - "$ROOT/bin/omarchy-agent-usage-kimi" <<'PY' -import importlib.machinery -import importlib.util -import io -import json -import os -import sys -import time -from contextlib import redirect_stdout -from pathlib import Path - -collector_path = str(Path(sys.argv[1])) -test_home = Path(os.environ["TEST_HOME"]) - -os.environ["HOME"] = str(test_home / "home") -os.environ["KIMI_CODE_HOME"] = str(test_home / "kimi-code") - -loader = importlib.machinery.SourceFileLoader("kimi_collector", collector_path) -spec = importlib.util.spec_from_loader(loader.name, loader) -scanner = importlib.util.module_from_spec(spec) -loader.exec_module(scanner) - -summary = {} - -# The env key wins without touching any credentials file. -os.environ["KIMI_API_KEY"] = "kimi_env" -summary["envKeyWins"] = scanner.api_key() == "kimi_env" -del os.environ["KIMI_API_KEY"] - -# The kimi-code store wins over the legacy kimi-cli one; legacy is the fallback. -code_dir = test_home / "kimi-code" / "credentials" -legacy_dir = test_home / "home" / ".kimi" / "credentials" -code_dir.mkdir(parents=True, exist_ok=True) -legacy_dir.mkdir(parents=True, exist_ok=True) -(legacy_dir / "kimi-code.json").write_text("{}") -del os.environ["KIMI_CODE_HOME"] -summary["legacyFallback"] = scanner.credentials_dir() == legacy_dir -os.environ["KIMI_CODE_HOME"] = str(test_home / "kimi-code") -(code_dir / "kimi-code.json").write_text("{}") -summary["codeStorePreferred"] = scanner.credentials_dir() == code_dir - -# A fresh token is used as-is; an expired one is refreshed under the lock and -# the rotated single-use refresh token is persisted back. -scanner.CREDENTIALS_DIR = code_dir -scanner.CREDENTIALS_PATH = code_dir / "kimi-code.json" - -fresh = {"access_token": "at_fresh", "refresh_token": "rt_1", "expires_at": time.time() + 3600} -scanner.CREDENTIALS_PATH.write_text(json.dumps(fresh)) -summary["freshTokenUsed"] = scanner.api_key() == "at_fresh" - -expired = {"access_token": "at_old", "refresh_token": "rt_1", "expires_at": time.time() - 10} -scanner.CREDENTIALS_PATH.write_text(json.dumps(expired)) -scanner.request_refresh = lambda token: { - "access_token": "at_new", - "refresh_token": "rt_2", - "expires_in": 3600, -} -granted = scanner.refresh_credentials() -persisted = json.loads(scanner.CREDENTIALS_PATH.read_text()) -summary["refreshRotates"] = ( - granted["access_token"] == "at_new" - and persisted["refresh_token"] == "rt_2" - and persisted["access_token"] == "at_new" -) - -# Endpoint trust: plain HTTP is refused off-loopback; tokens only go to kimi.com. -try: - scanner.trusted_url("http://evil.example.com/usages") - summary["httpRefused"] = False -except scanner.KimiError: - summary["httpRefused"] = True -summary["loopbackAllowed"] = scanner.trusted_url("http://127.0.0.1:9000/usages")[1] is False -summary["kimiTrusted"] = scanner.trusted_url("https://api.kimi.com/coding/v1/usages")[1] is True - -# The usage payload renders into display rows: the weekly quota, the 5-hour -# session window, used derived from remaining, and nanosecond stamps trimmed. -payload = { - "usage": {"limit": 100, "used": 33, "reset_at": "2026-08-23T11:52:13.042087123Z"}, - "limits": [ - { - "window": {"duration": 300, "timeUnit": "MINUTES"}, - "detail": {"limit": 50, "remaining": 17, "reset_at": "2026-08-17T16:52:13Z"}, - }, - { - "window": {"duration": 7, "timeUnit": "DAYS"}, - "detail": {"name": "Weekly quota", "limit": 200, "used": 50}, - }, - {"window": {"duration": 0}, "detail": {"limit": 0, "used": 0}}, - ], -} -limits = scanner.to_limits(payload) -summary["limits"] = limits - -os.environ["KIMI_API_KEY"] = "kimi_env" -scanner.quota = lambda key: payload -sys.argv = ["omarchy-agent-usage-kimi", "--limits-only"] -out = io.StringIO() -with redirect_stdout(out): - scanner.main() -summary["record"] = json.loads(out.getvalue()) -summary["quotaOnly"] = not summary["record"]["hasLocalStats"] and not summary["record"]["hasPromptStats"] - -print(json.dumps(summary, separators=(",", ":"))) -PY -) - -[[ $(jq -r '.envKeyWins' <<<"$result") == "true" ]] || - fail "Kimi collector prefers KIMI_API_KEY over stored credentials" "$result" -pass "Kimi collector prefers KIMI_API_KEY over stored credentials" - -[[ $(jq -r '"\(.codeStorePreferred):\(.legacyFallback)"' <<<"$result") == "true:true" ]] || - fail "Kimi collector prefers the kimi-code store and falls back to legacy kimi-cli" "$result" -pass "Kimi collector prefers the kimi-code store and falls back to legacy kimi-cli" - -[[ $(jq -r '"\(.freshTokenUsed):\(.refreshRotates)"' <<<"$result") == "true:true" ]] || - fail "Kimi collector refreshes expired tokens and persists the rotated pair" "$result" -pass "Kimi collector refreshes expired tokens and persists the rotated pair" - -[[ $(jq -r '"\(.httpRefused):\(.loopbackAllowed):\(.kimiTrusted)"' <<<"$result") == "true:true:true" ]] || - fail "Kimi collector only sends tokens to kimi.com over HTTPS" "$result" -pass "Kimi collector only sends tokens to kimi.com over HTTPS" - -[[ $(jq -c '.limits' <<<"$result") == '[{"label":"Weekly limit","percent":0.33,"resetsAt":"2026-08-23T11:52:13.042087Z"},{"label":"Session (5-hour)","percent":0.66,"resetsAt":"2026-08-17T16:52:13Z"},{"label":"Weekly quota","percent":0.25,"resetsAt":""}]' ]] || - fail "Kimi collector renders usage windows into display rows" "$result" -pass "Kimi collector renders usage windows into display rows" - -[[ $(jq -r '.record.ready | tostring' <<<"$result") == "true" && $(jq -r '.record.id' <<<"$result") == "kimi" ]] || - fail "Kimi collector prints the display-ready record contract" "$result" -pass "Kimi collector prints the display-ready record contract" - -[[ $(jq -r '.quotaOnly' <<<"$result") == "true" ]] || fail "Kimi quota is not local usage statistics" "$result" -pass "Kimi quota is not local usage statistics" diff --git a/test/shell.d/agent-usage-local-scanner-test.sh b/test/shell.d/agent-usage-local-scanner-test.sh new file mode 100755 index 00000000000..8843c4eea83 --- /dev/null +++ b/test/shell.d/agent-usage-local-scanner-test.sh @@ -0,0 +1,345 @@ +#!/bin/bash + +# The Local collector counts the turns that ran on hardware you own: its own +# reading of Codex rollouts and their config.toml, the self-hosted runtimes it +# knows by name, the providers local.json names, and the rows a provider +# collector claims, which it leaves out. How the shared pi and opencode readers +# work, and the unused-machine record, are held centrally by the shared and +# partition tests, and the scan cache by the contract check. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox +require_command rg + +python3 - <<'PY' +from agent_usage_fixtures import * + +STAMP = today_stamp() + + +def use_home(): + sandbox() + os.environ["CODEX_HOME"] = str(home() / ".codex") + + +# A machine that never ran a self-hosted model has no local stats to show. +use_home() +record = collect("local", "--force") +expect("Local with no usage says it has no local stats", [record["ready"], record["hasLocalStats"]], [False, False], record) + + +def rollouts(): + return Path(os.environ["CODEX_HOME"]) / "sessions" / datetime.now().strftime("%Y/%m/%d") + + +# A token_count line for a turn spending 1000 input and 100 output tokens, +# with the session's running total so far. +def token_count(total): + return {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": 1000, "output_tokens": 100, "cached_input_tokens": 0}, + "total_token_usage": {"input_tokens": total}, + }}} + + +def meta(provider): + return {"timestamp": STAMP, "type": "session_meta", "payload": {"model_provider": provider}} + + +def rollout(name, *lines): + path = rollouts() / f"rollout-{name}.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(line) + "\n" for line in lines)) + return path + + +# A prompt as Codex 0.163 writes it in a paginated rollout. +def typed(text="hello"): + item = {"type": "UserMessage", "id": "u1", "content": [{"type": "text", "text": text, "text_elements": []}]} + return {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "item_completed", "thread_id": "t1", "turn_id": "turn1", "item": item}} + + +def context(model): + return {"timestamp": STAMP, "type": "turn_context", "payload": {"model": model}} + + +def turn(name, provider, model): + return rollout(name, meta(provider), context(model), typed(), token_count(1000)) + + +def local_config(*providers): + write_json(config_home() / "omarchy" / "agents" / "local.json", {"providers": list(providers)}) + + +def local(*args): + return collect("local", *(args or ("--force",))) + + +def models(record): + return sorted(record["modelUsage"]) + + +# ------------------------------------------------------------------ Codex + +use_home() +(Path(os.environ["CODEX_HOME"]) / "config.toml").parent.mkdir(parents=True, exist_ok=True) +(Path(os.environ["CODEX_HOME"]) / "config.toml").write_text( + '[model_providers.homelab]\nbase_url = "http://server.home:11434/v1"\n\n' + '[model_providers.openrouter]\nbase_url = "https://openrouter.ai/api/v1"\n' +) +homelab = turn("homelab", "homelab", "qwen-net") # another machine on the LAN +ollama = turn("ollama", "ollama", "qwen-here") # this machine, a well-known runtime +turn("paid", "openrouter", "claude-paid") # a third-party API: neither collector +turn("sub", "openai", "gpt-sub") # the Codex subscription +turn("cloud", "ollama", "gpt-oss:120b-cloud") # Ollama Cloud, through the local daemon +# A quota update repeats the last turn's usage with an unchanged cumulative +# total; a real second turn moves it. +with ollama.open("a") as handle: + handle.write(json.dumps(token_count(1000)) + "\n" + json.dumps(token_count(2000)) + "\n") + # A line nested too deep to parse (a RecursionError, not a ValueError) + # costs only itself. + handle.write('{"type": "token_count", "info": ' + "[" * 100000 + "\n") +# A fork copies its parent's session_meta after its own; the first decides. +rollout("fork", meta("openai"), meta("ollama"), {"timestamp": STAMP, "type": "turn_context", "payload": {"model": "qwen-forked"}}, token_count(1000)) + +local_record = collect("local") +expect("Local counts Codex turns on LAN endpoints and known runtimes, nothing else", models(local_record), ["qwen-here", "qwen-net"], local_record) +expect("Local counts a repeated Codex token_count once, and each typed prompt once", + [local_record["todayTotalTokens"], local_record["modelUsage"]["qwen-here"]["inputTokens"], local_record["todayPrompts"], local_record["ready"]], + [3300, 2000, 2, True], local_record) + +# The Local and Codex collectors split the rollouts between them: the +# subscription's turns are Codex's alone, and no model shows in both. +codex_record = collect("codex", "--force") +expect("Codex keeps the subscription's own rollouts", models(codex_record), ["gpt-sub", "qwen-forked"], codex_record) +expect("Local and Codex never count the same rollout", set(models(local_record)) & set(models(codex_record)), set()) + +# The cloud-tagged turn Local leaves out is Ollama Cloud's, as Codex's `--oss` +# wrote it: counted there once, with the prompt that asked for it. +with FakeHttp(): + ollama_record = collect("ollama", "--force") +expect("Ollama counts the cloud-tagged Codex rollout Local leaves out", + [ollama_record["todayTokensByModel"], ollama_record["todayPrompts"], set(models(local_record)) & set(models(ollama_record))], + [{"gpt-oss": 1100}, 1, set()], ollama_record) + +# A rollout is replayed from the per-file cache while its size and mtime stand: +# change its contents but neither, and the old count stays. +stat = homelab.stat() +homelab.write_text(homelab.read_text().replace('"output_tokens": 100', '"output_tokens": 900')) +os.utime(homelab, ns=(stat.st_atime_ns, stat.st_mtime_ns)) +replayed = collect("local", "--limits-only") +for path in scan_cache_files("local"): + path.unlink() +expired = collect("local") +reread = local() +expect("Local replays an unchanged rollout from its file cache, after the scan cache too, and --force rereads it", + [record["modelUsage"]["qwen-net"]["outputTokens"] for record in (replayed, expired, reread)], [100, 100, 900]) + +# A provider the heuristics can't see is named in local.json, and takes effect +# without --force, since the classification is part of the scan key. +local_config("mystery") +turn("mystery", "mystery", "qwen-custom") +record = collect("local") +expect("Local counts the providers local.json names", record["modelUsage"].get("qwen-custom", {}).get("inputTokens"), 1000, record) + +# A Codex config that isn't UTF-8 names no endpoints, like a missing one, and +# the ones it named before stay remembered. +config = Path(os.environ["CODEX_HOME"]) / "config.toml" +saved = config.read_bytes() +config.write_bytes(b'a = "\xff\xfe"\n') +record = local() +expect("Local reads past a Codex config that isn't UTF-8", models(record), ["qwen-custom", "qwen-here", "qwen-net"], record) +config.write_bytes(saved) + +# A provider collector's id is never local, whatever local.json says. +local_config("openrouter", "openai") +record = collect("local") +expect("Local never counts a provider another collector counts", models(record), ["qwen-here", "qwen-net"], record) + +# A cache file of the wrong shape is a miss, and a cache that can't be written +# still leaves a record. +for path in scan_cache_files("local"): + path.write_text("[]") +record = collect("local", "--limits-only") +expect("Local rescans past a scan cache of the wrong shape", models(record), ["qwen-here", "qwen-net"], record) +cache_home = os.environ["XDG_CACHE_HOME"] +(home() / "not-a-dir").write_text("") +os.environ["XDG_CACHE_HOME"] = str(home() / "not-a-dir") +with contextlib.redirect_stderr(io.StringIO()): + record = collect("local") +os.environ["XDG_CACHE_HOME"] = cache_home +expect("Local still prints its record when the cache can't be written", models(record), ["qwen-here", "qwen-net"], record) + +# Prompts are the ones typed, not the requests answering them, each under the +# model that answered it even when a legacy rollout writes it before the +# turn_context: a tool call's two requests are one prompt, Codex's own context +# and a subagent's messages are none, and a cloud model's prompt is no more +# local than its tokens. A subagent naming no thread that started it starts no +# session of its own either. +use_home() +rollout("paginated", meta("ollama"), context("qwen-tools"), typed("run ls"), token_count(1000), + {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "item_completed", "item": {"type": "CommandExecution", "id": "c1"}}}, + token_count(2000), context("gpt-oss:120b-cloud"), typed("now in the cloud"), token_count(3000)) +rollout("legacy", meta("ollama"), + {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "user_message", "message": "", "kind": "user_instructions"}}, + {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "user_message", "message": "fix it"}}, context("qwen-legacy"), token_count(1000)) +rollout("subagent", {"timestamp": STAMP, "type": "session_meta", "payload": {"model_provider": "ollama", "source": {"subagent": "review"}}}, + context("qwen-review"), typed(), token_count(1000)) +record = local() +expect("Local counts the prompts typed, under the model that answered them", { + "prompts": record["todayPrompts"], "sessions": record["totalSessions"], "models": models(record), "ready": record["ready"], +}, {"prompts": 2, "sessions": 2, "models": ["qwen-legacy", "qwen-review", "qwen-tools"], "ready": True}, record) + +# A rollout records its provider's id, never the URL it reached, so a custom +# provider's endpoint is remembered from config.toml: its history stays local +# after the provider leaves the config, as in a real Codex 0.163 run on a +# provider since removed. A profile's .config.toml names providers too. +# One the config never named is still nobody's, and a config that moves an id +# to a hosted URL wins over what was remembered. +use_home() +codex_config = Path(os.environ["CODEX_HOME"]) / "config.toml" +codex_config.parent.mkdir(parents=True, exist_ok=True) +codex_config.write_text('model_provider = "mocklocal"\n\n[model_providers.mocklocal]\nname = "mock"\nbase_url = "http://127.0.0.1:18441/v1"\n' + '\n[model_providers.boxed]\nbase_url = "http://192.168.1.20:8000/v1"\n') +turn("mocklocal", "mocklocal", "gpt-5.1-codex") +turn("boxed", "boxed", "qwen-boxed") +turn("unseen", "unseen", "qwen-unseen") +(codex_config.parent / "work.config.toml").write_text('[model_providers.profiled]\nbase_url = "http://gpu-box:8080/v1"\n') +turn("profiled", "profiled", "qwen-profiled") +before = local() +codex_config.write_text('model = "gpt-5.1-codex"\n\n[model_providers.boxed]\nbase_url = "https://api.example.com/v1"\n') +after = local() +expect("Local keeps a custom Codex provider's history after it leaves config.toml", [models(before), models(after)], + [["gpt-5.1-codex", "qwen-boxed", "qwen-profiled"], ["gpt-5.1-codex", "qwen-profiled"]], after) + +# The base URL decides, not the id: a Codex entry under an id another +# collector claims is local's on this machine, as `kimchi codex` writes one +# for a self-hosted Kimchi gateway, and still the claiming collector's +# nowhere else (a real Codex 0.163 run served 10074 tokens there). +use_home() +codex_config = Path(os.environ["CODEX_HOME"]) / "config.toml" +codex_config.parent.mkdir(parents=True, exist_ok=True) +codex_config.write_text('model_provider = "kimchi"\nmodel = "kimi-k2-mock"\n\n[model_providers.kimchi]\nname = "Kimchi Gateway"\n' + 'base_url = "http://127.0.0.1:18491/openai/v1"\nwire_api = "responses"\n') +turn("kimchi-here", "kimchi", "kimi-k2-mock") +turn("kimchi-cloud", "kimchi", "glm-4.6:cloud") +record = local() +expect("Local counts a Codex entry on this machine whatever collector claims its id, but not a cloud-tagged model", + models(record), ["kimi-k2-mock"], record) +expect("and the claiming collector doesn't count it too", collect("kimchi", "--force").get("modelUsage") or {}, {}) + +# Only a turn with no time at all takes the file's day; a time of 0, false or +# "" can't be trusted, so that turn counts nowhere, as in the Codex collector. +use_home() +untimed = dict(token_count(9000)) +del untimed["timestamp"] +rollout("falsy", meta("ollama"), {"timestamp": STAMP, "type": "turn_context", "payload": {"model": "qwen-falsy"}}, + *({**token_count(1000 * n), "timestamp": falsy} for n, falsy in enumerate((0, False, ""), 1)), untimed) +record = local() +expect("Local counts a Codex turn with an untrusted time nowhere, and one with none on the file's day", + [record["todayTotalTokens"], record["modelUsage"].get("qwen-falsy", {}).get("inputTokens")], [1100, 1000], record) + +# ------------------------------------------------------------- pi and opencode + +# Runtimes known by name, whatever their case, are local; Ollama's cloud +# models and provider, another provider, and pi driving the Codex API are not. +use_home() +codex_api = pi_reply("ollama", 999, id="x", model="gpt-5", api="openai-codex-responses") +write_pi_session(home() / ".pi" / "agent" / "sessions" / "--work--" / "one.jsonl", [ + pi_reply("ollama", 100, id="a", model="qwen3:30b"), + pi_reply("LMStudio", 200, id="b", model="qwen3:8b"), + pi_reply("vllm", 300, id="c", model="llama-70b"), + pi_reply("llamacpp", 40, id="d", model="gemma"), + pi_reply("ollama", 999, id="e", model="glm-4.6:cloud"), + pi_reply("ollama-cloud", 999, id="f", model="qwen3:30b"), + pi_reply("anthropic", 999, id="g", model="claude"), + codex_api, +]) +# opencode: a known runtime, a provider local.json names, and a cloud model. +local_config("homelab") +opencode_db(message=[ + ("m1", "ses_1", oc_v1("ollama", {"input": 100, "output": 10, "reasoning": 5, "cache": {"read": 0, "write": 0}}, model="qwen-oc")), + ("m2", "ses_1", oc_v1("homelab", 50, model="qwen-cfg")), + ("m3", "ses_1", oc_v1("ollama", 999, model="kimi-k2:1t-cloud")), +]) +record = local() +expect("Local counts pi and opencode replies on self-hosted runtimes and configured providers", + [record["todayTotalTokens"], models(record)], [805, ["gemma", "llama-70b", "qwen-cfg", "qwen-oc", "qwen3:30b", "qwen3:8b"]], record) + +# A local.json naming ids the provider collectors claim leaves only their rows +# out: opencode's google provider and Vertex only on antigravity- models, which +# the agy collector counts, and pi's google rows, which agy never reads, stay +# local's, as does an id no collector matches, like minimax-coding. +use_home() +local_config("google", "google-antigravity", "google-vertex", "minimax-coding") +opencode_db(message=[ + ("m1", "ses_1", oc_v1("google", 10, model="gemini-oc")), + ("m2", "ses_1", oc_v1("google-vertex", 999, model="antigravity-gemini")), + ("m3", "ses_1", oc_v1("minimax-coding", 20, model="minimax-gateway")), +], session_message=[ + ("s1", "ses_2", "assistant", oc_v2("google-vertex", 30, model="gemini-vertex")), +]) +write_pi_session(home() / ".pi" / "agent" / "sessions" / "--claimed--" / "one.jsonl", [pi_reply("google", 40, id="p", model="antigravity-pi-gemini")]) +record = local() +expect("Local leaves out exactly the rows another collector claims", + [record["todayTotalTokens"], models(record)], [100, ["antigravity-pi-gemini", "gemini-oc", "gemini-vertex", "minimax-gateway"]], record) + +# Claude Code pointed at a self-hosted endpoint (Ollama's guide sets +# localhost:11434) ran its models here, but a cloud-tagged one is Ollama's +# and a Claude one went on to Anthropic; with no endpoint, nothing is local. +sandbox() +claude_code_session("box", [("qwen3-coder", 30), ("glm-4.6:cloud", 300), ("claude-sonnet-4-5", 40)]) +record = local() +expect("Local counts no Claude Code reply while Claude Code is sent to Anthropic", [record["todayTokensByModel"], record["ready"]], [{}, False], record) +claude_code_session("box", [("qwen3-coder", 30), ("glm-4.6:cloud", 300), ("claude-sonnet-4-5", 40)], base_url="http://localhost:11434") +record = local() +expect("Local counts the Claude Code replies a self-hosted endpoint ran itself", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"qwen3-coder": 30}, 1, 1], record) + +# A MiniMax Code custom provider on this machine or the LAN ran its models +# here; one at a gateway nothing can place is nobody's. +sandbox() +mcode_routed_session("box", "custom_provider:box", "qwen-box", 1387, base_url="http://127.0.0.1:18082/anthropic") +mcode_routed_session("lan", "custom_provider:lan", "qwen-lan", 30, base_url="http://gpu-box.lan:8080/v1") +mcode_routed_session("far", "custom_provider:far", "qwen-far", 5, base_url="https://gateway.example.com/v1") +record = local() +expect("Local counts MiniMax Code replies on a self-hosted custom provider, and none elsewhere", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"qwen-box": 1387, "qwen-lan": 30}, 2, 2], record) + +# ---------------------------------------------------------------- contract + +use_home() + + +def add_session(): + add_session.count = getattr(add_session, "count", 0) + 1 + write_pi_session(home() / ".pi" / "agent" / "sessions" / f"grown-{add_session.count}.jsonl", + [pi_reply("ollama", 10, id=f"grown-{add_session.count}", model="qwen3:8b")]) + + +check_scan_cache(Adapter("local", grow=add_session, pi_root=lambda: home() / ".pi" / "agent" / "sessions")) + +# A harness's route config can move its rows onto this machine without its +# store changing, so the cached scan is kept only while every harness still +# classifies the same way, and the match has to survive the cache's JSON. +import collectors as shared # noqa: E402 + +sandbox() +local = load_collector("local") +before = local.harness_match() +(home() / ".qwen").mkdir(parents=True, exist_ok=True) +(home() / ".qwen" / ".env").write_text("OPENAI_BASE_URL=http://127.0.0.1:8080/v1\n") +after = local.harness_match() +expect("a harness route change changes local's scan match, which survives JSON", + [before != after, json.loads(json.dumps(after)) == after], [True, True], after) +real_routes = shared.mcode_routes +shared.mcode_routes = lambda data_dir=None: {"mine": "local", "gateway": "openrouter"} +try: + routes = shared.mcode_scan_match() +finally: + shared.mcode_routes = real_routes +expect("MiniMax Code's route match survives the cache's JSON round trip", json.loads(json.dumps(routes)), routes, routes) +PY diff --git a/test/shell.d/agent-usage-minimax-scanner-test.sh b/test/shell.d/agent-usage-minimax-scanner-test.sh new file mode 100644 index 00000000000..35da76e07c8 --- /dev/null +++ b/test/shell.d/agent-usage-minimax-scanner-test.sh @@ -0,0 +1,206 @@ +#!/bin/bash + +# The MiniMax collector's local stats: MiniMax Code's own sessions (and where +# MINIMAX_DATA_DIR moves them), and pi, omp and opencode usage on MiniMax's +# provider ids. How pi and opencode stores are read (roots, forks, both +# opencode stores, OPENCODE_DB, opencode-v2.db, untrusted values) is held +# centrally by the shared and partition tests, the unused-machine record by +# the partition test, and the scan cache by the contract check. The limits +# are agent-usage-minimax-test.sh's. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + + + +def minimax(*args): + with FakeHttp(): + return collect("minimax", *(args or ("--force",))) + + +def mcode_session(root, name, lines): + path = Path(root) / "v2" / "sessions" / name / "messages.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join((line if isinstance(line, str) else json.dumps(line)) + "\n" for line in lines)) + return path + + +def mcode_reply(usage, timestamp=None, model="MiniMax-M2", provider=None): + message = {"role": "assistant", "timestamp": ms_in(0) if timestamp is None else timestamp, "model": model, "usage": usage} + if provider is not None: + message["provider"] = provider + return {"message": message} + + +def sessions(): + return home() / ".pi" / "agent" / "sessions" + + +# ---------------------------------------------------------- every source + +# MiniMax Code's own sessions: user messages are prompts, assistant usage is +# tokens, and a row with a garbage timestamp, token count or line drops on +# its own. Beside them, pi and omp sessions and opencode rows on any MiniMax provider +# id; other providers, and a gateway named like one, drop. +sandbox() +mcode_session(home() / ".minimax", "s1", [ + {"message": {"role": "user", "timestamp": ms_in(0), "content": "hi"}}, + mcode_reply({"input": 10, "output": 5, "cacheRead": 3, "cacheWrite": 2, "totalTokens": 20}), + mcode_reply({"input": 9999}, timestamp=True), + mcode_reply({"input": -4, "output": 1e300}), + "not json", + "[" * 100000, # nested too deep to parse: a RecursionError, not a ValueError +]) +mcode_session(home() / ".minimax", "s2", [ + {"message": {"role": "user", "timestamp": "2020-01-01T00:00:00Z"}}, + mcode_reply({"totalTokens": 7}, timestamp="2020-01-01T00:00:01Z"), +]) +stamp = iso_in(0) +parent = write_pi_session(sessions() / "p" / "a.jsonl", [ + pi_reply("minimax", id="m1", model="MiniMax-M2", usage={"input": 100, "output": 50}, timestamp=stamp), + pi_reply("openai", id="m2", model="gpt", usage={"input": 999}), +]) +write_pi_session(home() / ".omp" / "profiles" / "work" / "agent" / "sessions" / "p" / "b.jsonl", [ + pi_reply("minimax", id="m1", model="MiniMax-M2", usage={"input": 100, "output": 50}, timestamp=stamp), + pi_reply("minimax-cn", id="m3", model="MiniMax-M2", usage={"input": 1000}), +], parent=parent) +opencode_db(tables=("message",), message=[ + ("o1", "ses1", oc_v1("minimax-coding-plan", {"input": 7, "output": 2, "reasoning": 1, "cache": {"read": 0, "write": 0}}, model="MiniMax-M2")), + ("o2", "ses1", oc_v1("anthropic", {"input": 999}, model="claude")), + ("o3", "ses1", oc_v1("minimax-proxy", {"input": 999}, model="MiniMax-M2")), +]) + +# Today: MiniMax Code 20, pi 150 (the forked copy counts once) + 1000, +# opencode 10. Prompts: one MiniMax Code message today and one before, two +# pi turns, one opencode turn. +record = minimax() +expect("MiniMax counts MiniMax Code, pi/omp and opencode usage, prompts and sessions per source, skipping garbage rows", { + key: record[key] for key in ("todayTotalTokens", "todayPrompts", "totalPrompts", "totalSessions", "activeDays") +} | {"model": record["modelUsage"].get("MiniMax-M2")}, { + "todayTotalTokens": 1180, "todayPrompts": 4, "totalPrompts": 5, "totalSessions": 5, "activeDays": 2, + "model": {"inputTokens": 1124, "outputTokens": 58, "cacheReadInputTokens": 3, "cacheCreationInputTokens": 2}, +}, record) + +# In use here without a key: local stats show and the limits wait for one. +expect("MiniMax in use without a key waits for auth", + {key: record[key] for key in ("ready", "hasLocalStats", "hasPromptStats", "usageStatusText", "limits", "tierLabel")}, + {"ready": True, "hasLocalStats": True, "hasPromptStats": True, "usageStatusText": "Waiting for auth", "limits": [], "tierLabel": ""}, record) + +# MINIMAX_DATA_DIR moves MiniMax Code's data, ahead of MAVIS_DATA_DIR. +mcode_session(home() / "moved", "s", [mcode_reply({"input": 5})]) +os.environ.update(MINIMAX_DATA_DIR=str(home() / "moved"), MAVIS_DATA_DIR=str(home() / ".minimax")) +record = minimax() +os.environ.pop("MINIMAX_DATA_DIR") +os.environ.pop("MAVIS_DATA_DIR") +expect("MiniMax reads MiniMax Code's sessions from MINIMAX_DATA_DIR, ahead of MAVIS_DATA_DIR", record["todayTotalTokens"], 1165, record) + +# A scan that couldn't read one of MiniMax Code's files isn't kept for the +# next run. +if os.geteuid() != 0: + for path in scan_cache_files("minimax"): + path.unlink() + locked = home() / ".minimax" / "v2" / "sessions" / "s2" / "messages.jsonl" + locked.chmod(0) + try: + minimax() + finally: + locked.chmod(0o644) + record = minimax("--limits-only") + expect("MiniMax does not cache a scan that missed a MiniMax Code file", record["totalPrompts"], 5, record) + + # Nor one that couldn't list a session directory, whose files it never saw. + for path in scan_cache_files("minimax"): + path.unlink() + locked.parent.chmod(0) + try: + minimax() + finally: + locked.parent.chmod(0o755) + record = minimax("--limits-only") + expect("MiniMax does not cache a scan that couldn't list a MiniMax Code session directory", record["totalPrompts"], 5, record) +else: + print("ok - MiniMax does not cache a scan that missed a MiniMax Code file # SKIP root reads anything") + +# MiniMax Code 0.6 names each reply's route, as `minimax` (its Token Plan), +# `minimax_api`, or `custom_provider:` for a custom endpoint or a Z.AI or +# OpenRouter preset. Only MiniMax's own routes count, and a prompt goes with +# the route of the reply that answered it. +sandbox() +mcode_session(home() / ".minimax", "routed", [ + {"message": {"role": "user", "timestamp": ms_in(0), "content": [{"type": "text", "text": "one"}]}}, + mcode_reply({"input": 1234, "output": 51, "cacheRead": 71, "cacheWrite": 31, "totalTokens": 1387}, provider="custom_provider:mockanth"), + {"message": {"role": "user", "timestamp": ms_in(0), "content": [{"type": "text", "text": "two"}]}}, + mcode_reply({"input": 2234, "output": 52, "cacheRead": 72, "cacheWrite": 0, "totalTokens": 2358}, model="glm-5", provider="custom_provider:zai-coding-plan"), + {"message": {"role": "user", "timestamp": ms_in(0), "content": [{"type": "text", "text": "three"}]}}, + mcode_reply({"input": 10, "output": 5, "totalTokens": 15}, provider="minimax"), + {"message": {"role": "user", "timestamp": ms_in(0), "content": [{"type": "text", "text": "four"}]}}, + mcode_reply({"input": 20, "output": 5, "totalTokens": 25}, provider="minimax_api"), +]) +mcode_session(home() / ".minimax", "elsewhere", [ + {"message": {"role": "user", "timestamp": ms_in(0), "content": "unanswered, after a custom route"}}, + mcode_reply({"input": 999, "output": 1}, model="glm-mock", provider="custom_provider:mockoai"), + {"message": {"role": "user", "timestamp": ms_in(0), "content": "still waiting"}}, +]) +record = minimax() +expect("MiniMax counts only MiniMax Code replies on MiniMax's own routes, and the prompts they answered", + {key: record[key] for key in ("todayTotalTokens", "todayPrompts", "todaySessions", "todayTokensByModel")}, + {"todayTotalTokens": 40, "todayPrompts": 2, "todaySessions": 1, "todayTokensByModel": {"MiniMax-M2": 40}}, record) + +# Claude Code pointed at MiniMax's Anthropic API is MiniMax's; with no +# endpoint, only a MiniMax model's reply is. +sandbox() +claude_code_session("mm", [("MiniMax-M2", 300), ("claude-sonnet-4-5", 40)]) +record = minimax() +expect("MiniMax counts Claude Code's MiniMax replies, and no Claude one", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"MiniMax-M2": 300}, 1, 1], record) +claude_code_session("mm", [("MiniMax-M2", 300), ("claude-sonnet-4-5", 40)], base_url="https://api.minimax.io/anthropic") +record = minimax() +expect("MiniMax counts every Claude Code reply while Claude Code is pointed at MiniMax", + [record["todayTokensByModel"], record["todayPrompts"]], [{"MiniMax-M2": 300, "claude-sonnet-4-5": 40}, 2], record) + +# Nor is a scan that couldn't read one of Claude Code's transcripts kept for +# the next run. +if os.geteuid() != 0: + sandbox() + claude_code_session("mm", [("MiniMax-M2", 300)]) + locked = claude_code_session("late", [("MiniMax-M2", 40)]) + locked.chmod(0) + try: + minimax() + finally: + locked.chmod(0o644) + record = minimax("--limits-only") + expect("MiniMax does not cache a scan that missed a Claude Code transcript", record["todayTokensByModel"], {"MiniMax-M2": 340}, record) +else: + print("ok - MiniMax does not cache a scan that missed a Claude Code transcript # SKIP root reads anything") + +# A MiniMax Code custom provider is billed by whoever runs its baseURL in +# config.yaml: one at MiniMax's own host is MiniMax's, one at OpenRouter's +# isn't, and neither is one gone from the config. +sandbox() +mcode_routed_session("own", "custom_provider:mm", "MiniMax-M2", 300, base_url="https://api.minimax.io/anthropic") +mcode_routed_session("router", "custom_provider:router", "MiniMax-M2-or", 40, base_url="https://openrouter.ai/api/v1") +mcode_routed_session("gone", "custom_provider:gone", "MiniMax-M2-gone", 5) +record = minimax() +expect("MiniMax counts a MiniMax Code custom provider at MiniMax's host, and no other", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"MiniMax-M2": 300}, 1, 1], record) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_reply(): + add_reply.count = getattr(add_reply, "count", 0) + 1 + mcode_session(home() / ".minimax", f"grown-{add_reply.count}", [mcode_reply({"input": 5})]) + + +check_scan_cache(Adapter("minimax", grow=add_reply, pi_root=sessions)) +PY 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..aef22253d22 --- /dev/null +++ b/test/shell.d/agent-usage-minimax-test.sh @@ -0,0 +1,203 @@ +#!/bin/bash + +# The MiniMax collector's limits: the Token Plan quota answer and how its +# remaining percentages become used ones, the sign-ins mmx, opencode and pi +# keep, refusals MiniMax names in its answer, and its adapter for the +# contract checks. The unused-machine record is held centrally by the +# partition test, the limits cache's mechanics by the shared test, and the +# scan cache by the scanner test. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +REMAINS = "api.minimax.io/v1/token_plan/remains" + + +def quota(end_time, weekly_end_time): + return { + "model_remains": [ + {"model_name": "general", "end_time": end_time, "current_interval_remaining_percent": 10, "current_interval_status": 1, + "weekly_end_time": weekly_end_time, "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"}, + } + + +# Windows still open, so a kept copy outlives a failed check. +ANSWER = quota(ms_in(3600), ms_in(86400)) + + +def keyed(key="mmx_test"): + os.environ["MINIMAX_API_KEY"] = key + + +def minimax(answer=None, mode="up"): + with FakeHttp({REMAINS: answer}, mode=mode) as http: + record = collect("minimax", "--force") + return record, http + + +# ---------------------------------------------------------------- contract + +sandbox() +fresh = check_limits_contract(Adapter("minimax", plant=keyed, answer={REMAINS: ANSWER})) +expect("MiniMax's Token Plan and windows come from its quota answer", + [fresh["tierLabel"], [(entry["label"], round(entry["percent"], 2)) for entry in fresh["limits"]]], + ["Token Plan", [("Session (5-hour)", 0.9), ("Weekly (7-day)", 0.31)]], fresh) + +# An answer shaped like nothing expected is still a failed check: the kept +# limits stand. +record, _ = minimax(KeyError("surprise")) +check("MiniMax collector keeps the last limits past an unexpected answer", + record["usageStatusText"] == "" and record["limitsStale"] is True and record["limits"] == fresh["limits"] + and "KeyError" in record["authHelpText"], record) + +# A bad key MiniMax names in base_resp rather than with a status code. +record, _ = minimax({"base_resp": {"status_code": 1004, "status_msg": "login fail"}}) +expect("MiniMax collector reads an authentication failure in base_resp as a refused key", + [record["usageStatusText"], record["limitsStale"], bool(record.get("retryAdvised"))], ["Sign-in expired", True, False], record) + +# Another key never shows this one's numbers, and a failure with nothing kept +# says so; only one that reached no server asks for a retry. +keyed("other") +record, _ = minimax(mode=500) +expect("MiniMax collector keys its limits cache by credential and reports a failure with nothing kept", + [record["usageStatusText"], record["limits"], record["tierLabel"], bool(record.get("retryAdvised"))], + ["MiniMax limits unavailable", [], "", False], record) +record, _ = minimax(mode="down") +expect("MiniMax collector advises a retry only when the API can't be reached", + [record["usageStatusText"], record.get("retryAdvised")], ["MiniMax limits unavailable", True], record) + +# The same key in the other region is another account: its cache is kept per +# host, so the China host never shows the global allowance. +keyed() +os.environ["MINIMAX_REGION"] = "cn" +record, http = minimax(mode="down") +del os.environ["MINIMAX_REGION"] +expect("MiniMax collector keeps one key's limits per region", + [http.requests[0]["url"].startswith("https://api.minimaxi.com/"), record["limits"], record["usageStatusText"]], + [True, [], "MiniMax limits unavailable"], record) +record, _ = minimax(mode="down") +expect("MiniMax collector still keeps the global host's limits for that key", [record["limits"], record["limitsStale"]], [fresh["limits"], True], record) + +# ------------------------------------------------------------- the answer + +sandbox() +keyed() +record, http = minimax(quota(1782399600000, 1782691200000)) +expect("MiniMax remaining percentages become used fractions, with their resets normalized", + [(entry["label"], round(entry["percent"], 2), entry["resetsAt"]) for entry in record["limits"]], + [("Session (5-hour)", 0.9, "2026-06-25T15:00:00+00:00"), ("Weekly (7-day)", 0.31, "2026-06-29T00:00:00+00:00")], record) +expect("MiniMax collector asks with the key as a bearer", http.requests[0]["headers"].get("authorization"), "Bearer mmx_test") + +# Without a "general" row the first row that isn't a media model's is the plan. +record, _ = minimax({"model_remains": [ + {"model_name": "video", "current_interval_remaining_percent": 100}, + {"model_name": "Image-01", "current_interval_remaining_percent": 100}, + {"model_name": "MiniMax-M2", "current_interval_remaining_percent": 25}, +]}) +expect("MiniMax collector skips media quota rows when choosing the plan's row", [entry["percent"] for entry in record["limits"]], [0.75], record) + +# Each window's status outranks its percentage, as in mmx and omp: 2 is +# exhausted however much the stale percentage says remains, 3 is unlimited, +# with nothing to fill. A row in mmx's "not in plan" shape (both windows +# unlimited, nothing allotted) is never the plan's row. +record, _ = minimax({"model_remains": [ + {"model_name": "general", "current_interval_total_count": 0, "current_interval_remaining_percent": 100, "current_interval_status": 3, + "current_weekly_total_count": 0, "current_weekly_remaining_percent": 100, "current_weekly_status": 3}, + {"model_name": "MiniMax-M2", "current_interval_total_count": 1500, "current_interval_remaining_percent": 40, "current_interval_status": 2, + "current_weekly_total_count": 0, "current_weekly_remaining_percent": 100, "current_weekly_status": 3}, +], "base_resp": {"status_code": 0, "status_msg": "success"}}) +expect("MiniMax reads an exhausted window as spent, drops an unlimited one, and skips a model outside the plan", + [(entry["label"], entry["percent"]) for entry in record["limits"]], [("Session (5-hour)", 1.0)], record) + +# The bearer never follows a redirect to another host. +elsewhere, stolen = stub_server(lambda request: {"model_remains": []}) +origin, _ = stub_server(lambda request: (302, b"", {"Location": elsewhere + "/steal"})) +try: + load_collector("minimax").fetch_quota("secret", origin) + error = "" +except Exception as failure: + error = str(failure) +check("MiniMax collector refuses to follow a redirect with the bearer", stolen == [] and "redirected" in error, error) + +# ------------------------------------------------------------ credentials + +sandbox() +os.environ["MMX_CONFIG_DIR"] = str(home() / ".mmx") + + +def credentials(): + return list(load_collector("minimax").credentials()) + + +write_json(home() / ".mmx" / "config.json", {"api_key": "mmx_saved"}) +expect("MiniMax credentials use the key mmx saved", credentials(), ["mmx_saved", "https://api.minimax.io"]) + +write_json(home() / ".mmx" / "config.json", {"oauth": {"access_token": "mmx_oauth", "refresh_token": "refresh", "expires_at": "2099-01-01T00:00:00Z", "region": "cn"}}) +expect("MiniMax credentials use the OAuth session and region saved by mmx", credentials(), ["mmx_oauth", "https://api.minimaxi.com"]) + +# mmx 1.0's own order: base_url, then the OAuth resource_url, then the region, +# with MINIMAX_REGION over the saved one. +write_json(home() / ".mmx" / "config.json", {"base_url": "https://mmx.example", "oauth": {"access_token": "mmx_oauth", "resource_url": "https://resource.example"}}) +expect("MiniMax asks mmx's base_url ahead of its OAuth resource_url", credentials(), ["mmx_oauth", "https://mmx.example"]) +write_json(home() / ".mmx" / "config.json", {"api_key": "mmx_saved", "region": "cn"}) +os.environ["MINIMAX_REGION"] = "global" +expect("MiniMax lets MINIMAX_REGION move mmx's key off the region mmx saved", credentials(), ["mmx_saved", "https://api.minimax.io"]) +del os.environ["MINIMAX_REGION"] + +(home() / ".mmx" / "config.json").unlink() +opencode = opencode_auth(None, "minimax", "mmx_opencode") +expect("MiniMax credentials use opencode's login", credentials()[0], "mmx_opencode") + +# A China-region key pi saved answers on the China host. +opencode.unlink() +pi_auth(None, "minimax-cn", "mmx_pi_cn") +expect("MiniMax credentials use a China-region key pi saved, on the China host", credentials(), ["mmx_pi_cn", "https://api.minimaxi.com"]) + +# A key answers only on its own region's host, so mmx's region applies to +# mmx's key alone: a key from the environment, opencode or pi signs in where +# its own provider id says. +write_json(home() / ".mmx" / "config.json", {"region": "cn"}) +opencode_auth(None, "minimax", "mmx_opencode_global") +expect("MiniMax asks with a global key opencode saved on the global host, whatever region mmx names", + credentials(), ["mmx_opencode_global", "https://api.minimax.io"]) +write_json(home() / ".mmx" / "config.json", {"api_key": "mmx_saved", "region": "cn"}) +os.environ["MINIMAX_API_KEY"] = "mmx_env" +expect("MiniMax asks with an exported key on the global host, whatever region mmx names", credentials(), ["mmx_env", "https://api.minimax.io"]) +os.environ["MINIMAX_BASE_URL"] = "https://minimax.example" +expect("MiniMax asks MINIMAX_BASE_URL when it is set", credentials(), ["mmx_env", "https://minimax.example"]) +del os.environ["MINIMAX_API_KEY"], os.environ["MINIMAX_BASE_URL"] +expect("MiniMax asks with mmx's own key on the host mmx names", credentials(), ["mmx_saved", "https://api.minimaxi.com"]) + +# Wherever the host comes from, the key never goes over plain HTTP to another +# machine, nor to a URL that isn't one. +for name, setup in ( + ("MINIMAX_BASE_URL", lambda: os.environ.update(MINIMAX_API_KEY="mmx_env", MINIMAX_BASE_URL="http://evil.example")), + ("mmx's base_url", lambda: write_json(home() / ".mmx" / "config.json", {"api_key": "mmx_saved", "base_url": "http://evil.example"})), + ("mmx's resource_url", lambda: write_json(home() / ".mmx" / "config.json", {"oauth": {"access_token": "mmx_oauth", "resource_url": "not a url"}})), +): + setup() + record, http = minimax(ANSWER) + expect(f"MiniMax sends its key to no untrusted endpoint from {name}", + [http.count, record["usageStatusText"], "HTTPS" in record["authHelpText"]], [0, "MiniMax limits unavailable", True], record) + os.environ.pop("MINIMAX_API_KEY", None) + os.environ.pop("MINIMAX_BASE_URL", None) + +# A key saved under a provider id MiniMax doesn't bill belongs to a gateway +# in front of it, and is never sent to MiniMax as a bearer. +sandbox() +os.environ["MMX_CONFIG_DIR"] = str(home() / ".mmx") +opencode_auth(None, "minimax-coding", "gateway_key") +pi_auth(None, "minimax-token-plan", "gateway_key") +pi_auth(home() / ".omp" / "agent", "minimax-proxy", "gateway_key") +expect("MiniMax credentials take keys only under the provider ids MiniMax bills", credentials()[0], "") +PY diff --git a/test/shell.d/agent-usage-muse-dedup-test.sh b/test/shell.d/agent-usage-muse-dedup-test.sh new file mode 100755 index 00000000000..4a91e03cd21 --- /dev/null +++ b/test/shell.d/agent-usage-muse-dedup-test.sh @@ -0,0 +1,81 @@ +#!/bin/bash + +# Muse's native sessions repeat a completion under its record id: in retained +# frames, in copied sessions, and in the .msp-view directories that hold views +# of the same history. Each completion counts once, while records without a +# usable id count one by one. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +import shutil + +from agent_usage_fixtures import * + +sandbox() +# Only view directories inside the sessions tree are left out, so a data +# directory whose own parent looks like one is still read. +os.environ["MUSE_DATA_DIR"] = str(home() / ".msp-view-parent" / "muse") +root = Path(os.environ["MUSE_DATA_DIR"]) / "sessions" +VIEWS = (".msp-view", ".msp-view.saved", "2026/09/08/session/.msp-view-backup") + + +def event(event_id, tokens): + return { + "id": event_id, + "recorded_at": int(time.time() * 1_000_000), + "payload": {"kind": "run", "event": {"kind": "model_completed", "model": "muse-test", "usage": {"input_tokens": tokens}}}, + } + + +def write(relative, entries): + path = root / relative / "session.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(entry) + "\n" for entry in entries)) + + +first = event("first", 110) +child = event("child", 55) +# A record whose usage is malformed claims no id, so a later valid record +# with the same id still counts. +rejected = event("later-valid", 0) +rejected["payload"]["event"]["usage"] = [] +write("2026/09/08/session", [ + {"id": "first", "payload": {"kind": "metadata", "record": {"model_id": "muse-test"}}}, + first, first, + {"children": [{"record_json": json.dumps(first)}, {"record_json": json.dumps(child)}]}, + event(None, 7), event(None, 7), event([], 3), + rejected, event("later-valid", 5), +]) +write("2026/09/08/session/subagent/child", [child]) +write("2026/09/08/copied-session", [first]) +for view in VIEWS: + write(view, [event(None, 9000)]) + +# 110 + 55 + 7 + 7 + 3 + 5. +record = collect("muse", "--force") +expect("Muse counts unique completions across plain records, retained frames, and copied sessions", + [record["todayTotalTokens"], record["totalPrompts"], record["modelUsage"]["muse-test"]["inputTokens"]], [187, 6, 187], record) + +for view in VIEWS: + shutil.rmtree(root / view) +record = collect("muse", "--force") +expect("Muse leaves out view copies without dropping records with a missing or malformed id", + [record["todayTotalTokens"], record["totalPrompts"]], [187, 6], record) + +# A scan cached before completions were deduplicated is made again, even for +# --limits-only. +def predates_dedup(cached): + cached["museSchema"] = cached.get("museSchema", 0) - 1 + cached["stats"]["todayTotalTokens"] = 999999 + + +edit_scan_cache("muse", predates_dedup) +record = collect("muse", "--limits-only") +expect("Muse rescans a scan cached before completions were deduplicated", record["todayTotalTokens"], 187, record) +PY diff --git a/test/shell.d/agent-usage-muse-limits-test.sh b/test/shell.d/agent-usage-muse-limits-test.sh new file mode 100755 index 00000000000..1d4dccb5ac5 --- /dev/null +++ b/test/shell.d/agent-usage-muse-limits-test.sh @@ -0,0 +1,173 @@ +#!/bin/bash + +# The Muse collector's limits: the login the Muse CLI keeps, Muse's key +# endpoint and the subscription windows, plan and pay-as-you-go answers it +# gives. The unused-machine record is held centrally by the partition test, +# the limits cache by the contract checks, run here against a real stand-in +# for the key endpoint, and Muse's local stats by its scanner test. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +SUBSCRIBER = {"body": {"subs_tier_name": "Power Usage", "subs_usage": { + "window": {"used_percent": 8, "window_duration_mins": 300, "resets_at": 4102444800}, + "weekly": {"used_percent": 15, "resets_at": 4103049600}}}} +PAYG = {"body": {"api_key": "minted", "is_subs_active": False, "subs_tier_name": None}} +state = {"spec": SUBSCRIBER} + + +# A stand-in for Muse's key endpoint, answering the POST from state["spec"]: +# a body, or a status, or raw bytes. +def key_endpoint(request): + spec = state["spec"] + if spec.get("status", 200) != 200: + return spec["status"], b"" + return spec.get("raw") or spec.get("body", {}) + + +ENDPOINT, reqs = stub_server(key_endpoint) +ENDPOINT += "/key" + + +@contextlib.contextmanager +def serve(mode): + os.environ["MUSE_KEY_ENDPOINT"] = ENDPOINT + if mode in ("down", "timeout", "cut", "garbage"): + with FakeHttp(mode=mode): + yield + return + state["spec"] = SUBSCRIBER if mode == "up" else {"status": mode} + yield + + +def muse(spec=None, *args, fresh_cache=True): + """A run against the key endpoint answering `spec`, from no cache unless + fresh_cache is False.""" + if fresh_cache: + for path in limits_cache_files("muse"): + path.unlink() + os.environ["MUSE_KEY_ENDPOINT"] = ENDPOINT + if spec is not None: + state["spec"] = spec + return collect("muse", *(args or ("--force",))) + + +def auth_path(): + return config_home() / "muse" / "auth.json" + + +def signed_in(token="test-token"): + write_json(auth_path(), {"schema_version": 1, "providers": {"meta": {"access_token": token, "api_base_url": "https://api.meta.ai/v1"}}}) + + +def auth_state(): + return (auth_path().stat().st_mtime, auth_path().read_bytes(), sorted(path.name for path in auth_path().parent.iterdir())) + + +# ---------------------------------------------------------------- contract + +sandbox() +signed_in() +os.utime(auth_path(), (1767225600, 1767225600)) +auth_before = auth_state() + +fresh = check_limits_contract(Adapter("muse", serve=serve, count=lambda: len(reqs))) +expect("Muse reads the session and weekly windows, their resets, and the plan from the key endpoint", + [fresh["tierLabel"], [[entry["label"], entry["percent"], entry["resetsAt"]] for entry in fresh["limits"]]], + ["Power Usage", [["Session (5-hour)", 0.08, "2100-01-01T00:00:00+00:00"], ["Weekly (7-day)", 0.15, "2100-01-08T00:00:00+00:00"]]], fresh) +expect("Muse authenticates the probe with the saved token", reqs[-1]["headers"].get("Authorization"), "Bearer test-token") + +leaks = [str(path) for path in home().rglob("*") if path.is_file() and path != auth_path() and b"test-token" in path.read_bytes()] +expect("Muse keeps the token out of every cache and file of its own", leaks, []) + +# The plan is kept with the windows, so a reused or kept answer still names it. +with serve(500): + reused = collect("muse") + kept = collect("muse", "--force") +expect("a reused and a kept Muse answer keep the plan", [reused["tierLabel"], kept["tierLabel"], kept["limitsStale"]], ["Power Usage", "Power Usage", True], kept) + +# The login file belongs to the Muse CLI: every run so far only read it. +expect("Muse never writes the login file", auth_state(), auth_before) + +# MUSE_KEY_ENDPOINT moves the endpoint, but the token only goes over HTTPS or +# plain HTTP to loopback. +before = len(reqs) +for path in limits_cache_files("muse"): + path.unlink() +with FakeHttp({"/key": SUBSCRIBER["body"]}) as http: + os.environ["MUSE_KEY_ENDPOINT"] = "http://muse.example.com/key" + plain = collect("muse", "--force") +os.environ["MUSE_KEY_ENDPOINT"] = ENDPOINT +expect("the Muse token never goes to a plain HTTP endpoint off loopback", + [http.requests, len(reqs) - before, "HTTPS" in plain["authHelpText"]], [[], 0, True], plain) + +# ------------------------------------------------------- answers without windows + +# Pay-as-you-go: the key mints, but there is no subscription window to meter +# and no plan name to invent. Local metering only: not an error, not expired. +payg = muse(PAYG) +expect("Muse stays silent for a pay-as-you-go account", + [payg["limits"], payg["tierLabel"], payg["usageStatusText"], payg["authHelpText"], "retryAdvised" in payg], [[], "", "", "", False], payg) + +# An active plan without usage windows (the live Muse Code shape) keeps its +# name and shows local metering only. +plan = muse({"body": {"api_key": "minted", "is_subs_active": True, "subs_tier_name": "Muse Code Everyday Usage"}}) +expect("Muse stays silent for an active plan without usage windows", + [plan["limits"], plan["tierLabel"], plan["usageStatusText"], plan["authHelpText"], "retryAdvised" in plan], + [[], "Muse Code Everyday Usage", "", "", False], plan) +expect("a Muse sign-in shows before its first session, pay-as-you-go included", [payg["ready"], plan["ready"]], [True, True]) +expect("a Muse sign-in with no sessions here has no local stats", [payg["hasLocalStats"], plan["hasLocalStats"]], [False, False]) + +# A plan that moved to pay-as-you-go leaves its old windows behind: a failed +# check after that brings none of them back. The answer is reused like any +# other, so a panel flick doesn't mint a key again. +muse(SUBSCRIBER) +muse(PAYG, fresh_cache=False) +asked = len(reqs) +reused = collect("muse") +expect("a recent pay-as-you-go answer is reused", [len(reqs) - asked, reused["limits"]], [0, []], reused) +switched = muse({"status": 500}, fresh_cache=False) +expect("a former plan's windows stay out after it moved to pay-as-you-go", switched["limits"], [], switched) + +# Neither active nor pay-as-you-go and no windows: a warning, never invented +# meters. +unknown = muse({"body": {"api_key": "minted", "subs_tier_name": "Mystery"}}) +expect("Muse reports unavailable limits for an answer of unknown shape", unknown["usageStatusText"], "Muse limits unavailable", unknown) + +hours = muse({"body": {"subs_tier_name": "T", "subs_usage": {"window": {"used_percent": 5, "window_duration_mins": 120, "resets_at": 4102444800}, + "weekly": {"used_percent": 1, "resets_at": 4103049600}}}}) +expect("Muse labels a whole-hour session window in hours", hours["limits"][0]["label"], "Session (2-hour)", hours) + +# omp's schema for these windows takes resets_at as a string or a number, and +# reads a string as an ISO time. +iso = muse({"body": {"subs_tier_name": "T", "subs_usage": {"window": {"used_percent": 5, "resets_at": "2100-01-01T00:00:00Z"}, + "weekly": {"used_percent": 1, "resets_at": "2100-01-08T00:00:00+00:00"}}}}) +expect("Muse reads a reset given as an ISO time", [entry["resetsAt"] for entry in iso["limits"]], + ["2100-01-01T00:00:00+00:00", "2100-01-08T00:00:00+00:00"], iso) + +# ---------------------------------------------------------------- the login + +# A malformed login file is a missing credential, not a crash: the Muse CLI is +# set up here, so it waits for a sign-in. +for malformed in ("[1,2]", '"nope"', '{"providers":[]}', '{"providers":{"meta":7}}'): + auth_path().write_text(malformed) + garbled = muse(SUBSCRIBER) + expect(f"Muse waits for a sign-in past a malformed login file ({malformed})", + [garbled["limits"], garbled["usageStatusText"], garbled["ready"]], [[], "Waiting for auth", True], garbled) + +# MUSE_AUTH_PATH wins over a different login at the default place. +signed_in("default-token") +custom = home() / "custom auth.json" +write_json(custom, {"providers": {"meta": {"access_token": "custom-token"}}}) +os.environ["MUSE_AUTH_PATH"] = str(custom) +muse(SUBSCRIBER) +del os.environ["MUSE_AUTH_PATH"] +expect("Muse honors MUSE_AUTH_PATH over the default login", reqs[-1]["headers"].get("Authorization"), "Bearer custom-token") +PY diff --git a/test/shell.d/agent-usage-muse-scanner-test.sh b/test/shell.d/agent-usage-muse-scanner-test.sh new file mode 100755 index 00000000000..dc2b360d33d --- /dev/null +++ b/test/shell.d/agent-usage-muse-scanner-test.sh @@ -0,0 +1,389 @@ +#!/bin/bash + +# The Muse collector's local metering: native `muse` CLI sessions and the +# CLI's session index, plus its own filter over opencode and pi sessions (the +# meta provider, never Muse Spark through OpenCode Zen). How the shared pi and +# opencode readers work, and the unused-machine record, are held centrally by +# the shared and partition tests, and the scan cache by the contract check. +# Muse's limits have a test of their own. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox +require_command rg + +python3 - <<'PY' +import sqlite3 +from datetime import date, timedelta + +from agent_usage_fixtures import * + +NOW_US = int(time.time()) * 1_000_000 +NOW_MS = NOW_US // 1000 +OLD_US = NOW_US - 8 * 86400 * 1_000_000 +OLD_DATE = datetime.fromtimestamp(OLD_US / 1_000_000).astimezone().date().isoformat() + + +def muse_dir(): + return Path(os.environ["MUSE_DATA_DIR"]) + + +def new_muse_home(): + sandbox() + os.environ["MUSE_DATA_DIR"] = str(home() / "muse-data") + + +def completion(model=None, recorded_at=NOW_US, input=0, output=0, reasoning=0, cache_read=0, cache_write=0, **usage): + event = {"kind": "model_completed", "usage": {"input_tokens": input, "output_tokens": output, "reasoning_tokens": reasoning, + "cache_read_tokens": cache_read, "cache_write_tokens": cache_write, **usage}} + if model is not None: + event["model"] = model + line = {"payload": {"kind": "run", "event": event}} + if recorded_at is not None: + line["recorded_at"] = recorded_at + return line + + +def model_record(kind, model): + return {"recorded_at": NOW_US, "payload": {"kind": kind, "record": {"model_id": model}}} + + +def session(relative, *lines, root=None): + path = (root or muse_dir() / "sessions") / relative / "session.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join((line if isinstance(line, str) else json.dumps(line)) + "\n" for line in lines)) + return path + + +def oc_tokens(input, output=0, reasoning=0, read=0, write=0): + return {"input": input, "output": output, "reasoning": reasoning, "cache": {"read": read, "write": write}} + + +def muse(*args): + return collect("muse", *(args or ("--force",))) + + +def pick(record, *keys): + return {key: record[key] for key in keys} + + +# ---------------------------------------------------------- edge times + +# A zoneless time at the edge of the calendar, in a zone off UTC, dates +# nothing: that row goes undated, and the record still prints. +new_muse_home() +saved_tz = os.environ.get("TZ") +os.environ["TZ"] = "Asia/Tokyo" +time.tzset() +session("sess-edge", completion("muse-spark-1.3-contributor", "0001-01-01T00:00:00", input=10), completion("muse-spark-1.3-contributor", "9999-12-31T23:59:59", input=20)) +try: + record = muse() +finally: + if saved_tz is None: + os.environ.pop("TZ") + else: + os.environ["TZ"] = saved_tz + time.tzset() +expect("Muse reads zoneless times at year 1 and 9999 off UTC as undated", [record["ready"], record["modelUsage"]["muse-spark-1.3-contributor"]["inputTokens"]], [True, 30], record) + +# ---------------------------------------------------------- native and opencode + +# One main session with a completion and a subagent under it, one session from +# 8 days ago, and opencode rows: Meta's own (counted), Muse Spark through Zen +# (the opencode record's), another provider, and a prompt. +new_muse_home() +session("sess-A", + {"recorded_at": NOW_US, "payload": {"kind": "metadata", "record": {"model_id": "muse-spark-1.3-contributor", "provider_id": "meta"}}}, + completion("muse-spark-1.3-contributor", input=1000, output=100, reasoning=20, cache_read=600)) +session("sess-A/subagent/sub-1", completion("muse-spark-1.3-contributor", input=500, output=50, cache_read=100)) +session("sess-B", completion("muse-spark-1.2-contributor", OLD_US, input=2000, output=200)) +opencode_db(message=[ + ("op-1", "op-1", oc_v1("meta", oc_tokens(300, 60, 10, 90, 7), model="muse-spark-1.2-contributor", created=NOW_MS)), + ("op-2", "op-2", oc_v1("opencode", oc_tokens(100, 20, read=30), model="muse-spark-1.3-contributor-free", created=NOW_MS)), + ("op-3", "op-3", oc_v1("openai", oc_tokens(999, 999), model="gpt-5.4", created=NOW_MS)), + ("op-4", "op-4", oc_v1("opencode", oc_tokens(0), model="muse-spark-1.3-contributor-free", created=NOW_MS, role="user")), +]) +record = muse() +# Native input includes cached tokens, so uncached input is split out: sess-A +# (1000-600) + 100 out + 600 cached = 1100; the subagent (500-100) + 50 + 100 +# = 550. opencode input is already cache-exclusive: 300 + 70 + 90 + 7 = 467. +expect("Muse sums native sessions and Meta's opencode rows, with metering only", pick(record, "ready", "hasPromptStats", "limits", "todayTotalTokens", "todayPrompts") | {"balance": "balance" in record}, + {"ready": True, "hasPromptStats": True, "limits": [], "todayTotalTokens": 2117, "todayPrompts": 3, "balance": False}, record) +expect("Muse rolls subagents into their parent session", + pick(record, "todaySessions", "totalSessions", "totalPrompts"), {"todaySessions": 2, "totalSessions": 3, "totalPrompts": 4}, record) +expect("Muse splits cached tokens out of native input without counting reasoning twice, and keeps opencode input as it is", { + "today13": record["todayTokensByModel"]["muse-spark-1.3-contributor"], + "m13": record["modelUsage"]["muse-spark-1.3-contributor"], + "m12input": record["modelUsage"]["muse-spark-1.2-contributor"]["inputTokens"], +}, { + "today13": 1650, + "m13": {"cacheCreationInputTokens": 0, "cacheReadInputTokens": 700, "inputTokens": 800, "outputTokens": 150}, + "m12input": 2300, +}, record) +check("Muse keeps active dates from before the recent week", OLD_DATE in record["activeDates"], record["activeDates"]) + +# A symlink to a real session would count it twice if followed, a dangling one +# must not break the scan, and a session reached through a symlinked directory +# must not escape the tree. +outside = session("evil", completion("evil-model", input=100, output=10), root=home() / "outside") +(muse_dir() / "sessions" / "sess-link").mkdir() +(muse_dir() / "sessions" / "sess-link" / "session.jsonl").symlink_to(muse_dir() / "sessions" / "sess-A" / "session.jsonl") +(muse_dir() / "sessions" / "sess-dangling").mkdir() +(muse_dir() / "sessions" / "sess-dangling" / "session.jsonl").symlink_to(muse_dir() / "sessions" / "nope" / "session.jsonl") +(muse_dir() / "sessions" / "sess-evildir").symlink_to(outside.parent) +record = muse() +expect("Muse skips symlinked session files and sessions outside the tree", + [record["todayTotalTokens"], "evil-model" in record["modelUsage"]], [2117, False], record) + +# A model id too long for the record fails its row closed, named on the +# completion or on the metadata before it. +big_model = "x" * 500 +session("sess-huge", completion(big_model, input=100, output=10)) +session("sess-huge-metadata", model_record("metadata", big_model), completion(input=100, output=10)) +record = muse() +expect("Muse skips completions under an oversized model name", [record["todayTotalTokens"], len(record["modelUsage"])], [2117, 2], record) + +# Many distinct models stay exact: every bucket is kept and the totals +# reconcile with the all-time total (today's 2817 and sess-B's 2200). +session("sess-many", *(completion("cap-model-%02d" % i, input=10) for i in range(70))) +record = muse() +expect("Muse keeps exact totals and every bucket with many models", + [record["todayTotalTokens"], len(record["modelUsage"]), sum(sum(bucket.values()) for bucket in record["modelUsage"].values())], + [2817, 72, 5017], record) + +# Negative counts can't take usage away. +session("sess-negative", completion("negative-model", input=-100, output=10, reasoning=-5, cache_read=-20, cache_write=-30)) +record = muse() +expect("Muse clamps negative token fields", record["modelUsage"]["negative-model"], + {"cacheCreationInputTokens": 0, "cacheReadInputTokens": 0, "inputTokens": 0, "outputTokens": 10}, record) + +# ---------------------------------------------------------------- sources + +# Without MUSE_DATA_DIR the native sessions are under XDG_DATA_HOME. +sandbox() +session("xdg-session", completion("muse-spark-1.3", input=10, output=2, reasoning=1, cache_read=4), root=data_home() / "muse" / "sessions") +record = muse() +expect("Muse follows XDG_DATA_HOME", record["todayTotalTokens"], 12, record) + +# A completion without a model takes the run_model's, a retained frame's +# record_json child reads like a plain event, an ISO stamp lands on today, and +# a missing recorded_at falls back on the file's time. +new_muse_home() +child = completion("model-via-child", input=20, output=2) +session("sess-events", + model_record("run_model", "model-via-run-model"), + completion(input=100, output=10), + completion("model-explicit", datetime.now().astimezone().isoformat(timespec="seconds"), input=50, output=5), + {"children": [{"record_json": json.dumps(child)}]}, + completion("model-via-mtime", None, input=7)) +record = muse() +expect("Muse counts every event shape in one session", pick(record, "todayTotalTokens", "todayPrompts", "todaySessions") | { + "runModel": record["modelUsage"]["model-via-run-model"]["inputTokens"] + record["modelUsage"]["model-via-run-model"]["outputTokens"], + "child": record["modelUsage"]["model-via-child"]["inputTokens"] + record["modelUsage"]["model-via-child"]["outputTokens"], + "today": sorted(record["todayTokensByModel"]), +}, { + "todayTotalTokens": 194, "todayPrompts": 4, "todaySessions": 1, "runModel": 110, "child": 22, + "today": ["model-explicit", "model-via-child", "model-via-mtime", "model-via-run-model"], +}, record) + +# Current logs carry a completion as record_type "event" with no payload +# kind, and count cached tokens inside input_tokens; a cache figure larger +# than its input never claims more than the input. +new_muse_home() + + +def current(input, output, cache_read=None, cache_write=None, offset=0): + usage = {"input_tokens": input, "output_tokens": output} + if cache_read is not None: + usage.update(cache_read_tokens=cache_read, cache_write_tokens=cache_write) + return {"schema_version": 1, "recorded_at": NOW_US + offset, "record_type": "event", "payload_type": "runtime.session", + "payload": {"event": {"kind": "model_completed", "model": "muse-wonder-1.0", "usage": usage}}} + + +session("2026/10/12/demo", + current(38219, 290, 0, 0), + current(42901, 357, 38132, 0, 1000), + current(10, 1, 50, 7, 2000), + {"payload": {"event": {"kind": "model_completed", "model": "muse", "usage": {"estimated_tokens": 5}}}}, + {"payload": {"event": {"kind": "tool_completed", "usage": {"input_tokens": 999}}}}) +record = muse() +expect("Muse reads current event logs and clamps cache tokens to the input", pick(record, "totalPrompts", "totalSessions", "todayPrompts", "todayTotalTokens") | {"m": record["modelUsage"]["muse-wonder-1.0"]}, + {"totalPrompts": 3, "totalSessions": 1, "todayPrompts": 3, "todayTotalTokens": 81778, + "m": {"cacheCreationInputTokens": 0, "cacheReadInputTokens": 38142, "inputTokens": 42988, "outputTokens": 648}}, record) + +# The CLI's session index holds the real prompt and session counts: a session +# nobody prompted isn't one, and an undated row counts on no day. +index = sqlite3.connect(muse_dir() / "session-index.db") +index.execute("CREATE TABLE sessions (session_id TEXT PRIMARY KEY, prompt_count INTEGER NOT NULL DEFAULT 0, created_at_us INTEGER)") +index.executemany("INSERT INTO sessions VALUES (?, ?, ?)", [("demo", 2, NOW_US), ("earlier", 5, OLD_US), ("opened", 0, NOW_US), ("undated", 1, "soon")]) +index.commit() +index.close() +record = muse() +expect("Muse takes prompt and session counts from the session index", + pick(record, "totalPrompts", "totalSessions", "todayPrompts", "todaySessions", "todayTotalTokens") | {"old": OLD_DATE in record["activeDates"]}, + {"totalPrompts": 8, "totalSessions": 3, "todayPrompts": 2, "todaySessions": 1, "todayTotalTokens": 81778, "old": True}, record) + +# A native session the index has no row for yet counts its completions as +# prompts, as it would with no index. +fresh = session("2026/10/12/fresh", current(10, 1)) +record = muse() +fresh.unlink() +expect("Muse counts a native session the index hasn't seen yet", + pick(record, "totalPrompts", "totalSessions", "todayPrompts", "todaySessions", "todayTotalTokens"), + {"totalPrompts": 9, "totalSessions": 4, "todayPrompts": 3, "todaySessions": 2, "todayTotalTokens": 81789}, record) + +# A session carried on from an earlier day counts today what its prompt count +# grew by since the last scan before today, not its whole count. +prompts_files = list(cache_root().glob("muse-prompts-*.json")) +check("Muse keeps the session prompt counts it saw", prompts_files, sorted(path.name for path in cache_root().iterdir())) +kept = json.loads(prompts_files[0].read_text()) +kept.update(day=(date.today() - timedelta(days=1)).isoformat(), last={"earlier": 3}) +prompts_files[0].write_text(json.dumps(kept)) +record = muse() +expect("Muse counts today the prompts a session from an earlier day grew by", + pick(record, "totalPrompts", "todayPrompts", "todaySessions"), {"totalPrompts": 8, "todayPrompts": 4, "todaySessions": 2}, record) +index = sqlite3.connect(muse_dir() / "session-index.db") +index.execute("UPDATE sessions SET prompt_count = 7 WHERE session_id = 'earlier'") +index.commit() +index.close() +record = muse() +expect("Muse keeps counting today from the count before today", + pick(record, "totalPrompts", "todayPrompts", "todaySessions"), {"totalPrompts": 10, "todayPrompts": 6, "todaySessions": 2}, record) + +# After a gap of days, what a count grew by since the last scan wasn't all +# typed today, so the counts start over from now, as on a first scan. +kept = json.loads(prompts_files[0].read_text()) +kept.update(day=(date.today() - timedelta(days=3)).isoformat(), last={"earlier": 3}) +prompts_files[0].write_text(json.dumps(kept)) +record = muse() +expect("Muse puts nothing a session grew by over a gap of days on today", + pick(record, "totalPrompts", "todayPrompts", "todaySessions"), {"totalPrompts": 10, "todayPrompts": 2, "todaySessions": 1}, record) + +# An index that's there but unreadable falls back on completions, and that +# scan isn't cached. +(muse_dir() / "session-index.db").write_text("not a database") +for path in scan_cache_files("muse"): + path.unlink() +record = collect("muse") +expect("Muse counts completions past an unreadable session index", pick(record, "totalPrompts", "totalSessions"), {"totalPrompts": 3, "totalSessions": 1}, record) +check("Muse doesn't cache a scan past an unreadable session index", not scan_cache_files("muse"), scan_cache_files("muse")) + +# A corrupt opencode database can't take the native sessions with it, and +# that incomplete scan is served but never cached. +new_muse_home() +session("sess-native", completion("native-model", input=30, output=3)) +(data_home() / "opencode").mkdir(parents=True) +(data_home() / "opencode" / "opencode.db").write_text("not a database") +record = muse() +again = collect("muse") +expect("Muse serves native stats past a corrupt opencode database, uncached", + [record["todayTotalTokens"], record["ready"], again["todayTotalTokens"], scan_cache_files("muse")], [33, True, 33, []], record) + +# An unreadable native session leaves the scan incomplete: the sessions that +# could be read are served, and that partial scan is never cached. +if os.geteuid() != 0: + new_muse_home() + session("sess-readable", completion("native-model", input=30, output=3)) + unreadable = session("sess-unreadable", completion("native-model", input=500, output=50)) + unreadable.chmod(0) + record = muse() + again = collect("muse") + unreadable.chmod(0o644) + expect("Muse serves readable native sessions past an unreadable one, uncached", + [record["todayTotalTokens"], again["todayTotalTokens"], scan_cache_files("muse")], [33, 33, []], record) + record = collect("muse") + expect("Muse counts the session once it can be read", record["todayTotalTokens"], 583, record) + +# A microsecond epoch written as a string lands on its real day; a boolean +# stamp falls back on the file's time rather than 1970; an opencode time that +# isn't an object reads as no time (today), tokens that aren't an object skip +# the row, and a cache that isn't one keeps the row's input. Today: 5 + 44 + 11. +# The database has the bare shape of older opencode builds, with no id column. +new_muse_home() +session("sess-odd", completion("string-ts-model", str(OLD_US), input=60, output=6), completion("bool-ts-model", True, input=5)) +(data_home() / "opencode").mkdir(parents=True) +bare = sqlite3.connect(data_home() / "opencode" / "opencode.db") +bare.execute("CREATE TABLE message (session_id TEXT, data TEXT)") +bare.executemany("INSERT INTO message VALUES (?, ?)", [ + ("odd-1", json.dumps({"role": "assistant", "providerID": "meta", "modelID": "odd-time-model", "tokens": oc_tokens(40, 4), "time": 12345})), + ("odd-2", json.dumps({"role": "assistant", "providerID": "meta", "modelID": "odd-tokens-model", "tokens": [1, 2], "time": {"created": NOW_MS}})), + ("odd-3", json.dumps({"role": "assistant", "providerID": "meta", "modelID": "odd-cache-model", "tokens": {"input": 11, "output": 0, "reasoning": 0, "cache": "nope"}, "time": {"created": NOW_MS}})), +]) +bare.commit() +bare.close() +record = muse() +usage = record["modelUsage"] +expect("Muse survives odd timestamp and token shapes without misfiling them", { + "today": record["todayTotalTokens"], + "prompts": record["totalPrompts"], + "stringEpoch": OLD_DATE in record["activeDates"], + "epoch": bool({"1969-12-31", "1970-01-01"} & set(record["activeDates"])), + "oddTime": [usage["odd-time-model"]["inputTokens"], usage["odd-time-model"]["outputTokens"]], + "oddTokens": "odd-tokens-model" in usage, + "oddCache": [usage["odd-cache-model"]["inputTokens"], usage["odd-cache-model"]["outputTokens"]], +}, {"today": 60, "prompts": 4, "stringEpoch": True, "epoch": False, "oddTime": [40, 4], "oddTokens": False, "oddCache": [11, 0]}, record) + +# pi sessions count on Meta's own provider: not Muse Spark through Zen, and +# not a prompt. +new_muse_home() +prompt = pi_reply("meta", 5000, id="user", model="muse-spark") +prompt["message"]["role"] = "user" +write_pi_session(home() / ".pi" / "agent" / "sessions" / "p" / "a.jsonl", [ + prompt, + pi_reply("meta", 10, id="meta", model="muse-spark"), + pi_reply("opencode", 5000, id="zen", model="muse-spark"), +]) +record = muse() +expect("Muse counts Meta's pi replies alone", [record["todayTotalTokens"], record["totalPrompts"]], [10, 1], record) + +# ---------------------------------------------------------- child agents + +# Muse credits each billed run to its owner in the session's log, as `muse +# exec` writes it: the main agent, and the reminder children that watch a +# prompt. A child with a log of its own is counted from its completions; one +# without (a skill reminder) is known only from its attribution, which counts +# once however many copies of it the log holds, and answers no prompt. +new_muse_home() + + +def attribution(owner_id, owner_type, input, cached, output, usage_id): + record = {"goal_attribution": {"mode": "none"}, "owner": {"owner_id": owner_id, "owner_type": owner_type, "requester_kind": "reminder", + "run_id": "run-1", "session_id": "run-1"}, "quantity": {"cached_tokens": cached, "input_tokens": input, "main_llm_steps": 0, + "output_tokens": output, "reasoning_tokens": 5, "reported": True, "unit": "tokens"}, "schema_version": 1, + "usage_family": "provider", "usage_id": usage_id} + return {"schema_version": 1, "id": "rec-" + usage_id, "recorded_at": NOW_US, "record_type": "event", "payload_type": "runtime.session", + "payload": {"event": {"kind": "goal_usage_attribution", "record": record}, "kind": "run", "run_id": "run-1"}} + + +skill_reminder = attribution("child-skill", "reminder_child", 4000, 400, 40, "usage-skill") +session("sess-C", + {"recorded_at": NOW_US, "payload": {"kind": "metadata", "record": {"model_id": "muse-spark-1.3-contributor", "provider_id": "meta"}}}, + completion("muse-spark-1.3-contributor", input=1000, output=10, cache_read=100), + attribution("main-root", "main_root", 1000, 100, 10, "usage-main"), + attribution("child-logged", "reminder_child", 2000, 200, 20, "usage-logged"), + skill_reminder, + {"id": "frame-1", "children": [{"record_json": json.dumps(skill_reminder)}]}) +session("sess-C/subagent/child-logged", + completion("muse-spark-1.3-contributor", input=2000, output=20, cache_read=200), + attribution("child-logged", "reminder_child", 2000, 200, 20, "usage-logged-call")) +record = muse() +# Main 1010 + the logged child 2020 + the skill reminder 4040. +expect("Muse counts a child agent's calls known only from the session's attribution, once and with no prompt", + pick(record, "todayTotalTokens", "todayPrompts", "todaySessions") | {"cached": record["modelUsage"]["muse-spark-1.3-contributor"]["cacheReadInputTokens"]}, + {"todayTotalTokens": 7070, "todayPrompts": 2, "todaySessions": 1, "cached": 700}, record) + +# ---------------------------------------------------------------- contract + +new_muse_home() + + +def add_session(): + add_session.count = getattr(add_session, "count", 0) + 1 + session(f"grown-{add_session.count}", completion("muse-spark", input=10)) + + +check_scan_cache(Adapter("muse", grow=add_session, pi_root=lambda: home() / ".pi" / "agent" / "sessions")) +PY 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..5280e9e3e0d --- /dev/null +++ b/test/shell.d/agent-usage-ollama-scanner-test.sh @@ -0,0 +1,279 @@ +#!/bin/bash + +# The Ollama Cloud collector's own parts: the ollama-cloud filter over pi and +# omp sessions with model tags merged, the configured model prefix, Ollama's +# balance and plan answers, and its key. Roots, forks, cloud-tagged models +# through the local daemon, both opencode stores, OPENCODE_DB and untrusted +# values are held by the shared and partition tests, as is the +# unused-machine record; the limits and scan caches go through the contract +# checks. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +# The checks below compare today's totals with rows stamped now. +freeze_today() +BALANCE = "ollama.com/api/balance" +ME = "ollama.com/api/me" +SESSION_RESET = iso_in(3 * 3600) +WEEKLY_RESET = iso_in(4 * 86400) +# A legacy plan's answer, as docs.ollama.com/api/balance documents it. +ROUTES = { + BALANCE: { + "included": { + "session": {"remaining_percent": 75, "resets_at": SESSION_RESET}, + "weekly": {"remaining_percent": 40, "resets_at": WEEKLY_RESET}, + }, + "purchased": {"balance_usd": 25}, + }, + ME: {"Plan": "pro"}, +} + + +def ollama(*args, routes=None, mode="up"): + with FakeHttp(ROUTES if routes is None else routes, mode=mode) as http: + record = collect("ollama", *(args or ("--force",))) + ollama.requests = http.count + return record + + +def sessions(): + return home() / ".pi" / "agent" / "sessions" + + +def prefix(value): + if value is None: + (config_home() / "omarchy" / "agents" / "ollama.json").unlink() + else: + write_json(config_home() / "omarchy" / "agents" / "ollama.json", {"modelPrefix": value}) + + +# ------------------------------------------------------------ local stats + +# pi and omp sessions on the ollama-cloud provider are the local stats; other +# providers and prompts in the same files stay out, and model version tags +# merge. +sandbox() +write_pi_session(sessions() / "project" / "pi.jsonl", [ + pi_reply("ollama-cloud", id="pi-1", model="deepseek-v4-pro:0813", usage={"input": 10, "output": 4, "cacheRead": 3, "cacheWrite": 2, "totalTokens": 19}), + pi_reply("ollama-cloud", id="pi-2", model="deepseek-v4-pro:0731", usage={"input": 5, "output": 1, "cacheRead": 0, "cacheWrite": 0, "totalTokens": 6}), + pi_reply("anthropic", id="pi-3", model="claude-test", usage={"input": 999, "output": 999}), + {"type": "message", "id": "pi-4", "timestamp": iso_in(0), "message": {"role": "user", "provider": "ollama-cloud", "model": "deepseek-v4-pro:0813", "usage": {"input": 999, "output": 999}}}, + pi_reply("ollama-cloud", id="pi-5", model="deepseek-v4-pro:0813", usage={"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": 0}), +]) +write_pi_session(home() / ".omp" / "agent" / "sessions" / "project" / "omp.jsonl", [ + pi_reply("ollama-cloud", id="omp-1", model="kimi-k3", usage={"input": 20, "output": 5, "cacheRead": 4, "cacheWrite": 1, "totalTokens": 30}), +]) +record = ollama() +expect("Ollama counts ollama-cloud usage from pi and omp, merging model tags", [record["name"], record["todayTotalTokens"], record["modelUsage"]], ["Ollama Cloud", 55, { + "deepseek-v4-pro": {"inputTokens": 15, "outputTokens": 5, "cacheReadInputTokens": 3, "cacheCreationInputTokens": 2}, + "kimi-k3": {"inputTokens": 20, "outputTokens": 5, "cacheReadInputTokens": 4, "cacheCreationInputTokens": 1}, +}], record) + +# In use here, but no key to ask for the account's limits. +expect("Ollama in use without a key waits for auth", [record["ready"], record["usageStatusText"], record["limitsStale"], ollama.requests], + [True, "Waiting for auth", False, 0], record) + +# A configured model prefix restricts the local stats to matching models, and +# a scan only answers for the filter it applied: the same filter reuses it, +# another scans again. +prefix("kimi") +expect("Ollama honors the configured model prefix", list(ollama()["modelUsage"]), ["kimi-k3"]) +expect("Ollama reuses a scan made with the same model prefix", list(ollama("--limits-only")["modelUsage"]), ["kimi-k3"]) +prefix(None) +expect("Ollama scans again when the model prefix changes", sorted(ollama("--limits-only")["modelUsage"]), ["deepseek-v4-pro", "kimi-k3"]) + +# Claude Code pointed at ollama.com is Ollama's whatever model answered, and +# a cloud-tagged model is Ollama's wherever Claude Code sent it; a model the +# local daemon ran itself is not. +sandbox() +claude_code_session("cloud", [("glm-4.6:cloud", 300), ("qwen3-coder", 30), ("claude-sonnet-4-5", 40)], base_url="http://localhost:11434") +record = ollama() +expect("Ollama counts Claude Code's cloud-tagged replies through the local daemon, and nothing it ran itself", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"glm-4.6": 300}, 1, 1], record) +claude_code_session("cloud", [("kimi-k2.7-code:cloud", 300), ("claude-sonnet-4-5", 40)], base_url="https://ollama.com") +record = ollama() +expect("Ollama counts every Claude Code reply while Claude Code is pointed at ollama.com", + [record["todayTokensByModel"], record["todayPrompts"]], [{"kimi-k2.7-code": 300, "claude-sonnet-4-5": 40}, 2], record) + +# Codex's `--oss` rollouts name the local daemon's `ollama` provider: a +# cloud-tagged model ran on Ollama's servers and counts here, with its typed +# prompt, while one the daemon ran itself is the local collector's. A +# model_providers entry named `ollama-cloud` is Ollama Cloud whatever the +# model. A built-in `openai` rollout is the Codex collector's. +sandbox() +os.environ["CODEX_HOME"] = str(home() / ".codex") +day_dir = home() / ".codex" / "sessions" / datetime.now().strftime("%Y/%m/%d") +day_dir.mkdir(parents=True) +STAMP = today_stamp() + + +def codex_rollout(name, provider, model, tokens): + lines = [ + {"timestamp": STAMP, "type": "session_meta", "payload": {"model_provider": provider}}, + {"timestamp": STAMP, "type": "turn_context", "payload": {"model": model}}, + {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "user_message", "message": "go"}}, + {"timestamp": STAMP, "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": tokens, "cached_input_tokens": 10, "output_tokens": 5}, + "total_token_usage": {"input_tokens": tokens}}}}, + ] + (day_dir / f"rollout-{name}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + + +codex_rollout("cloud", "ollama", "gpt-oss:120b-cloud", 6280) +codex_rollout("here", "ollama", "gpt-oss:20b", 4280) +codex_rollout("hosted", "ollama-cloud", "qwen3-coder:480b", 980) +codex_rollout("sub", "openai", "gpt-5.1-codex", 777) +record = ollama() +expect("Ollama counts Codex rollouts on its cloud models and its own provider, prompts as typed", + [record["todayTokensByModel"], record["modelUsage"].get("gpt-oss"), record["todayPrompts"], record["todaySessions"]], + [{"gpt-oss": 6285, "qwen3-coder": 985}, {"inputTokens": 6270, "outputTokens": 5, "cacheReadInputTokens": 10, "cacheCreationInputTokens": 0}, 2, 2], record) +write_json(config_home() / "omarchy" / "agents" / "ollama.json", {"modelPrefix": "qwen"}) +record = ollama() +expect("Ollama's model prefix filters Codex rollouts too", record["todayTokensByModel"], {"qwen3-coder": 985}, record) +del os.environ["CODEX_HOME"] + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_reply(): + add_reply.count = getattr(add_reply, "count", 0) + 1 + write_pi_session(sessions() / "grown" / f"{add_reply.count}.jsonl", [pi_reply("ollama-cloud", 5, id=f"grown-{add_reply.count}")]) + + +# A probe asks for the balance and then the plan. +adapter = Adapter("ollama", plant=lambda: os.environ.update(OLLAMA_API_KEY="key"), answer=ROUTES, requests_per_probe=2, + grow=add_reply, pi_root=sessions) +fresh = check_limits_contract(adapter) +expect("Ollama's plan and a legacy plan's session and weekly use and resets come from its answers", { + "tierLabel": fresh["tierLabel"], "limits": fresh["limits"], +}, { + "tierLabel": "Pro", "limits": [ + {"label": "Session (5-hour)", "percent": 0.25, "resetsAt": SESSION_RESET}, + {"label": "Weekly (7-day)", "percent": 0.6, "resetsAt": WEEKLY_RESET}, + ], +}, fresh) +check_scan_cache(adapter) + +# ---------------------------------------------------------------- the key + +# A rejected key is a sign-in to redo, and says so. +record = ollama(mode=401) +check("Ollama explains a rejected key", "rejected" in record["authHelpText"], record) + +# A plan lookup that fails keeps the plan last named. +record = ollama(routes={**ROUTES, ME: urllib.error.URLError(OSError("plan lookup timed out"))}) +expect("Ollama keeps the cached plan when the plan lookup fails", [record["tierLabel"], record["limitsStale"]], ["Pro", False], record) + +# An answer never rolls back one another run stamped while this one waited +# on the network. +def later_run(request): + edit_limits_cache("ollama", fetched_ago=-60) + return ROUTES[BALANCE] + + +ollama(routes={**ROUTES, BALANCE: later_run}) +expect("Ollama never overwrites a newer cached answer with its own", + [json.loads(path.read_text())["fetchedAtMs"] > ms_in(30) for path in limits_cache_files("ollama")], [True]) +edit_limits_cache("ollama", fetched_ago=0) + +# Another key never shows this one's numbers or plan. +os.environ["OLLAMA_API_KEY"] = "other-key" +record = ollama(mode="down") +expect("Ollama keeps each key's limits and plan apart", [record["limits"], record["tierLabel"], record["usageStatusText"]], + [[], "", "Ollama limits unavailable"], record) + +# Ollama's UserResponse struct tags the plan `plan` (api/types.go), as the +# server re-encodes it; the capitalized spelling still reads. +record = ollama(routes={**ROUTES, ME: {"id": "00000000-0000-0000-0000-000000000000", "email": "a@example.com", "name": "a", "plan": "pro"}}) +expect("Ollama reads the plan under its JSON tag's lowercase spelling", record["tierLabel"], "Pro", record) + +# ---------------------------------------------------------- the answer + +# With nothing kept, a failed check says the limits are unavailable; only an +# unreachable endpoint asks for an early retry. +sandbox() +os.environ["OLLAMA_API_KEY"] = "key" +record = ollama(mode="down") +expect("Ollama offline with nothing kept says its limits are unavailable and advises a retry", + [record["usageStatusText"], record.get("retryAdvised")], ["Ollama limits unavailable", True], record) +record = ollama(mode=500) +expect("Ollama advises no retry for a server's error status", [record["usageStatusText"], record.get("retryAdvised")], + ["Ollama limits unavailable", None], record) + +# A current plan's answer, as documented: the included credits are a monthly +# meter of what is spent, resetting when the period ends. +record = ollama(routes={BALANCE: { + "included": {"balance_usd": 72.5, "allowance_usd": 100, "period": {"from": iso_in(-86400), "until": WEEKLY_RESET}}, + "purchased": {"balance_usd": 25}, +}, ME: {"Plan": "pro"}}) +expect("Ollama meters a current plan's included credits until the period ends", [record["limits"], record["usageStatusText"]], + [[{"label": "Included credits", "title": "Monthly", "percent": 0.275, "resetsAt": WEEKLY_RESET}], ""], record) +expect("Ollama shows the purchased credits left as a balance, with no funded figure to alarm on", record.get("balance"), + {"remaining": 25.0, "funded": 0, "spent": 0, "currency": "USD", "estimated": False}, record) + +# The balance is kept with the limits: a later failed check shows both, stale. +record = ollama(mode=500) +expect("Ollama keeps the purchased balance stale past a failed check", + [record.get("balance", {}).get("remaining"), len(record["limits"]), record["limitsStale"]], [25.0, 1, True], record) + +# An answer with no usable windows explains itself instead of printing none. +sandbox() +os.environ["OLLAMA_API_KEY"] = "key" +record = ollama(routes={BALANCE: {"included": {"session": {"remaining_percent": "unknown"}, "weekly": None}, "purchased": {"balance_usd": 0}}, + ME: {"Plan": "pro"}}) +expect("Ollama reports a balance answer without usable limits", [record["limits"], record["usageStatusText"], "no limits" in record["authHelpText"]], + [[], "Ollama limits unavailable", True], record) + +# Values out of range or of the wrong type skip that window, never the rest: +# a negative remainder is dropped, one past 100 reads as untouched, and a +# reset that won't parse as none. +record = ollama(routes={BALANCE: {"included": { + "session": {"remaining_percent": -20, "resets_at": SESSION_RESET}, + "weekly": {"remaining_percent": 140, "resets_at": "not a time"}, +}}, ME: {"Plan": "pro"}}) +expect("Ollama skips a window it can't read and keeps the rest in range", + record["limits"], [{"label": "Weekly (7-day)", "percent": 0.0, "resetsAt": ""}], record) + +# Without OLLAMA_API_KEY, the key pi or omp signed in with stands in, pi's +# first, and goes out as a Bearer token on both requests. A pi `$VAR` or +# `!command` entry is pi's to resolve, never a key to send. +sandbox() +pi_auth(home() / ".omp" / "agent", "ollama-cloud", "omp-key") +pi_auth(None, "ollama-cloud", "pi-key") +with FakeHttp(ROUTES) as http: + record = collect("ollama", "--force") +expect("Ollama asks with pi's saved key before omp's, as a Bearer token", + [sorted({request["headers"].get("authorization") for request in http.requests}), len(http.requests), len(record["limits"])], + [["Bearer pi-key"], 2, 2], record) +pi_auth(None, "ollama-cloud", "$OLLAMA_KEY") +expect("Ollama passes over pi's $VAR entry to omp's key", load_collector("ollama").api_key(), "omp-key") +pi_auth(home() / ".omp" / "agent", "ollama-cloud", "!pass show ollama") +expect("Ollama never sends a pi $VAR or !command entry as a key", load_collector("ollama").api_key(), "") +# A MiniMax Code custom provider on the local daemon that answered with a +# cloud-tagged model ran on Ollama's servers, so it is Ollama's, prompt and +# all, and no local work. +sandbox() +mcode_routed_session("box", "custom_provider:box", "gpt-oss:120b-cloud", 500, base_url="http://127.0.0.1:11434/v1") +record = ollama() +expect("Ollama counts a MiniMax Code cloud-tagged model sent through the local daemon", + [record["todayTokensByModel"], record["todayPrompts"]], [{"gpt-oss": 500}, 1], record) + +# A MiniMax Code prompt takes the model of the reply that answered it, so the +# model prefix keeps it with its reply. +sandbox() +mcode_routed_session("oc", "custom_provider:oc", "claude-sonnet-4-5", 40, base_url="https://ollama.com/v1") +prefix("claude") +record = ollama() +expect("Ollama's model prefix keeps a MiniMax Code prompt with the reply that answered it", + [record["todayTokensByModel"], record["todayPrompts"]], [{"claude-sonnet-4-5": 40}, 1], record) +PY diff --git a/test/shell.d/agent-usage-opencode-limits-test.sh b/test/shell.d/agent-usage-opencode-limits-test.sh new file mode 100644 index 00000000000..4a386b98fbc --- /dev/null +++ b/test/shell.d/agent-usage-opencode-limits-test.sh @@ -0,0 +1,262 @@ +#!/bin/bash + +# The OpenCode collector's limits: the Go plan's windows from the Zen gateway +# in both shapes it answers, a key without a Go plan, every credential store +# a key can come from, and its adapter for the contract checks. The +# unused-machine record is held centrally by the partition test, the limits +# cache's mechanics by the shared test, and the scan cache by the scanner +# test. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +ZEN = "opencode.ai/zen/go/v1/usage" +CONSOLE = "opencode.ai/inference/go/v1/usage" +USAGE = {"usage": { + "rolling": {"percent": 16, "resetsAt": "2999-01-01T00:00:00.000Z"}, + "weekly": {"percent": 6, "resetsAt": "2999-01-05T00:00:00.000Z"}, + "monthly": {"percent": 29, "resetsAt": "2999-01-11T15:38:40.000Z"}, +}} + + +def keyed(key="zen_test"): + os.environ["OPENCODE_API_KEY"] = key + + +def opencode(answer=None, mode="up", *args): + with FakeHttp({ZEN: answer}, mode=mode) as http: + record = collect("opencode", *(args or ("--force",))) + return record, http + + +def windows(record): + return [(entry["label"], entry["percent"]) for entry in record["limits"]] + + +def candidates(): + return [candidate.token for candidate in load_collector("opencode").credential_candidates()] + + +# ---------------------------------------------------------------- contract + +sandbox() +fresh = check_limits_contract(Adapter("opencode", plant=keyed, answer={ZEN: USAGE})) +expect("OpenCode collector maps every Zen usage window, under the Go plan", + [windows(fresh), fresh["tierLabel"]], [[("Rolling (5h)", 0.16), ("Weekly (7-day)", 0.06), ("Monthly", 0.29)], "Go"], fresh) + +# An answer (limits, or no plan) never rolls back one another run stamped +# while this one waited on the network. +for later_answer in (USAGE, 403): + def later_run(request, later_answer=later_answer): + edit_limits_cache("opencode", fetched_ago=-60) + return later_answer + opencode(later_run) + expect(f"OpenCode collector never overwrites a newer cached answer with its own ({'limits' if later_answer == USAGE else 'no plan'})", + [json.loads(path.read_text())["fetchedAtMs"] > ms_in(30) for path in limits_cache_files("opencode")], [True]) + edit_limits_cache("opencode", fetched_ago=0) +opencode(USAGE) + +# A failed check keeps the plan along with the windows. +record, _ = opencode(mode=500) +expect("OpenCode collector keeps the Go plan with the kept windows", [record["tierLabel"], record["limitsStale"]], ["Go", True], record) + +# Another sign-in never shows this one's numbers. +keyed("zen_other") +record, _ = opencode(mode="down") +expect("OpenCode collector caches limits per credential", [record["limits"], record["usageStatusText"]], [[], "OpenCode limits unavailable"], record) + +# ------------------------------------------------------------- the answer + +sandbox() +keyed() + +# Another deployment spells the windows Usage, and a rate-limited +# window is spent whatever it reads. +record, _ = opencode({ + "rollingUsage": {"usagePercent": 40, "resetInSec": 3600, "status": "rate-limited"}, + "weeklyUsage": {"usagePercent": "12.5", "resetInSec": 0}, +}) +expect("OpenCode collector reads the alternate payload shape", + [windows(record), record["limits"][0]["resetsAt"] != "", record["limits"][1]["resetsAt"]], + [[("Rolling (5h)", 1.0), ("Weekly (7-day)", 0.125)], True, ""], record) + +# resetsAt is an ISO time or nothing, so an epoch, a list or garbage falls back +# to resetInSec, and a reset further out than a date goes reads as none. +for odd in (1792700000, ["x"], "garbage"): + record, _ = opencode({"usage": { + "rolling": {"percent": 10, "resetsAt": odd, "resetInSec": 3600}, + "weekly": {"percent": 20, "resetsAt": odd}, + "monthly": {"percent": 30, "resetInSec": 1e14}, + }}) + expect(f"OpenCode collector shows only an ISO reset time (resetsAt {odd!r})", + [windows(record), record["limits"][0]["resetsAt"][:4].isdigit() and "T" in record["limits"][0]["resetsAt"], + record["limits"][1]["resetsAt"], record["limits"][2]["resetsAt"]], + [[("Rolling (5h)", 0.1), ("Weekly (7-day)", 0.2), ("Monthly", 0.3)], True, "", ""], record) + +# With nothing kept, each way a check fails says what happened. +for path in limits_cache_files("opencode"): + path.unlink() +cases = { + "a payload without windows": ({"usage": {}}, "up", "OpenCode limits unavailable", "no quota windows", False), + "a refused key": (None, 401, "Sign-in expired", "rejected the saved key", False), + "a non-auth status": (None, 500, "OpenCode limits unavailable", "returned status 500", False), + "no server answering": (None, "down", "OpenCode limits unavailable", "Couldn't reach", True), + "a timeout": (None, "timeout", "OpenCode limits unavailable", "Couldn't reach", True), + # A server that answered with garbage was reached; retrying soon won't help. + "an unparseable answer": (None, "garbage", "OpenCode limits unavailable", "returned no data", False), + "an answer that breaks the reader": (RuntimeError("broken"), "up", "OpenCode limits unavailable", "returned no data", False), +} +for name, (answer, mode, status, help_text, retry) in cases.items(): + record, _ = opencode(answer, mode) + check(f"OpenCode collector explains {name} with nothing kept, and calls nothing stale", + record["limits"] == [] and record["limitsStale"] is False and record["usageStatusText"] == status + and help_text in record["authHelpText"] and bool(record.get("retryAdvised")) is retry, record) + +# A Zen pay-as-you-go key has no Go plan behind it. That is no trouble, and +# the answer is remembered so the next runs don't ask again. +record, http = opencode(None, 403) +check("OpenCode collector stays silent for a key without a Go plan", + record["limits"] == [] and record["usageStatusText"] == "" and record["authHelpText"] == "" and record["tierLabel"] == "" + and record["limitsStale"] is False and http.count == 1, record) +check("an OpenCode key with no sessions here has no local stats", + record["ready"] is True and record["totalSessions"] == 0 and record["hasLocalStats"] is False, record) +record, http = opencode(None, 403, "--limits-only") +check("OpenCode collector doesn't ask again about a key without a Go plan", + record["usageStatusText"] == "" and record["limits"] == [] and http.count == 0, record) +record, _ = opencode({"usage": {"weekly": {"percent": 10, "resetsAt": "2999-01-01T00:00:00.000Z"}}}) +expect("OpenCode collector --force asks again after a Go plan appears", [windows(record), record["tierLabel"]], [[("Weekly (7-day)", 0.1)], "Go"], record) +record, _ = opencode(None, 403) +expect("OpenCode collector drops a cancelled plan's windows", + [record["limits"], record["usageStatusText"], record["tierLabel"]], [[], "", ""], record) + +# Two ways in, each answering for itself: OPENCODE_API_KEY is asked first. +sandbox() +os.environ.update(OPENCODE_API_KEY="key_first", OPENCODE_GO_API_KEY="key_second") + + +def by_key(**answers): + return lambda request: answers[request.get_header("Authorization").removeprefix("Bearer ")] + + +# Every way in refused keeps the open windows of whichever refusal has them, +# not just the first one asked, which has none. +opencode(by_key(key_first=401, key_second=USAGE)) +record, _ = opencode(by_key(key_first=401, key_second=401)) +expect("OpenCode collector keeps the windows of a later refused key when the first has none", + [windows(record), record["limitsStale"], record["usageStatusText"]], + [[("Rolling (5h)", 0.16), ("Weekly (7-day)", 0.06), ("Monthly", 0.29)], True, "Sign-in expired"], record) + +# A key that says it has no Go plan makes an outage on the other one, with +# nothing kept for it, no news. +for path in limits_cache_files("opencode"): + path.unlink() +record, _ = opencode(by_key(key_first=403, key_second=500)) +expect("OpenCode collector lets a no-plan key quiet another key's outage", + [record["limits"], record["usageStatusText"], record["authHelpText"]], [[], "", ""], record) +del os.environ["OPENCODE_API_KEY"], os.environ["OPENCODE_GO_API_KEY"] + +# ------------------------------------------------------------ credentials + +# The key can come from opencode's own credential stores or, when a machine +# codes through pi, from pi's. +sandbox() +auth = data_home() / "opencode" / "auth.json" +write_json(auth, {"opencode": {"type": "api_key", "key": "from_opencode"}}) +expect("OpenCode collector reads opencode's own credential", candidates(), ["from_opencode"]) + +# A custom provider merely named like opencode holds a key for somewhere else. +write_json(auth, {"opencode-local": {"type": "api_key", "key": "gateway_secret"}}) +check("OpenCode collector never sends a look-alike provider's key", "gateway_secret" not in candidates()) + +# An auth.json entry may be the key itself, and the Go key has an +# environment variable of its own beside OPENCODE_API_KEY. +write_json(auth, {"opencode-go": " plain_key "}) +expect("OpenCode collector reads a plain-string auth.json key", candidates(), ["plain_key"]) +os.environ["OPENCODE_GO_API_KEY"] = "go_env" +expect("OpenCode collector reads OPENCODE_GO_API_KEY first", candidates(), ["go_env", "plain_key"]) + +sandbox() +pi_auth(None, "opencode-go", "from_pi") +expect("OpenCode collector reads pi's credential as a last resort", candidates(), ["from_pi"]) +pi_auth(home() / ".omp" / "agent", "opencode", "from_omp") +expect("OpenCode collector reads omp's credential after pi's", candidates(), ["from_pi", "from_omp"]) +os.environ["PI_CODING_AGENT_DIR"] = str(home() / "pi-moved") +pi_auth(home() / "pi-moved", "opencode", "from_moved") +expect("OpenCode collector reads pi's credential where PI_CODING_AGENT_DIR moves it", candidates(), ["from_moved", "from_omp"]) +pi_auth(home() / "pi-moved", "opencode", "!pass show opencode") +pi_auth(home() / ".omp" / "agent", "opencode", "$OPENCODE_KEY") +expect("OpenCode collector never sends a pi !command or $VAR entry as a key", candidates(), []) + +# OpenCode V2 keeps its credentials in the database: a Go key first, and the +# Console session last, asked on Console's endpoint with its workspace. +sandbox() +db = data_home() / "opencode" / "opencode.db" +db.parent.mkdir(parents=True) +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE credential (id text, integration_id text, active integer, time_updated integer, value text)") +conn.executemany("INSERT INTO credential VALUES (?, ?, ?, ?, ?)", [ + ("c1", "opencode-go", 1, 2, json.dumps({"type": "key", "key": "store_key"})), + ("c2", "opencode-go", 0, 3, json.dumps({"type": "key", "key": "inactive_key"})), + ("c3", "opencode", 1, 1, json.dumps({"type": "oauth", "access": "console_token", "expires": ms_in(60), "metadata": {"orgID": "org_1"}})), +]) +conn.commit() +conn.close() +found = load_collector("opencode").credential_candidates() +expect("OpenCode collector reads V2 credentials, the Console session last", + [(c.token, c.url.endswith("/inference/go/v1/usage"), c.headers) for c in found], + [("store_key", False, {}), ("console_token", True, {"x-opencode-org-id": "org_1"})]) + +# One value nested too deeply to parse (a RecursionError, not a ValueError) +# costs only its own entry: the database's other key and pi's still count. +DEEP = "[" * 100000 + "]" * 100000 +(data_home() / "opencode" / "auth.json").write_text(DEEP) +pi_auth(None, "opencode", "from_pi") +conn = sqlite3.connect(db) +conn.execute("INSERT INTO credential VALUES ('c4', 'opencode-go', 1, 9, ?)", (DEEP,)) +conn.commit() +conn.close() +with FakeHttp({ZEN: USAGE, CONSOLE: USAGE}): + record = collect("opencode", "--force") +expect("OpenCode collector skips a credential nested too deeply to parse", + [[c.token for c in load_collector("opencode").credential_candidates()], windows(record)], + [["store_key", "from_pi", "console_token"], [("Rolling (5h)", 0.16), ("Weekly (7-day)", 0.06), ("Monthly", 0.29)]], record) +(data_home() / "opencode" / "auth.json").unlink() +(home() / ".pi" / "agent" / "auth.json").unlink() +conn = sqlite3.connect(db) +conn.execute("DELETE FROM credential WHERE id = 'c4'") +conn.commit() +conn.close() + + +def refuse_key(request): + if "store_key" in request.get_header("Authorization"): + return 401 + return {"usage": {"weekly": {"percent": 30, "resetsAt": "2999-01-01T00:00:00.000Z"}}} + + +with FakeHttp({ZEN: refuse_key, CONSOLE: refuse_key}) as http: + record = collect("opencode", "--force") +asked = http.requests[-1] +expect("OpenCode collector hands a refused key over to the Console session", + [windows(record), asked["url"], asked["headers"].get("authorization"), asked["headers"].get("x-opencode-org-id")], + [[("Weekly (7-day)", 0.3)], "https://opencode.ai/inference/go/v1/usage", "Bearer console_token", "org_1"], record) + +# Every way in refused, with the Console session expired, says so. +conn = sqlite3.connect(db) +conn.execute("UPDATE credential SET value = ? WHERE id = 'c3'", (json.dumps({"type": "oauth", "access": "console_token", "expires": 1000}),)) +conn.commit() +conn.close() +for path in limits_cache_files("opencode"): + path.unlink() +with FakeHttp(mode=401): + record = collect("opencode", "--force") +check("OpenCode collector names an expired Console session", + record["usageStatusText"] == "Sign-in expired" and "Console session has expired" in record["authHelpText"] and record["limits"] == [], record) +PY diff --git a/test/shell.d/agent-usage-opencode-scanner-test.sh b/test/shell.d/agent-usage-opencode-scanner-test.sh new file mode 100644 index 00000000000..8c4de60f151 --- /dev/null +++ b/test/shell.d/agent-usage-opencode-scanner-test.sh @@ -0,0 +1,112 @@ +#!/bin/bash + +# The OpenCode collector's local stats: pi, omp and opencode usage on the Zen +# and Go providers only, with the thinking suffix stripped from model names. +# How pi and opencode stores are read (roots, forks, both opencode stores, +# OPENCODE_DB, untrusted values) is held centrally by the shared and +# partition tests, the unused-machine record by the partition test, and the +# scan cache by the contract check. The limits are agent-usage-opencode-limits-test.sh's. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + + +def opencode(*args): + with FakeHttp(): + return collect("opencode", *(args or ("--force",))) + + +def sessions(): + return home() / ".pi" / "agent" / "sessions" + + +# ------------------------------------------------------------- pi and omp + +# pi and omp can spend an OpenCode subscription without opencode's CLI ever +# running, so its own database cannot be the only source. A custom +# "opencode-local" gateway is somewhere else, and a user line is no reply. +sandbox() +write_pi_session(sessions() / "project" / "pi.jsonl", [ + pi_reply("opencode", id="pi-1", model="deepseek-v4.1-flash:high", usage={"input": 100, "output": 15, "reasoning": 5, "cacheRead": 20, "cacheWrite": 3}), + pi_reply("opencode-local", id="pi-2", model="gpt-nope", usage={"input": 999, "output": 999}), + {"type": "message", "id": "pi-3", "timestamp": iso_in(0), "message": {"role": "user", "provider": "opencode", "model": "deepseek-v4.1-flash", "usage": {"input": 999, "output": 999}}}, +]) +write_pi_session(home() / ".omp" / "agent" / "sessions" / "project" / "omp.jsonl", [ + pi_reply("opencode-go", id="omp-1", model="minimax-m3", usage={"input": 50, "output": 5, "reasoningTokens": 1, "cacheRead": 5, "cacheWrite": 0}), + pi_reply("openai-codex", id="other-1", model="gpt-nope", usage={"input": 999, "output": 999}), +]) +record = opencode() +expect("OpenCode counts its Zen and Go usage from pi and omp, reasoning already in output, stripping the thinking suffix", { + "name": record["name"], "today": record["todayTotalTokens"], "models": record["modelUsage"], +}, { + "name": "OpenCode", "today": 198, "models": { + "deepseek-v4.1-flash": {"inputTokens": 100, "outputTokens": 15, "cacheReadInputTokens": 20, "cacheCreationInputTokens": 3}, + "minimax-m3": {"inputTokens": 50, "outputTokens": 5, "cacheReadInputTokens": 5, "cacheCreationInputTokens": 0}, + }, +}, record) + +# ---------------------------------------------------------------- opencode + +# A subscription burned entirely through opencode has no pi session files; +# usage comes from opencode's own database, filtered to its providers. +sandbox() +db = opencode_db(tables=("message",), message=[ + ("msg_1", "ses_1", oc_v1("opencode", {"input": 80, "output": 40, "reasoning": 5, "cache": {"read": 30, "write": 0}}, model="deepseek-v4-flash")), + ("msg_2", "ses_1", oc_v1("opencode-go", {"input": 10, "output": 2}, model="minimax-m3")), + ("msg_3", "ses_1", oc_v1("anthropic", {"input": 999, "output": 999}, model="claude-opus-5")), + ("msg_7", "ses_1", oc_v1("openai", {"input": 999, "output": 999}, model="gpt-5.5")), + ("msg_4", "ses_1", oc_v1("opencode-local", {"input": 999, "output": 999}, model="gpt-nope")), + ("msg_5", "ses_1", oc_v1("opencode", {"input": 999, "output": 999}, model="deepseek-v4-flash", role="user")), + ("msg_6", "ses_1", '["not","an","object"]'), +]) +record = opencode() +expect("OpenCode counts opencode's own rows on its providers, leaving Anthropic, OpenAI, look-alike gateways, prompts and malformed rows out", { + "today": record["todayTotalTokens"], "sessions": record["totalSessions"], "models": record["modelUsage"], +}, { + "today": 167, "sessions": 1, "models": { + "deepseek-v4-flash": {"inputTokens": 80, "outputTokens": 45, "cacheReadInputTokens": 30, "cacheCreationInputTokens": 0}, + "minimax-m3": {"inputTokens": 10, "outputTokens": 2, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, + }, +}, record) + +# Valid JSON with a nested field of the wrong shape (a numeric time, a list of +# tokens or of cache counts, a list for a pi message or its usage) costs that +# field or that row, never the record. +sandbox() +write_pi_session(sessions() / "project" / "pi.jsonl", [ + {"type": "message", "id": "pi-odd", "timestamp": iso_in(0), "message": ["provider", "opencode"]}, + {"type": "message", "id": "pi-usage", "timestamp": iso_in(0), "message": {"role": "assistant", "provider": "opencode", "model": "kimi-k3", "usage": [1, 2]}}, +]) +odd_db = data_home() / "opencode" / "opencode.db" +odd_db.parent.mkdir(parents=True) +conn = sqlite3.connect(odd_db) +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, data text NOT NULL)") +conn.executemany("INSERT INTO message VALUES (?, 'ses_odd', ?)", [(row_id, json.dumps(data)) for row_id, data in ( + ("odd_time", {"role": "assistant", "providerID": "opencode", "modelID": "kimi-k3", "tokens": {"input": 40, "output": 0}, "time": 12345}), + ("odd_tokens", {"role": "assistant", "providerID": "opencode", "modelID": "kimi-k3", "tokens": [1, 2], "time": {"created": ms_in(0)}}), + ("odd_cache", {"role": "assistant", "providerID": "opencode", "modelID": "kimi-k3", "tokens": {"input": 2, "cache": [9]}, "time": {"created": ms_in(0)}}), +)]) +conn.commit() +conn.close() +record = opencode() +expect("OpenCode survives malformed nested fields", [record["totalPrompts"], record["modelUsage"]["kimi-k3"]["inputTokens"]], [2, 42], record) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_reply(): + add_reply.count = getattr(add_reply, "count", 0) + 1 + write_pi_session(sessions() / "grown" / f"{add_reply.count}.jsonl", [pi_reply("opencode-go", 5, id=f"grown-{add_reply.count}")]) + + +check_scan_cache(Adapter("opencode", grow=add_reply, pi_root=sessions)) +PY 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..a12067cff19 --- /dev/null +++ b/test/shell.d/agent-usage-openrouter-scanner-test.sh @@ -0,0 +1,373 @@ +#!/bin/bash + +# The OpenRouter collector's own sources: the key from OPENROUTER_API_KEY, +# opencode's or pi's or omp's saved login, or its config; the /key meter and +# /credits balance, each kept on its own; and with a management key the +# account's 30-day /activity history in place of the local scan. The +# unused-machine record and which local rows count as OpenRouter's are held +# centrally by the partition test, the pi and opencode readers by the shared +# test, and the limits and scan caches by the contract checks. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +# 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 + +python3 - <<'PY' +from agent_usage_fixtures import * + +API = "openrouter.ai/api/v1/" +TODAY = datetime.now().astimezone() +METER = {"data": {"limit": 10, "limit_remaining": 4, "limit_reset": "monthly"}} +CREDITS = {"data": {"total_credits": 10, "total_usage": 1}} +ACTIVITY = {"data": [{"date": TODAY.strftime("%Y-%m-%d"), "model": "a/b", "prompt_tokens": 10, "completion_tokens": 5}]} +ROUTES = {API + "activity": ACTIVITY, API + "key": METER, API + "credits": CREDITS} +DOWN = urllib.error.URLError(OSError("Name or service not known")) + + +def openrouter(*args, routes=None, mode="up", force=True, **env): + saved = dict(os.environ) + os.environ.update(env) + try: + with FakeHttp(ROUTES if routes is None else routes, mode=mode) as http: + record = collect("openrouter", *(args or (("--force",) if force else ()))) + finally: + os.environ.clear() + os.environ.update(saved) + openrouter.asked = sorted(f"/{request['url'].rsplit('/', 1)[1]}={request['headers']['authorization'][7:]}" for request in http.requests) + return record + + +def with_routes(**answers): + return dict(ROUTES, **{API + name: answer for name, answer in answers.items()}) + + +def signed_in(): + os.environ["OPENROUTER_API_KEY"] = "sk-test" + + +def picked(record, *keys): + return [record.get(key) for key in keys] + + +def remaining(record): + return (record.get("balance") or {}).get("remaining") + + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_session(): + add_session.count = getattr(add_session, "count", 0) + 1 + write_pi_session(home() / ".pi" / "agent" / "sessions" / "--grown--" / f"{add_session.count}.jsonl", + [pi_reply("openrouter", 1, id=f"g-{add_session.count}", model="anthropic/claude-test")]) + + +adapter = Adapter("openrouter", plant=signed_in, answer=ROUTES, requests_per_probe=2, grow=add_session, pi_root=lambda: home() / ".pi" / "agent" / "sessions") +fresh = check_limits_contract(adapter) +expect("an inference key's credit limit drains toward empty, named for its period, beside the balance", + [fresh["limits"], remaining(fresh), fresh["tierLabel"], fresh["hasPromptStats"]], + [[{"label": "Monthly key limit", "title": "Monthly key limit", "percent": 0.6, "resetsAt": ""}], 9.0, "Prepaid", True], fresh) +check_scan_cache(adapter) +del os.environ["OPENROUTER_API_KEY"] + +# --------------------------------------------------------------- the answer + +module = load_collector("openrouter") + +# The local scan counts OpenRouter's assistant rows under the model's own +# name, without its provider prefix. +sandbox() +opencode_db(message=[("m1", "ses_a", oc_v1("openrouter", 100, model="openai/gpt-4o-mini"))]) +record = openrouter(OPENROUTER_API_KEY="sk-test") +expect("a model is named without its provider prefix", sorted(record["modelUsage"]), ["gpt-4o-mini"], record) + +# /activity folds prompt tokens into input and completion tokens into output, +# one bucket per model. Reasoning is part of the completion count (OpenRouter's +# AI SDK provider reads text = completion - reasoning), so it is never added +# again; the first row is the spec's own ActivityItem example. A row naming +# reasoning alone still counts it as output. +activity = module.summarize_activity([ + {"byok_usage_inference": 0.012, "cached_tokens": 10, "completion_tokens": 125, "date": TODAY.strftime("%Y-%m-%d"), + "endpoint_id": "550e8400-e29b-41d4-a716-446655440000", "model": "openai/gpt-4.1", "model_permaslug": "openai/gpt-4.1-2025-04-14", + "prompt_tokens": 50, "provider_name": "OpenAI", "reasoning_tokens": 25, "requests": 5, "usage": 0.015}, + {"date": TODAY.strftime("%Y-%m-%d"), "model": "openai/gpt-4o-mini", "prompt_tokens": 5, "completion_tokens": 0, "reasoning_tokens": 4}, + {"date": "2020-01-01", "model": "x/y-model", "prompt_tokens": 7, "completion_tokens": 3, "reasoning_tokens": 0}, + "garbage-row", +], TODAY) +expect("activity folds prompt and completion tokens per model, with its active dates", + [activity["todayTotalTokens"], sorted(activity["modelUsage"]), len(activity["activeDates"])], [184, ["gpt-4.1-2025-04-14", "gpt-4o-mini", "y-model"], 2], activity) +expect("activity counts reasoning inside the completion tokens, once, as output", + [activity["modelUsage"]["gpt-4.1-2025-04-14"]["inputTokens"], activity["modelUsage"]["gpt-4.1-2025-04-14"]["outputTokens"], + activity["modelUsage"]["gpt-4o-mini"]["outputTokens"]], [50, 125, 4], activity) +expect("a key without a credit limit has no meter, and one with a limit names its period, not a reset time", + [module.key_limit_meter({"limit": None, "limit_remaining": None}), module.key_limit_meter({"limit": 10, "limit_remaining": 4, "limit_reset": "weekly"})], + [[], [{"label": "Weekly key limit", "title": "Weekly key limit", "percent": 0.6, "resetsAt": ""}]]) + +# A read that stalls on its body times out the read itself: a transport +# failure like any other. +class StalledBody: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self, *args): + raise TimeoutError("The read operation timed out") + + +real_urlopen = urllib.request.urlopen +urllib.request.urlopen = lambda request, timeout=None: StalledBody() +try: + module.api_get("/key", "sk-test") + stalled = "no error" +except module.OpenRouterTransportError: + stalled = "transport" +finally: + urllib.request.urlopen = real_urlopen +expect("a body that stalls is a transport failure", stalled, "transport") + +# --------------------------------------------------------------- the keys + +# A management key brings the account's 30-day history in place of the +# local scan: account-wide numbers with no prompt counts. +sandbox() +record = openrouter(OPENROUTER_API_KEY="sk-test", OPENROUTER_MANAGEMENT_KEY="sk-mgmt") +expect("a management key's history is account-wide, with no prompt counts, beside the balance", + picked(record, "scope", "hasPromptStats", "todayTotalTokens") + [remaining(record), openrouter.asked], + ["account", False, 15, 9.0, ["/activity=sk-mgmt", "/credits=sk-mgmt", "/key=sk-test"]], record) + +# A management key alone still collects the history and the balance, with no +# inference key's meter to show. +sandbox() +record = openrouter(OPENROUTER_MANAGEMENT_KEY="sk-mgmt", routes=with_routes(credits={"data": {"total_credits": 10, "total_usage": 4}})) +expect("a management key alone collects history and balance", + picked(record, "ready", "scope", "limits", "usageStatusText") + [remaining(record), openrouter.asked], + [True, "account", [], "", 6.0, ["/activity=sk-mgmt", "/credits=sk-mgmt"]], record) + +# The history is reused like a scan, and when /activity can't be reached, +# today's last answer stands in for it. +sandbox() +os.environ["OPENROUTER_MANAGEMENT_KEY"] = "sk-mgmt" +openrouter(OPENROUTER_API_KEY="sk-test") +reused = openrouter(OPENROUTER_API_KEY="sk-test", force=False) +expect("the activity history is reused", [[ask for ask in openrouter.asked if ask.startswith("/activity")], reused["scope"]], [[], "account"], reused) +with contextlib.redirect_stderr(io.StringIO()): + kept = openrouter(OPENROUTER_API_KEY="sk-test", routes=with_routes(activity=DOWN)) +expect("the activity history is kept past an unreachable check", picked(kept, "scope", "todayTotalTokens", "retryAdvised"), ["account", 15, True], kept) +del os.environ["OPENROUTER_MANAGEMENT_KEY"] + +# A key nothing was spent through, with no credits and no limit, is an +# unused machine: ready=false and no status from outside the vocabulary. +sandbox() +record = openrouter(OPENROUTER_API_KEY="sk-test", routes=with_routes(key={"data": {"limit": None}}, credits={"data": {"total_credits": 0, "total_usage": 0}})) +expect("a key with nothing to show is an unused machine", picked(record, "ready", "usageStatusText", "authHelpText"), [False, "", ""], record) + +# Without any key, what pi, omp and opencode spent through OpenRouter here +# still shows, waiting for a key for the balance and meter. +write_pi_session(home() / ".pi" / "agent" / "sessions" / "--p--" / "s.jsonl", [pi_reply("openrouter", 100, id="a1", model="anthropic/claude-test")]) +record = openrouter() +expect("local usage without a key makes the record ready and waits for auth", + picked(record, "ready", "usageStatusText", "tierLabel", "limits", "balance", "hasLocalStats", "totalPrompts") + [openrouter.asked], + [True, "Waiting for auth", "", [], None, True, 1, []], record) +expect("and says how to add one", record["authHelpText"].startswith("Set OPENROUTER_API_KEY"), True, record) + +# The key pi or omp saved with /login is used, a `$VAR` or `!command` entry +# is left to pi, and PI_CODING_AGENT_DIR moves pi's. +sandbox() +pi_auth(None, "openrouter", "sk-pi") +pi_auth(home() / ".omp" / "agent", "openrouter", "sk-omp") +pi_auth(home() / "moved", "openrouter", "$OPENROUTER_KEY") +saved = [module.api_key({})] +os.environ["PI_CODING_AGENT_DIR"] = str(home() / "moved") +saved.append(module.api_key({})) +del os.environ["PI_CODING_AGENT_DIR"] +record = openrouter() +expect("pi's and omp's saved OpenRouter key is used", saved + [openrouter.asked, record["tierLabel"]], + ["sk-pi", "sk-omp", ["/credits=sk-pi", "/key=sk-pi"], "Prepaid"], record) +opencode_auth(None, "openrouter", "sk-opencode") +record = openrouter() +expect("opencode's saved OpenRouter key comes before pi's", openrouter.asked, ["/credits=sk-opencode", "/key=sk-opencode"], record) + +# A "management" key that /activity refuses (an inference key, or one +# revoked) falls back to the local scan: machine numbers with prompt counts, +# and no sign-in asked for, since the inference key still answers, and reads +# the balance /credits refuses the other key. +sandbox() +write_pi_session(home() / ".pi" / "agent" / "sessions" / "--p--" / "s.jsonl", [pi_reply("openrouter", 100, id="a1", model="anthropic/claude-test")]) +for code in (401, 403): + def refused_manager(request, code=code): + if request.get_header("Authorization") == "Bearer sk-refused": + return FakeHttp.refusal(request.full_url, code, {"error": {"code": code, "message": "refused"}}) + return CREDITS + record = openrouter(OPENROUTER_API_KEY="sk-test", OPENROUTER_MANAGEMENT_KEY="sk-refused", routes=with_routes(activity=code, credits=refused_manager)) + expect(f"a management key /activity refuses ({code}) falls back to the local scan", + picked(record, "scope", "hasPromptStats", "totalPrompts", "todayTotalTokens", "usageStatusText", "retryAdvised") + [remaining(record), openrouter.asked], + [None, True, 1, 100, "", None, 9.0, ["/activity=sk-refused", "/credits=sk-refused", "/credits=sk-test", "/key=sk-test"]], record) + +# ---------------------------------------------------------- failed checks + +# Each key keeps its own meter: another key whose check fails never shows +# this one's numbers. +sandbox() +os.environ["OPENROUTER_API_KEY"] = "sk-test" +openrouter() +other = openrouter(mode=500, OPENROUTER_API_KEY="sk-other") +expect("another key never shows this key's meter, and says what failed", + picked(other, "limits", "usageStatusText", "retryAdvised", "authHelpText") + [len(limits_cache_files("openrouter"))], + [[], "OpenRouter limits unavailable", None, "OpenRouter's API returned status 500.", 1], other) + +# A refused key (403 as well as 401) is a sign-in to redo, never a raw +# status code, and not worth a sooner retry. +for code in (401, 403): + openrouter() + refused = openrouter(mode=code) + expect(f"a {code} keeps the last meter and asks to sign in again", + picked(refused, "usageStatusText", "limitsStale", "retryAdvised", "ready", "limits") + [refused["authHelpText"].startswith("OpenRouter refused the API key. ")], + ["Sign-in expired", True, None, True, fresh["limits"], True], refused) + +# OpenRouter documents /credits as management-key only, answering an +# inference key with its 403 example. That key is healthy: the meter shows, +# with no balance and nothing to sign in to again, and the refusal is kept, so +# later refreshes don't ask again. A management key, when there is one, reads +# the balance instead, the only /credits request. +def management_only(request): + if request.get_header("Authorization") == "Bearer sk-mgmt": + return CREDITS + return FakeHttp.refusal(request.full_url, 403, {"error": {"code": 403, "message": "Only management keys can perform this operation"}}) + + +sandbox() +record = openrouter(OPENROUTER_API_KEY="sk-test", routes=with_routes(credits=management_only)) +expect("an inference key /credits refuses still shows its meter, with no balance and no sign-in asked", + picked(record, "usageStatusText", "authHelpText", "limitsStale", "balance", "retryAdvised") + [len(record["limits"]), openrouter.asked], + ["", "", False, None, None, 1, ["/credits=sk-test", "/key=sk-test"]], record) +record = openrouter(OPENROUTER_API_KEY="sk-test", routes=with_routes(credits=management_only)) +expect("an inference key /credits refused isn't asked again on the next refresh", + picked(record, "usageStatusText", "limitsStale", "balance") + [len(record["limits"]), openrouter.asked], + ["", False, None, 1, ["/key=sk-test"]], record) +record = openrouter(OPENROUTER_API_KEY="sk-test", OPENROUTER_MANAGEMENT_KEY="sk-mgmt", + routes=with_routes(credits=management_only)) +expect("a management key reads the balance an inference key is refused", + picked(record, "usageStatusText", "limitsStale") + [remaining(record), openrouter.asked], + ["", False, 9.0, ["/activity=sk-mgmt", "/credits=sk-mgmt", "/key=sk-test"]], record) + +# A management key revoked after that (OpenRouter's 401 "User not found.") +# is a sign-in to redo, with its last good balance kept, stale, whether the +# inference key is refused /credits afresh or by its kept refusal. +def revoked_manager(request): + if request.get_header("Authorization") == "Bearer sk-mgmt": + return FakeHttp.refusal(request.full_url, 401, {"error": {"code": 401, "message": "User not found."}}) + return management_only(request) + + +sandbox() +openrouter(OPENROUTER_API_KEY="sk-test", OPENROUTER_MANAGEMENT_KEY="sk-mgmt", routes=with_routes(credits=management_only)) +for attempt in ("afresh", "by its kept refusal"): + record = openrouter(OPENROUTER_API_KEY="sk-test", OPENROUTER_MANAGEMENT_KEY="sk-mgmt", + routes=with_routes(activity=401, credits=revoked_manager)) + expect(f"a revoked management key asks to sign in again and keeps its balance, the inference key refused {attempt}", + picked(record, "usageStatusText", "limitsStale") + [remaining(record)], ["Sign-in expired", True, 9.0], record) + # The inference key still answers /key, so the help names the management + # key to replace, not OPENROUTER_API_KEY. + expect(f"a revoked management key's help names the management key, the inference key refused {attempt}", + [record["authHelpText"].startswith("OpenRouter refused the management key. "), "OPENROUTER_MANAGEMENT_KEY" in record["authHelpText"]], + [True, True], record) + +# A management key alone, revoked, names that key too; with the inference +# key refused as well, the help names the inference key. +record = openrouter(OPENROUTER_MANAGEMENT_KEY="sk-mgmt", routes=with_routes(activity=401, credits=revoked_manager)) +expect("a revoked management key with no inference key names the management key", + [record["usageStatusText"], record["authHelpText"].startswith("OpenRouter refused the management key. ")], ["Sign-in expired", True], record) +record = openrouter(OPENROUTER_API_KEY="sk-test", OPENROUTER_MANAGEMENT_KEY="sk-mgmt", routes=with_routes(activity=401, key=401, credits=revoked_manager)) +expect("both keys refused names the API key", + [record["usageStatusText"], record["authHelpText"].startswith("OpenRouter refused the API key. ")], ["Sign-in expired", True], record) + +# A key /credits answers, where OpenRouter still lets one read it, in a fresh +# home, since the refusal above is kept for sk-test. +sandbox() +os.environ["OPENROUTER_API_KEY"] = "sk-test" + +# A failed /credits keeps the last good balance, marked stale; with none +# kept, the balance says it is unavailable. +openrouter() +record = openrouter(routes=with_routes(credits=502)) +expect("a failed balance check keeps the last balance, stale", [remaining(record)] + picked(record, "limitsStale", "usageStatusText"), [9.0, True, ""], record) +for path in cache_root().glob("openrouter-credits-*.json"): + path.unlink() +record = openrouter(routes=with_routes(credits=502)) +expect("a failed balance check with nothing kept says so", picked(record, "balance", "limitsStale", "usageStatusText"), [None, True, "Balance unavailable"], record) + +# A /credits answer without both totals is a failed check, not an empty +# balance: it never overwrites the last good one, so a later failure still +# has it to keep. +openrouter() +record = openrouter(routes=with_routes(credits={"data": {"total_usage": 1}})) +expect("an incomplete balance answer keeps the last balance, stale", + [remaining(record)] + picked(record, "limitsStale", "usageStatusText", "retryAdvised"), [9.0, True, "", None], record) +record = openrouter(routes=with_routes(credits=502)) +expect("an incomplete balance answer never replaces the cached balance", [remaining(record), record["limitsStale"]], [9.0, True], record) + +# A /key answer without a readable allowance is a failed check too, not an +# unlimited key: the last meter stays, stale, and survives a later failure. +for incomplete in ({"data": {}}, {}, {"data": {"limit": 10}}, {"data": {"limit": 10, "limit_remaining": "lots"}}): + openrouter() + record = openrouter(routes=with_routes(key=incomplete)) + expect(f"an incomplete /key answer {incomplete} keeps the last meter, stale", + picked(record, "limits", "limitsStale", "usageStatusText"), [fresh["limits"], True, ""], record) + record = openrouter(routes=with_routes(key=502)) + expect(f"an incomplete /key answer {incomplete} never replaces the cached meter", picked(record, "limits", "limitsStale"), [fresh["limits"], True], record) + +# An explicitly unlimited key is a good answer that clears the meter. +openrouter() +record = openrouter(routes=with_routes(key={"data": {"limit": None, "limit_remaining": None}})) +expect("an unlimited key replaces the meter with none, fresh", picked(record, "limits", "limitsStale", "usageStatusText"), [[], False, ""], record) +# Claude Code pointed at OpenRouter's Anthropic-compatible API is +# OpenRouter's, counted from this machine even with no key. +sandbox() +claude_code_session("or", [("anthropic/claude-sonnet-4.5", 70), ("claude-sonnet-4-5", 40)]) +record = openrouter() +expect("OpenRouter counts no Claude Code reply while Claude Code is sent to Anthropic", [record["todayTokensByModel"], record["ready"]], [{}, False], record) +claude_code_session("or", [("anthropic/claude-sonnet-4.5", 70), ("claude-sonnet-4-5", 40)], base_url="https://openrouter.ai/api") +record = openrouter() +expect("OpenRouter counts every Claude Code reply while Claude Code is pointed at it", + [record["todayTokensByModel"], record["todayPrompts"], record["usageStatusText"]], [{"claude-sonnet-4.5": 70, "claude-sonnet-4-5": 40}, 2, "Waiting for auth"], record) + +# OpenRouter usage that other harnesses keep: DeepSeek Harness's pi-ai +# `openrouter` route (its DeepSeek routes are DeepSeek's), a Kimchi session +# run with `--provider openrouter` (its gateway routes are Kimchi's), and a +# MiniMax Code custom provider at openrouter.ai (one elsewhere isn't). +sandbox() +os.environ["DSH_HOME"] = str(home() / "dsh") +dsh_log = home() / "dsh" / "sessions" / "--home-me--" / "s1" / "session.v4.jsonl" +dsh_log.parent.mkdir(parents=True) +dsh_log.write_text("".join(json.dumps(line) + "\n" for line in ( + {"type": "session", "version": 4, "id": "s1", "createdAt": ms_in(0), "isSeeded": False}, + {"type": "user/message", "seq": 0, "time": ms_in(0), "data": {"source": {"kind": "user"}, "role": "user", "id": "u1"}}, + {"type": "request/header", "seq": 1, "time": ms_in(0), "data": {"header": {"config": {"provider": "openrouter", "model": "z-ai/glm-mock"}}}}, + {"type": "assistant/message", "seq": 2, "time": ms_in(0), "data": {"turn": 1, "step": 1, "usage": {"inputTokens": 5234, "outputTokens": 55, "cacheReadTokens": 75}}}, + {"type": "request/header", "seq": 3, "time": ms_in(0), "data": {"header": {"config": {"provider": "deepseek-official", "model": "deepseek-flash"}}}}, + {"type": "assistant/message", "seq": 4, "time": ms_in(0), "data": {"turn": 2, "step": 1, "usage": {"inputTokens": 999}}}, +))) +write_pi_session(home() / ".config" / "kimchi" / "harness" / "sessions" / "--home-me--" / "k.jsonl", [ + pi_reply("openrouter", id="k1", model="z-ai/glm-or-mock", usage={"input": 5234, "output": 55, "cacheRead": 75, "totalTokens": 5364}), + pi_reply("kimchi-dev/zai", id="k2", model="glm-mock", usage={"input": 777}), +]) +mcode_routed_session("m1", "custom_provider:router", "z-ai/glm-mcode", 2358, base_url="https://openrouter.ai/api/v1") +mcode_routed_session("m2", "custom_provider:acme", "acme", 444, base_url="https://llm.acme.example/v1") +record = openrouter() +expect("OpenRouter counts dsh's openrouter route, Kimchi's --provider openrouter, and MiniMax Code at openrouter.ai, nothing else of theirs", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], + [{"glm-mock": 5364, "glm-or-mock": 5364, "glm-mcode": 2358}, 3, 3], record) +del os.environ["DSH_HOME"] + +PY diff --git a/test/shell.d/agent-usage-partition-test.sh b/test/shell.d/agent-usage-partition-test.sh new file mode 100644 index 00000000000..0a16bf8b500 --- /dev/null +++ b/test/shell.d/agent-usage-partition-test.sh @@ -0,0 +1,2651 @@ +#!/bin/bash + +# Tokens are counted once, under the provider that bills them: pi, omp and +# opencode sessions are split by provider into the matching collector, and +# the `local` collector is the exact complement for self-hosted runtimes. Each +# collector's own test checks its own filter; this one runs every collector +# against one machine holding usage on every provider any of them claims, and +# checks the filters add up to a partition: each provider counted by exactly +# one collector, an unclaimed provider by none, and no copy counted twice. +# Every collector reads those shared sources the same way, so the fixture +# also holds each case where the copies once drifted: forks, compactions, a +# relocated opencode database, a moved pi tree, and a root nested in another. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command python3 +require_command rg +require_command timeout + +COLLECTOR_TIMEOUT=${AGENT_USAGE_PARTITION_TIMEOUT:-60} +TIMEOUT=$(command -v timeout) + +# Every fixture row is today's, and the collectors run as their own processes, +# where no frozen clock reaches. A run that crossed local midnight would see +# its rows land on yesterday, so one that would start that close waits for the +# new day instead. +seconds_to_midnight=$(( $(date -d 'tomorrow 00:00' +%s) - $(date +%s) )) +if (( seconds_to_midnight < 120 )); then + sleep $(( seconds_to_midnight + 1 )) +fi + +WORK=$(mktemp -d) +trap 'rm -rf "$WORK"' EXIT + +collectors=() +for collector in "$ROOT"/bin/omarchy-agent-usage-*; do + [[ ${collector##*/omarchy-agent-usage-} == "update" ]] && continue + collectors+=("${collector##*/omarchy-agent-usage-}") +done +(( ${#collectors[@]} > 0 )) || fail "agent usage collectors are found under bin/" + +# Only what the collectors need to read their sources: Python and rg. No +# agent CLI (codex, gh, omp, secret-tool, mise) can be found, so none is +# spawned, and nothing holds a credential. +mkdir -p "$WORK/bin" +ln -s "$(command -v python3)" "$WORK/bin/python3" +ln -s "$(command -v rg)" "$WORK/bin/rg" + +# Each run gets the environment below and nothing else: no API key from the +# developer's shell leaks in, and every HTTP(S) request goes to a proxy on a +# port nothing listens on, so a collector that does reach for the network +# fails at once instead of talking to a real service. +run_collector() { + local home="$1" id="$2" out="$3" status=0 + shift 3 + env -i \ + HOME="$home" \ + PATH="$WORK/bin" \ + LANG=C.UTF-8 \ + ${TZ+TZ="$TZ"} \ + OMARCHY_PATH="$ROOT" \ + OMARCHY_LANGUAGE_FILE="$OMARCHY_LANGUAGE_FILE" \ + XDG_CONFIG_HOME="$home/.config" \ + XDG_DATA_HOME="$home/.local/share" \ + XDG_STATE_HOME="$home/.local/state" \ + XDG_CACHE_HOME="$home/.cache" \ + XDG_RUNTIME_DIR="$home/.run" \ + http_proxy=http://127.0.0.1:9 https_proxy=http://127.0.0.1:9 \ + HTTP_PROXY=http://127.0.0.1:9 HTTPS_PROXY=http://127.0.0.1:9 \ + ALL_PROXY=http://127.0.0.1:9 \ + "$@" \ + "$TIMEOUT" --kill-after=5 "$COLLECTOR_TIMEOUT" "$ROOT/bin/omarchy-agent-usage-$id" --force \ + >"$out.json" 2>"$out.err" || status=$? + echo "$status" >"$out.status" +} + +# Runs every collector against one home, in parallel, into a results dir. +run_all() { + local home="$1" results="$2" id + shift 2 + mkdir -p "$results" "$home/.cache" "$home/.run" + for id in "${collectors[@]}"; do + run_collector "$home" "$id" "$results/$id" "$@" & + done + wait +} + +# ------------------------------------------------------------------ fixture + +# One machine with usage on every provider id a collector claims, in every +# store pi, omp and opencode keep: +# +# - pi and omp's default session roots, a profiles//agent/sessions root +# under each, and a pi tree PI_CODING_AGENT_DIR moves elsewhere, read as +# well as the default one. +# - A forked pi session repeating its parent's lines with their ids and +# timestamps, which count once. +# - A pi profile symlinked into the default root, so its sessions sit inside +# another root: its lines carry no id, so only dropping the nested root +# keeps them from counting twice. +# - Usage kept only as totalTokens, which still counts, and a line dated at +# the epoch's zero, which counts nowhere. +# - A compaction's usage, which names no provider and bills the model of the +# reply before it. +# - opencode's database with rows only in the original `message` store, rows +# only in the V2 `session_message` store, rows migrated into both under one +# id (counted once), compactions (counted like replies), thinking tokens +# (counted as output), a forked session whose copies of its source's +# messages count nowhere while its own reply does, and a row dated at zero, +# which counts nowhere either. +# - Codex rollouts on the built-in provider, on `--oss`'s local daemon (a +# cloud-tagged model and a local one), on a provider named `ollama-cloud`, +# on LM Studio, and on a custom provider nothing can place. +# - DeepSeek Harness logs on its DeepSeek routes, pi-ai's openrouter route +# and a hand-declared gateway; Kimchi harness sessions on its gateway, +# `--provider openrouter` and a custom provider; and MiniMax Code replies +# on its own route and on custom providers at MiniMax, OpenRouter, Z.ai, +# this machine, an unplaceable gateway, and one gone from its config. +# +# Every row carries a unique tag (zqNNN) in its model name and a +# recognizable token total, 100 * provider + source, so a collector's +# modelUsage says exactly which rows it counted and how many times. +FIXTURE_HOME="$WORK/home" +mkdir -p "$FIXTURE_HOME" + +python3 - "$FIXTURE_HOME" "$WORK/cells.json" <<'PY' +import json +import sqlite3 +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +cells_path = Path(sys.argv[2]) + +# (provider, model stem, api, owner on pi/omp, owner on opencode). An owner of +# None means nobody may count it; "-" means the harness never records it. +# +# - pi's `openai` is an OpenAI API key, which no collector bills; opencode's +# `openai` is the ChatGPT sign-in the Codex collector reads. +# - pi records an Antigravity sign-in as `google-antigravity`; opencode's +# Antigravity plugin records a `google` provider (`google`, `google-vertex`) +# with an `antigravity-` model. +PROVIDERS = [ + ("anthropic", "claude-pt", "", "claude", "claude"), + ("openai-codex", "gpt-pt", "openai-codex-responses", "codex", "-"), + ("openai", "gpt-pt", "openai-codex-responses", "codex", "-"), + ("openai", "gpt-pt", "", "-", "codex"), + ("opencode", "zen-pt", "", "opencode", "opencode"), + ("opencode-go", "go-pt", "", "opencode", "opencode"), + ("ollama-cloud", "cloud-pt", "", "ollama", "ollama"), + ("ollama", "tagged-pt@:cloud", "", "ollama", "ollama"), + ("ollama", "tagged-pt@:120b-cloud", "", "ollama", "ollama"), + ("ollama", "plain-pt@:8b", "", "local", "local"), + ("lmstudio", "plain-pt", "", "local", "local"), + ("deepseek", "deepseek-pt", "", "deepseek", "deepseek"), + ("zai", "glm-pt", "", "zai", "zai"), + ("zhipu", "glm-pt", "", "zai", "zai"), + ("minimax", "minimax-pt", "", "minimax", "minimax"), + ("minimax-cn", "minimax-pt", "", "minimax", "minimax"), + ("minimax-coding-plan", "minimax-pt", "", "minimax", "minimax"), + ("minimax-cn-coding-plan", "minimax-pt", "", "minimax", "minimax"), + ("openrouter", "router-pt", "", "openrouter", "openrouter"), + ("google-antigravity", "antigravity-pt", "", "agy", "-"), + ("google", "antigravity-pt", "", "-", "agy"), + ("google-vertex", "antigravity-pt", "", "-", "agy"), + ("devin", "devin-pt", "", "devin", "devin"), + ("synthetic", "synthetic-pt", "", "synthetic", "synthetic"), + ("mistral", "devstral-pt", "", "vibe", "vibe"), + ("meta", "muse-spark-pt", "", "muse", "muse"), + ("alibaba-coding-plan", "qwen-pt", "", "qwen", "qwen"), + ("alibaba-coding-plan-cn", "qwen-pt", "", "qwen", "qwen"), + ("alibaba-token-plan", "qwen-pt", "", "qwen", "qwen"), + ("alibaba-token-plan-cn", "qwen-pt", "", "qwen", "qwen"), + ("qwen-token-plan", "qwen-pt", "", "qwen", "qwen"), + ("qwen-token-plan-cn", "qwen-pt", "", "qwen", "qwen"), + ("qwen-token-plan-individual", "qwen-pt", "", "qwen", "qwen"), + ("qwen-portal", "qwen-pt", "", "qwen", "qwen"), + ("kimi-coding", "kimi-pt", "", "kimi", "kimi"), + ("kimi-code-plan-global", "kimi-pt", "", "kimi", "kimi"), + ("kimi-code-plan-cn", "kimi-pt", "", "kimi", "kimi"), + # Moonshot's pay-as-you-go API counts under Kimi, as Claude Code's turns + # to it do; Alibaba's pay-as-you-go keys carry no plan. + ("alibaba", "qwen-pt", "", None, None), + ("moonshotai", "kimi-pt", "", "kimi", "kimi"), + ("moonshotai-cn", "kimi-pt", "", "kimi", "kimi"), + ("acme-unclaimed", "acme-pt", "", None, None), +] + +PI_SOURCES = [ + ("pi", home / ".pi/agent/sessions/project"), + ("omp", home / ".omp/agent/sessions/project"), + ("pi-profile", home / ".pi/profiles/work/agent/sessions/project"), + ("omp-profile", home / ".omp/profiles/work/agent/sessions/project"), + # PI_CODING_AGENT_DIR=$HOME/pi-moved on every fixture run. + ("pi-moved", home / "pi-moved/sessions/project"), + # ~/.pi/profiles/nest links here, inside the default root. + ("pi-nested", home / ".pi/agent/sessions/nest-profile/agent/sessions/project"), + ("omp-total", home / ".omp/agent/sessions/total"), + ("omp-zero-stamp", home / ".omp/agent/sessions/zero"), + ("pi-compaction", home / ".pi/agent/sessions/compacted"), +] +OC_SOURCES = ["opencode-message", "opencode-session_message", "opencode-migrated", "opencode-compaction", + "opencode-reasoning", "opencode-fork-copy", "opencode-fork-reply", "opencode-zero-stamp", + "opencode-message-fork-copy", "opencode-message-fork-reply"] + +now = datetime.now(timezone.utc) +stamp = now.isoformat(timespec="milliseconds").replace("+00:00", "Z") +now_ms = int(time.time() * 1000) - 1000 +# The forked session began a minute ago: its copies of the source session's +# messages predate it, its own reply doesn't. +fork_ms = now_ms - 60_000 + +cells = [] +lines = {name: [] for name, _ in PI_SOURCES} +message_rows, session_rows = [], [] + + +def model_for(stem, tag): + # A stem of "name@:tag" keeps the Ollama tag after the unique marker, so a + # collector that strips tags still sees which row it counted. + base, _, suffix = stem.partition("@") + return f"{base}-{tag}{suffix}" + + +for p, (provider, stem, api, pi_owner, oc_owner) in enumerate(PROVIDERS, 1): + for s, (source, _) in enumerate(PI_SOURCES, 1): + if pi_owner == "-": + continue + tag = f"zq{len(cells):03d}" + tokens = 100 * p + s + usage = {"input": tokens, "output": 0, "cacheRead": 0, "cacheWrite": 0} + if source == "omp-total": + usage = {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": tokens} + message = {"role": "assistant", "provider": provider, "model": model_for(stem, tag), "usage": usage} + if api: + message["api"] = api + line = {"type": "message", "id": f"{source}-{tag}", "timestamp": 0 if source == "omp-zero-stamp" else stamp, "message": message} + if source == "pi-nested": + del line["id"] + if source == "pi-compaction": + # A compaction names no provider: it ran on the model of the reply + # before it, here one that spent nothing itself. + lines[source].append({**line, "id": f"{source}-{tag}-reply", "message": {**message, "usage": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0}}}) + line = {"type": "compaction", "id": f"{source}-{tag}", "timestamp": stamp, "summary": "s", "firstKeptEntryId": "x", "tokensBefore": tokens, "usage": usage} + lines[source].append(line) + owner = None if source == "omp-zero-stamp" else pi_owner + cells.append({"tag": tag, "provider": provider, "api": api, "source": source, "tokens": tokens, "owner": owner}) + for s, source in enumerate(OC_SOURCES, len(PI_SOURCES) + 1): + if oc_owner == "-": + continue + tag = f"zq{len(cells):03d}" + tokens = 100 * p + s + model = model_for(stem, tag) + counts = {"input": tokens, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}} + if source == "opencode-reasoning": + counts = {"input": tokens - 7, "output": 0, "reasoning": 7, "cache": {"read": 0, "write": 0}} + created = 0 if source == "opencode-zero-stamp" else now_ms + v1 = json.dumps({"role": "assistant", "providerID": provider, "modelID": model, "tokens": counts, "time": {"created": created}}) + v2 = json.dumps({"model": {"providerID": provider, "id": model}, "tokens": counts, "time": {"created": now_ms}}) + message_id = f"msg_{tag}" + if source in ("opencode-message", "opencode-migrated", "opencode-reasoning", "opencode-zero-stamp"): + message_rows.append((message_id, f"ses_{source}", now_ms, now_ms, v1)) + if source in ("opencode-session_message", "opencode-migrated"): + session_rows.append((message_id, f"ses_{source}", "assistant", now_ms, v2)) + if source == "opencode-compaction": + session_rows.append((message_id, f"ses_{source}", "compaction", now_ms, v2)) + if source == "opencode-fork-copy": + session_rows.append((message_id, "ses_fork", "assistant", fork_ms - 5_000, v2)) + if source == "opencode-fork-reply": + session_rows.append((message_id, "ses_fork", "assistant", now_ms, v2)) + if source == "opencode-message-fork-copy": + # As opencode 1.19.3's `run --fork` writes one into the original store: + # a plain new session, the copy under a new id, its time the source's. + message_rows.append((message_id, "ses_message_fork", now_ms, now_ms, json.dumps({**json.loads(v1), "time": {"created": fork_ms - 5_000}}))) + if source == "opencode-message-fork-reply": + message_rows.append((message_id, "ses_message_fork", now_ms, now_ms, v1)) + owner = None if source in ("opencode-fork-copy", "opencode-zero-stamp", "opencode-message-fork-copy") else oc_owner + cells.append({"tag": tag, "provider": provider, "api": "", "source": source, "tokens": tokens, "owner": owner}) + +for source, directory in PI_SOURCES: + directory.mkdir(parents=True, exist_ok=True) + with (directory / "session.jsonl").open("w") as handle: + for line in lines[source]: + handle.write(json.dumps(line) + "\n") + +# A fork of the default pi session, in another project: its parent's lines +# again, under the same ids and timestamps. +fork = home / ".pi/agent/sessions/forked" +fork.mkdir(parents=True) +with (fork / "session.jsonl").open("w") as handle: + for line in lines["pi"]: + handle.write(json.dumps(line) + "\n") + +# Relative, so a copy of this home keeps pointing into itself. +(home / ".pi/profiles/nest").symlink_to("../agent/sessions/nest-profile") + +# The schema the opencode collector's own tests build: the original store and +# the V2 one side by side, as a database opencode has migrated looks. +db = home / ".local/share/opencode/opencode.db" +db.parent.mkdir(parents=True, exist_ok=True) +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL, time_created integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE TABLE session_v2 (id text PRIMARY KEY, fork_session_id text, time_created integer NOT NULL)") +conn.execute("INSERT INTO session_v2 VALUES ('ses_fork', 'ses_opencode-session_message', ?)", (fork_ms,)) +conn.execute("CREATE TABLE session (id text PRIMARY KEY, parent_id text, title text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL)") +conn.execute("INSERT INTO session VALUES ('ses_message_fork', NULL, 'Mock reply (fork #1)', ?, ?)", (fork_ms, fork_ms)) +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", message_rows) +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, ?, ?)", session_rows) +conn.commit() +conn.close() + + +# Codex rollouts, as Codex 0.163 writes them: the built-in `openai` provider +# is the Codex collector's, `--oss` on the local daemon's `ollama` provider is +# the local collector's unless the model is cloud-tagged (Ollama Cloud's), a +# provider named `ollama-cloud` is Ollama Cloud's, LM Studio is local, and a +# custom provider nothing can place is nobody's. +CODEX_ROLLOUTS = [ + ("openai", "gpt-codex@", "codex"), + ("ollama", "gpt-oss@:120b-cloud", "ollama"), + ("ollama-cloud", "qwen-codex@", "ollama"), + ("ollama", "plain-codex@:20b", "local"), + ("lmstudio", "plain-codex@", "local"), + ("acme-unclaimed", "acme-codex@", None), +] +rollouts = home / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") +rollouts.mkdir(parents=True) +for n, (provider, stem, owner) in enumerate(CODEX_ROLLOUTS): + tag = f"zq{len(cells):03d}" + tokens = 80000 + 100 * n + base, _, suffix = stem.partition("@") + lines = [ + {"timestamp": stamp, "type": "session_meta", "payload": {"model_provider": provider}}, + {"timestamp": stamp, "type": "turn_context", "payload": {"model": f"{base}-{tag}{suffix}"}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "user_message", "message": "go"}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": tokens, "output_tokens": 0}, "total_token_usage": {"input_tokens": tokens}}}}, + ] + (rollouts / f"rollout-{tag}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + cells.append({"tag": tag, "provider": provider, "api": "", "source": "codex-rollout", "tokens": tokens, "owner": owner}) + + +# DeepSeek Harness session logs, one per route, as dsh 0.2 writes them (left +# uncompressed): its DeepSeek routes and pi-ai's `deepseek` are DeepSeek's, +# pi-ai's `openrouter` OpenRouter's, and a hand-declared gateway nobody's. +DSH_ROUTES = [ + ("deepseek-official", "deepseek-dsh@", "deepseek"), + ("deepseek", "deepseek-piai@", "deepseek"), + ("openrouter", "router-dsh@", "openrouter"), + ("acme-gateway", "acme-dsh@", None), +] +for n, (provider, stem, owner) in enumerate(DSH_ROUTES): + tag = f"zq{len(cells):03d}" + tokens = 81000 + 100 * n + log = home / f".dsh/sessions/--home-me--/session-{tag}/session.v4.jsonl" + log.parent.mkdir(parents=True) + log.write_text("".join(json.dumps(record) + "\n" for record in ( + {"type": "session", "version": 4, "id": f"dsh-{tag}", "createdAt": now_ms, "isSeeded": False}, + {"type": "request/header", "seq": 1, "time": now_ms, "data": {"header": {"config": {"provider": provider, "model": model_for(stem, tag)}}}}, + {"type": "assistant/message", "seq": 2, "time": now_ms, "data": {"turn": 1, "step": 1, "usage": {"inputTokens": tokens}}}, + ))) + cells.append({"tag": tag, "provider": provider, "api": "", "source": "dsh", "tokens": tokens, "owner": owner}) + +# Kimchi's harness sessions, in pi's format: its own gateway routes are +# Kimchi's, a session run with `--provider openrouter` OpenRouter's, and one +# on a custom provider nobody's. +kimchi_lines = [] +for n, (provider, stem, owner) in enumerate([ + ("kimchi-dev/zai", "glm-kimchi@", "kimchi"), + ("openrouter", "router-kimchi@", "openrouter"), + ("my-gateway", "acme-kimchi@", None), +]): + tag = f"zq{len(cells):03d}" + tokens = 82000 + 100 * n + kimchi_lines.append({"type": "message", "id": f"kimchi-{tag}", "timestamp": stamp, + "message": {"role": "assistant", "provider": provider, "model": model_for(stem, tag), "usage": {"input": tokens}}}) + cells.append({"tag": tag, "provider": provider, "api": "", "source": "kimchi-harness", "tokens": tokens, "owner": owner}) +kimchi_sessions = home / ".config/kimchi/harness/sessions/--home-me--" +kimchi_sessions.mkdir(parents=True) +(kimchi_sessions / "session.jsonl").write_text("".join(json.dumps(line) + "\n" for line in kimchi_lines)) + +# MiniMax Code sessions, one per route, and the custom providers its +# config.yaml declares: its own routes are MiniMax's, a custom provider +# whatever its baseURL's host bills (MiniMax, OpenRouter, Z.ai, this +# machine), and one at a gateway nothing can place, or gone from the config, +# nobody's. +MCODE_ROUTES = [ + ("minimax", "MiniMax-mcode@", "minimax"), + ("custom_provider:mm", "MiniMax-custom@", "minimax"), + ("custom_provider:router", "router-mcode@", "openrouter"), + ("custom_provider:zed", "glm-mcode@", "zai"), + ("custom_provider:box", "plain-mcode@", "local"), + ("custom_provider:acme", "acme-mcode@", None), + ("custom_provider:gone", "acme-gone@", None), +] +(home / ".minimax").mkdir() +(home / ".minimax/config.yaml").write_text("""logLevel: info +custom_provider: + mm: + name: mm + kind: custom + options: + apiKey: sk-placeholder + baseURL: https://api.minimax.io/anthropic + router: + name: router + api: openai-completions + options: + apiKey: sk-placeholder + baseURL: https://openrouter.ai/api/v1 + zed: + options: + baseURL: "https://api.z.ai/api/coding/paas/v4" + box: + options: + baseURL: http://127.0.0.1:8080/v1 # llama.cpp + acme: + options: + baseURL: https://llm.acme.example/v1 +""") +for n, (provider, stem, owner) in enumerate(MCODE_ROUTES): + tag = f"zq{len(cells):03d}" + tokens = 83000 + 100 * n + session = home / f".minimax/v2/sessions/2026/01/01/{tag}/messages.jsonl" + session.parent.mkdir(parents=True) + session.write_text("".join(json.dumps(line) + "\n" for line in ( + {"message": {"role": "user", "timestamp": now_ms, "content": "go"}}, + {"message": {"role": "assistant", "timestamp": now_ms, "provider": provider, "model": model_for(stem, tag), "usage": {"input": tokens}}}, + ))) + cells.append({"tag": tag, "provider": provider, "api": "", "source": "mcode", "tokens": tokens, "owner": owner}) + + +# Codex 0.163 rollouts with a remote compaction, as it writes one: a +# token_usage_record for every response, the compaction's included, where the +# token_count after the compaction spends nothing and the next one counts only +# the reply. A record repeated under the same id and running total is one +# response. +CODEX_RECORD_ROLLOUTS = [ + ("openai", "gpt-record@", "codex"), + ("ollama-cloud", "qwen-record@", "ollama"), + ("lmstudio", "plain-record@", "local"), +] +for n, (provider, stem, owner) in enumerate(CODEX_RECORD_ROLLOUTS): + tag = f"zq{len(cells):03d}" + compaction, reply = 84000 + 100 * n, 1000 + n + first = {"input_tokens": compaction, "output_tokens": 0} + thread = {"input_tokens": compaction + reply, "output_tokens": 0} + usage_record = {"timestamp": stamp, "type": "token_usage_record", "payload": { + "turn_id": "turn-1", "response_id": "resp_mock_0002", "usage": {"input_tokens": reply, "output_tokens": 0}, "thread_token_usage": thread}} + lines = [ + {"timestamp": stamp, "type": "session_meta", "payload": {"model_provider": provider, "source": "exec"}}, + {"timestamp": stamp, "type": "turn_context", "payload": {"model": model_for(stem, tag)}}, + {"timestamp": stamp, "type": "token_usage_record", "payload": { + "turn_id": "turn-1", "response_id": "resp_mock_0001", "usage": first, "thread_token_usage": first}}, + {"timestamp": stamp, "type": "compacted", "payload": {"message": ""}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": 0, "output_tokens": 0}, "total_token_usage": {}}}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "item_completed", "item": {"type": "UserMessage", "id": "u1", "content": []}}}, + usage_record, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": reply, "output_tokens": 0}, "total_token_usage": {"input_tokens": reply, "output_tokens": 0}}}}, + usage_record, + ] + (rollouts / f"rollout-{tag}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + cells.append({"tag": tag, "provider": provider, "api": "", "source": "codex-record-rollout", "tokens": compaction + reply, "owner": owner}) + +cells_path.write_text(json.dumps(cells)) +PY + +EMPTY_HOME="$WORK/empty" +mkdir -p "$EMPTY_HOME" + +# The empty home runs in a zone whose day differs from the system zone's, so +# a collector dating its record by /etc/localtime rather than the caller's TZ +# shows. +14 and -12 hours are never both on the system's day; the days on +# either side of the run cover one that straddles that zone's midnight. +EMPTY_TZ=Pacific/Kiritimati +if [[ $(TZ=$EMPTY_TZ date +%F) == "$(env -u TZ date +%F)" ]]; then + EMPTY_TZ=Etc/GMT+12 +fi +TZ=$EMPTY_TZ date +%F >"$WORK/empty-days" +TZ=$EMPTY_TZ run_all "$EMPTY_HOME" "$WORK/empty-results" +TZ=$EMPTY_TZ date +%F >>"$WORK/empty-days" +run_all "$FIXTURE_HOME" "$WORK/results" PI_CODING_AGENT_DIR="$FIXTURE_HOME/pi-moved" + +# The same machine with opencode's database moved by OPENCODE_DB. A relative +# name resolves beside the default one, as opencode itself resolves it. +RELOCATED_HOME="$WORK/relocated" +cp -a "$FIXTURE_HOME" "$RELOCATED_HOME" +rm -rf "$RELOCATED_HOME/.cache" +mkdir -p "$RELOCATED_HOME/.local/share/opencode/moved" +mv "$RELOCATED_HOME/.local/share/opencode/opencode.db" "$RELOCATED_HOME/.local/share/opencode/moved/elsewhere.db" +run_all "$RELOCATED_HOME" "$WORK/relocated-results" PI_CODING_AGENT_DIR="$RELOCATED_HOME/pi-moved" OPENCODE_DB=moved/elsewhere.db + +# A local.json naming every claimed provider (and a self-hosted box of its +# own) must not pull any of them into the local collector. Fireworks counts +# its whole account rather than these rows, so no cell owns its pi turns, but +# naming fireworks or fireworks-ai in local.json must not count them twice. +CONFIGURED_HOME="$WORK/configured" +cp -a "$FIXTURE_HOME" "$CONFIGURED_HOME" +rm -rf "$CONFIGURED_HOME/.cache" +mkdir -p "$CONFIGURED_HOME/.config/omarchy/agents" "$CONFIGURED_HOME/.pi/agent/sessions/fireworks" +python3 - "$WORK/cells.json" "$CONFIGURED_HOME/.config/omarchy/agents/local.json" "$CONFIGURED_HOME/.pi/agent/sessions/fireworks/session.jsonl" <<'PY' +import json +import sys +from datetime import datetime, timezone + +cells = json.load(open(sys.argv[1])) +providers = sorted({cell["provider"] for cell in cells if cell["owner"] and cell["owner"] != "local"} | {"fireworks", "fireworks-ai"}) +json.dump({"providers": providers}, open(sys.argv[2], "w")) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +with open(sys.argv[3], "w") as handle: + for provider in ("fireworks", "fireworks-ai"): + message = {"role": "assistant", "provider": provider, "model": f"{provider}-account-pt", "usage": {"input": 77, "output": 0, "cacheRead": 0, "cacheWrite": 0}} + handle.write(json.dumps({"type": "message", "id": f"fw-{provider}", "timestamp": stamp, "message": message}) + "\n") +PY +mkdir -p "$CONFIGURED_HOME/.cache" "$CONFIGURED_HOME/.run" +run_collector "$CONFIGURED_HOME" local "$WORK/configured-local" PI_CODING_AGENT_DIR="$CONFIGURED_HOME/pi-moved" + +# OpenRouter reads no local usage at all without a key, so the keyless run +# above can only show it missing everything. A placeholder key, with every +# request going to the dead proxy, shows what its own filter counts. +run_collector "$FIXTURE_HOME" openrouter "$WORK/keyed-openrouter" OPENROUTER_API_KEY=sk-or-placeholder PI_CODING_AGENT_DIR="$FIXTURE_HOME/pi-moved" + + +# Claude Code fronts whatever ANTHROPIC_BASE_URL names, and its transcripts +# record only the model that answered each reply. One home holds a reply by +# each kind of model (and a fork's copy of one), and every collector runs +# against it once per endpoint Claude Code can be sent to; each reply must +# land in exactly the collector that bills it there, or in none. +CLAUDE_HOME="$WORK/claude-code" +mkdir -p "$CLAUDE_HOME" +python3 - "$CLAUDE_HOME" "$WORK/claude-cells.json" "$WORK/claude-endpoints" <<'PY' +import json +import re +import sys +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +# (model, owner with no endpoint, owner behind a self-hosted one, owner +# behind a gateway nothing can place). A provider's own host owns them all. +MODELS = [ + ("claude-sonnet-zq900", "claude", "claude", "claude"), + ("glm-zq901", "zai", "local", None), + ("MiniMax-zq902", "minimax", "local", None), + ("deepseek-zq903", "deepseek", "local", None), + ("kimi-zq904", "kimi", "local", None), + ("moonshot-zq905", "kimi", "local", None), + ("gpt-oss-zq906:120b-cloud", "ollama", "ollama", "ollama"), + ("qwen3-zq907", None, "local", None), +] +ENDPOINTS = { + "": 1, "https://api.z.ai/api/anthropic": "zai", "https://ollama.com": "ollama", "https://openrouter.ai/api": "openrouter", + "https://api.minimax.io/anthropic": "minimax", "https://api.deepseek.com/anthropic": "deepseek", + "https://api.kimi.com/coding/": "kimi", "http://localhost:11434": 2, "https://llm-proxy.example.com/anthropic": 3, +} +projects = home / ".claude/projects/project" +projects.mkdir(parents=True) +lines, cells = [], [] +for n, (model, *owners) in enumerate(MODELS): + tokens = 90000 + 100 * n + lines.append({"type": "user", "uuid": f"prompt-{n}", "sessionId": "main", "timestamp": stamp, "message": {"role": "user", "content": "go"}}) + lines.append({"type": "assistant", "sessionId": "main", "timestamp": stamp, "requestId": f"req-{n}", + "message": {"id": f"msg-{n}", "role": "assistant", "model": model, "usage": {"input_tokens": tokens, "output_tokens": 0}}}) + for index, (url, owner) in enumerate(ENDPOINTS.items()): + cells.append({"endpoint": index, "url": url, "tag": re.search(r"zq\d{3}", model).group(0), "model": model, "tokens": tokens, + "owner": owners[owner - 1] if isinstance(owner, int) else owner}) +(projects / "main.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) +# A fork copies the first reply into a session of its own. +(projects / "fork.jsonl").write_text(json.dumps({**lines[1], "sessionId": "fork"}) + "\n") +Path(sys.argv[2]).write_text(json.dumps(cells)) +Path(sys.argv[3]).write_text("".join(url + "\n" for url in ENDPOINTS)) +PY +mapfile -t claude_endpoints <"$WORK/claude-endpoints" +for endpoint in "${!claude_endpoints[@]}"; do + url=${claude_endpoints[endpoint]} + if [[ -n $url ]]; then + run_all "$CLAUDE_HOME" "$WORK/claude-results-$endpoint" ANTHROPIC_BASE_URL="$url" + else + run_all "$CLAUDE_HOME" "$WORK/claude-results-$endpoint" + fi +done + +# ----------------------------------------------------------------- analysis + +python3 - "$WORK" "${collectors[@]}" >"$WORK/report.tsv" <<'PY' +import json +import re +import sys +from pathlib import Path + +work = Path(sys.argv[1]) +collectors = sys.argv[2:] +cells = {cell["tag"]: cell for cell in json.load(open(work / "cells.json"))} +empty_days = set((work / "empty-days").read_text().split()) + +STATS = ["todayPrompts", "todaySessions", "todayTotalTokens", "todayTokensByModel", "recentDays", + "totalPrompts", "totalSessions", "activeDays", "activeDates", "modelUsage"] +# The fields every record carries, per the contract's record table and the +# reference collector. retryAdvised, hasPromptStats, scope, balance and +# accounts are conditional and checked only for shape when present. +REQUIRED = ["schemaVersion", "id", "name", "updatedAt", "ready", "tierLabel", "usageStatusText", + "authHelpText", "limits", "limitsStale", "limitsFetchedAt", "hasLocalStats"] + STATS +BUCKET = ["inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens"] + +results = [] + + +def report(check, ok, detail=""): + results.append((check, ok, detail)) + + +def load(prefix): + status = (work / f"{prefix}.status").read_text().strip() + raw = (work / f"{prefix}.json").read_text() + err = (work / f"{prefix}.err").read_text().strip() + if status != "0": + return None, f"exited {status}{' (timed out)' if status in ('124', '137') else ''}: {err[-400:]}" + try: + record = json.loads(raw) + except ValueError: + return None, f"printed no JSON record: {raw[:200]!r} {err[-200:]}" + if not isinstance(record, dict): + return None, f"printed a non-object: {raw[:200]!r}" + return record, "" + + +def shape_problems(cid, record): + problems = [f"missing {key}" for key in REQUIRED if key not in record] + if record.get("schemaVersion") != 1: + problems.append(f"schemaVersion {record.get('schemaVersion')!r}") + if record.get("id") != cid: + problems.append(f"id {record.get('id')!r}") + for key in ("ready", "limitsStale", "hasLocalStats"): + if key in record and not isinstance(record[key], bool): + problems.append(f"{key} is not a boolean") + for key in ("retryAdvised", "hasPromptStats"): + if key in record and not isinstance(record[key], bool): + problems.append(f"{key} is not a boolean") + if "scope" in record and record["scope"] != "account": + problems.append(f"scope {record['scope']!r}") + for key in ("name", "updatedAt", "tierLabel", "usageStatusText", "authHelpText"): + if key in record and not isinstance(record[key], str): + problems.append(f"{key} is not a string") + if "limits" in record and not isinstance(record["limits"], list): + problems.append("limits is not a list") + if "recentDays" in record and (not isinstance(record["recentDays"], list) or len(record["recentDays"]) != 7): + problems.append("recentDays is not seven days") + if "modelUsage" in record and isinstance(record["modelUsage"], dict): + for model, bucket in record["modelUsage"].items(): + if not isinstance(bucket, dict) or any(key not in bucket for key in BUCKET): + problems.append(f"modelUsage[{model}] lacks the token buckets") + elif "modelUsage" in record: + problems.append("modelUsage is not an object") + return problems + + +# 1. Every collector in an empty home: the unused-machine record. +for cid in collectors: + record, error = load(f"empty-results/{cid}") + if record is None: + report("empty-home", False, f"{cid}: {error}") + continue + problems = shape_problems(cid, record) + if record.get("ready") is not False: + problems.append(f"ready {record.get('ready')!r}") + if record.get("usageStatusText") != "": + problems.append(f"usageStatusText {record.get('usageStatusText')!r} (authHelpText {record.get('authHelpText')!r})") + elif record.get("authHelpText") != "": + problems.append(f"authHelpText {record.get('authHelpText')!r} with no status") + if record.get("limits"): + problems.append("limits not empty") + if record.get("tierLabel") != "": + problems.append(f"tierLabel {record.get('tierLabel')!r}") + if "balance" in record: + problems.append("has a balance") + if any(record.get(key) for key in ("todayPrompts", "todaySessions", "todayTotalTokens", "totalPrompts", "totalSessions", "activeDays")) \ + or record.get("todayTokensByModel") or record.get("modelUsage") or record.get("activeDates"): + problems.append("stats not empty") + if record.get("retryAdvised"): + problems.append("retryAdvised with nothing to retry") + days = record.get("recentDays") or [{}] + if days[-1].get("date") not in empty_days: + problems.append(f"recentDays ends {days[-1].get('date')!r}, not the caller's TZ day ({', '.join(sorted(empty_days))})") + report("empty-home", not problems, f"{cid}: {'; '.join(problems)}" if problems else cid) + + +def counted(record): + """tag -> tokens a record counted for that fixture row, and keys matching none.""" + per_tag, strays = {}, [] + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in BUCKET) if isinstance(bucket, dict) else 0 + tags = re.findall(r"zq\d{3}", model) + if len(tags) != 1 or tags[0] not in cells: + if total: + strays.append(f"{model}={total}") + continue + per_tag[tags[0]] = per_tag.get(tags[0], 0) + total + return per_tag, strays + + +# 2. The partition. +counts = {tag: {} for tag in cells} +collector_totals = {} +fixture_counts = {} +for cid in collectors: + record, error = load(f"results/{cid}") + if record is None: + report("fixture-run", False, f"{cid}: {error}") + continue + problems = shape_problems(cid, record) + report("fixture-shape", not problems, f"{cid}: {'; '.join(problems)}" if problems else cid) + per_tag, strays = counted(record) + fixture_counts[cid] = per_tag + for tag, total in per_tag.items(): + counts[tag][cid] = total + if strays: + report("attribution", False, f"{cid}: modelUsage keys that name no fixture row: {', '.join(strays)}") + today = record.get("todayTotalTokens") or 0 + collector_totals[cid] = today + model_sum = sum(per_tag.values()) + if today != model_sum: + report("today-matches-models", False, f"{cid}: todayTotalTokens {today} but modelUsage over fixture rows sums to {model_sum}") + else: + report("today-matches-models", True, cid) + +claimed_total = 0 +for tag, cell in sorted(cells.items()): + owner, tokens = cell["owner"], cell["tokens"] + label = f"{cell['provider']}{' api=' + cell['api'] if cell['api'] else ''} via {cell['source']} ({tag}, {tokens} tokens)" + seen = counts[tag] + if owner: + claimed_total += tokens + wrong = [] + for cid, total in sorted(seen.items()): + times = total / tokens + if cid != owner: + wrong.append(f"{cid} counted it" + (f" {times:g}x" if times != 1 else "")) + elif total != tokens: + wrong.append(f"{cid} counted it {times:g}x") + if owner is None: + report("unclaimed", not wrong, f"{label}: {'; '.join(wrong)}" if wrong else label) + continue + if owner not in seen: + wrong.insert(0, f"{owner} missed it") + report("partition", not wrong, f"{label}, expected {owner}: {'; '.join(wrong)}" if wrong else label) + +grand = sum(collector_totals.values()) +report("sum", grand == claimed_total, f"claimed providers total {claimed_total}, collectors' todayTotalTokens sum to {grand}" + + "".join(f"\n {cid}={total}" for cid, total in sorted(collector_totals.items()) if total)) + +# 3. The database moved by OPENCODE_DB counts exactly the rows it did in place. +for cid in collectors: + record, error = load(f"relocated-results/{cid}") + if record is None: + report("opencode-db-env", False, f"{cid}: {error}") + continue + if cid not in fixture_counts: + continue + per_tag, _ = counted(record) + before = fixture_counts[cid] + changed = [] + for tag in sorted(set(per_tag) | set(before)): + if per_tag.get(tag, 0) != before.get(tag, 0): + cell = cells[tag] + changed.append(f"{cell['provider']} via {cell['source']} ({tag}) {before.get(tag, 0)} -> {per_tag.get(tag, 0)}") + report("opencode-db-env", not changed, f"{cid}: {'; '.join(changed)}" if changed else cid) + +# 4. A local.json naming every claimed provider. +record, error = load("configured-local") +if record is None: + report("local-config", False, f"local: {error}") +else: + per_tag, strays = counted(record) + leaked = sorted({f"{cells[tag]['provider']}" for tag in per_tag if cells[tag]["owner"] != "local"} | + {stray.partition("-account-pt=")[0] for stray in strays if "-account-pt=" in stray}) + missed = sorted({cells[tag]["provider"] + " via " + cells[tag]["source"] for tag in cells if cells[tag]["owner"] == "local" and tag not in per_tag}) + problems = ([f"counts claimed providers it was configured with: {', '.join(leaked)}"] if leaked else []) + \ + ([f"missed {', '.join(missed)}"] if missed else []) + report("local-config", not problems, f"local: {'; '.join(problems)}" if problems else "local") + +# 5. OpenRouter with a key: its own filter, once each, and nothing else. +record, error = load("keyed-openrouter") +if record is None: + report("openrouter-keyed", False, f"openrouter: {error}") +else: + per_tag, strays = counted(record) + problems = [f"stray {', '.join(strays)}"] if strays else [] + for tag, cell in sorted(cells.items()): + total = per_tag.get(tag, 0) + label = f"{cell['provider']} via {cell['source']} ({tag})" + if cell["owner"] == "openrouter" and total != cell["tokens"]: + problems.append(f"{label} counted {total / cell['tokens']:g}x") + elif cell["owner"] != "openrouter" and total: + problems.append(f"{label} counted, owned by {cell['owner']}") + report("openrouter-keyed", not problems, f"openrouter: {'; '.join(problems)}" if problems else "openrouter") + +# 6. Claude Code's replies, behind every endpoint it can be sent to. +claude_cells = json.load(open(work / "claude-cells.json")) +for endpoint in sorted({cell["endpoint"] for cell in claude_cells}): + cells_here = {cell["tag"]: cell for cell in claude_cells if cell["endpoint"] == endpoint} + url = next(iter(cells_here.values()))["url"] or "Anthropic (no ANTHROPIC_BASE_URL)" + seen = {tag: {} for tag in cells_here} + for cid in collectors: + record, error = load(f"claude-results-{endpoint}/{cid}") + if record is None: + report("claude-code", False, f"{url}: {cid}: {error}") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in BUCKET) if isinstance(bucket, dict) else 0 + for tag in re.findall(r"zq\d{3}", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total + for tag, cell in sorted(cells_here.items()): + owner, tokens = cell["owner"], cell["tokens"] + wrong = [f"{cid} counted it {total / tokens:g}x" for cid, total in sorted(seen[tag].items()) if cid != owner or total != tokens] + if owner and owner not in seen[tag]: + wrong.insert(0, f"{owner} missed it") + label = f"{cell['model']} sent to {url}, expected {owner or 'nobody'}" + report("claude-code", not wrong, f"{label}: {'; '.join(wrong)}" if wrong else label) + +for check, ok, detail in results: + print(f"{check}\t{'ok' if ok else 'fail'}\t{detail}".replace("\n", "\\n")) +PY + +# ------------------------------------------------------------------ verdict + +declare -A DESCRIPTIONS=( + [empty-home]="every collector prints the unused-machine record in an empty home: ready false, no plan, status or help, no limits, no balance, empty stats, the full field set, dated by the caller's TZ" + [fixture-run]="every collector prints a record against the fixture home, offline and in time" + [fixture-shape]="every collector's fixture record has the contract's full field set" + [attribution]="every token a collector counts comes from a fixture row" + [today-matches-models]="each collector's todayTotalTokens equals the fixture rows in its modelUsage" + [partition]="each claimed provider's tokens are counted by exactly its collector, once, from every root and store, forks and nesting included" + [unclaimed]="a provider no collector claims, a fork's copies of its source's messages, and rows dated at zero are counted by nobody" + [sum]="the claimed providers' tokens equal the collectors' todayTotalTokens summed" + [opencode-db-env]="with opencode's database moved by a relative OPENCODE_DB, every collector counts the same rows" + [local-config]="a local.json naming claimed providers doesn't pull them into the local collector" + [openrouter-keyed]="with a key, the OpenRouter collector counts exactly the OpenRouter rows, once each" + [claude-code]="each Claude Code reply is counted once, by the collector that bills it behind every endpoint Claude Code can be sent to, or by none" +) +order=(empty-home fixture-run fixture-shape attribution today-matches-models partition unclaimed sum opencode-db-env local-config openrouter-keyed claude-code) + +failed=() +details="" +for check in "${order[@]}"; do + lines=$(awk -F'\t' -v c="$check" '$1 == c && $2 == "fail" { print " " $3 }' "$WORK/report.tsv") + if [[ -z $lines ]]; then + pass "${DESCRIPTIONS[$check]}" + else + failed+=("$check") + details+="${DESCRIPTIONS[$check]}:"$'\n'"${lines//\\n/$'\n'}"$'\n' + fi +done + +if (( ${#failed[@]} > 0 )); then + fail "agent usage collectors partition tokens by provider (failing: ${failed[*]})" "$details" +fi + +# ------------------------------------------------------- harness provider ids + +# pi, omp and opencode each name a plan's provider their own way, so a plan +# a collector meters can reach it under several ids. One home holds a reply +# on each id below, written as pi 1.4.0, omp 18.12.6 and opencode 1.19.3 +# wrote them for that provider, and each must be counted once, by its owner +# (Moonshot's pay-as-you-go API by Kimi). +IDS_HOME="$WORK/provider-ids" +mkdir -p "$IDS_HOME" +python3 - "$IDS_HOME" "$WORK/ids-cells.json" <<'PY' +import json +import sqlite3 +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +now_ms = int(time.time() * 1000) - 1000 +# (harness, provider, owner) +IDS = [ + ("pi", "zai-coding-cn", "zai"), + ("omp", "zhipu-coding-plan", "zai"), + ("omp", "kimi-code", "kimi"), + ("omp", "muse-code", "muse"), + ("omp", "moonshot", "kimi"), + ("opencode", "zai-coding-plan", "zai"), + ("opencode", "zhipuai-coding-plan", "zai"), + ("opencode", "zhipuai", "zai"), +] +cells = [] +rows = [] +for n, (harness, provider, owner) in enumerate(IDS): + tag = f"zq{700 + n}" + model = f"ids-{tag}" + tokens = 70000 + 100 * n + usage = {"input": tokens, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": tokens, "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "total": 0}} + if harness == "opencode": + rows.append((f"msg_{tag}", f"ses_{tag}", now_ms, now_ms, json.dumps({ + "parentID": f"msg_{tag}_u", "role": "assistant", "mode": "build", "agent": "build", "cost": 0, + "tokens": {"total": tokens, "input": tokens, "output": 0, "reasoning": 0, "cache": {"write": 0, "read": 0}}, + "modelID": model, "providerID": provider, "time": {"created": now_ms, "completed": now_ms}, "finish": "stop"}))) + elif harness == "omp": + runtime = {"version": "18.12.6"} + lines = [ + {"type": "session", "version": 3, "id": f"s-{tag}", "timestamp": stamp, "cwd": "/home/me"}, + {"type": "model_change", "id": f"mc-{tag}", "parentId": None, "timestamp": stamp, "runtime": runtime, "model": f"{provider}/{model}", "resolvedModelIsFallback": False}, + {"type": "message", "id": f"u-{tag}", "parentId": f"mc-{tag}", "timestamp": stamp, "runtime": runtime, "message": {"role": "user", "content": [{"type": "text", "text": "Say hi"}], "attribution": "user", "timestamp": now_ms}}, + {"type": "message", "id": f"a-{tag}", "parentId": f"u-{tag}", "timestamp": stamp, "runtime": runtime, "message": { + "role": "assistant", "content": [{"type": "text", "text": "hi"}], "api": "openai-completions", "provider": provider, "model": model, + "usage": usage, "stopReason": "stop", "timestamp": now_ms}}, + ] + path = home / f".omp/agent/sessions/-/{tag}.jsonl" + else: + lines = [ + {"type": "session", "version": 3, "id": f"s-{tag}", "timestamp": stamp, "cwd": "/home/me"}, + {"type": "model_change", "id": f"mc-{tag}", "parentId": None, "timestamp": stamp, "provider": provider, "modelId": model}, + {"type": "message", "id": f"u-{tag}", "parentId": f"mc-{tag}", "timestamp": stamp, "message": {"role": "user", "content": [{"type": "text", "text": "Say hi"}], "timestamp": now_ms}}, + {"type": "message", "id": f"a-{tag}", "parentId": f"u-{tag}", "timestamp": stamp, "message": { + "role": "assistant", "content": [{"type": "text", "text": "hi"}], "api": "openai-completions", "provider": provider, "model": model, + "usage": usage, "stopReason": "stop", "timestamp": now_ms}}, + ] + path = home / f".pi/agent/sessions/--home-me--/{tag}.jsonl" + if harness != "opencode": + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(line) + "\n" for line in lines)) + cells.append({"tag": tag, "harness": harness, "provider": provider, "tokens": tokens, "owner": owner}) + +db = home / ".local/share/opencode/opencode.db" +db.parent.mkdir(parents=True) +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", rows) +conn.commit() +conn.close() +Path(sys.argv[2]).write_text(json.dumps(cells)) +PY +run_all "$IDS_HOME" "$WORK/ids-results" + +python3 - "$WORK" "${collectors[@]}" >"$WORK/ids-report.txt" <<'PY' +import json +import re +import sys +from pathlib import Path + +work = Path(sys.argv[1]) +cells = {cell["tag"]: cell for cell in json.load(open(work / "ids-cells.json"))} +seen = {tag: {} for tag in cells} +for cid in sys.argv[2:]: + try: + record = json.loads((work / f"ids-results/{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no record: {(work / f'ids-results/{cid}.err').read_text()[-300:]}") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zq\d{3}", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total +for tag, cell in sorted(cells.items()): + owner, tokens = cell["owner"], cell["tokens"] + wrong = [f"{cid} counted it {total / tokens:g}x" for cid, total in sorted(seen[tag].items()) if cid != owner or total != tokens] + if owner and owner not in seen[tag]: + wrong.insert(0, f"{owner} missed it") + if wrong: + print(f"{cell['harness']}'s {cell['provider']} ({tag}), expected {owner or 'nobody'}: {'; '.join(wrong)}") +PY +if [[ -s $WORK/ids-report.txt ]]; then + fail "each harness's own id for a plan is counted once, by the collector that meters it" "$(cat "$WORK/ids-report.txt")" +fi +pass "each harness's own id for a plan is counted once, by the collector that meters it" + +# ----------------------------------------------------- harness stores' routes + +# Kimi Code, Qwen Code, Hermes and Mistral Vibe keep stores of their own and +# record where each request went: each request on a route counts once, by +# the harness's own collector when its service billed it, by `local` on a +# server you run, and by nobody on any other provider or a gateway, since no +# other collector reads these stores. +HARNESS_HOME="$WORK/harnesses" +mkdir -p "$HARNESS_HOME" +python3 - "$HARNESS_HOME" "$WORK/harness-cells.json" <<'PY' +import json +import sqlite3 +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +now = time.time() +now_ms = int(now * 1000) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +cells = [] + + +def cell(source, route, owner): + tag = f"zq{900 + len(cells)}" + tokens = 91000 + 100 * len(cells) + cells.append({"tag": tag, "source": source, "route": route, "tokens": tokens, "owner": owner}) + return tag, tokens + + +# Kimi Code: a model alias per route in config.toml, one usage.record each. +KIMI_ROUTES = [ + ("plan", 'type = "kimi"\nbase_url = "https://api.kimi.com/coding/v1"', "kimi"), + ("box", 'type = "kimi"\nbase_url = "http://192.168.1.20:8000/v1"', "local"), + ("anth", 'type = "anthropic"', None), + ("acme", 'type = "openai"\nbase_url = "https://llm.acme.example/v1"', None), +] +config, wire = [], [] +for name, provider, owner in KIMI_ROUTES: + tag, tokens = cell("kimi-code", name, owner) + config.append(f'[providers.{name}]\n{provider}\napi_key = "k"\n\n[models."{name}/m-{tag}"]\nprovider = "{name}"\nmodel = "m-{tag}"\nmax_context_size = 1000\n') + wire.append({"type": "usage.record", "agentId": "main", "model": f"{name}/m-{tag}", "usage": {"inputOther": tokens, "output": 0, "inputCacheRead": 0, "inputCacheCreation": 0}, "usageScope": "turn", "time": now_ms}) +(home / ".kimi-code").mkdir() +(home / ".kimi-code" / "config.toml").write_text("\n".join(config)) +path = home / ".kimi-code" / "sessions" / "wd_p_1" / "session_1" / "agents" / "main" / "wire.jsonl" +path.parent.mkdir(parents=True) +path.write_text("".join(json.dumps(line) + "\n" for line in wire)) + +# Qwen Code: settings.json's modelProviders say where each model goes. +QWEN_ROUTES = [ + ("openai", "https://coding-intl.dashscope.aliyuncs.com/v1", "qwen"), + ("openai", None, None), + ("openai", "http://localhost:1234/v1", "local"), + ("openai", "https://openrouter.ai/api/v1", None), + ("anthropic", None, None), +] +providers, lines = {}, [] +for auth, url, owner in QWEN_ROUTES: + tag, tokens = cell("qwen-code", f"{auth} {url}", owner) + if url: + providers.setdefault(auth, []).append({"id": f"q-{tag}", "baseUrl": url}) + lines.append({"uuid": tag, "sessionId": "s", "timestamp": stamp, "type": "system", "subtype": "ui_telemetry", "systemPayload": {"uiEvent": { + "event.name": "qwen-code.api_response", "event.timestamp": stamp, "response_id": tag, "model": f"q-{tag}", "auth_type": auth, + "input_token_count": tokens, "output_token_count": 0, "cached_content_token_count": 0}}}) +(home / ".qwen" / "projects" / "p" / "chats").mkdir(parents=True) +(home / ".qwen" / "settings.json").write_text(json.dumps({"modelProviders": providers})) +(home / ".qwen" / "projects" / "p" / "chats" / "s.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + +# Hermes: each session and request names the base URL it was billed through. +(home / ".hermes").mkdir() +conn = sqlite3.connect(home / ".hermes" / "state.db") +conn.executescript(""" +CREATE TABLE sessions (id TEXT PRIMARY KEY, model TEXT, started_at REAL NOT NULL, billing_provider TEXT, billing_base_url TEXT); +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, billing_provider TEXT NOT NULL DEFAULT '', billing_base_url TEXT NOT NULL DEFAULT '', + 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); +""") +HERMES_ROUTES = [ + ("nous", "https://inference-api.nousresearch.com/v1", "hermes"), + ("custom", "http://127.0.0.1:8080/v1", "local"), + ("openrouter", "https://openrouter.ai/api/v1", None), +] +for provider, url, owner in HERMES_ROUTES: + tag, tokens = cell("hermes", provider, owner) + conn.execute("INSERT INTO sessions VALUES (?, ?, ?, ?, ?)", (tag, f"h-{tag}", now, provider, url)) + conn.execute("INSERT INTO session_model_usage (session_id, model, billing_provider, billing_base_url, input_tokens) VALUES (?, ?, ?, ?, ?)", (tag, f"h-{tag}", provider, url, tokens)) +conn.commit() +conn.close() + +# Mistral Vibe: a legacy session per provider, its config dumped beside it. +VIBE_ROUTES = [ + ("mistral", {"name": "mistral", "backend": "mistral"}, "vibe"), + ("llamacpp", {"name": "llamacpp", "api_base": "http://127.0.0.1:8080/v1", "backend": "generic"}, "local"), + ("proxy", {"name": "proxy", "api_base": "https://llm.acme.example/v1", "backend": "generic"}, None), +] +for n, (provider, entry, owner) in enumerate(VIBE_ROUTES): + tag, tokens = cell("vibe", provider, owner) + session = home / ".vibe" / "logs" / "session" / f"session_20260101_12000{n}_{tag}" + session.mkdir(parents=True) + (session / "meta.json").write_text(json.dumps({"session_id": session.name, "start_time": stamp, "config": { + "active_model": f"v-{tag}", "models": {f"v-{tag}": {"name": f"v-{tag}", "provider": provider, "alias": f"v-{tag}"}}, "providers": [entry]}, + "stats": {"session_prompt_tokens": tokens, "session_cached_tokens": 0, "session_completion_tokens": 0}})) + +Path(sys.argv[2]).write_text(json.dumps(cells)) +PY +run_all "$HARNESS_HOME" "$WORK/harness-results" + +harness_report=$(python3 - "$WORK" "${collectors[@]}" <<'PY' +import json +import re +import sys +from pathlib import Path + +work = Path(sys.argv[1]) +cells = {cell["tag"]: cell for cell in json.load(open(work / "harness-cells.json"))} +seen = {tag: {} for tag in cells} +for cid in sys.argv[2:]: + try: + record = json.loads((work / "harness-results" / f"{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no JSON record") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zq9\d\d", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total +for tag, cell in sorted(cells.items()): + owner, tokens = cell["owner"], cell["tokens"] + wrong = [f"{cid} counted it {total / tokens:g}x" for cid, total in sorted(seen[tag].items()) if cid != owner or total != tokens] + if owner and owner not in seen[tag]: + wrong.insert(0, f"{owner} missed it") + if wrong: + print(f"{cell['source']} on {cell['route']} ({tag}), expected {owner or 'nobody'}: {'; '.join(wrong)}") +PY +) +if [[ -z $harness_report ]]; then + pass "each Kimi Code, Qwen Code, Hermes and Vibe request is counted once, by the collector that bills its recorded route, or by none" +else + fail "each Kimi Code, Qwen Code, Hermes and Vibe request is counted once, by the collector that bills its recorded route, or by none" "$harness_report" +fi + +# ------------------------------------------- Claude Code sent to Model Studio + +# Alibaba's Model Studio documents pointing Claude Code at its +# Anthropic-compatible API (international, mainland and the Coding Plan's), +# the hosts Qwen Code's replies are billed under qwen from. Every reply in +# the Claude Code home above, whatever model answered, is then qwen's, once, +# and no other collector's. +ALIBABA_ENDPOINTS=( + https://dashscope-intl.aliyuncs.com/apps/anthropic + https://dashscope.aliyuncs.com/apps/anthropic + https://coding.dashscope.aliyuncs.com/apps/anthropic +) +for endpoint in "${!ALIBABA_ENDPOINTS[@]}"; do + run_all "$CLAUDE_HOME" "$WORK/alibaba-results-$endpoint" ANTHROPIC_BASE_URL="${ALIBABA_ENDPOINTS[endpoint]}" +done + +alibaba_report=$(python3 - "$WORK" "${ALIBABA_ENDPOINTS[@]}" -- "${collectors[@]}" <<'PY' +import json +import re +import sys +from pathlib import Path + +work = Path(sys.argv[1]) +split = sys.argv.index("--") +endpoints, collectors = sys.argv[2:split], sys.argv[split + 1:] +tokens = {cell["tag"]: cell["tokens"] for cell in json.load(open(work / "claude-cells.json"))} +for index, url in enumerate(endpoints): + seen = {tag: {} for tag in tokens} + for cid in collectors: + try: + record = json.loads((work / f"alibaba-results-{index}" / f"{cid}.json").read_text()) + except ValueError: + print(f"{url}: {cid} printed no JSON record") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zq\d{3}", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total + for tag, expected in sorted(tokens.items()): + wrong = [f"{cid} counted it {total / expected:g}x" for cid, total in sorted(seen[tag].items()) if cid != "qwen" or total != expected] + if "qwen" not in seen[tag]: + wrong.insert(0, "qwen missed it") + if wrong: + print(f"{tag} sent to {url}, expected qwen: {'; '.join(wrong)}") +PY +) +if [[ -z $alibaba_report ]]; then + pass "each Claude Code reply sent to Model Studio's Anthropic API is counted once, by qwen" +else + fail "each Claude Code reply sent to Model Studio's Anthropic API is counted once, by qwen" "$alibaba_report" +fi + +# ------------------------------------------------- local.json and claimed ids + +# The local collector drops every id another collector claims from what +# local.json names (CLAIMED_PROVIDERS), so the claimer counts its rows and +# local never does. That list is kept by hand: each collector's own provider +# ids (its PROVIDERS, PROVIDER or NATIVE_PROVIDER, the exact id in its pi/omp +# pattern and its opencode prefilter) must all be in it. +claimed_report=$(python3 - "$ROOT" <<'PY' +import ast +import re +import sys +from pathlib import Path + +bin_dir = Path(sys.argv[1]) / "bin" + +def literals(node): + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return [node.value] + if isinstance(node, (ast.Tuple, ast.List, ast.Set)): + return [value for item in node.elts for value in literals(item)] + return [] + +def assigned(tree, names): + return [node.value for node in tree.body if isinstance(node, ast.Assign) and any(getattr(target, "id", "") in names for target in node.targets)] + +local = ast.parse((bin_dir / "omarchy-agent-usage-local").read_text()) +claimed = {pid for value in assigned(local, ("CLAIMED_PROVIDERS",)) for pid in literals(value)} +missing = [] if claimed else ["CLAIMED_PROVIDERS is not found in omarchy-agent-usage-local"] +for path in sorted(bin_dir.glob("omarchy-agent-usage-*")): + cid = path.name.removeprefix("omarchy-agent-usage-") + if cid in ("local", "update"): + continue + tree = ast.parse(path.read_text()) + ids = {pid for value in assigned(tree, ("PROVIDERS", "PROVIDER", "NATIVE_PROVIDER")) for pid in literals(value)} + for node in ast.walk(tree): + if isinstance(node, ast.Constant) and isinstance(node.value, str): + ids.update(re.findall(r'"provider"\\s\*:\\s\*"([a-z0-9-]+)"', node.value)) + if isinstance(node, ast.Call) and getattr(node.func, "id", "") == "scan_opencode" and node.args: + ids.update(value[:-1] for value in literals(node.args[0]) if value.endswith('"')) + missing += [f"{pid} ({cid})" for pid in sorted(ids - claimed)] +if missing: + print("CLAIMED_PROVIDERS lacks " + ", ".join(missing)) +PY +) +if [[ -z $claimed_report ]]; then + pass "local's CLAIMED_PROVIDERS holds every provider id another collector claims" +else + fail "local's CLAIMED_PROVIDERS holds every provider id another collector claims" "$claimed_report" +fi + +# A local.json naming each harness's own id for a plan leaves every row to +# the collector that meters it: counted once, never by local too. +CLAIMED_IDS_HOME="$WORK/provider-ids-configured" +cp -a "$IDS_HOME" "$CLAIMED_IDS_HOME" +rm -rf "$CLAIMED_IDS_HOME/.cache" +mkdir -p "$CLAIMED_IDS_HOME/.config/omarchy/agents" +python3 - "$WORK/ids-cells.json" "$CLAIMED_IDS_HOME/.config/omarchy/agents/local.json" <<'PY' +import json +import sys + +cells = json.load(open(sys.argv[1])) +json.dump({"providers": sorted({cell["provider"] for cell in cells})}, open(sys.argv[2], "w")) +PY +mkdir -p "$CLAIMED_IDS_HOME/.cache" "$CLAIMED_IDS_HOME/.run" "$WORK/ids-configured-results" +for id in local kimi muse zai; do + run_collector "$CLAIMED_IDS_HOME" "$id" "$WORK/ids-configured-results/$id" & +done +wait + +claimed_ids_report=$(python3 - "$WORK" local kimi muse zai <<'PY' +import json +import re +import sys +from pathlib import Path + +work = Path(sys.argv[1]) +cells = {cell["tag"]: cell for cell in json.load(open(work / "ids-cells.json"))} +seen = {tag: {} for tag in cells} +for cid in sys.argv[2:]: + try: + record = json.loads((work / f"ids-configured-results/{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no record: {(work / f'ids-configured-results/{cid}.err').read_text()[-300:]}") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zq\d{3}", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total +for tag, cell in sorted(cells.items()): + owner, tokens = cell["owner"], cell["tokens"] + wrong = [f"{cid} counted it {total / tokens:g}x" for cid, total in sorted(seen[tag].items()) if cid != owner or total != tokens] + if owner not in seen[tag]: + wrong.insert(0, f"{owner} missed it") + if wrong: + print(f"{cell['harness']}'s {cell['provider']} ({tag}), expected {owner}: {'; '.join(wrong)}") +PY +) +if [[ -z $claimed_ids_report ]]; then + pass "a local.json naming a plan's harness ids leaves each row to the collector that meters it, counted once" +else + fail "a local.json naming a plan's harness ids leaves each row to the collector that meters it, counted once" "$claimed_ids_report" +fi + +# ------------------------------------------------- Claude Code cost-state + +# A cost-state line's totals hold the requests no reply line records (a +# /compact, the session title): what it holds beyond the replies counts once, +# under the collector that bills that model behind the endpoint, as the +# replies do. +COST_HOME="$WORK/claude-cost" +mkdir -p "$COST_HOME" +python3 - "$COST_HOME" "$WORK/cost-cells.json" <<'PY' +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +# (model, reply tokens, beyond them, owner with no endpoint, owner behind Z.ai) +MODELS = [("claude-sonnet-zq920", 1387, 3001387, "claude", "zai"), ("glm-zq921", 2390, 24387, "zai", "zai"), ("claude-haiku-zq922", 0, 917, "claude", "zai")] +lines, usage = [], {} +for n, (model, replied, beyond, *_) in enumerate(MODELS): + lines.append({"type": "user", "uuid": f"cost-prompt-{n}", "sessionId": "cost", "timestamp": stamp, "message": {"role": "user", "content": "go"}}) + if replied: + lines.append({"type": "assistant", "sessionId": "cost", "timestamp": stamp, "requestId": f"req-cost-{n}", + "message": {"id": f"msg-cost-{n}", "role": "assistant", "model": model, "usage": {"input_tokens": replied, "output_tokens": 0}}}) + usage[model + "[1m]"] = {"inputTokens": replied + beyond, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0, "costUSD": 0} +lines.append({"type": "cost-state", "sessionId": "cost", "startTime": 1792669864166, "modelUsage": usage}) +projects = home / ".claude/projects/project" +projects.mkdir(parents=True) +(projects / "cost.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) +Path(sys.argv[2]).write_text(json.dumps([{"tag": model.split("-")[-1], "tokens": replied + beyond, "owners": owners} for model, replied, beyond, *owners in MODELS])) +PY +run_all "$COST_HOME" "$WORK/cost-results-0" +run_all "$COST_HOME" "$WORK/cost-results-1" ANTHROPIC_BASE_URL="https://api.z.ai/api/anthropic" + +cost_report=$(python3 - "$WORK" "${collectors[@]}" <<'PY' +import json +import re +import sys +from pathlib import Path + +work = Path(sys.argv[1]) +cells = json.load(open(work / "cost-cells.json")) +for endpoint in (0, 1): + seen = {cell["tag"]: {} for cell in cells} + for cid in sys.argv[2:]: + try: + record = json.loads((work / f"cost-results-{endpoint}" / f"{cid}.json").read_text()) + except ValueError: + print(f"endpoint {endpoint}: {cid} printed no JSON record") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zq\d{3}", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total + for cell in cells: + owner = cell["owners"][endpoint] + wrong = [f"{cid} counted {total} of {cell['tokens']}" for cid, total in sorted(seen[cell["tag"]].items()) if cid != owner or total != cell["tokens"]] + if owner not in seen[cell["tag"]]: + wrong.insert(0, f"{owner} missed it") + if wrong: + print(f"{cell['tag']} at endpoint {endpoint}, expected {owner}: {'; '.join(wrong)}") +PY +) +if [[ -z $cost_report ]]; then + pass "what a Claude Code cost-state holds beyond the replies is counted once, by the collector that bills its model" +else + fail "what a Claude Code cost-state holds beyond the replies is counted once, by the collector that bills its model" "$cost_report" +fi + +# Codex rollouts on a model_providers entry, as Codex 0.163 writes them: the +# config.toml base_url says who served them, so one at a provider's host is +# that provider's collector's whatever the entry's id (Z.ai, OpenRouter, +# DeepSeek, Moonshot, Model Studio, MiniMax, ollama.com), one on the user's +# own hardware is local's, and one at a host nothing can place nobody's. A +# provider gone from the config keeps the base_url it was last seen with. +CODEX_ROUTES_HOME="$WORK/codex-routes" +mkdir -p "$CODEX_ROUTES_HOME/.codex" "$CODEX_ROUTES_HOME/.local/state/omarchy/agents" +cat >"$CODEX_ROUTES_HOME/.codex/config.toml" <<'TOML' +model = "gpt-5.1-codex" +[model_providers.zai] +name = "zai" +base_url = "https://api.z.ai/api/coding/paas/v4" +wire_api = "responses" +[model_providers.openrouter] +name = "openrouter" +base_url = "https://openrouter.ai/api/v1" +wire_api = "responses" +[model_providers.dsk] +base_url = "https://api.deepseek.com/v1" +[model_providers.moon] +base_url = "https://api.moonshot.ai/v1" +[model_providers.dash] +base_url = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1" +[model_providers.mm] +base_url = "https://api.minimax.io/v1" +[model_providers.ollama] +base_url = "https://ollama.com/v1" +[model_providers.box] +base_url = "http://192.168.1.20:8080/v1" +[model_providers.acme] +base_url = "https://llm.acme.example/v1" +TOML +python3 - "$CODEX_ROUTES_HOME" "$WORK/codex-routes-cells.json" <<'PY' +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +(home / ".local/state/omarchy/agents/local-codex-providers.json").write_text(json.dumps( + {"homes": {str(home / ".codex"): {"gone": "https://api.z.ai/api/paas/v4"}}})) +ROUTES = [ + ("openai", "gpt-5.1-codex", "codex"), + ("zai", "glm-4.6", "zai"), + ("openrouter", "z-ai/glm-4.6", "openrouter"), + ("dsk", "deepseek-chat", "deepseek"), + ("moon", "kimi-k2-0905-preview", "kimi"), + ("dash", "qwen3-coder-plus", "qwen"), + ("mm", "MiniMax-M2", "minimax"), + ("ollama", "gpt-oss:120b", "ollama"), + ("box", "qwen3-coder:30b", "local"), + ("acme", "acme-coder", None), + ("gone", "glm-4.5", "zai"), +] +rollouts = home / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") +rollouts.mkdir(parents=True) +cells = [] +for n, (provider, model, owner) in enumerate(ROUTES): + tag = f"zr{n:03d}" + base, colon, suffix = model.partition(":") + name = f"{base}-{tag}{colon}{suffix}" + thread = f"01a16422-7868-7633-a566-{n:012d}" + first = {"input_tokens": 8234 + n, "cached_input_tokens": 71, "cache_write_input_tokens": 31, "output_tokens": 51, "reasoning_output_tokens": 0, "total_tokens": 8285 + n} + second = {"input_tokens": 9234 + n, "cached_input_tokens": 72, "cache_write_input_tokens": 32, "output_tokens": 52, "reasoning_output_tokens": 0, "total_tokens": 9286 + n} + total = {key: first[key] + second[key] for key in first} + lines = [ + {"timestamp": stamp, "type": "session_meta", "payload": {"session_id": thread, "id": thread, "originator": "codex_exec", "cli_version": "0.163.0", "source": "exec", "thread_source": "user", "model_provider": provider}}, + {"timestamp": stamp, "type": "turn_context", "payload": {"model": name}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "item_completed", "thread_id": thread, "item": {"type": "UserMessage", "id": f"u{n}", "content": [{"type": "text", "text": "go"}]}}}, + {"timestamp": stamp, "type": "token_usage_record", "payload": {"thread_id": thread, "response_id": "resp_mock_0001", "usage": first, "thread_token_usage": first}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": first, "last_token_usage": first}}}, + {"timestamp": stamp, "type": "token_usage_record", "payload": {"thread_id": thread, "response_id": "resp_mock_0002", "usage": second, "thread_token_usage": total}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": total, "last_token_usage": second}}}, + ] + (rollouts / f"rollout-{tag}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + cells.append({"tag": tag, "provider": provider, "tokens": total["input_tokens"] + total["output_tokens"], "owner": owner}) +json.dump(cells, open(sys.argv[2], "w")) +PY + +codex_routes_report() { + python3 - "$WORK/codex-routes-cells.json" "$@" <<'PY' +import json +import re +import sys +from pathlib import Path + +cells = {cell["tag"]: cell for cell in json.load(open(sys.argv[1]))} +results = Path(sys.argv[2]) +seen = {tag: {} for tag in cells} +for cid in sys.argv[3:]: + try: + record = json.loads((results / f"{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no record: {(results / f'{cid}.err').read_text()[-300:]}") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zr\d{3}", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total +for tag, cell in sorted(cells.items()): + owner, tokens = cell["owner"], cell["tokens"] + wrong = [f"{cid} counted it {total / tokens:g}x" for cid, total in sorted(seen[tag].items()) if cid != owner or total != tokens] + if owner is not None and owner not in seen[tag]: + wrong.insert(0, f"{owner} missed it") + if wrong: + print(f"Codex provider {cell['provider']} ({tag}), expected {owner}: {'; '.join(wrong)}") +PY +} + +run_all "$CODEX_ROUTES_HOME" "$WORK/codex-routes-results" +codex_routes=$(codex_routes_report "$WORK/codex-routes-results" "${collectors[@]}") +if [[ -z $codex_routes ]]; then + pass "a Codex model_providers entry counts once, under the collector its base_url's host bills" +else + fail "a Codex model_providers entry counts once, under the collector its base_url's host bills" "$codex_routes" +fi + +# Moving a provider's base_url is a new scan: a cached one that counted it +# under Z.ai no longer stands once config.toml sends it elsewhere. +mkdir -p "$WORK/codex-routes-cached" +run_cached_zai() { + env -i HOME="$CODEX_ROUTES_HOME" PATH="$WORK/bin" LANG=C.UTF-8 ${TZ+TZ="$TZ"} OMARCHY_PATH="$ROOT" OMARCHY_LANGUAGE_FILE="$OMARCHY_LANGUAGE_FILE" \ + XDG_CONFIG_HOME="$CODEX_ROUTES_HOME/.config" XDG_DATA_HOME="$CODEX_ROUTES_HOME/.local/share" XDG_STATE_HOME="$CODEX_ROUTES_HOME/.local/state" \ + XDG_CACHE_HOME="$CODEX_ROUTES_HOME/.cache" XDG_RUNTIME_DIR="$CODEX_ROUTES_HOME/.run" \ + http_proxy=http://127.0.0.1:9 https_proxy=http://127.0.0.1:9 HTTP_PROXY=http://127.0.0.1:9 HTTPS_PROXY=http://127.0.0.1:9 ALL_PROXY=http://127.0.0.1:9 \ + "$TIMEOUT" --kill-after=5 "$COLLECTOR_TIMEOUT" "$ROOT/bin/omarchy-agent-usage-zai" --limits-only \ + >"$WORK/codex-routes-cached/zai.json" 2>"$WORK/codex-routes-cached/zai.err" || true +} +run_cached_zai +sed -i 's#https://api.z.ai/api/coding/paas/v4#https://llm.acme.example/zai/v1#' "$CODEX_ROUTES_HOME/.codex/config.toml" +run_cached_zai +if codex_cached=$(python3 - "$WORK/codex-routes-cached/zai.json" <<'PY' +import json +import sys + +try: + usage = json.load(open(sys.argv[1])).get("modelUsage") or {} +except ValueError: + usage = {"no record": {}} +models = sorted(model for model in usage if "zr" in model or model == "no record") +if models != ["glm-4.5-zr010"]: + print(f"zai counted {models}, expected only glm-4.5-zr010") +PY +) && [[ -z $codex_cached ]]; then + pass "moving a Codex provider's base_url off Z.ai rescans rather than keeping a cached count" +else + fail "moving a Codex provider's base_url off Z.ai rescans rather than keeping a cached count" "$codex_cached" +fi + +# --force reads every Codex rollout again in each collector that bills one, +# even one rewritten to the same size and mtime its per-file cache keeps. +python3 - "$CODEX_ROUTES_HOME" <<'PY' +import os +import sys +from pathlib import Path + +for rollout in sorted((Path(sys.argv[1]) / ".codex" / "sessions").rglob("rollout-zr*.jsonl")): + stat = rollout.stat() + rollout.write_text(rollout.read_text().replace('"output_tokens": 51,', '"output_tokens": 91,').replace('"output_tokens": 103,', '"output_tokens": 143,')) + os.utime(rollout, ns=(stat.st_atime_ns, stat.st_mtime_ns)) +PY +run_all "$CODEX_ROUTES_HOME" "$WORK/codex-routes-forced" +if codex_forced=$(python3 - "$WORK/codex-routes-cells.json" "$WORK/codex-routes-forced" "${collectors[@]}" <<'PY' +import json +import re +import sys +from pathlib import Path + +# Z.ai's coding endpoint moved off its host above, so that rollout is nobody's now. +cells = {cell["tag"]: {**cell, "tokens": cell["tokens"] + 40, "owner": None if cell["tag"] == "zr001" else cell["owner"]} for cell in json.load(open(sys.argv[1]))} +routed = {"zai", "openrouter", "deepseek", "kimi", "qwen", "minimax", "ollama"} +for cid in sys.argv[3:]: + if cid not in routed: + continue + try: + record = json.loads((Path(sys.argv[2]) / f"{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no record") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zr\d{3}", model): + if tag in cells and cells[tag]["owner"] == cid and total != cells[tag]["tokens"]: + print(f"{cid} counted {tag} as {total}, expected the rewritten {cells[tag]['tokens']}") +PY +) && [[ -z $codex_forced ]]; then + pass "--force rereads a Codex rollout rewritten to the same size and mtime in every collector that bills one" +else + fail "--force rereads a Codex rollout rewritten to the same size and mtime in every collector that bills one" "$codex_forced" +fi + +# ------------------------------------- Claude Code and Codex sent to Synthetic + +# Synthetic documents Claude Code on its Anthropic-compatible API +# (ANTHROPIC_BASE_URL=https://api.synthetic.new/anthropic, with hf:org/model +# ids and syn: aliases as the models) and an OpenAI-compatible one at +# /openai/v1. Behind that base URL every reply is synthetic's, a Claude-named +# one too; with no override an hf: or syn: model still is, and a Codex +# model_providers entry at the host is, whatever its id. The lines are shaped +# as Claude Code 2.1.308 and Codex 0.163 write them. +SYNTHETIC_HOME="$WORK/synthetic-routes" +mkdir -p "$SYNTHETIC_HOME/.codex" +cat >"$SYNTHETIC_HOME/.codex/config.toml" <<'TOML' +[model_providers.syn] +name = "syn" +base_url = "https://api.synthetic.new/openai/v1" +TOML +python3 - "$SYNTHETIC_HOME" "$WORK/synthetic-cells.json" <<'PY' +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +# (model, owner with no override, owner behind Synthetic's base URL). +MODELS = [ + ("hf:moonshotai/Kimi-K2.5-zs900", "synthetic", "synthetic"), + ("syn:large:text-zs901", "synthetic", "synthetic"), + ("hf:zai-org/GLM-4.7-Flash-zs902", "synthetic", "synthetic"), + ("claude-sonnet-4-5-zs903", "claude", "synthetic"), + ("glm-4.6-zs904", "zai", "synthetic"), +] +session = "1baa699a-5508-4dc9-bbe8-9b8b236c5300" +lines, cells = [], [] +for n, (model, bare, routed) in enumerate(MODELS): + tokens = 8001234 + 1000 * n + prompt = f"894c905e-cd07-4593-946b-c63e9a80{n:04d}" + lines.append({"parentUuid": None, "isSidechain": False, "promptId": f"dcabb8e9-c1ea-4e48-a2ce-7a27caa1{n:04d}", "type": "user", "uuid": prompt, + "timestamp": stamp, "origin": {"kind": "human"}, "promptSource": "typed", "userType": "external", "entrypoint": "cli", + "sessionId": session, "version": "2.1.308", "message": {"role": "user", "content": "go"}}) + lines.append({"parentUuid": prompt, "isSidechain": False, "requestRef": f"78b6e39e-b112-4b92-a1cb-2489e2f2{n:04d}", "requestedModel": model, + "type": "assistant", "uuid": f"da577048-3a3d-4f5e-aed2-0df4c92b{n:04d}", "timestamp": stamp, "userType": "external", + "entrypoint": "cli", "sessionId": session, "version": "2.1.308", + "message": {"id": f"msg_syn_{n:04d}", "type": "message", "role": "assistant", "model": model, "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", "usage": {"input_tokens": tokens, "cache_creation_input_tokens": 31, "cache_read_input_tokens": 71, + "output_tokens": 51, "service_tier": "standard"}}}) + cells.append({"tag": f"zs{900 + n}", "tokens": tokens + 31 + 71 + 51, "owners": [bare, routed]}) +projects = home / ".claude/projects/-home-dhh-proj" +projects.mkdir(parents=True) +(projects / f"{session}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + +thread = "01a16375-5c70-7982-8b22-06dc854c1a00" +usage = {"input_tokens": 8234, "cached_input_tokens": 71, "cache_write_input_tokens": 31, "output_tokens": 51, "reasoning_output_tokens": 0, "total_tokens": 8285} +rollouts = home / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") +rollouts.mkdir(parents=True) +(rollouts / f"rollout-zs910-{thread}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in [ + {"timestamp": stamp, "type": "session_meta", "payload": {"session_id": thread, "id": thread, "originator": "codex_exec", "cli_version": "0.163.0", "source": "exec", "thread_source": "user", "model_provider": "syn"}}, + {"timestamp": stamp, "type": "turn_context", "payload": {"model": "hf:zai-org/GLM-5.3-zs910"}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "item_completed", "thread_id": thread, "item": {"type": "UserMessage", "id": "u0", "content": [{"type": "text", "text": "go"}]}}}, + {"timestamp": stamp, "type": "token_usage_record", "payload": {"thread_id": thread, "response_id": "resp_syn_0001", "usage": usage, "thread_token_usage": usage}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": usage, "last_token_usage": usage}}}, +])) +cells.append({"tag": "zs910", "tokens": usage["input_tokens"] + usage["output_tokens"], "owners": ["synthetic", "synthetic"]}) +json.dump(cells, open(sys.argv[2], "w")) +PY +run_all "$SYNTHETIC_HOME" "$WORK/synthetic-results-0" +run_all "$SYNTHETIC_HOME" "$WORK/synthetic-results-1" ANTHROPIC_BASE_URL="https://api.synthetic.new/anthropic" + +synthetic_report=$(python3 - "$WORK" "${collectors[@]}" <<'PY' +import json +import re +import sys +from pathlib import Path + +work = Path(sys.argv[1]) +cells = json.load(open(work / "synthetic-cells.json")) +for run, label in enumerate(("with no override", "sent to https://api.synthetic.new/anthropic")): + seen = {cell["tag"]: {} for cell in cells} + for cid in sys.argv[2:]: + try: + record = json.loads((work / f"synthetic-results-{run}" / f"{cid}.json").read_text()) + except ValueError: + print(f"{label}: {cid} printed no JSON record") + continue + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zs9\d\d", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total + for cell in cells: + tag, owner, tokens = cell["tag"], cell["owners"][run], cell["tokens"] + wrong = [f"{cid} counted it {total / tokens:g}x" for cid, total in sorted(seen[tag].items()) if cid != owner or total != tokens] + if owner not in seen[tag]: + wrong.insert(0, f"{owner} missed it") + if wrong: + print(f"{tag} {label}, expected {owner}: {'; '.join(wrong)}") +PY +) +if [[ -z $synthetic_report ]]; then + pass "Claude Code and Codex set up as Synthetic documents count once, under synthetic" +else + fail "Claude Code and Codex set up as Synthetic documents count once, under synthetic" "$synthetic_report" +fi + +# ----------------------------------------------- prompts no reply answered + +# A typed prompt whose request failed, or that got no reply at all, still +# counts once, under the collector billing the model the session was set to +# use for it, and holds its session in use: pi and omp write the failed reply +# with its provider and no usage, opencode the user row with its model. The +# lines are shaped as pi 1.4.0, omp 18.12.6 and opencode 1.19.3 wrote them +# against a server refusing the request with a 400. A prompt typed at an +# untrusted time counts nowhere. +UNANSWERED_HOME="$WORK/unanswered" +python3 - "$UNANSWERED_HOME" <<'PY' +import json +import sqlite3 +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +now_ms = int(time.time() * 1000) - 1000 +runtime = {"version": "18.12.6"} +zero = {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": 0, "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "total": 0}} + + +def spent(n): + return {"input": 121101 + n, "output": 1251, "cacheRead": 1271, "cacheWrite": 0, "reasoning": 1204, "totalTokens": 123623 + n, "cost": zero["cost"]} + + +def user(n, text, at=None, omp=False): + message = {"role": "user", "content": [{"type": "text", "text": text}], "timestamp": now_ms if at is None else at} + if omp: + message["attribution"] = "user" + return {"type": "message", "id": f"u{n}", "timestamp": stamp if at is None else at, "message": message} + + +def reply(n, provider, model, usage): + message = {"role": "assistant", "content": [{"type": "text", "text": "" if usage is zero else "hi"}], "api": "openai-completions", + "provider": provider, "model": model, "usage": usage, "stopReason": "error" if usage is zero else "stop", "timestamp": now_ms} + if usage is zero: + message["errorMessage"] = '400: {"message":"refused","type":"invalid_request_error"}' + return {"type": "message", "id": f"a{n}", "timestamp": stamp, "message": message} + + +def write(path, lines): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(line) + "\n" for line in lines)) + + +pi = home / ".pi/agent/sessions/--home-me--" +omp = home / ".omp/agent/sessions/-" +# pi: a failed prompt, then an answered one, on Synthetic. +write(pi / "failed-then-answered.jsonl", [ + {"type": "session", "version": 3, "id": "s1", "timestamp": stamp, "cwd": "/home/me"}, + {"type": "model_change", "id": "m1", "parentId": None, "timestamp": stamp, "provider": "synthetic", "modelId": "e2e-failpi"}, + user(11, "FAILME"), reply(11, "synthetic", "e2e-failpi", zero), + user(12, "Say hi"), reply(12, "synthetic", "e2e-failpi", spent(1)), +]) +# pi: a prompt failed on Z.ai, then the session switched to DeepSeek, which +# answered the next one: one prompt each. +write(pi / "failed-then-switched.jsonl", [ + {"type": "session", "version": 3, "id": "s2", "timestamp": stamp, "cwd": "/home/me"}, + {"type": "model_change", "id": "m2", "parentId": None, "timestamp": stamp, "provider": "zai", "modelId": "e2e-failzai"}, + user(21, "FAILME"), reply(21, "zai", "e2e-failzai", zero), + {"type": "model_change", "id": "m2b", "parentId": "a21", "timestamp": stamp, "provider": "deepseek", "modelId": "e2e-failswitch"}, + user(22, "Say hi"), reply(22, "deepseek", "e2e-failswitch", spent(2)), +]) +# pi: a failed prompt typed at an untrusted time, the session's only one. +write(pi / "untimed.jsonl", [ + {"type": "session", "version": 3, "id": "s3", "timestamp": stamp, "cwd": "/home/me"}, + {"type": "model_change", "id": "m3", "parentId": None, "timestamp": stamp, "provider": "zai", "modelId": "e2e-untimed"}, + user(31, "FAILME", at=0), reply(31, "zai", "e2e-untimed", zero), +]) +# omp: a failed prompt, then an answered one, on Z.ai, with a handoff +# compaction that records no usage. +write(omp / "failed-then-answered.jsonl", [ + {"type": "session", "version": 3, "id": "s4", "timestamp": stamp, "cwd": "/home/me"}, + {"type": "model_change", "id": "m4", "parentId": None, "timestamp": stamp, "runtime": runtime, "model": "zai/e2e-failomp", "resolvedModelIsFallback": False}, + user(41, "FAILME", omp=True), reply(41, "zai", "e2e-failomp", zero), + user(42, "Say hi", omp=True), reply(42, "zai", "e2e-failomp", spent(4)), + {"type": "compaction", "id": "c4", "parentId": "a42", "timestamp": stamp, "runtime": runtime, "summary": "s", "firstKeptEntryId": "a42", "tokensBefore": 123929, "tokensAfter": 2054, "method": "handoff"}, +]) +# omp: a prompt on DeepSeek that got no reply at all (the process died), so +# only the model switch says where it was going. +write(omp / "no-reply.jsonl", [ + {"type": "session", "version": 3, "id": "s5", "timestamp": stamp, "cwd": "/home/me"}, + {"type": "model_change", "id": "m5", "parentId": None, "timestamp": stamp, "runtime": runtime, "model": "deepseek/e2e-noreply", "resolvedModelIsFallback": False}, + user(51, "hello?", omp=True), +]) + +db = home / ".local/share/opencode/opencode.db" +db.parent.mkdir(parents=True) +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE TABLE session (id text PRIMARY KEY, parent_id text, title text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL)") +rows = [] + + +def oc_user(session, n, provider, model): + rows.append((f"msg_{session}_u{n}", session, now_ms + n, {"role": "user", "time": {"created": now_ms + n}, "agent": "build", "model": {"providerID": provider, "modelID": model}, "summary": {"diffs": []}})) + + +def oc_step(session, n, provider, model, tokens): + data = {"parentID": f"msg_{session}_u{n}", "role": "assistant", "mode": "build", "agent": "build", "path": {"cwd": "/home/me", "root": "/"}, "cost": 0, + "tokens": tokens, "modelID": model, "providerID": provider, "time": {"created": now_ms + n + 1, "completed": now_ms + n + 2}} + if tokens["input"]: + data["finish"] = "stop" + else: + data["error"] = {"name": "APIError", "data": {"message": "refused", "statusCode": 400, "isRetryable": False}} + rows.append((f"msg_{session}_a{n}", session, now_ms + n + 1, data)) + + +failed = {"input": 0, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}} +answered = {"total": 123827, "input": 121301, "output": 47, "reasoning": 1206, "cache": {"write": 0, "read": 1273}} +# A failed prompt, then `-c` and an answered one, on Synthetic. +oc_user("ses_failoc", 1, "synthetic", "oc-failoc") +oc_step("ses_failoc", 1, "synthetic", "oc-failoc", failed) +oc_user("ses_failoc", 3, "synthetic", "oc-failoc") +oc_step("ses_failoc", 3, "synthetic", "oc-failoc", answered) +# A session holding only a failed prompt, on Z.ai. +oc_user("ses_failzai", 1, "zai", "oc-failzai") +oc_step("ses_failzai", 1, "zai", "oc-failzai", failed) +for session in ("ses_failoc", "ses_failzai"): + conn.execute("INSERT INTO session VALUES (?, NULL, 'FAILME', ?, ?)", (session, now_ms - 1000, now_ms - 1000)) +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [(mid, session, at, at, json.dumps(data)) for mid, session, at, data in rows]) +conn.commit() +conn.close() +PY +run_all "$UNANSWERED_HOME" "$WORK/unanswered-results" + +unanswered_report=$(python3 - "$WORK/unanswered-results" "${collectors[@]}" <<'PY' +import json +import sys +from pathlib import Path + +results = Path(sys.argv[1]) +# (prompts, sessions) each collector must count; every other one counts none. +expected = {"synthetic": (4, 2), "zai": (4, 3), "deepseek": (2, 2)} +for cid in sys.argv[2:]: + try: + record = json.loads((results / f"{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no JSON record") + continue + want = expected.get(cid, (0, 0)) + got = (record.get("totalPrompts"), record.get("totalSessions")) + today = (record.get("todayPrompts"), record.get("todaySessions")) + if got != want or today != want: + print(f"{cid}: prompts and sessions {got} (today {today}), expected {want}") +PY +) +if [[ -z $unanswered_report ]]; then + pass "a prompt no reply answered in pi, omp or opencode counts once, under the model its session was set to use" +else + fail "a prompt no reply answered in pi, omp or opencode counts once, under the model its session was set to use" "$unanswered_report" +fi + +# ------------------------------ Codex's built-in provider and its own threads + +# Codex 0.163 sends its built-in `openai` provider wherever config.toml's +# openai_base_url says (OPENAI_BASE_URL does nothing), and its rollouts still +# record only `"model_provider":"openai"`. Behind a provider's host every +# built-in rollout is that provider's; behind a self-hosted URL a model not +# OpenAI's is local's; behind a host nothing can place it counts nowhere, +# and OpenAI's own models stay codex's behind either. Dropping the key puts +# them back under codex. A guardian thread counts under the session that +# started it, in every collector. The lines are shaped as Codex 0.160 and +# 0.163 write them. +CODEX_NATIVE_HOME="$WORK/codex-native-routes" +mkdir -p "$CODEX_NATIVE_HOME/.codex" +python3 - "$CODEX_NATIVE_HOME" <<'PY' +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +rollouts = home / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") +rollouts.mkdir(parents=True) + + +def rollout(tag, thread, meta, model, n, typed=True): + usage = {"input_tokens": 9234 + n, "cached_input_tokens": 71, "cache_write_input_tokens": 31, "output_tokens": 51, "reasoning_output_tokens": 0, "total_tokens": 9285 + n} + turn = f"01a1647f-6fd1-7341-86cf-b00a9315{n:04d}" + item = {"type": "UserMessage", "id": f"01a1647f-7085-7231-a532-dda74a01{n:04d}", "content": [{"type": "text", "text": "go", "text_elements": []}]} if typed else {"type": "Reasoning", "id": f"rs_{n:04d}", "summary_text": [], "raw_content": []} + lines = [ + {"timestamp": stamp, "ordinal": 0, "type": "session_meta", "payload": {"session_id": thread, "id": thread, "cli_version": "0.163.0", "history_mode": "paginated", **meta}}, + {"timestamp": stamp, "ordinal": 5, "type": "turn_context", "payload": {"model": model, "turn_id": turn}}, + {"timestamp": stamp, "ordinal": 8, "type": "event_msg", "payload": {"type": "item_completed", "thread_id": thread, "turn_id": turn, "item": item}}, + {"timestamp": stamp, "ordinal": 10, "type": "token_usage_record", "payload": {"thread_id": thread, "turn_id": turn, "session_id": meta.get("session_id", thread), "response_id": f"resp_zo_{n:04d}", "usage": usage, "turn_token_usage": usage, "thread_token_usage": usage}}, + {"timestamp": stamp, "ordinal": 13, "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": usage, "last_token_usage": usage, "model_context_window": 258400}}}, + ] + (rollouts / f"rollout-{tag}-{thread}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + + +exec_meta = {"originator": "codex_exec", "source": "exec", "thread_source": "user", "model_provider": "openai"} +rollout("zo920", "01a1647f-6fb5-7081-a481-257bbb530920", exec_meta, "hf:zai-org/GLM-4.6-zo920", 0) +rollout("zo921", "01a1647f-6fb5-7081-a481-257bbb530921", exec_meta, "gpt-5.1-codex-zo921", 1) +parent = "01a116c6-3c98-7931-bbd3-71e4e2ed0922" +rollout("zo922", parent, {"originator": "codex-tui", "source": "cli", "thread_source": "user", "model_provider": "syn"}, "gpt-6.1-sol-zo922", 2) +rollout("zo923", "01a1284c-6fa4-7c13-8d29-e2d8b5660923", + {"session_id": parent, "parent_thread_id": parent, "originator": "codex-tui", "source": {"subagent": {"other": "guardian"}}, "thread_source": "guardian_review", "model_provider": "syn"}, + "codex-auto-review-zo923", 3, typed=False) +PY +# (openai_base_url, then each tag's owner, then collector=prompts/sessions). +CODEX_NATIVE_RUNS=( + "https://api.synthetic.new/openai/v1|zo920=synthetic zo921=synthetic zo922=synthetic zo923=synthetic|synthetic=3/3 codex=0/0 local=0/0" + "http://127.0.0.1:18493/v1|zo920=local zo921=codex zo922=synthetic zo923=synthetic|synthetic=1/1 codex=1/1 local=1/1" + "https://llm-proxy.example.com/v1|zo920= zo921=codex zo922=synthetic zo923=synthetic|synthetic=1/1 codex=1/1 local=0/0" + "|zo920=codex zo921=codex zo922=synthetic zo923=synthetic|synthetic=1/1 codex=2/2 local=0/0" +) +codex_native_report="" +for run in "${!CODEX_NATIVE_RUNS[@]}"; do + IFS='|' read -r base owners counts <<<"${CODEX_NATIVE_RUNS[$run]}" + { + if [[ -n $base ]]; then + printf 'openai_base_url = "%s"\n\n' "$base" + fi + printf '[model_providers.syn]\nname = "syn"\nbase_url = "https://api.synthetic.new/openai/v1"\n' + } >"$CODEX_NATIVE_HOME/.codex/config.toml" + run_all "$CODEX_NATIVE_HOME" "$WORK/codex-native-results-$run" + report=$(python3 - "$WORK/codex-native-results-$run" "${base:-no openai_base_url}" "$owners" "$counts" "${collectors[@]}" <<'PY' +import json +import re +import sys +from pathlib import Path + +results, label = Path(sys.argv[1]), sys.argv[2] +owners = dict(pair.split("=") for pair in sys.argv[3].split()) +counts = {cid: tuple(int(n) for n in want.split("/")) for cid, want in (pair.split("=") for pair in sys.argv[4].split())} +seen = {tag: {} for tag in owners} +for cid in sys.argv[5:]: + try: + record = json.loads((results / f"{cid}.json").read_text()) + except ValueError: + print(f"{label}: {cid} printed no JSON record") + continue + if cid in counts and (record.get("totalPrompts"), record.get("totalSessions")) != counts[cid]: + print(f"{label}: {cid} counted prompts and sessions {(record.get('totalPrompts'), record.get('totalSessions'))}, expected {counts[cid]}") + for model, bucket in (record.get("modelUsage") or {}).items(): + total = sum(bucket.get(key) or 0 for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")) + for tag in re.findall(r"zo92\d", model): + if tag in seen and total: + seen[tag][cid] = seen[tag].get(cid, 0) + total +for tag, owner in owners.items(): + wrong = [f"{cid} counted it" for cid in sorted(seen[tag]) if cid != owner] + if owner and owner not in seen[tag]: + wrong.insert(0, f"{owner} missed it") + if wrong: + print(f"{label}: {tag}, expected {owner or 'nobody'}: {'; '.join(wrong)}") +PY +) + if [[ -n $report ]]; then + codex_native_report+="$report"$'\n' + fi +done +if [[ -z $codex_native_report ]]; then + pass "Codex's built-in provider counts where openai_base_url sends it, and a guardian thread under its parent's session" +else + fail "Codex's built-in provider counts where openai_base_url sends it, and a guardian thread under its parent's session" "$codex_native_report" +fi + +# Kimchi's gateway bills what reaches it through any agent, so its usage is +# Kimchi's alone: pi and omp sessions on its provider ids (kimchi, kimchi-dev, +# kimchi-dev/, kimchi-experimental) in pi's and omp's trees and its +# own harness, opencode's `kimchi` provider, and the Claude Code and Codex +# setups `kimchi claudecode` and `kimchi codex` write, sent to llm.kimchi.dev +# and llm.eu.kimchi.dev. Model names that would otherwise say Kimi, Z.ai, +# MiniMax, DeepSeek or Anthropic never take a turn the gateway served, and a +# lookalike id (kimchi-proxy) counts nowhere. +KIMCHI_HOME="$WORK/kimchi-gateway" +python3 - "$KIMCHI_HOME" <<'PY' +import json +import sqlite3 +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +now_ms = int(time.time() * 1000) - 1000 + + +def write(path, lines): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(line) + "\n" for line in lines)) + + +# pi 1.4.0 and omp 18.12.6 lines: a model switch, a typed prompt, its reply. +def pi_session(path, provider, model, tokens, omp=False): + switch = {"type": "model_change", "id": "m0", "parentId": None, "timestamp": stamp} + switch.update({"runtime": {"version": "18.12.6"}, "model": f"{provider}/{model}"} if omp else {"provider": provider, "modelId": model}) + user = {"role": "user", "content": [{"type": "text", "text": "Say hi"}], "timestamp": now_ms} + if omp: + user["attribution"] = "user" + write(path, [ + {"type": "session", "version": 3, "id": path.stem, "timestamp": stamp, "cwd": "/home/me"}, + switch, + {"type": "message", "id": f"u-{path.stem}", "parentId": "m0", "timestamp": stamp, "message": user}, + {"type": "message", "id": f"a-{path.stem}", "parentId": f"u-{path.stem}", "timestamp": stamp, "message": { + "role": "assistant", "content": [{"type": "text", "text": "hi"}], "api": "openai-completions", "provider": provider, "model": model, + "usage": {"input": tokens, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": tokens}, "stopReason": "stop", "timestamp": now_ms}}, + ]) + + +pi = home / ".pi/agent/sessions/--home-me--" +omp = home / ".omp/agent/sessions/-" +harness = home / ".config/kimchi/harness/sessions/--home-me--" +pi_session(pi / "pi-kimchi.jsonl", "kimchi", "zqk01-minimax-m3", 101) +pi_session(pi / "pi-kimchi-dev.jsonl", "kimchi-dev", "zqk02-kimi-k2.6", 102) +pi_session(pi / "pi-lookalike.jsonl", "kimchi-proxy", "zqk00-glm-5", 999) +pi_session(omp / "omp-kimchi.jsonl", "kimchi", "zqk03-kimi-k2.5", 103, omp=True) +pi_session(omp / "omp-upstream.jsonl", "kimchi-dev/zai", "zqk04-glm-5.2", 104, omp=True) +pi_session(omp / "omp-experimental.jsonl", "kimchi-experimental", "zqk05-deepseek-v4", 105, omp=True) +pi_session(harness / "harness-kimchi.jsonl", "kimchi", "zqk06-glm-5.2", 106) + +# opencode 1.19.3 rows on the provider Kimchi's opencode setup adds. +db = home / ".local/share/opencode/opencode.db" +db.parent.mkdir(parents=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)") +user = {"role": "user", "time": {"created": now_ms}, "agent": "build", "model": {"providerID": "kimchi", "modelID": "zqk07-minimax-m3"}} +reply = {"parentID": "msg_k_u", "role": "assistant", "mode": "build", "agent": "build", "cost": 0, "modelID": "zqk07-minimax-m3", "providerID": "kimchi", + "tokens": {"total": 107, "input": 107, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, "time": {"created": now_ms + 1, "completed": now_ms + 2}, "finish": "stop"} +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [("msg_k_u", "ses_k", now_ms, now_ms, json.dumps(user)), ("msg_k_a", "ses_k", now_ms + 1, now_ms + 1, json.dumps(reply))]) +conn.commit() +conn.close() + +# Claude Code with the settings.json env `kimchi claudecode` writes, its +# replies on a Claude model and on models whose names say another provider. +lines = [] +for n, (model, tokens) in enumerate((("claude-sonnet-4-5-zqk08", 108), ("kimi-k2.6-zqk09", 109), ("glm-5.2-zqk10", 110), ("MiniMax-M3-zqk11", 111), ("deepseek-v4-flash-zqk12", 112))): + lines.append({"type": "user", "uuid": f"kc-prompt-{n}", "sessionId": "kc", "timestamp": stamp, "message": {"role": "user", "content": "go"}}) + lines.append({"type": "assistant", "sessionId": "kc", "timestamp": stamp, "requestId": f"kc-req-{n}", + "message": {"id": f"kc-msg-{n}", "role": "assistant", "model": model, "usage": {"input_tokens": tokens, "output_tokens": 0}}}) +write(home / ".claude/projects/project/kc.jsonl", lines) +(home / ".claude/settings.json").write_text(json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://llm.eu.kimchi.dev/anthropic", "ANTHROPIC_API_KEY": ""}})) + +# Codex as `kimchi codex` configures it. +(home / ".codex").mkdir() +(home / ".codex/config.toml").write_text('model_provider = "kimchi"\nmodel = "zqk13-minimax-m3"\n\n[model_providers.kimchi]\nname = "Kimchi Gateway"\nbase_url = "https://llm.kimchi.dev/openai/v1"\nwire_api = "responses"\n') +write(home / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") / "rollout-kimchi.jsonl", [ + {"timestamp": stamp, "type": "session_meta", "payload": {"model_provider": "kimchi"}}, + {"timestamp": stamp, "type": "turn_context", "payload": {"model": "zqk13-minimax-m3"}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "user_message", "message": "go"}}, + {"timestamp": stamp, "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": 113, "cached_input_tokens": 0, "output_tokens": 0}, "total_token_usage": {"input_tokens": 113}}}}, +]) +PY +run_all "$KIMCHI_HOME" "$WORK/kimchi-gateway-results" + +kimchi_gateway_report=$(python3 - "$WORK/kimchi-gateway-results" "${collectors[@]}" <<'PY' +import json +import re +import sys +from pathlib import Path + +results = Path(sys.argv[1]) +# Every zqkNN row once, under kimchi, and zqk00 under nobody: seven pi-format +# sessions, opencode's, Claude Code's and Codex's, one prompt per reply. +expected = {f"zqk{n:02d}": 100 + n for n in range(1, 14)} +seen = {} +for cid in sys.argv[2:]: + try: + record = json.loads((results / f"{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no JSON record") + continue + for model, tokens in (record.get("todayTokensByModel") or {}).items(): + tag = re.search(r"zqk\d\d", model) + if tag: + seen.setdefault(tag.group(0), []).append((cid, tokens)) + if cid != "kimchi" and record.get("todayPrompts"): + print(f"{cid} counted {record['todayPrompts']} of Kimchi's prompts") + if cid == "kimchi" and (record.get("todayPrompts"), record.get("todaySessions")) != (13, 9): + print(f"kimchi: prompts and sessions {(record.get('todayPrompts'), record.get('todaySessions'))}, expected (13, 9)") +for tag in sorted(set(expected) | set(seen)): + want = [("kimchi", expected[tag])] if tag in expected else [] + if seen.get(tag, []) != want: + print(f"{tag}: counted {seen.get(tag, [])}, expected {want}") +PY +) +if [[ -z $kimchi_gateway_report ]]; then + pass "usage Kimchi's gateway billed through pi, omp, opencode, Claude Code or Codex counts once, under kimchi" +else + fail "usage Kimchi's gateway billed through pi, omp, opencode, Claude Code or Codex counts once, under kimchi" "$kimchi_gateway_report" +fi + +# ------------------------------------------- requests that never got a reply + +# Shapes pi 1.4.0 and omp 18.12.6 wrote against a mock server. A prompt +# whose agent was killed mid-request leaves no assistant line, so it books +# on the model the session was on, never on the next provider's reply, and a +# fork copying it books it once. omp writes a reply that overflowed the +# model's context window twice under new ids, one response served once. A +# prompt omp refused before sending (Muse Code's invalid credential, an +# errorId and no duration) counts nowhere; one a server refused (a duration +# and a 400) counts. +STRANDED_HOME="$WORK/stranded" +python3 - "$STRANDED_HOME" <<'PY' +import json +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +now_ms = int(time.time() * 1000) - 1000 +cost = {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "total": 0} +zero = {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": 0, "cost": cost} + + +def spent(n): + return {"input": 141601 + n, "output": 1456, "cacheRead": 1476, "cacheWrite": 0, "totalTokens": 144533 + n, "reasoningTokens": 1409, "cost": cost} + + +def user(n, omp=False): + message = {"role": "user", "content": [{"type": "text", "text": "Say hi"}], "timestamp": now_ms} + if omp: + message["attribution"] = "user" + return {"type": "message", "id": f"u{n}", "parentId": None, "timestamp": stamp, "message": message} + + +def reply(n, provider, model, usage, **extra): + message = {"role": "assistant", "content": [{"type": "text", "text": "hi"}], "api": "openai-completions", "provider": provider, "model": model, + "usage": usage, "stopReason": "stop", "timestamp": now_ms, **extra} + return {"type": "message", "id": f"a{n}", "parentId": f"u{n}", "timestamp": stamp, "message": message} + + +def switch(n, provider, model, omp=False): + if omp: + return {"type": "model_change", "id": f"m{n}", "timestamp": stamp, "model": f"{provider}/{model}"} + return {"type": "model_change", "id": f"m{n}", "timestamp": stamp, "provider": provider, "modelId": model} + + +def write(path, lines): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(line) + "\n" for line in lines)) + + +pi = home / ".pi/agent/sessions/--home-me--" +omp = home / ".omp/agent/sessions/-" +# pi: killed waiting on Z.ai, then switched to DeepSeek, which answered. +write(pi / "hang-then-switched.jsonl", [ + {"type": "session", "version": 3, "id": "s1", "timestamp": stamp, "cwd": "/home/me"}, + switch(1, "zai", "e2e-hangpi2"), user(11), + switch(12, "deepseek", "e2e-hangpi2ds"), user(12), reply(12, "deepseek", "e2e-hangpi2ds", spent(12), responseId="chatcmpl-12"), +]) +# pi: killed waiting on Z.ai, then `--fork` onto DeepSeek, which answered. +parent = [ + {"type": "session", "version": 3, "id": "s2", "timestamp": stamp, "cwd": "/home/me"}, + switch(2, "zai", "e2e-hangfork"), user(21), +] +write(pi / "hang-parent.jsonl", parent) +write(pi / "hang-fork.jsonl", [ + {"type": "session", "version": 3, "id": "s3", "timestamp": stamp, "cwd": "/home/me", "parentSession": str(pi / "hang-parent.jsonl")}, + *parent[1:], + switch(22, "deepseek", "e2e-hangforkds"), user(22), reply(22, "deepseek", "e2e-hangforkds", spent(22), responseId="chatcmpl-22"), +]) +# omp: one response past the declared context window, written twice. +twice = reply(31, "deepseek", "e2e-dupb", spent(31), responseId="chatcmpl-31", duration=3.1, ttft=1.2, completedAt=now_ms) +again = json.loads(json.dumps(twice)) +again["id"] = "a31b" +again["message"]["errorId"] = 0 +write(omp / "overflow.jsonl", [ + {"type": "session", "id": "s4", "timestamp": stamp, "cwd": "/home/me"}, + switch(3, "deepseek", "e2e-dupb", omp=True), user(31, omp=True), twice, again, +]) +# omp: Muse Code's credential refused before sending, twice. +refusal = {"stopReason": "error", "errorMessage": "Muse Code credential is invalid; sign in again", "errorId": 0, "completedAt": now_ms} +muse = [reply(n, "muse-code", "e2e-musecode-omp", zero) for n in (41, 42)] +for line in muse: + line["message"].update(refusal) +write(omp / "refused-locally.jsonl", [ + {"type": "session", "id": "s5", "timestamp": stamp, "cwd": "/home/me"}, + switch(4, "muse-code", "e2e-musecode-omp", omp=True), user(41, omp=True), muse[0], user(42, omp=True), muse[1], +]) +# omp: a 400 from Z.ai's server. +served = reply(51, "zai", "e2e-failomp", zero) +served["message"].update({"stopReason": "error", "errorMessage": "400 mock refused FAILME", "errorStatus": 400, "errorId": 400, "duration": 3.2, "completedAt": now_ms}) +write(omp / "refused-by-server.jsonl", [ + {"type": "session", "id": "s6", "timestamp": stamp, "cwd": "/home/me"}, + switch(5, "zai", "e2e-failomp", omp=True), user(51, omp=True), served, +]) +PY +run_all "$STRANDED_HOME" "$WORK/stranded-results" + +stranded_report=$(python3 - "$WORK/stranded-results" "${collectors[@]}" <<'PY' +import json +import sys +from pathlib import Path + +results = Path(sys.argv[1]) +# (prompts, sessions, tokens) each collector must count; every other one none. +expected = {"zai": (3, 3, 0), "deepseek": (3, 3, (144533 + 12) + (144533 + 22) + (144533 + 31))} +for cid in sys.argv[2:]: + try: + record = json.loads((results / f"{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no JSON record") + continue + want = expected.get(cid, (0, 0, 0)) + got = (record.get("todayPrompts"), record.get("todaySessions"), record.get("todayTotalTokens")) + if got != want: + print(f"{cid}: prompts, sessions and tokens {got}, expected {want}") + if cid == "muse" and record.get("hasLocalStats"): + print("muse: in use for prompts that never left the machine") +PY +) +if [[ -z $stranded_report ]]; then + pass "a prompt killed mid-request, a reply omp wrote twice and a request omp never sent each count once or nowhere" +else + fail "a prompt killed mid-request, a reply omp wrote twice and a request omp never sent each count once or nowhere" "$stranded_report" +fi + +# A collector reads only its own share of opencode's `message` store. Finding +# the first step answering each prompt across providers once read every +# assistant row in the database on each collector's rebuild, so a collector +# with no rows there paid for everyone's history. Counted in SQLite VM steps, +# not seconds, so a slow machine can't flake it: a collector with one session +# among a thousand others' costs no more than a few gated scans of the table, +# and still books its prompt only where the first step is its own. +OC_SCALE_HOME="$WORK/oc-scale-home" +mkdir -p "$OC_SCALE_HOME" +oc_scale_report=$(python3 - "$ROOT" "$OC_SCALE_HOME" <<'PY' +import json +import os +import sqlite3 +import sys +import time +from pathlib import Path + +root, home = Path(sys.argv[1]), Path(sys.argv[2]) +os.environ["HOME"] = str(home) +os.environ["XDG_DATA_HOME"] = str(home / ".local/share") +os.environ.pop("OPENCODE_DB", None) +sys.path.insert(0, str(root / "shell" / "plugins" / "agents")) +import collectors + +db = home / ".local/share/opencode/opencode.db" +db.parent.mkdir(parents=True, exist_ok=True) +conn = sqlite3.connect(db) +# The schema opencode 1.x creates, its session index included. +conn.execute("CREATE TABLE session (id text PRIMARY KEY, parent_id text, title text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL)") +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE INDEX message_session_time_created_id_idx ON message (session_id, time_created, id)") +now_ms = int(time.time() * 1000) - 60000 +tokens = {"total": 1320, "input": 1200, "output": 20, "reasoning": 0, "cache": {"write": 0, "read": 100}} +rows = [] + + +def prompt(session, n, providers): + # A typed prompt and the tool loop answering it, one step per provider. + at = now_ms + n * 100 + rows.append((f"msg_{session}_u{n}", session, at, {"role": "user", "time": {"created": at}, "agent": "build", "model": {"providerID": providers[0], "modelID": "m"}, "summary": {"diffs": []}})) + for step, provider in enumerate(providers, 1): + rows.append((f"msg_{session}_a{n}_{step}", session, at + step, {"parentID": f"msg_{session}_u{n}", "role": "assistant", "mode": "build", "agent": "build", "path": {"cwd": "/home/me", "root": "/"}, "cost": 0, "tokens": tokens, "modelID": "m", "providerID": provider, "time": {"created": at + step, "completed": at + step + 1}, "finish": "tool-calls"})) + + +sessions = [f"ses_other{s:04d}" for s in range(1000)] +for session in sessions: + for n in range(3): + prompt(session, n, ["anthropic"] * 4) +# The collector's one session: a prompt it answered first, and one whose +# first step was another provider's. +sessions.append("ses_mine") +prompt("ses_mine", 0, ["devin", "devin"]) +prompt("ses_mine", 1, ["anthropic", "devin"]) +for session in sessions: + conn.execute("INSERT INTO session VALUES (?, NULL, 't', ?, ?)", (session, now_ms - 1000, now_ms - 1000)) +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [(mid, session, at, at, json.dumps(data)) for mid, session, at, data in rows]) +conn.commit() +conn.close() + +steps = [0] +connect = sqlite3.connect + + +def counted(*args, **kwargs): + handle = connect(*args, **kwargs) + handle.set_progress_handler(lambda: steps.__setitem__(0, steps[0] + 1), 100) + return handle + + +def cost(run): + steps[0] = 0 + sqlite3.connect = counted + try: + result = run() + finally: + sqlite3.connect = connect + return steps[0], result + + +def gated_scan(): + handle = sqlite3.connect(db.as_uri() + "?mode=ro", uri=True) + try: + return handle.execute("SELECT m.id, m.data FROM message m WHERE m.data LIKE '%\"role\"%:%\"assistant\"%' AND m.data LIKE ?", ('%"providerID"%:%"devin"%',)).fetchall() + finally: + handle.close() + + +scan_cost, (found, complete) = cost(lambda: collectors.scan_opencode(['devin"'], lambda provider, model, entry: provider == "devin")) +gate_cost, _ = cost(gated_scan) +problems = [] +if not complete: + problems.append("the scan did not complete") +if (len(found), sum(row.prompts for row in found)) != (3, 1): + problems.append(f"devin's rows {len(found)} with {sum(row.prompts for row in found)} prompts, expected 3 with 1") +if scan_cost > 6 * gate_cost: + problems.append(f"scan_opencode took {scan_cost} progress ticks, over 6x the {gate_cost} of one gated scan") +print("\n".join(problems)) +PY +) +if [[ -z $oc_scale_report ]]; then + pass "a collector's opencode scan costs a few gated scans of the store, not a pass over every provider's steps" +else + fail "a collector's opencode scan costs a few gated scans of the store, not a pass over every provider's steps" "$oc_scale_report" +fi + +# ------------------- Codex profiles and a claimed id on the user's own hardware + +# Codex 0.163 applies a .config.toml profile only to a run given +# --profile, and that run's rollout still records only +# `"model_provider":"openai"`, so a profile's openai_base_url says nothing +# about plain `codex` runs: here it moves only the rollouts on the model the +# profile sets (zp941), while ChatGPT-plan rollouts (zp940), and a model no +# profile ties (zp943), stay codex's. A profile setting an OpenAI model can't +# be told from a plain run and moves nothing. A model_providers entry whose +# id another collector claims but whose base_url is this machine, as a +# self-hosted Kimchi gateway's `kimchi codex` setup writes it, is local's: +# the base URL decides (zp942). The lines are shaped as Codex 0.163 wrote +# them in real `codex exec` runs. +CODEX_PROFILE_HOME="$WORK/codex-profiles" +mkdir -p "$CODEX_PROFILE_HOME/.codex" +python3 - "$CODEX_PROFILE_HOME" <<'PY' +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +rollouts = home / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") +rollouts.mkdir(parents=True) + + +def rollout(tag, provider, model, n): + thread = f"01a164a4-8b38-7013-9a23-709319c3{n:04d}" + turn = f"01a164a4-8b8b-71f1-b648-20be1fa2{n:04d}" + usage = {"input_tokens": 18234 + n, "cached_input_tokens": 71, "cache_write_input_tokens": 31, "output_tokens": 51, "reasoning_output_tokens": 0, "total_tokens": 18285 + n} + item = {"type": "UserMessage", "id": f"01a164a4-8c01-7231-a532-dda74a01{n:04d}", "content": [{"type": "text", "text": "go", "text_elements": []}]} + lines = [ + {"timestamp": stamp, "ordinal": 0, "type": "session_meta", "payload": {"session_id": thread, "id": thread, "cli_version": "0.163.0", "history_mode": "paginated", "originator": "codex_exec", "source": "exec", "thread_source": "user", "model_provider": provider}}, + {"timestamp": stamp, "ordinal": 5, "type": "turn_context", "payload": {"model": model, "turn_id": turn}}, + {"timestamp": stamp, "ordinal": 8, "type": "event_msg", "payload": {"type": "item_completed", "thread_id": thread, "turn_id": turn, "item": item}}, + {"timestamp": stamp, "ordinal": 10, "type": "token_usage_record", "payload": {"thread_id": thread, "turn_id": turn, "session_id": thread, "response_id": f"resp_zp_{n:04d}", "usage": usage, "turn_token_usage": usage, "thread_token_usage": usage}}, + ] + (rollouts / f"rollout-2026-10-22T13-24-27-{thread}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + + +rollout("zp940", "openai", "gpt-5.1-codex-zp940", 940) +rollout("zp941", "openai", "hf:zai-org/GLM-4.6-zp941", 941) +rollout("zp942", "kimchi", "kimi-k2-zp942", 942) +rollout("zp943", "openai", "qwen3-coder-zp943", 943) +codex = home / ".codex" +(codex / "config.toml").write_text( + 'model = "gpt-5.1-codex-zp940"\n\n' + '[model_providers.kimchi]\nname = "Kimchi Gateway"\nbase_url = "http://127.0.0.1:18491/openai/v1"\nwire_api = "responses"\n' +) +(codex / "syn.config.toml").write_text('model = "hf:zai-org/GLM-4.6-zp941"\nopenai_base_url = "https://api.synthetic.new/openai/v1"\n') +(codex / "or.config.toml").write_text('model = "gpt-5.1-codex-zp940"\nopenai_base_url = "https://openrouter.ai/api/v1"\n') +PY +run_all "$CODEX_PROFILE_HOME" "$WORK/codex-profile-results" +codex_profile_report=$(python3 - "$WORK/codex-profile-results" "${collectors[@]}" <<'PY' +import json +import re +import sys +from pathlib import Path + +results = Path(sys.argv[1]) +owners = {"zp940": "codex", "zp941": "synthetic", "zp942": "local", "zp943": "codex"} +counts = {"codex": (2, 2), "synthetic": (1, 1), "local": (1, 1), "openrouter": (0, 0), "kimchi": (0, 0)} +seen = {tag: set() for tag in owners} +for cid in sys.argv[2:]: + try: + record = json.loads((results / f"{cid}.json").read_text()) + except ValueError: + print(f"{cid} printed no JSON record") + continue + if cid in counts and (record.get("totalPrompts"), record.get("totalSessions")) != counts[cid]: + print(f"{cid} counted prompts and sessions {(record.get('totalPrompts'), record.get('totalSessions'))}, expected {counts[cid]}") + for model, bucket in (record.get("modelUsage") or {}).items(): + if any(bucket.get(key) for key in ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")): + for tag in re.findall(r"zp94\d", model): + seen[tag].add(cid) +for tag, owner in owners.items(): + if seen[tag] != {owner}: + print(f"{tag}: counted by {sorted(seen[tag]) or 'nobody'}, expected {owner}") +PY +) +if [[ -z $codex_profile_report ]]; then + pass "a Codex profile's openai_base_url moves only the rollouts on its model, and a claimed id on this machine is local's" +else + fail "a Codex profile's openai_base_url moves only the rollouts on its model, and a claimed id on this machine is local's" "$codex_profile_report" +fi + +# Kimchi 1.18 names each gateway route kimchi-dev/, and the kimchi +# collector claims every such id by its prefix. A local.json naming one must +# not make local count those rows as well, while a self-hosted id beside it +# still counts. The replies are pi's own shape on Kimchi's gateway. +PREFIX_HOME="$WORK/kimchi-prefix" +mkdir -p "$PREFIX_HOME/.config/omarchy/agents" "$PREFIX_HOME/.pi/agent/sessions/-" "$PREFIX_HOME/.cache" "$PREFIX_HOME/.run" +echo '{"providers": ["Kimchi-Dev/MoonshotAI", "my-llama"]}' >"$PREFIX_HOME/.config/omarchy/agents/local.json" +python3 - "$PREFIX_HOME/.pi/agent/sessions/-/session.jsonl" <<'PY' +import json +import sys +from datetime import datetime, timezone + +now = datetime.now(timezone.utc) +stamp = now.isoformat(timespec="milliseconds").replace("+00:00", "Z") +ms = int(now.timestamp() * 1000) +lines = [ + {"type": "session", "version": 3, "id": "01a16377-aaaa-72fe-95e0-9faf99826e64", "timestamp": stamp, "cwd": "/home/me"}, + {"type": "model_change", "id": "29d6576e", "parentId": None, "timestamp": stamp, "model": "kimchi/kimi-k2.5"}, +] +for index, (provider, model, tokens) in enumerate((("kimchi-dev/moonshotai", "kimi-k2.5", 4201), ("my-llama", "llama-3.3", 3101))): + lines.append({"type": "message", "id": f"u{index}", "timestamp": stamp, "message": {"role": "user", "content": [{"type": "text", "text": "Say hi"}], "attribution": "user", "timestamp": ms}}) + usage = {"input": tokens, "output": 82, "cacheRead": 102, "cacheWrite": 0, "totalTokens": tokens + 184, "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "total": 0}} + lines.append({"type": "message", "id": f"a{index}", "parentId": f"u{index}", "timestamp": stamp, "message": {"role": "assistant", "content": [{"type": "text", "text": "Hi."}], "api": "openai-completions", "provider": provider, "model": model, "usage": usage, "stopReason": "stop", "timestamp": ms, "responseId": f"chatcmpl-{index}"}}) +with open(sys.argv[1], "w") as handle: + handle.writelines(json.dumps(line) + "\n" for line in lines) +PY +run_collector "$PREFIX_HOME" local "$WORK/prefix-local" +run_collector "$PREFIX_HOME" kimchi "$WORK/prefix-kimchi" +prefix_report=$(python3 - "$WORK/prefix-local.json" "$WORK/prefix-kimchi.json" <<'PY' +import json +import sys + +problems = [] +for path, expected in zip(sys.argv[1:], ({"llama-3.3"}, {"kimi-k2.5"})): + try: + models = set(json.load(open(path)).get("todayTokensByModel") or {}) + except (OSError, ValueError) as error: + models = {f"unreadable: {error}"} + if models != expected: + problems.append(f"{path.rsplit('/', 1)[-1]} counts {sorted(models)}, expected {sorted(expected)}") +print("\n".join(problems)) +PY +) +if [[ -z $prefix_report ]]; then + pass "local leaves a configured kimchi-dev/ id to the kimchi collector" +else + fail "local leaves a configured kimchi-dev/ id to the kimchi collector" "$prefix_report" +fi + +# ------------------------------- A Codex rollout switching models mid-thread + +# Behind a self-hosted openai_base_url, one built-in rollout can switch from +# an OpenAI model to a local one between turns (`/model`). Each typed prompt +# goes with the model that answered it, so codex books the first and local +# the second, and neither books the other's. +CODEX_MIXED_HOME="$WORK/codex-mixed-models" +mkdir -p "$CODEX_MIXED_HOME/.codex" +printf 'openai_base_url = "http://127.0.0.1:18495/v1"\n' >"$CODEX_MIXED_HOME/.codex/config.toml" +python3 - "$CODEX_MIXED_HOME" <<'PY' +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +rollouts = home / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") +rollouts.mkdir(parents=True) +thread = "01a164b2-4c11-7a20-9d31-5e6f7a8b0950" +lines = [{"timestamp": stamp, "ordinal": 0, "type": "session_meta", "payload": {"session_id": thread, "id": thread, "cli_version": "0.163.0", "history_mode": "paginated", "originator": "codex-tui", "source": "cli", "thread_source": "user", "model_provider": "openai"}}] +total = {"input_tokens": 0, "cached_input_tokens": 0, "cache_write_input_tokens": 0, "output_tokens": 0, "reasoning_output_tokens": 0, "total_tokens": 0} +for n, model in enumerate(("gpt-5.1-codex-zm950", "qwen3-coder-zm951")): + turn = f"01a164b2-4c8b-71f1-b648-20be1fa2{n:04d}" + usage = {"input_tokens": 12234 + n, "cached_input_tokens": 71, "cache_write_input_tokens": 31, "output_tokens": 51, "reasoning_output_tokens": 0, "total_tokens": 12285 + n} + total = {key: total[key] + usage[key] for key in total} + item = {"type": "UserMessage", "id": f"01a164b2-4d01-7231-a532-dda74a01{n:04d}", "content": [{"type": "text", "text": "go", "text_elements": []}]} + lines += [ + {"timestamp": stamp, "ordinal": 10 * n + 1, "type": "turn_context", "payload": {"model": model, "turn_id": turn}}, + {"timestamp": stamp, "ordinal": 10 * n + 2, "type": "event_msg", "payload": {"type": "item_completed", "thread_id": thread, "turn_id": turn, "item": item}}, + {"timestamp": stamp, "ordinal": 10 * n + 3, "type": "token_usage_record", "payload": {"thread_id": thread, "turn_id": turn, "session_id": thread, "response_id": f"resp_zm_{n:04d}", "usage": usage, "turn_token_usage": usage, "thread_token_usage": total}}, + {"timestamp": stamp, "ordinal": 10 * n + 4, "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": total, "last_token_usage": usage, "model_context_window": 258400}}}, + ] +(rollouts / f"rollout-2026-10-22T14-02-11-{thread}.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) +PY +run_collector "$CODEX_MIXED_HOME" codex "$WORK/codex-mixed-codex" +run_collector "$CODEX_MIXED_HOME" local "$WORK/codex-mixed-local" +codex_mixed_report=$(python3 - "$WORK/codex-mixed-codex.json" "$WORK/codex-mixed-local.json" <<'PY' +import json +import sys + +problems = [] +for path, model in zip(sys.argv[1:], ("gpt-5.1-codex-zm950", "qwen3-coder-zm951")): + name = path.rsplit("/", 1)[-1] + try: + record = json.load(open(path)) + except (OSError, ValueError) as error: + problems.append(f"{name} is unreadable: {error}") + continue + models = set(record.get("todayTokensByModel") or {}) + if models != {model}: + problems.append(f"{name} counts {sorted(models)}, expected {[model]}") + if (record.get("totalPrompts"), record.get("totalSessions")) != (1, 1): + problems.append(f"{name} counted prompts and sessions {(record.get('totalPrompts'), record.get('totalSessions'))}, expected (1, 1)") +print("\n".join(problems)) +PY +) +if [[ -z $codex_mixed_report ]]; then + pass "a Codex rollout switching from an OpenAI model to a local one books each prompt once, with the model that answered it" +else + fail "a Codex rollout switching from an OpenAI model to a local one books each prompt once, with the model that answered it" "$codex_mixed_report" +fi + +# `codex exec -m resume` and /model switch a thread's model with a +# thread_settings_applied event, as Codex 0.163 writes it, and the request a +# switch-time compaction sends goes out on the new model before that turn's +# turn_context names it. Behind a self-hosted openai_base_url the compaction +# is local's, as is a prompt typed after the switch that no request answered. +CODEX_SWITCH_HOME="$WORK/codex-switched-models" +mkdir -p "$CODEX_SWITCH_HOME/.codex" +printf 'openai_base_url = "http://127.0.0.1:18496/v1"\n' >"$CODEX_SWITCH_HOME/.codex/config.toml" +python3 - "$CODEX_SWITCH_HOME" <<'PY' +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +rollouts = home / ".codex/sessions" / datetime.now().strftime("%Y/%m/%d") +rollouts.mkdir(parents=True) + + +def context(model): + return lambda thread, n: {"type": "turn_context", "payload": {"model": model, "turn_id": f"turn{n}"}} + + +def switch(model): + return lambda thread, n: {"type": "event_msg", "payload": {"type": "thread_settings_applied", "thread_id": thread, "thread_settings": { + "model": model, "model_provider_id": "openai", "approval_policy": "never", "approvals_reviewer": "auto_review", + "collaboration_mode": {"mode": "default", "settings": {"model": model, "reasoning_effort": None, "developer_instructions": None}}, + "cwd": str(home), "disabled_plugin_ids": [], "runtime_workspace_roots": [str(home)]}}} + + +def typed(thread, n): + item = {"type": "UserMessage", "id": f"01a164cc-7301-7231-a532-dda74a01{n:04d}", "content": [{"type": "text", "text": "go", "text_elements": []}]} + return {"type": "event_msg", "payload": {"type": "item_completed", "thread_id": thread, "turn_id": f"turn{n}", "item": item}} + + +def compacted(thread, n): + return {"type": "compacted", "payload": {"message": "", "replacement_history": []}} + + +def write(thread, steps): + """A rollout of `steps`: line makers, and the input tokens of each request + a token_usage_record names.""" + total = {"input_tokens": 0, "cached_input_tokens": 0, "cache_write_input_tokens": 0, "output_tokens": 0, "reasoning_output_tokens": 0, "total_tokens": 0} + lines = [{"type": "session_meta", "payload": {"session_id": thread, "id": thread, "cli_version": "0.163.0", "history_mode": "paginated", "originator": "codex_exec", "source": "exec", "thread_source": "user", "model_provider": "openai"}}] + for n, step in enumerate(steps): + if isinstance(step, int): + usage = {"input_tokens": step, "cached_input_tokens": 71, "cache_write_input_tokens": 31, "output_tokens": 51, "reasoning_output_tokens": 0, "total_tokens": step + 51} + total = {key: total[key] + usage[key] for key in total} + lines.append({"type": "token_usage_record", "payload": {"thread_id": thread, "turn_id": f"turn{n}", "session_id": thread, "response_id": f"resp_zs_{n:04d}", "usage": usage, "turn_token_usage": usage, "thread_token_usage": total}}) + else: + lines.append(step(thread, n)) + (rollouts / f"rollout-2026-10-22T14-07-55-{thread}.jsonl").write_text("".join(json.dumps({"timestamp": stamp, **line}) + "\n" for line in lines)) + + +local = "hf:zai-org/GLM-4.6-zs961" +write("01a164cc-54d9-7f81-80d2-0cbf0cd6a960", [ + context("gpt-5.1-codex-zs960"), switch("gpt-5.1-codex-zs960"), typed, 301234, + switch(local), switch(local), 2234, compacted, context(local), switch(local), typed, 3234, +]) +write("01a164cc-54d9-7f81-80d2-0cbf0cd6a962", [context("gpt-5.1-codex-zs962"), typed, 1234, switch(local), typed]) +PY +run_collector "$CODEX_SWITCH_HOME" codex "$WORK/codex-switched-codex" +run_collector "$CODEX_SWITCH_HOME" local "$WORK/codex-switched-local" +codex_switched_report=$(python3 - "$WORK/codex-switched-codex.json" "$WORK/codex-switched-local.json" <<'PY' +import json +import sys + +expected = { + "codex-switched-codex.json": ({"gpt-5.1-codex-zs960": 301285, "gpt-5.1-codex-zs962": 1285}, 2, 2), + "codex-switched-local.json": ({"hf:zai-org/GLM-4.6-zs961": 2285 + 3285}, 2, 2), +} +problems = [] +for path in sys.argv[1:]: + name = path.rsplit("/", 1)[-1] + try: + record = json.load(open(path)) + except (OSError, ValueError) as error: + problems.append(f"{name} is unreadable: {error}") + continue + got = (record.get("todayTokensByModel"), record.get("totalPrompts"), record.get("totalSessions")) + if got != expected[name]: + problems.append(f"{name} counts tokens, prompts and sessions {got}, expected {expected[name]}") +print("\n".join(problems)) +PY +) +if [[ -z $codex_switched_report ]]; then + pass "a Codex model switch books the requests and prompts after thread_settings_applied to the model switched to" +else + fail "a Codex model switch books the requests and prompts after thread_settings_applied to the model switched to" "$codex_switched_report" +fi + +# --------------------- A pi or omp prompt a fork answered on another provider + +# A fork copies its parent's messages with their ids and timestamps, so a +# prompt the parent's provider failed on and a fork answered after switching +# provider (`/model`) is one typed message. It counts once, under the +# collector of the reply that answered it, never also under the parent's +# provider, even when the fork holds no line of the parent's provider at all. +FORKED_HOME="$WORK/pi-forked-provider" +python3 - "$FORKED_HOME" <<'PY' +import json +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +home = Path(sys.argv[1]) +now = datetime.now(timezone.utc) + + +def stamp(offset): + return (now + timedelta(seconds=offset)).isoformat(timespec="milliseconds").replace("+00:00", "Z") + + +def millis(offset): + return int((now + timedelta(seconds=offset)).timestamp() * 1000) + + +def user(entry_id, offset): + return {"type": "message", "id": entry_id, "parentId": None, "timestamp": stamp(offset), + "message": {"role": "user", "content": [{"type": "text", "text": "go"}], "timestamp": millis(offset)}} + + +def reply(entry_id, provider, model, offset, tokens, **extra): + usage = {"input": tokens, "output": 1 if tokens else 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": tokens + 1 if tokens else 0, + "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "total": 0}} + message = {"role": "assistant", "content": [], "api": "openai-completions", "provider": provider, "model": model, "usage": usage, + "stopReason": "stop" if tokens else "error", "timestamp": millis(offset), **extra} + if not tokens: + message["errorMessage"] = "400 status code (no body)" + return {"type": "message", "id": entry_id, "parentId": None, "timestamp": stamp(offset), "message": message} + + +def write(path, header, lines): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(line) + "\n" for line in [header, *lines])) + + +# pi 1.4.0: kimchi failed on the prompt, and the fork answered it on openrouter. +pi = home / ".pi/agent/sessions/--home-me--" +parent = [{"type": "model_change", "id": "zf1", "parentId": None, "timestamp": stamp(0), "provider": "kimchi", "modelId": "kimi-k2.5-zf950"}, + user("zf2", 1), reply("zf3", "kimchi", "kimi-k2.5-zf950", 2, 0)] +write(pi / "parent.jsonl", {"type": "session", "version": 3, "id": "zf-parent", "timestamp": stamp(0), "cwd": "/home/me"}, parent) +write(pi / "fork.jsonl", {"type": "session", "version": 3, "id": "zf-fork", "timestamp": stamp(3), "cwd": "/home/me", "parentSession": str(pi / "parent.jsonl")}, + parent + [{"type": "model_change", "id": "zf4", "parentId": None, "timestamp": stamp(3), "provider": "openrouter", "modelId": "z-ai/glm-zf951"}, + reply("zf5", "openrouter", "z-ai/glm-zf951", 4, 40)]) +# omp 18.12: zai failed on the prompt, and a fork cut before that reply +# answered it on synthetic, so the fork holds no zai line. +omp = home / ".omp/agent/sessions/--home-me--" +prompt = user("zf6", 1) +write(omp / "parent.jsonl", {"type": "session", "version": 3, "id": "zf-oparent", "timestamp": stamp(0), "cwd": "/home/me"}, + [prompt, reply("zf7", "zai", "glm-zf952", 2, 0, errorStatus=400, duration=310)]) +write(omp / "fork.jsonl", {"type": "session", "version": 3, "id": "zf-ofork", "timestamp": stamp(3), "cwd": "/home/me", "parentSession": str(omp / "parent.jsonl")}, + [prompt, {"type": "model_change", "id": "zf8", "parentId": None, "timestamp": stamp(3), "model": "synthetic/hf:zai-org/GLM-zf953"}, + reply("zf9", "synthetic", "hf:zai-org/GLM-zf953", 4, 50, duration=820)]) +PY +for id in kimchi openrouter zai synthetic; do + run_collector "$FORKED_HOME" "$id" "$WORK/pi-forked-$id" +done +forked_report=$(python3 - "$WORK" <<'PY' +import json +import sys + +expected = {"kimchi": (0, 0), "openrouter": (1, 41), "zai": (0, 0), "synthetic": (1, 51)} +problems = [] +for id, want in expected.items(): + try: + record = json.load(open(f"{sys.argv[1]}/pi-forked-{id}.json")) + except (OSError, ValueError) as error: + problems.append(f"{id} is unreadable: {error}") + continue + got = (record.get("totalPrompts") or 0, record.get("todayTotalTokens") or 0) + if got != want: + problems.append(f"{id} counted (prompts, tokens) {got}, expected {want}") +print("\n".join(problems)) +PY +) +if [[ -z $forked_report ]]; then + pass "a pi or omp prompt a fork answered on another provider counts once, under the provider that answered it" +else + fail "a pi or omp prompt a fork answered on another provider counts once, under the provider that answered it" "$forked_report" +fi diff --git a/test/shell.d/agent-usage-pi-scan-test.sh b/test/shell.d/agent-usage-pi-scan-test.sh new file mode 100755 index 00000000000..dd8e50483ca --- /dev/null +++ b/test/shell.d/agent-usage-pi-scan-test.sh @@ -0,0 +1,215 @@ +#!/bin/bash + +# Every collector reads the whole pi/omp history through scan_pi_sessions, +# each for its own providers, so what one scan parses is paid once per +# collector on every refresh. A file holding none of a collector's lines +# must cost it nothing beyond rg's search: its prompts, tool results and +# compactions are never parsed, and no file is read again for the model a +# summary ran on. A count, the same on every machine, rather than a time. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command python3 +require_command rg + +WORK=$(mktemp -d) +trap 'rm -rf "$WORK"' EXIT + +WORK="$WORK" python3 - <<'PY' +import json +import os +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +import collectors # noqa: E402 + +work = Path(os.environ["WORK"]) +home = work / "home" +home.mkdir() +os.environ["HOME"] = str(home) +for name in ("PI_CODING_AGENT_DIR", "PI_CODING_AGENT_SESSION_DIR"): + os.environ.pop(name, None) + + +def check(description, condition, detail=""): + if condition: + print("ok - " + description) + return + print("not ok - " + description, file=sys.stderr) + if detail: + print(detail, file=sys.stderr) + sys.exit(1) + + +stamp = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +runtime = {"version": "18.12.6"} + + +def usage(tokens): + return {"input": tokens, "output": 10, "cacheRead": 0, "cacheWrite": 0, "totalTokens": tokens + 10, "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "total": 0}} + + +def pi_session(provider, model, n, tokens): + """A tool-loop session as pi 1.4.0 writes one, with a pasted image and a + compaction, on one provider.""" + lines = [ + {"type": "session", "version": 3, "id": f"s{n}", "timestamp": stamp, "cwd": "/proj"}, + {"type": "model_change", "id": f"mc{n}", "parentId": None, "timestamp": stamp, "provider": provider, "modelId": model}, + {"type": "message", "id": f"u{n}", "timestamp": stamp, "message": {"role": "user", "content": [{"type": "text", "text": "read it"}, {"type": "image", "mimeType": "image/png", "data": "iVBORw0KGgo" * 2000}], "timestamp": 1}}, + {"type": "message", "id": f"a{n}", "timestamp": stamp, "message": {"role": "assistant", "content": [{"type": "toolCall", "id": f"c{n}", "name": "read", "arguments": {"path": "/x"}}], + "api": "openai-completions", "provider": provider, "model": model, "usage": usage(tokens), "stopReason": "toolUse", "timestamp": 2}}, + {"type": "message", "id": f"r{n}", "timestamp": stamp, "message": {"role": "toolResult", "toolCallId": f"c{n}", "toolName": "read", "content": [{"type": "text", "text": "x" * 4000}], "isError": False, "timestamp": 3}}, + {"type": "message", "id": f"b{n}", "timestamp": stamp, "message": {"role": "assistant", "content": [{"type": "text", "text": "done"}], + "api": "openai-completions", "provider": provider, "model": model, "usage": usage(tokens + 1), "stopReason": "stop", "timestamp": 4}}, + {"type": "compaction", "id": f"cp{n}", "parentId": f"b{n}", "timestamp": stamp, "summary": "s", "firstKeptEntryId": f"u{n}", "tokensBefore": 40000, "usage": usage(tokens + 2)}, + ] + return lines + + +def write(path, lines): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(line) + "\n" for line in lines)) + + +sessions = home / ".pi" / "agent" / "sessions" / "--proj--" +write(sessions / "mine.jsonl", pi_session("demo", "demo-model", 0, 1000)) +for n in range(1, 41): + write(sessions / f"other-{n}.jsonl", pi_session(f"other{n % 4}", "other-model", n, 50000)) +# As omp 18.12.6 writes one: the switch names provider/model, no field +# `pattern` sees, and the only usage is a compaction on that model. +write(home / ".omp" / "agent" / "sessions" / "-" / "omp-compacted.jsonl", [ + {"type": "session", "version": 3, "id": "o1", "timestamp": stamp, "cwd": "/proj"}, + {"type": "model_change", "id": "omc", "parentId": None, "timestamp": stamp, "runtime": runtime, "model": "demo/omp-model", "resolvedModelIsFallback": False}, + {"type": "message", "id": "ou", "timestamp": stamp, "runtime": runtime, "message": {"role": "user", "content": [{"type": "text", "text": "compact"}], "attribution": "user", "timestamp": 1}}, + {"type": "compaction", "id": "ocp", "timestamp": stamp, "runtime": runtime, "summary": "s", "firstKeptEntryId": "ou", "tokensBefore": 900, "usage": usage(700)}, +]) + +parsed = {"files": set(), "lines": 0} +pi_rows = collectors._pi_rows + + +def counting_rows(path, lines, accept): + parsed["files"].add(Path(path).name) + parsed["lines"] += len(lines) + return pi_rows(path, lines, accept) + + +collectors._pi_rows = counting_rows + + +def scan(provider): + for key in parsed: + parsed[key] = set() if key != "lines" else 0 + pattern = r'"provider"\s*:\s*"' + provider + '"' + return collectors.scan_pi_sessions(pattern, lambda name, model, entry: name == provider) + + +rows, complete = scan("nobody") +check("a scan for a provider with no sessions parses no line", + complete and not rows and not parsed["files"] and parsed["lines"] == 0, json.dumps({**parsed, "files": sorted(parsed["files"])})) + +rows, complete = scan("demo") +check("a scan parses only the files holding its own lines, the omp switch's included", + complete and parsed["files"] == {"mine.jsonl", "omp-compacted.jsonl"}, json.dumps(sorted(parsed["files"]))) +# The omp prompt no reply answered counts on the switch's model. +check("its own replies, prompts and compactions still count", + sorted((row.model, row.input_tokens, row.prompts) for row in rows) + == [("demo-model", 1000, 1), ("demo-model", 1001, 0), ("demo-model", 1002, 0), ("omp-model", 0, 1), ("omp-model", 700, 0)], + json.dumps([row._asdict() for row in rows])) + +rows, complete = scan("other1") +check("another provider's scan parses only its own ten files", + complete and len(parsed["files"]) == 10 and all(name.startswith("other-") for name in parsed["files"]) and len(rows) == 30, + json.dumps({"files": sorted(parsed["files"]), "rows": len(rows)})) +# A prompt goes with the reply that answered it, even one rg's pattern keeps +# out of the parse: another provider's reply after a handoff answered it, and +# the next reply here answers none. +turns = work / "turns" +write(turns / "handoff.jsonl", [ + {"type": "model_change", "id": "hm", "timestamp": stamp, "provider": "demo", "modelId": "demo-model"}, + {"type": "message", "id": "hu", "timestamp": stamp, "message": {"role": "user", "content": "ask the other one"}}, + {"type": "message", "id": "ho", "timestamp": stamp, "message": {"role": "assistant", "provider": "elsewhere", "model": "e", "usage": usage(300)}}, + {"type": "message", "id": "hd", "timestamp": stamp, "message": {"role": "assistant", "provider": "demo", "model": "demo-model", "usage": usage(30)}}, +]) +rows, _ = collectors.scan_pi_sessions(r'"provider"\s*:\s*"demo"', lambda name, model, entry: name == "demo", roots=[turns]) +check("a prompt another provider answered is not booked on the next reply here", + [(row.input_tokens, row.prompts) for row in rows] == [(30, 0)], json.dumps([row._asdict() for row in rows])) + +# A prompt counts on the day it was typed, not the day its reply came, and +# one whose time can't be trusted counts nowhere. +yesterday = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat(timespec="milliseconds").replace("+00:00", "Z") +dated = work / "dated" +write(dated / "midnight.jsonl", [ + {"type": "model_change", "id": "dm", "timestamp": stamp, "provider": "demo", "modelId": "demo-model"}, + {"type": "message", "id": "du", "timestamp": yesterday, "message": {"role": "user", "content": "late"}}, + {"type": "message", "id": "da", "timestamp": stamp, "message": {"role": "assistant", "provider": "demo", "model": "demo-model", "usage": usage(40)}}, + {"type": "message", "id": "dx", "timestamp": 0, "message": {"role": "user", "content": "untimed"}}, + {"type": "message", "id": "db", "timestamp": stamp, "message": {"role": "assistant", "provider": "demo", "model": "demo-model", "usage": usage(50)}}, +]) +rows, _ = collectors.scan_pi_sessions(r'"provider"\s*:\s*"demo"', lambda name, model, entry: name == "demo", roots=[dated]) +today = collectors.local_day(stamp) +check("a prompt counts on its own day, and an untimed one not at all", + sorted((row.day, row.input_tokens, row.prompts) for row in rows) + == sorted([(collectors.local_day(yesterday), 0, 1), (today, 40, 0), (today, 50, 0)]), + json.dumps([row._asdict() for row in rows])) + + +def turn(n, provider, model): + return [{"type": "message", "id": f"fu{n}", "timestamp": stamp, "message": {"role": "user", "content": "go", "timestamp": 1}}, + {"type": "message", "id": f"fa{n}", "timestamp": stamp, "message": {"role": "assistant", "provider": provider, "model": model, "usage": usage(100), "timestamp": 2}}] + + +# A fork copies its parent's whole history, as pi --fork does. A collector +# that owns one fork's own reply reads the other files of the fork tree +# only when they hold a copy of a prompt it booked: here none does. +forks = work / "forks" +history = [{"type": "model_change", "id": "fm", "timestamp": stamp, "provider": "busy", "modelId": "busy-model"}] +for n in range(20): + history += turn(n, "busy", "busy-model") +write(forks / "parent.jsonl", [{"type": "session", "version": 3, "id": "fp", "timestamp": stamp, "cwd": "/proj"}, *history]) +for n in range(6): + extra = [] + if n == 5: + extra = [{"type": "model_change", "id": "fk", "timestamp": stamp, "provider": "late", "modelId": "late-model"}, *turn("k", "late", "late-model")] + write(forks / f"fork-{n}.jsonl", [{"type": "session", "version": 3, "id": f"ff{n}", "timestamp": stamp, "cwd": "/proj", "parentSession": str(forks / "parent.jsonl")}, *history, *extra]) +for key in parsed: + parsed[key] = set() if key != "lines" else 0 +rows, complete = collectors.scan_pi_sessions(r'"provider"\s*:\s*"late"', lambda name, model, entry: name == "late", roots=[forks]) +check("a fork's own reply parses no other file of its fork tree that holds none of its prompts", + complete and parsed["files"] == {"fork-5.jsonl"} and [(row.input_tokens, row.prompts) for row in rows] == [(100, 1)], + json.dumps({"files": sorted(parsed["files"]), "rows": [row._asdict() for row in rows]})) + +# pi deletes a session from /resume and leaves its forks naming it. The +# fork's copies then count under the fork, and the scan is complete. +orphan = work / "orphan" +write(orphan / "fork.jsonl", [{"type": "session", "version": 3, "id": "of", "timestamp": stamp, "cwd": "/proj", "parentSession": str(orphan / "deleted.jsonl")}, + *history, {"type": "model_change", "id": "ok", "timestamp": stamp, "provider": "late", "modelId": "late-model"}, *turn("k", "late", "late-model")]) +for provider, want in (("busy", (2000, 20)), ("late", (100, 1))): + rows, complete = collectors.scan_pi_sessions(r'"provider"\s*:\s*"' + provider + '"', lambda name, model, entry: name == provider, roots=[orphan]) + check(f"a fork whose parent was deleted is read in full for {provider}", + complete and (sum(row.input_tokens for row in rows), sum(row.prompts for row in rows)) == want, + json.dumps({"complete": complete, "rows": [row._asdict() for row in rows]})) + +# Two forks of that deleted session copied its last prompt, still +# unanswered, and answered it on different providers: the gone parent still +# ties them, so the prompt counts once, under one collector. +siblings = work / "siblings" +copied = [{"type": "model_change", "id": "sm", "timestamp": stamp, "provider": "busy", "modelId": "busy-model"}, + {"type": "message", "id": "sq", "timestamp": stamp, "message": {"role": "user", "content": "go", "timestamp": 1}}] +for name, provider in (("a", "alpha"), ("b", "beta")): + write(siblings / f"fork-{name}.jsonl", [{"type": "session", "version": 3, "id": f"s{name}", "timestamp": stamp, "cwd": "/proj", "parentSession": str(siblings / "deleted.jsonl")}, + *copied, {"type": "model_change", "id": f"sc{name}", "timestamp": stamp, "provider": provider, "modelId": f"{provider}-model"}, + {"type": "message", "id": f"sr{name}", "timestamp": stamp, "message": {"role": "assistant", "provider": provider, "model": f"{provider}-model", "usage": usage(100), "timestamp": 2}}]) +prompts, completes = 0, [] +for provider in ("alpha", "beta"): + rows, complete = collectors.scan_pi_sessions(r'"provider"\s*:\s*"' + provider + '"', lambda name, model, entry, provider=provider: name == provider, roots=[siblings]) + prompts += sum(row.prompts for row in rows) + completes.append(complete) +check("a prompt two forks of a deleted session copied and answered on different providers counts once", + all(completes) and prompts == 1, json.dumps({"prompts": prompts, "complete": completes})) +PY diff --git a/test/shell.d/agent-usage-qwen-test.sh b/test/shell.d/agent-usage-qwen-test.sh new file mode 100755 index 00000000000..94479064938 --- /dev/null +++ b/test/shell.d/agent-usage-qwen-test.sh @@ -0,0 +1,409 @@ +#!/bin/bash + +# The Qwen Code collector's own sources: its chat transcripts, read the way +# Qwen Code's own statistics read them, where QWEN_HOME, QWEN_RUNTIME_DIR and +# settings.json put them, the per-transcript index, and the pi, omp and +# opencode turns on Alibaba's Qwen plan provider ids. The unused-machine +# record and exactly-once counting across collectors are held by the +# partition test, and the scan cache by the contract check. Neither Qwen Code +# nor Alibaba publishes a way to read a plan's quota, so there is no limits +# contract to hold. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +# The checks below compare today's totals with rows stamped NOW. +freeze_today() +NOW = fixture_now().isoformat(timespec="milliseconds").replace("+00:00", "Z") +TWO_DAYS_AGO = (fixture_now() - timedelta(days=2)).isoformat(timespec="milliseconds").replace("+00:00", "Z") + + +def qwen(*args): + return collect("qwen", *(args or ("--force",))) + + +def prompt(uuid, session, timestamp, subtype=None): + line = {"uuid": uuid, "parentUuid": None, "sessionId": session, "timestamp": timestamp, "type": "user", "message": {"role": "user", "parts": [{"text": "hi"}]}} + if subtype: + line["subtype"] = subtype + return line + + +# One billed answer, as Qwen Code logs it: a ui_telemetry api_response event. +def response(uuid, session, timestamp, model, auth, input, output, cached, thoughts, response_id="", total=0): + event = {"event.name": "qwen-code.api_response", "event.timestamp": timestamp, "response_id": response_id, "model": model, "auth_type": auth, + "input_token_count": input, "output_token_count": output, "cached_content_token_count": cached, "thoughts_token_count": thoughts, "total_token_count": total} + return {"uuid": uuid, "sessionId": session, "timestamp": timestamp, "type": "system", "subtype": "ui_telemetry", "systemPayload": {"uiEvent": event}} + + +def transcript(path, lines, mode="w"): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open(mode) as handle: + for line in lines: + handle.write((line if isinstance(line, str) else json.dumps(line)) + "\n") + return path + + +def chats(project="p"): + return home() / ".qwen" / "projects" / project / "chats" + + +# Qwen Code sends an OpenAI-compatible reply to OPENAI_BASE_URL, else to +# OpenAI's API, its default. The checks that aren't about routes run Qwen Code +# pointed at Model Studio, as from the shell it was started in. +def studio(): + sandbox() + os.environ["OPENAI_BASE_URL"] = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1" + + +# A machine that never ran Qwen Code has no local stats to show. +sandbox() +record = qwen() +expect("Qwen with no usage says it has no local stats", [record["ready"], record["hasLocalStats"]], [False, False], record) + +# ------------------------------------------------------------- transcripts + +studio() +transcript(chats("-home-me-app") / "s1.jsonl", [ + prompt("u1", "s1", NOW), + response("r1", "s1", NOW, "qwen3-coder-plus", "qwen-oauth", 1000, 200, 600, 50, "resp-1"), + prompt("u2", "s1", NOW, "mid_turn_user_message"), + # Goal and notification records are not typed prompts. + prompt("u3", "s1", NOW, "goal_runtime"), + # A provider that reports only the cached part reported a fully cached prompt. + response("r3", "s1", NOW, "qwen3-coder-plus", "openai", 0, 5, 80, 0, "resp-3"), + # Gemini's API is Google's, not a Qwen plan's, and so is the prompt it + # answers. + prompt("u5", "s1", NOW), + response("r2", "s1", NOW, "gemini-2.5-pro", "gemini", 300, 40, 0, 10, "resp-2"), + {"uuid": "x", "type": "system", "subtype": "ui_telemetry", "systemPayload": {"uiEvent": {"event.name": "qwen-code.tool_call", "model": "m"}}}, + {"uuid": "y", "type": "assistant", "model": "qwen3-coder-plus", "usageMetadata": {"promptTokenCount": 999}}, +]) +earlier = [prompt("u4", "s2", TWO_DAYS_AGO), response("r4", "s2", TWO_DAYS_AGO, "qwen3-coder-plus", "openai", 70, 7, 0, 3, "resp-4")] +transcript(chats("-home-me-app") / "s2.jsonl", earlier) +# A branched session carries the records it was branched from. +transcript(chats("-home-me-app") / "s3.jsonl", earlier) +transcript(chats("-home-me-app") / "s1.ledger.jsonl", [{"uuid": "l", "type": "user", "sessionId": "s1", "timestamp": NOW}]) + +record = qwen() +two_days_ago = (fixture_now().astimezone() - timedelta(days=2)).strftime("%Y-%m-%d") +expect("Qwen counts typed prompts and billed responses once each, booked on the day each arrived", { + "name": record["name"], + "counts": {key: record[key] for key in ("ready", "hasPromptStats", "todayPrompts", "todaySessions", "todayTotalTokens", "totalPrompts", "totalSessions", "activeDays")}, + "two days ago": next(day["messageCount"] for day in record["recentDays"] if day["date"] == two_days_ago), +}, { + "name": "Qwen Code", + "counts": {"ready": True, "hasPromptStats": True, "todayPrompts": 2, "todaySessions": 1, "todayTotalTokens": 1285, "totalPrompts": 3, "totalSessions": 2, "activeDays": 2}, + "two days ago": 77, +}, record) +expect("Qwen reads tokens the way Qwen Code's own statistics do", record["modelUsage"], { + "qwen3-coder-plus": {"inputTokens": 470, "outputTokens": 212, "cacheReadInputTokens": 680, "cacheCreationInputTokens": 0}, +}, record) + +# Qwen Code 0.26.0's records through the Responses API, whose output_tokens +# already include reasoning_tokens, with and without a total, and a record +# kept without its auth type, with a total: thinking the total already holds +# in the output counts once. +studio() +transcript(chats() / "r.jsonl", [ + prompt("u1", "r", NOW), + response("r1", "r", NOW, "qwen-resp-model", "openai-responses", 5234, 55, 75, 5, "resp_mock5", total=5289), + response("r2", "r", NOW, "qwen-resp-model", "openai-responses", 6234, 56, 76, 5, "resp_mock6"), + response("r3", "r", NOW, "qwen3-coder-plus", "", 300, 40, 0, 10, "resp-g1", total=350), + response("r4", "r", NOW, "qwen3-coder-plus", "", 300, 40, 0, 10, "resp-g2", total=340), +]) +record = qwen() +expect("Qwen counts a Responses API reply's reasoning once, and other thinking only outside its total", { + model: usage["outputTokens"] for model, usage in record["modelUsage"].items() +}, {"qwen-resp-model": 111, "qwen3-coder-plus": 90}, record) + +# Qwen Code 0.26.0's records from two sessions against a server that recycles +# response ids: each got chatcmpl-mock1208 and chatcmpl-mock1209 for replies +# of the same size, a minute apart. Both count. --fork-session copies a +# record verbatim, its uuid included, under the new session with a +# forkedFrom, and that copy counts once, in the session that wrote it, though +# the fork's file name sorts first; the fork counts for the prompt typed into +# it. +studio() +first = [ + prompt("8b0e5c1a-6f4d-4a0e-9a51-3c0f1b6c2d10", "f73baca5-d3d4-4eb9-8008-20dc4951f613", NOW), + response("8f4c3a4e-5876-4c73-9356-6507b21f70ea", "f73baca5-d3d4-4eb9-8008-20dc4951f613", NOW, "qwen-sp-model", "openai", 9312, 58, 78, 0, "chatcmpl-mock1208", total=9370), + response("28ac483b-350e-45aa-9286-486a3964925a", "f73baca5-d3d4-4eb9-8008-20dc4951f613", NOW, "qwen-sp-model", "openai", 10313, 59, 79, 0, "chatcmpl-mock1209", total=10372), +] +transcript(chats("-home-me-proj-qwen10") / "f73baca5-d3d4-4eb9-8008-20dc4951f613.jsonl", first) +transcript(chats("-home-me-proj-qwen11") / "61cea1b3-8fa7-433b-b69c-cce255da2433.jsonl", [ + prompt("4e2b7d90-1c3a-4f55-8e6b-7a9d0c1e2f34", "61cea1b3-8fa7-433b-b69c-cce255da2433", NOW), + response("0c2fb936-82a9-4a54-9d94-7c043a11b674", "61cea1b3-8fa7-433b-b69c-cce255da2433", NOW, "qwen-sp-model", "openai", 9312, 58, 78, 0, "chatcmpl-mock1208", total=9370), + response("d7f750d2-72ec-47ef-82a1-9e1400e2c289", "61cea1b3-8fa7-433b-b69c-cce255da2433", NOW, "qwen-sp-model", "openai", 10313, 59, 79, 0, "chatcmpl-mock1209", total=10372), +]) +branched = [{**line, "sessionId": "9d1e4b27-3f6a-4c08-b5e2-0a7c8d9e1f23", "forkedFrom": {"sessionId": line["sessionId"], "messageUuid": line["uuid"]}} for line in first] +transcript(chats("-home-me-proj-qwen10") / "9d1e4b27-3f6a-4c08-b5e2-0a7c8d9e1f23.jsonl", [ + *branched, + prompt("6bcf502c-f7ae-46ec-882c-721ccc6ca9a3", "9d1e4b27-3f6a-4c08-b5e2-0a7c8d9e1f23", NOW), + response("4b804879-1b4d-45b5-a2d6-81ef1ee79f4b", "9d1e4b27-3f6a-4c08-b5e2-0a7c8d9e1f23", NOW, "qwen-sp-model", "openai", 4314, 46, 0, 0, "chatcmpl-mock1210", total=4360), +]) +record = qwen() +expect("Qwen counts different replies a server gave the same response id, and a forked copy once, in its source session", { + key: record[key] for key in ("todayTotalTokens", "todayPrompts", "todaySessions") +}, {"todayTotalTokens": 2 * (9370 + 10372) + 4360, "todayPrompts": 3, "todaySessions": 3}, record) + +# ------------------------------------------------------------------ routes + +# A reply counts where it was billed. Qwen Code records only its auth type, +# and settings.json's modelProviders the base URL each model is sent to: +# Alibaba's hosts are Qwen's, OpenRouter's or a self-hosted server's are +# theirs, and Anthropic's API, which Qwen Code reaches with --auth-type +# anthropic, is no Qwen plan, nor OpenAI's, where Qwen Code sends a model no +# route or OPENAI_BASE_URL places. Each prompt goes with the reply that +# answered it. Shaped as Qwen Code 0.26.0 wrote them against a mock of each +# API. +sandbox() +write_json(home() / ".qwen" / "settings.json", {"modelProviders": {"openai": [ + {"id": "qwen3-coder-plus", "baseUrl": "https://coding-intl.dashscope.aliyuncs.com/v1", "envKey": "BAILIAN_CODING_PLAN_API_KEY"}, + {"id": "qwen/qwen3-coder", "baseUrl": "https://openrouter.ai/api/v1", "envKey": "OPENROUTER_API_KEY"}, + {"id": "local-qwen", "baseUrl": "http://127.0.0.1:18083/v1"}, +]}}) +transcript(chats() / "routes.jsonl", [ + prompt("u1", "routes", NOW), + response("r1", "routes", NOW, "qwen3-coder-plus", "openai", 1234, 51, 71, 5, "chatcmpl-mock1", total=1285), + prompt("u2", "routes", NOW), + response("r2", "routes", NOW, "claude-mock-model", "anthropic", 7234, 57, 77, 0, "msg_mock7", total=7291), + prompt("u3", "routes", NOW), + response("r3", "routes", NOW, "qwen/qwen3-coder", "openai", 5000, 50, 0, 0, "chatcmpl-or"), + prompt("u4", "routes", NOW), + response("r4", "routes", NOW, "local-qwen", "openai", 3000, 30, 0, 0, "chatcmpl-local"), + prompt("u5", "routes", NOW), + response("r5", "routes", NOW, "qwen-unlisted", "openai", 100, 10, 0, 0, "chatcmpl-default"), +]) +record = qwen() +expect("Qwen counts the replies Alibaba billed, with their prompts, and leaves other providers' out", + [sorted(record["modelUsage"]), record["todayPrompts"], record["todayTotalTokens"]], + [["qwen3-coder-plus"], 1, 1285], record) +local = collect("local", "--force") +expect("Local counts the Qwen Code replies sent to a self-hosted server", + [local["todayTokensByModel"], local["todayPrompts"]], [{"local-qwen": 3030}, 1], local) + + +# A request no reply answered, as Qwen Code logs it: a ui_telemetry api_error +# event naming the model and auth type it was sent with. +def failed(uuid, session, timestamp, model, auth, execution): + event = {"event.name": "qwen-code.api_error", "event.timestamp": timestamp, "response_id": "", "execution_id": execution, "model": model, "duration_ms": 2990, + "prompt_id": f"{session}########0", "auth_type": auth, "error_message": "Connection error. (cause: connect ECONNREFUSED 127.0.0.1:18499)", "error_type": "APIConnectionError"} + return {"uuid": uuid, "sessionId": session, "timestamp": timestamp, "type": "system", "subtype": "ui_telemetry", "systemPayload": {"uiEvent": event}} + + +# A prompt no reply answered goes where its requests were sent, else where +# settings.json's selected model goes: a server you run that refused them all +# is local's, not Qwen's. Shaped as Qwen Code 0.26.0 wrote a run against a +# dead port on a modelProviders route. +sandbox() +write_json(home() / ".qwen" / "settings.json", { + "modelProviders": {"openai": [{"id": "qwen-dead-model", "name": "dead", "baseUrl": "http://127.0.0.1:18499/v1", "envKey": "MOCK_QWEN_KEY"}, + {"id": "qwen3-coder-plus", "baseUrl": "https://coding-intl.dashscope.aliyuncs.com/v1"}]}, + "security": {"auth": {"selectedType": "openai"}}, + "model": {"name": "qwen-dead-model"}, +}) +transcript(chats() / "93e3e301.jsonl", [ + prompt("f5316beb", "93e3e301", NOW), + failed("1e40dc03", "93e3e301", NOW, "qwen-dead-model", "openai", "cb1ab8b5"), + failed("a76f1d2f", "93e3e301", NOW, "qwen-dead-model", "openai", "0746ab80"), +]) +transcript(chats() / "plan.jsonl", [prompt("p1", "plan", NOW), failed("p2", "plan", NOW, "qwen3-coder-plus", "openai", "e1")]) +# Typed, and stopped before any request: the selected model's route decides. +transcript(chats() / "bare.jsonl", [prompt("b1", "bare", NOW)]) +records = {cid: collect(cid, "--force") for cid in ("qwen", "local")} +expect("a Qwen Code prompt no reply answered counts where its route goes", + {cid: [record["todayPrompts"], record["todaySessions"], record["todayTotalTokens"]] for cid, record in records.items()}, + {"qwen": [1, 1, 0], "local": [2, 2, 0]}, records) + +# A scan answers only for the routes it was made under: a route moved after +# it is reflected at once, even on a --limits-only run. As Qwen Code 0.26.0 +# wrote a run against a server of the user's own. +def env_run(): + sandbox() + transcript(chats() / "env.jsonl", [ + prompt("54ccea9a", "env", NOW), + response("e1", "env", NOW, "qwen-env-model", "openai", 6309, 55, 75, 0, "chatcmpl-mock405", total=6364), + response("e2", "env", NOW, "qwen-env-model", "openai", 7310, 56, 76, 0, "chatcmpl-mock406", total=7366), + ]) + + +env_run() +write_json(home() / ".qwen" / "settings.json", {"modelProviders": {"openai": [{"id": "qwen-env-model", "baseUrl": "http://127.0.0.1:18491/v1"}]}}) +expect("Qwen leaves out a reply its route sends to a server of the user's own", collect("qwen")["todayTotalTokens"], 0) +write_json(home() / ".qwen" / "settings.json", {"modelProviders": {"openai": [{"id": "qwen-env-model", "baseUrl": "https://coding-intl.dashscope.aliyuncs.com/v1"}]}}) +expect("a Qwen scan made before a modelProviders route moved is not reused after it", collect("qwen", "--limits-only")["todayTotalTokens"], 13730) + +# A model no modelProviders route names is sent where OPENAI_BASE_URL (or +# ANTHROPIC_BASE_URL) says, which Qwen Code takes from the first .env file +# setting it (its home's, then ~/.env), else from settings.json's env, else +# from the security.auth.baseUrl /auth saved. +env_run() +ENV_ROUTES = [ + ("~/.qwen/.env", lambda: (home() / ".qwen" / ".env").write_text("OPENAI_API_KEY=sk-dummy\nexport OPENAI_BASE_URL=\"http://127.0.0.1:18491/v1\" # mock\nOPENAI_MODEL=qwen-env-model\n"), "local"), + ("~/.env", lambda: (home() / ".env").write_text("OPENAI_BASE_URL=http://127.0.0.1:18491/v1\n"), "local"), + ("settings.json env", lambda: write_json(home() / ".qwen" / "settings.json", {"env": {"OPENAI_BASE_URL": "http://192.168.1.20:8000/v1"}}), "local"), + ("security.auth.baseUrl", lambda: write_json(home() / ".qwen" / "settings.json", {"security": {"auth": {"selectedType": "openai", "baseUrl": "http://127.0.0.1:18491/v1"}}}), "local"), + ("settings.json env, OpenRouter", lambda: write_json(home() / ".qwen" / "settings.json", {"env": {"OPENAI_BASE_URL": "https://openrouter.ai/api/v1"}}), None), + ("settings.json env, Model Studio", lambda: write_json(home() / ".qwen" / "settings.json", {"env": {"OPENAI_BASE_URL": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"}}), "qwen"), + # The first .env file wins over settings.json's env. + ("~/.qwen/.env over settings.json env", lambda: ((home() / ".qwen" / ".env").write_text("OPENAI_BASE_URL=http://127.0.0.1:18491/v1\n"), + write_json(home() / ".qwen" / "settings.json", {"env": {"OPENAI_BASE_URL": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"}})), "local"), +] + + +def clear_routes(): + for path in (home() / ".qwen" / ".env", home() / ".env", home() / ".qwen" / "settings.json"): + path.unlink(missing_ok=True) + + +for where, configure, owner in ENV_ROUTES: + clear_routes() + configure() + records = {cid: collect(cid, "--force") for cid in ("qwen", "local")} + expect(f"a Qwen Code reply sent where {where} says counts under {owner or 'nobody'}", + {cid: [record["todayTotalTokens"], record["todayPrompts"]] for cid, record in records.items()}, + {cid: [13730, 1] if cid == owner else [0, 0] for cid in records}, records) + +# A .env route set after a scan is reflected at once too. +clear_routes() +expect("Qwen leaves out an OpenAI-compatible reply with no route, which went to OpenAI's API", collect("qwen")["todayTotalTokens"], 0) +(home() / ".qwen" / ".env").write_text("OPENAI_BASE_URL=https://dashscope-intl.aliyuncs.com/compatible-mode/v1\n") +expect("a Qwen scan made before a .env route is not reused after it", collect("qwen", "--limits-only")["todayTotalTokens"], 13730) + +# --------------------------------------------------------- untrusted fields + +studio() +transcript(chats() / "s.jsonl", [ + prompt("bad1", "s", "yesterday-ish"), + prompt("bad2", "s", "2999-01-01T00:00:00Z"), + response("bad3", "s", NOW, "m", "openai", "100", True, float("nan"), -5, "a"), + response("bad4", "s", NOW, "m", "openai", 1e30, 0, 0, 0, "b"), + response("ok1", "s", NOW, "m", "openai", 10, 2, 0, 0, "c"), + prompt("bad5", "s", "\u00b2"), + # A line nested deeper than the parser will go. + "[" * 100000 + '"user"', +]) +with (chats() / "s.jsonl").open("a") as torn: + torn.write('{"uuid":"z","type":"user","timestamp":') +record = qwen() +expect("Qwen skips bad times, bad token fields, a line nested too deep, and a torn last line", + [record["todayPrompts"], record["todayTotalTokens"]], [0, 12], record) + +# A record with no time is dated once, by the transcript's time as it is +# read, so it isn't today's again on every day after. +studio() +undated = transcript(chats() / "u.jsonl", [prompt("nt1", "u", None), response("nt2", "u", None, "m", "openai", 10, 2, 0, 0, "nt")]) +three_days_ago = time.time() - 3 * 86400 +os.utime(undated, (three_days_ago, three_days_ago)) +record = qwen() +expect("an undated record is dated by its transcript, once", + [record["todayPrompts"], record["todayTotalTokens"], record["totalPrompts"], record["activeDates"]], + [0, 0, 1, [datetime.fromtimestamp(three_days_ago).strftime("%Y-%m-%d")]], record) + +# ---------------------------------------------------------------- locations + +studio() +transcript(home() / "runtime" / "projects" / "p" / "chats" / "s.jsonl", [response("a", "s", NOW, "m", "openai", 1, 0, 0, 0, "ra")]) +transcript(home() / "qhome" / "projects" / "p" / "chats" / "s.jsonl", [response("b", "s", NOW, "m", "openai", 10, 0, 0, 0, "rb")]) +transcript(chats() / "s.jsonl", [response("c", "s", NOW, "m", "openai", 100, 0, 0, 0, "rc")]) +os.environ.update(QWEN_HOME=str(home() / "qhome"), QWEN_RUNTIME_DIR=str(home() / "runtime")) +record = qwen() +del os.environ["QWEN_HOME"], os.environ["QWEN_RUNTIME_DIR"] +expect("Qwen honors QWEN_HOME and QWEN_RUNTIME_DIR", record["todayTotalTokens"], 11, record) + +write_json(home() / ".qwen" / "settings.json", {"advanced": {"runtimeOutputDir": "~/out"}}) +transcript(home() / "out" / "projects" / "p" / "chats" / "s.jsonl", [response("d", "s", NOW, "m", "openai", 1000, 0, 0, 0, "rd")]) +record = qwen() +expect("Qwen reads the runtimeOutputDir from settings.json", record["todayTotalTokens"], 1100, record) + +# A runtimeOutputDir no path can name is no place to look. +write_json(home() / ".qwen" / "settings.json", {"advanced": {"runtimeOutputDir": "/\u0000x"}}) +record = qwen() +expect("Qwen ignores a runtimeOutputDir holding a NUL", record["todayTotalTokens"], 100, record) + +# ------------------------------------------------- Qwen plans in pi and opencode + +# pi, omp and opencode turns on Alibaba's Qwen plans count beside Qwen Code's +# own transcripts: every plan id once, and nothing on a pay-as-you-go key, a +# custom gateway that only starts with a plan's name, or a bare `qwen`. +studio() +PLAN_IDS = ["alibaba-coding-plan", "alibaba-coding-plan-cn", "alibaba-token-plan", "alibaba-token-plan-cn", + "qwen-token-plan", "qwen-token-plan-cn", "qwen-token-plan-individual", "qwen-portal"] +OTHER_IDS = ["alibaba", "alibaba-cn", "alibaba-coding-plan-proxy", "qwen", "dashscope"] +transcript(chats() / "s.jsonl", [prompt("u", "s", NOW), response("r", "s", NOW, "qwen3-coder-plus", "openai", 1, 0, 0, 0, "rq")]) +write_pi_session(home() / ".pi" / "agent" / "sessions" / "p" / "pi.jsonl", + [pi_reply(provider, 10 ** index, id=f"pi-{provider}", model=f"pi/{provider}") for index, provider in enumerate(PLAN_IDS[:4])] + + [pi_reply(provider, 7, id=f"pi-{provider}", model="other") for provider in OTHER_IDS]) +write_pi_session(home() / ".omp" / "agent" / "sessions" / "p" / "omp.jsonl", + [pi_reply(provider, 10 ** index, id=f"omp-{provider}", model=f"omp/{provider}") for index, provider in enumerate(PLAN_IDS[4:], 4)]) +opencode_db(message=[(f"m-{provider}", "oc1", oc_v1(provider, 3, model=f"oc-{provider}")) for provider in PLAN_IDS + OTHER_IDS], + session_message=[("v2-coding", "oc2", "assistant", oc_v2("alibaba-coding-plan", 5, model="qwen3-coder-next"))]) +record = qwen() +expect("Qwen counts every Qwen plan id from pi, omp and opencode once, beside its own transcripts", { + "tokens": record["todayTotalTokens"], + "models": sorted(record["modelUsage"]), + "prompts": record["todayPrompts"], + "sessions": record["todaySessions"], +}, { + "tokens": 1 + sum(10 ** index for index in range(len(PLAN_IDS))) + 3 * len(PLAN_IDS) + 5, + "models": sorted(["qwen3-coder-plus", "qwen3-coder-next"] + PLAN_IDS + [f"oc-{provider}" for provider in PLAN_IDS]), + "prompts": 1 + len(PLAN_IDS) + len(PLAN_IDS) + 1, + "sessions": 5, +}, record) +# Qwen plan usage alone makes the tab: there is no quota to wait on a +# sign-in for, so the status stays empty and no meters show. +sandbox() +opencode_db(session_message=[("only", "oc", "assistant", oc_v2("qwen-portal", 5))]) +record = qwen() +expect("Qwen shows Qwen plan usage without Qwen Code transcripts or a sign-in", + [record["ready"], record["usageStatusText"], record["authHelpText"], record["limits"], record["todayTotalTokens"]], + [True, "", "", [], 5], record) + +# ---------------------------------------------------------------- contract + +studio() + + +def add_response(): + add_response.count = getattr(add_response, "count", 0) + 1 + transcript(chats() / "s.jsonl", [response(f"g{add_response.count}", "s", fixture_now().isoformat(), "m", "openai", 5, 0, 0, 0, f"rg{add_response.count}")], mode="a") + + +check_scan_cache(Adapter("qwen", grow=add_response)) + +# The transcript index reads only what was appended, and starts over when the +# file is replaced, even by a longer one. +others = qwen()["todayTotalTokens"] +transcript(chats() / "t.jsonl", [response("a", "t", NOW, "m", "openai", 10, 0, 0, 0, "ra"), response("b", "t", NOW, "m", "openai", 5, 0, 0, 0, "rb")]) +expect("Qwen reads a new transcript", qwen()["todayTotalTokens"] - others, 15) +transcript(chats() / "t.jsonl", [response("c", "t", NOW, "m", "openai", 20, 0, 0, 0, "rc")], mode="a") +expect("Qwen picks up lines appended to a transcript", qwen()["todayTotalTokens"] - others, 35) +replacement = transcript(chats() / "t.new", [response(name, "t", NOW, "m", "openai", tokens, 0, 0, 0, f"r{name}") for name, tokens in zip("defg", (1, 2, 3, 4))]) +replacement.rename(chats() / "t.jsonl") +expect("Qwen rereads a transcript that was replaced", qwen()["todayTotalTokens"] - others, 10) + +# A scan that couldn't read every transcript isn't cached. +transcript(chats() / "s.jsonl", [response("h", "s", NOW, "m", "openai", 5, 0, 0, 0, "rh")], mode="a") +(chats() / "s.jsonl").chmod(0) +try: + if os.access(chats() / "s.jsonl", os.R_OK): + print("ok - Qwen does not cache a scan that could not read every transcript # SKIP running as root") + else: + for path in scan_cache_files("qwen"): + path.unlink() + with contextlib.redirect_stderr(io.StringIO()): + qwen() + check("Qwen does not cache a scan that could not read every transcript", not scan_cache_files("qwen"), + ", ".join(path.name for path in scan_cache_files("qwen"))) +finally: + (chats() / "s.jsonl").chmod(0o644) +PY diff --git a/test/shell.d/agent-usage-shared-test.sh b/test/shell.d/agent-usage-shared-test.sh new file mode 100644 index 00000000000..1f40dea5212 --- /dev/null +++ b/test/shell.d/agent-usage-shared-test.sh @@ -0,0 +1,1344 @@ +#!/bin/bash + +# The shared collector module, shell/plugins/agents/collectors.py, holds the +# parts of the collector contract every collector shares. This test holds the +# module's pieces to the contract directly, so a collector built on them +# inherits it, and a collector's own test covers only its own sources, +# credentials, and answer. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command python3 +require_command rg + +WORK=$(mktemp -d) +trap 'rm -rf "$WORK"' EXIT + +WORK="$WORK" python3 - <<'PY' +import contextlib +import json +import os +import sqlite3 +import sys +import time +from datetime import datetime, timedelta, timezone +from pathlib import Path + +sys.dont_write_bytecode = True +sys.path.insert(0, os.path.join(os.environ["OMARCHY_PATH"], "shell", "plugins", "agents")) +import collectors # noqa: E402 + +work = Path(os.environ["WORK"]) +home = work / "home" +home.mkdir() +os.environ["HOME"] = str(home) +os.environ["XDG_CACHE_HOME"] = str(work / "cache") +os.environ["XDG_DATA_HOME"] = str(home / ".local" / "share") +for name in ("PI_CODING_AGENT_DIR", "PI_CODING_AGENT_SESSION_DIR", "OPENCODE_DB"): + os.environ.pop(name, None) + + +def check(description, condition, detail=""): + if condition: + print("ok - " + description) + return + print("not ok - " + description, file=sys.stderr) + if detail: + print(detail, file=sys.stderr) + sys.exit(1) + + +# "today" is held to the day this run began, so one that straddles local +# midnight still dates its rows, stats and scans by one day. +sys.path.insert(0, os.path.join(os.environ["ROOT"], "test", "shell.d")) +from agent_usage_fixtures import fixture_now, freeze_today # noqa: E402 + +freeze_today() +today = collectors.today() +stamp = fixture_now().isoformat(timespec="milliseconds").replace("+00:00", "Z") + +# ------------------------------------------------------------- untrusted values + +check("local_day reads a row with no time as today's", collectors.local_day(None) == today) +check( + "local_day skips zero, negative, huge, NaN, boolean, and garbage times", + [collectors.local_day(value) for value in (0, -5, 1e300, float("nan"), True, "not a date", "", [1])] == [None] * 8, +) +check( + "local_day keeps the last half second of a day on that day, in seconds and milliseconds", + (lambda midnight: [collectors.local_day(midnight - 0.4), collectors.local_day((midnight - 0.4) * 1000)] + == [(datetime.fromtimestamp(midnight) - timedelta(days=1)).strftime("%Y-%m-%d")] * 2)( + datetime.combine(datetime.now().date(), datetime.min.time()).timestamp()), +) +check( + "local_day skips a damaged stamp dated before 2000 or after 2100", + [collectors.local_day(value) for value in (1, 86400, "1999-12-31T12:00:00", "2102-06-01T00:00:00Z")] == [None] * 4, +) +check( + "local_day skips digit strings that aren't a time: other scripts' digits and thousands of digits", + [collectors.local_day(value) for value in ("\u0661\u0667\u0660\u0660\u0660\u0660\u0660\u0660\u0660\u0660", "9" * 5000)] == [None] * 2, +) +check("local_day reads epoch milliseconds, seconds, and digit strings alike", len({ + collectors.local_day(1_700_000_000_000), + collectors.local_day(1_700_000_000), + collectors.local_day("1700000000000"), +}) == 1) +check( + "token counts skip booleans, strings, NaN, negatives, and absurd values", + [collectors.token_count(value) for value in (True, "12", float("nan"), -1, 1e300, 12)] == [0, 0, 0, 0, 0, 12], +) +check( + "usage_tokens takes the first spelling present", + collectors.usage_tokens({"input_tokens": 9, "input": 3}, "input", "inputTokens", "input_tokens") == 3, +) + +# ------------------------------------------------------------------ time + +real_time = time.time +time.time = lambda: 1_000.0009 +try: + check("stamps are floored, never rounded", collectors.now_ms() == 1_000_000) +finally: + time.time = real_time + +now = time.time() +check("the age guard takes a stamp younger than the limit", collectors.age_below(now - 5, 15)) +check("the age guard refuses a stamp as old as the limit", not collectors.age_below(now - 15, 15)) +check("the age guard refuses a stamp from the future", not collectors.age_below(now + 60, 15)) +check("a zero limit reuses nothing", not collectors.age_below(now, 0)) +check("an untrusted stamp reads as never", [collectors.stamp(value) for value in (True, "5", float("inf"), -3, 7.9)] == [0, 0, 0, 0, 7]) +check( + "the scan reuse windows are 20, 900, and 0 seconds", + (collectors.scan_max_age(False, False), collectors.scan_max_age(False, True), collectors.scan_max_age(True, True)) == (20, 900, 0), +) + +# ------------------------------------------------------------------- files + +# Collectors overlap in practice: the update command backgrounds one per agent +# while the panel refreshes on its own. Writers aiming at one cache file must +# all land, not trip over a shared temp path. +import threading # noqa: E402 + +race = work / "race" / "cache.json" +race.parent.mkdir() +failures = [] +barrier = threading.Barrier(8) + + +def hammer(writer): + barrier.wait() + for round_ in range(25): + try: + collectors.write_json(race, {"writer": writer, "round": round_}) + except Exception as error: + failures.append(repr(error)) + + +writers = [threading.Thread(target=hammer, args=(writer,)) for writer in range(8)] +for thread in writers: + thread.start() +for thread in writers: + thread.join() +check( + "concurrent writers to one file all land, leaving one intact file and no temp files", + not failures and isinstance(json.loads(race.read_text()).get("writer"), int) and [path.name for path in race.parent.iterdir()] == ["cache.json"], + json.dumps({"failures": failures[:3], "files": sorted(path.name for path in race.parent.iterdir())}), +) +check("written files stay readable to others", race.stat().st_mode & 0o777 == 0o644, oct(race.stat().st_mode)) + +# ------------------------------------------------------------------ record + +record = collectors.base_record("demo", "Demo") +check( + "the unused-machine record is unready, with no status, limits, or balance", + {key: record[key] for key in ("schemaVersion", "id", "name", "ready", "hasLocalStats", "hasPromptStats", "tierLabel", "usageStatusText", "authHelpText", "limits", "limitsStale", "limitsFetchedAt")} + == {"schemaVersion": 1, "id": "demo", "name": "Demo", "ready": False, "hasLocalStats": False, "hasPromptStats": True, "tierLabel": "", "usageStatusText": "", "authHelpText": "", "limits": [], "limitsStale": False, "limitsFetchedAt": 0} + and "balance" not in record, + json.dumps(record), +) +check( + "the unused-machine record carries the full, zeroed stats block", + [day["messageCount"] for day in record["recentDays"]] == [0] * 7 + and record["recentDays"][-1]["date"] == today + and all(record[key] in (0, {}, []) for key in ("todayPrompts", "todaySessions", "todayTotalTokens", "todayTokensByModel", "totalPrompts", "totalSessions", "activeDays", "activeDates", "modelUsage")), + json.dumps(record), +) + +stats = collectors.Stats() +stats.add(today, "s1", "m", 1, 2, 3, 4) +stats.add(today, "s1", "m", 10, 0, 0, 0) +stats.add("2000-01-01", "s2", "n", 5, 0, 0, 0) +snapshot = stats.snapshot() +check( + "the stats accumulator turns replies into the record's stats", + (snapshot["todayPrompts"], snapshot["todaySessions"], snapshot["todayTotalTokens"], snapshot["totalPrompts"], snapshot["totalSessions"], snapshot["activeDates"], snapshot["recentDays"][-1]["messageCount"], snapshot["todayTokensByModel"], snapshot["modelUsage"]["m"]) + == (2, 1, 20, 3, 2, ["2000-01-01", today], 20, {"m": 20}, {"inputTokens": 11, "outputTokens": 2, "cacheReadInputTokens": 3, "cacheCreationInputTokens": 4}), + json.dumps(snapshot), +) + +stats = collectors.Stats() +# A store that books a prompt apart from the tokens it spent. +stats.add_prompts(today, "s1") +stats.add(today, "s1", "m", 10, 5, 0, 0, prompts=0) +# Running totals covering three prompts, and a line with no tokens at all. +stats.add(today, "s2", "n", 4, 0, 0, 0, prompts=3) +stats.add(today, "s3", "z", 0, 0, 0, 0, prompts=0) +# A typed prompt booked on its own line, with no tokens and no model. +stats.add(today, "s4", "", 0, 0, 0, 0) +snapshot = stats.snapshot() +check( + "prompts count apart from tokens, and tokens alone are no prompt", + (snapshot["todayPrompts"], snapshot["totalPrompts"], snapshot["todaySessions"], snapshot["todayTotalTokens"], sorted(snapshot["modelUsage"]), snapshot["todayTokensByModel"]) + == (5, 5, 3, 19, ["m", "n"], {"m": 15, "n": 4}), + json.dumps(snapshot), +) +stats = collectors.Stats() +stats.add(None, None, "m", 7, 0, 0, 0, prompts=0) +stats.add_prompts(None, "s1", 2) +stats.add_prompts("2000-01-01", None) +stats.add_prompts(today, "s2", 0) +snapshot = stats.snapshot() +check( + "undated tokens with no session count toward the totals alone", + (snapshot["modelUsage"]["m"]["inputTokens"], snapshot["todayTotalTokens"], snapshot["totalPrompts"], snapshot["todayPrompts"], snapshot["totalSessions"], snapshot["todaySessions"], snapshot["activeDates"]) + == (7, 0, 3, 0, 2, 1, ["2000-01-01", today]), + json.dumps(snapshot), +) +from datetime import date # noqa: E402 +stats = collectors.Stats(today=date(2020, 3, 1)) +stats.add("2020-03-01", None, "m", 5, 0, 0, 0, prompts=0) +stats.add("2020-02-24", None, "m", 1, 0, 0, 0, prompts=0) +snapshot = stats.snapshot() +check( + "an injected today dates today's fields and the recent week", + (snapshot["todayTotalTokens"], [day["date"] for day in snapshot["recentDays"]][0], snapshot["recentDays"][0]["messageCount"], snapshot["recentDays"][-1]["date"]) + == (5, "2020-02-24", 1, "2020-03-01") + and collectors.Stats(today="2020-03-01").today == "2020-03-01", + json.dumps(snapshot), +) + +# ------------------------------------------------------------- scan cache + +calls = [] + + +def scan(): + calls.append(1) + return {**collectors.empty_stats(), "totalPrompts": len(calls)}, True + + +first = collectors.cached_scan("demo", ["a"], 20, scan) +again = collectors.cached_scan("demo", ["a"], 20, scan) +forced = collectors.cached_scan("demo", ["a"], 0, scan) +check("a scan is reused within its window and rebuilt at zero", (first["totalPrompts"], again["totalPrompts"], forced["totalPrompts"]) == (1, 1, 2)) +cache_file, _ = collectors.scan_cache_paths("demo", ["a"]) +cached = json.loads(cache_file.read_text()) +check("the scan cache records its day", cached["scanDate"] == today and cached["schemaVersion"] == 2) +cached["scanDate"] = "2000-01-01" +cache_file.write_text(json.dumps(cached)) +check("a scan from another day is a miss", collectors.cached_scan("demo", ["a"], 900, scan)["totalPrompts"] == 3) +check("a scan with other options is a miss", collectors.cached_scan("demo", ["a"], 900, scan, match={"modelPrefix": "kimi"})["totalPrompts"] == 4) +check("a scan with the same options is reused", collectors.cached_scan("demo", ["a"], 900, scan, match={"modelPrefix": "kimi"})["totalPrompts"] == 4) +collectors.cached_scan("demo", ["b"], 0, lambda: ({**collectors.empty_stats(), "totalPrompts": 99}, False)) +check("an incomplete scan is never cached", collectors.cached_scan("demo", ["b"], 900, scan)["totalPrompts"] == 5) +partial = collectors.cached_scan("demo", ["c"], 900, lambda: ({**collectors.empty_stats(), "totalPrompts": 7}, False), with_complete=True) +whole = collectors.cached_scan("demo", ["b"], 900, scan, with_complete=True) +check("a caller can ask whether the scan was complete, a cached one being so", partial[0]["totalPrompts"] == 7 and partial[1] is False and whole == (whole[0], True) and whole[0]["totalPrompts"] == 5) + +# A scan is rebuilt under its lock, so a run that starts while another scans +# waits for it, then takes its result rather than reading every source again. +import subprocess # noqa: E402 + +HOLDER = ''' +import json, os, sys, time, fcntl +fd = os.open(sys.argv[1], os.O_WRONLY | os.O_CREAT, 0o600) +fcntl.flock(fd, fcntl.LOCK_EX) +print("locked", flush=True) +time.sleep(1.0) +envelope = json.loads(sys.argv[3]) +envelope["scannedAtMs"] = int(time.time() * 1000) +with open(sys.argv[2], "w") as handle: + json.dump(envelope, handle) +''' +locked_cache, locked_lock = collectors.scan_cache_paths("demo", ["locked"]) +envelope = {"schemaVersion": collectors.SCAN_CACHE_VERSION, "scanDate": today, "stats": {**collectors.empty_stats(), "totalPrompts": 42}} +holder = subprocess.Popen([sys.executable, "-c", HOLDER, str(locked_lock), str(locked_cache), json.dumps(envelope)], stdout=subprocess.PIPE, text=True) +holder.stdout.readline() +scans_before = len(calls) +started = time.time() +waited_for = collectors.cached_scan("demo", ["locked"], 20, scan) +waited = time.time() - started +holder.wait() +check( + "a run that starts during a scan waits on its lock, then takes that scan rather than scanning again", + waited >= 0.8 and waited_for["totalPrompts"] == 42 and len(calls) == scans_before, + f"waited {waited:.2f}s, got {waited_for['totalPrompts']}, scanned {len(calls) - scans_before} more times", +) + +# ------------------------------------------------------------------ limits + +cache = collectors.LimitsCache("demo", "who@example.com") +check("each credential gets its own limits cache", cache.path.name.startswith("demo-limits-") and cache.path != collectors.LimitsCache("demo", "other").path) +past = (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat() +future = (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat() +fetched = cache.save([ + {"label": "Session", "percent": 0.8, "resetsAt": past}, + {"label": "Weekly", "percent": 0.3, "resetsAt": future}, + {"label": "Monthly", "percent": 0.5, "resetsAt": ""}, +]) +kept = collectors.LimitsCache("demo", "who@example.com") +check( + "a kept window whose reset passed reads 0 with no reset; others stay as they were", + [(entry["percent"], entry["resetsAt"]) for entry in kept.limits] == [(0.0, ""), (0.3, future), (0.5, "")], + json.dumps(kept.limits), +) +check("a fresh answer stands in for a probe, but not past --force", kept.fresh(False) and not kept.fresh(True)) +interval = collectors.LimitsCache("demo", "interval") +interval.save([{"label": "Weekly", "percent": 0.1, "resetsAt": future}]) +aged = json.loads(interval.path.read_text()) +aged["fetchedAtMs"] = collectors.now_ms() - 14_000 +interval.path.write_text(json.dumps(aged)) +young = collectors.LimitsCache("demo", "interval").fresh(False) +aged["fetchedAtMs"] = collectors.now_ms() - 16_000 +interval.path.write_text(json.dumps(aged)) +check( + "an answer stands in for a probe for 15 seconds", + collectors.PROBE_MIN_INTERVAL_SECONDS == 15 and young and not collectors.LimitsCache("demo", "interval").fresh(False), +) +check("the stale fallback carries when the limits were measured", kept.kept() == (kept.limits, fetched)) +os.utime(kept.path, (0, 0)) +check("the limits cache is read whatever its file time", collectors.LimitsCache("demo", "who@example.com").fetched_at == fetched) +data = json.loads(kept.path.read_text()) +data["fetchedAtMs"] = collectors.now_ms() + 60_000 +kept.path.write_text(json.dumps(data)) +future_stamped = collectors.LimitsCache("demo", "who@example.com") +check("a future-stamped answer is kept but never trusted as fresh", future_stamped.limits and not future_stamped.fresh(False)) +check("no answer at all has nothing to fall back on", collectors.LimitsCache("demo", "nobody").kept() == ([], 0)) +check("an identity holding a lone surrogate still gets its own limits cache", collectors.LimitsCache("demo", "who\udc80").path != cache.path) + +credits = collectors.LimitsCache("demo", "who@example.com", kind="credits") +check( + "another kind of answer is kept beside the limits under the same key", + credits.path.name == cache.path.name.replace("demo-limits-", "demo-credits-", 1), +) +credits.save([], balance={"remaining": 3}) +credits = collectors.LimitsCache("demo", "who@example.com", kind="credits") +check( + "an answer without limits stands in for a probe only when the caller allows it", + not credits.fresh(False) and credits.fresh(False, allow_empty=True) and not credits.fresh(True, allow_empty=True) + and not credits.fresh(False, seconds=0, allow_empty=True) and credits.data["balance"] == {"remaining": 3}, +) +check("no answer at all never stands in, even allowing an empty one", not collectors.LimitsCache("demo", "nobody").fresh(False, allow_empty=True)) +newer = json.loads(credits.path.read_text()) +newer["fetchedAtMs"] = collectors.now_ms() + 60_000 +newer["balance"] = {"remaining": 9} +credits.path.write_text(json.dumps(newer)) +older = collectors.LimitsCache("demo", "who@example.com", kind="credits") +kept_stamp = older.save([], only_if_newer=True, balance={"remaining": 1}) +check( + "an answer saved only if newer never rolls back one another run stamped later, but serves this run", + json.loads(credits.path.read_text())["balance"] == {"remaining": 9} and older.data["balance"] == {"remaining": 1} and older.fetched_at == kept_stamp, +) +older.save([], balance={"remaining": 2}) +check("a plain save replaces whatever was kept", json.loads(credits.path.read_text())["balance"] == {"remaining": 2}) +ahead = json.loads(credits.path.read_text()) +ahead["fetchedAtMs"] = collectors.now_ms() + 3_600_000 +credits.path.write_text(json.dumps(ahead)) +older.save([], only_if_newer=True, balance={"remaining": 3}) +check( + "a stamp an hour ahead, from before the clock went back, is replaced by an answer saved only if newer", + json.loads(credits.path.read_text())["balance"] == {"remaining": 3}, +) +named = collectors.LimitsCache.at(collectors.cache_root() / "demo-limits.json") +named.save([{"label": "Weekly", "percent": 0.1, "resetsAt": future}]) +check( + "a cache opened at a name of its own keeps and reads like any other", + named.path.name == "demo-limits.json" and named.lock_path.name == "demo-limits.lock" + and collectors.LimitsCache.at(named.path).fresh(False) and collectors.LimitsCache.at(named.path).limits == named.limits, +) + +import io # noqa: E402 + +printed = io.StringIO() +with contextlib.redirect_stdout(printed): + collectors.emit({"balance": {"remaining": float("inf"), "spent": [float("nan"), 1.5]}, "ready": True}) +check( + "emit prints a number that overflowed as 0, so the record stays JSON the panel can parse", + json.loads(printed.getvalue()) == {"balance": {"remaining": 0, "spent": [0, 1.5]}, "ready": True}, + printed.getvalue(), +) + +# ------------------------------------------------------------------- http + +import io +import urllib.error +import urllib.request + +real_urlopen = urllib.request.urlopen + + +def failing(error): + def open_(request, timeout=None): + raise error + return open_ + + +def outcome(opener): + urllib.request.urlopen = opener + try: + return collectors.open_json(urllib.request.Request("https://example.invalid")) + except collectors.HttpError as error: + return (error.status, error.transport, error.invalid) + finally: + urllib.request.urlopen = real_urlopen + + +check("open_json returns the JSON answer", outcome(lambda request, timeout=None: io.BytesIO(b'{"a":1}')) == {"a": 1}) +check("an HTTP status is a refusal, not a transport failure", outcome(failing(urllib.error.HTTPError("u", 401, "no", {}, None))) == (401, False, False)) +check("no route, no DNS, and a timeout are transport failures", {outcome(failing(error)) for error in (urllib.error.URLError("no route"), TimeoutError("slow"), ConnectionResetError())} == {(0, True, False)}) +check("a body that isn't JSON is an invalid answer", outcome(lambda request, timeout=None: io.BytesIO(b"")) == (0, False, True)) +check("a body nested deeper than the parser goes is an invalid answer", outcome(lambda request, timeout=None: io.BytesIO(b"[" * 200_000)) == (0, False, True)) + +import http.client # noqa: E402 + + +def error_of(opener): + urllib.request.urlopen = opener + try: + collectors.open_json(urllib.request.Request("https://example.invalid")) + except collectors.HttpError as error: + return error + finally: + urllib.request.urlopen = real_urlopen + + +class CutShort(io.BytesIO): + def read(self, *args): + raise http.client.IncompleteRead(b'{"a"', 10) + + +cut = error_of(lambda request, timeout=None: CutShort()) +check("a reply cut short is an answer, never a transport failure", (cut.status, cut.transport, cut.invalid, cut.cut_short) == (0, False, True, True)) +garbled = error_of(failing(http.client.BadStatusLine("HTTP/9 ???"))) +check("a garbled reply is an answer too", (garbled.transport, garbled.cut_short) == (False, True)) +dropped = error_of(failing(http.client.RemoteDisconnected("closed"))) +check("a server hanging up without a reply is a dropped connection, so transport", (dropped.transport, dropped.cut_short) == (True, False)) +reset = error_of(failing(ConnectionResetError())) +check("a reset connection is transport, not a reply cut short", (reset.transport, reset.cut_short) == (True, False)) +refused = error_of(failing(urllib.error.HTTPError("u", 429, "slow down", {"retry-after": "30"}, None))) +check("a refusal carries its own headers", refused.status == 429 and refused.headers.get("retry-after") == "30" and not refused.cut_short) +# Over a real socket: a server that stops short of its Content-Length and +# hangs up hands a sized read only what arrived, which is cut short, not an +# invalid answer. +import socket # noqa: E402 + + +def short_server(): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + + def serve(): + conn, _ = listener.accept() + with conn: + conn.recv(65536) + conn.sendall(b'HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 100\r\n\r\n{"a": 1, "b"') + listener.close() + + threading.Thread(target=serve, daemon=True).start() + return f"http://127.0.0.1:{listener.getsockname()[1]}/" + + +os.environ["no_proxy"] = "*" +try: + collectors.open_json(urllib.request.Request(short_server()), timeout=5) + short = None +except collectors.HttpError as error: + short = error +check("a body short of its Content-Length over a socket is cut short", short is not None and (short.transport, short.invalid, short.cut_short) == (False, True, True), repr(short and vars(short))) + + +def drip_server(): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + + def serve(): + conn, _ = listener.accept() + with conn: + conn.recv(65536) + conn.sendall(b'HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 1000\r\n\r\n') + with contextlib.suppress(OSError): + for _ in range(100): + conn.sendall(b" ") + time.sleep(0.1) + listener.close() + + threading.Thread(target=serve, daemon=True).start() + return f"http://127.0.0.1:{listener.getsockname()[1]}/" + + +began = time.monotonic() +try: + collectors.open_json(urllib.request.Request(drip_server()), timeout=0.3) + dripped = None +except collectors.HttpError as error: + dripped = error +check( + "a reply dripping a byte at a time runs out of time as one cut short instead of holding the collector", + dripped is not None and dripped.cut_short and not dripped.transport and time.monotonic() - began < 5, + repr(dripped and vars(dripped)), +) +bad_key = urllib.request.Request("https://example.invalid", headers={"Authorization": "Bearer sk-secret\nX: y"}) +urllib.request.urlopen = real_urlopen +try: + collectors.open_json(bad_key) + header = None +except collectors.HttpError as error: + header = error +check( + "a key http.client won't send is an invalid answer that doesn't quote the key", + header is not None and header.invalid and header.__cause__ is None and header.__suppress_context__ and "sk-secret" not in str(header), + repr(header), +) +check( + "a credential goes over HTTPS, or plain HTTP to loopback, and never to a URL with credentials in it", + [collectors.trusted_endpoint(url) for url in ("https://example.com/x", "http://127.0.0.1:9/x", "http://localhost/x", "http://[::1]:9/x", + "http://example.com/x", "ftp://example.com", "https://u:p@example.com", "http://[bad", "")] + == [True, True, True, True, False, False, False, False, False], +) +handlers = [handler for handler in urllib.request._opener.handlers if isinstance(handler, urllib.request.HTTPRedirectHandler)] +check("requests refuse redirects", handlers and handlers[0].redirect_request(None, None, 302, "Found", {}, "https://elsewhere") is None) + +# --------------------------------------------------------- pi/omp sessions + + +def write_session(path, lines, header=None): + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w") as handle: + if header is not None: + handle.write(json.dumps(header) + "\n") + for line in lines: + handle.write(json.dumps(line) + "\n") + + +def reply(provider, tokens, message_id=None, timestamp=stamp, model="m", usage=None): + line = {"type": "message", "timestamp": timestamp, "message": {"role": "assistant", "provider": provider, "model": model, "usage": usage or {"input": tokens}}} + if message_id: + line["id"] = message_id + return line + + +sessions = home / ".pi" / "agent" / "sessions" +moved = work / "moved" +alone = work / "alone" +(home / ".omp" / "profiles" / "work").mkdir(parents=True) +(home / ".pi" / "profiles").mkdir(parents=True) +# A profile linked into the default root sits inside it, so it is dropped. +(sessions / "nested").mkdir(parents=True) +(home / ".pi" / "profiles" / "nest").symlink_to(sessions / "nested") +os.environ["PI_CODING_AGENT_DIR"] = str(moved) +os.environ["PI_CODING_AGENT_SESSION_DIR"] = str(alone) +roots = collectors.pi_session_roots() +check( + "session roots are the defaults, each profile, and both overrides, nested ones dropped", + roots == [ + sessions.resolve(), + (home / ".omp" / "agent" / "sessions").resolve(), + (home / ".omp" / "profiles" / "work" / "agent" / "sessions").resolve(), + (moved / "sessions").resolve(), + alone.resolve(), + ], + "\n".join(map(str, roots)), +) +os.environ["PI_CODING_AGENT_DIR"] = str(home / ".pi" / "agent") +check("an override that aliases a default root is read once", collectors.pi_session_roots().count(sessions.resolve()) == 1) +os.environ.pop("PI_CODING_AGENT_DIR") + +# pi's SettingsManager.getSessionDir() reads `sessionDir` from the agent +# tree's settings.json, expanding a leading ~, and new sessions go there. +(home / ".pi" / "agent").mkdir(parents=True, exist_ok=True) +(home / ".pi" / "agent" / "settings.json").write_text(json.dumps({"sessionDir": "~/pi-sessions"})) +(home / ".omp" / "profiles" / "work" / "agent").mkdir(parents=True, exist_ok=True) +(home / ".omp" / "profiles" / "work" / "agent" / "settings.json").write_text(json.dumps({"sessionDir": "relative/sessions"})) +roots = collectors.pi_session_roots() +check( + "a settings.json sessionDir is read as a root, a relative one skipped", + (home / "pi-sessions").resolve() in roots and not any(root.name == "sessions" and root.parent.name == "relative" for root in roots), + "\n".join(map(str, roots)), +) +(home / ".pi" / "agent" / "settings.json").unlink() +(home / ".omp" / "profiles" / "work" / "agent" / "settings.json").unlink() + +parent = sessions / "p" / "parent.jsonl" +write_session(parent, [ + reply("demo", 100, "m1"), + reply("demo", 1), + reply("demo", 1), + reply("other", 5000, "o1"), + reply("demo", 7, "z", timestamp=0), + reply("demo", 0, "t", usage={"input": 0, "totalTokens": 40}), + reply("demo", 0, "r", usage={"input": 2, "output": 3, "reasoning": 4, "cacheRead": 5, "cache_write": 6}), +]) +# The fork is read first by name, and repeats the parent's message. +write_session(sessions / "a" / "fork.jsonl", [reply("demo", 100, "m1"), reply("demo", 10, "m2")], header={"type": "session", "parentSession": str(parent)}) +write_session(sessions / "a" / "copy-only.jsonl", [reply("demo", 100, "m1")], header={"type": "session", "parentSession": str(parent)}) +write_session(alone / "x.jsonl", [reply("demo", 1000, "m3")]) +# A $HOME that is a git checkout ignoring everything still gets read. +(home / ".git").mkdir() +(home / ".gitignore").write_text("*\n") + +rows, complete = collectors.scan_pi_sessions(r'"provider"\s*:\s*"(demo|other)"', lambda provider, model, entry: provider == "demo") +by_tokens = sorted(row.input_tokens + row.output_tokens + row.cache_read + row.cache_write for row in rows) +check("a scan that read everything says so", complete) +check( + "pi replies count once by (id, timestamp), id-less lines always, with totals, thinking, and garbage handled", + by_tokens == [1, 1, 10, 16, 40, 100, 1000], + json.dumps([row._asdict() for row in rows]), +) +reasoning = next(row for row in rows if row.input_tokens == 2) +check("thinking is already in output and never added again", (reasoning.output_tokens, reasoning.cache_read, reasoning.cache_write) == (3, 5, 6)) +# As pi 1.4.0 and omp 18.12.6 wrote them for a reply the server billed as 54 +# and 84 completion tokens, 7 and 37 of them reasoning: totalTokens leaves the +# thinking out because output already holds it. +thinking = work / "thinking" / "sessions" +write_session(thinking / "pi.jsonl", [reply("demo", 0, "pi1", usage={"input": 1401, "output": 54, "cacheRead": 74, "cacheWrite": 0, "reasoning": 7, "totalTokens": 1529})]) +write_session(thinking / "omp.jsonl", [reply("demo", 0, "omp1", usage={"input": 4401, "output": 84, "cacheRead": 104, "cacheWrite": 0, "totalTokens": 4589, "reasoningTokens": 37})]) +thinking_rows, _ = collectors.scan_pi_sessions(r'"provider"\s*:\s*"demo"', lambda provider, model, entry: True, roots=[thinking]) +check( + "pi's and omp's output already holds their thinking", + sorted((row.output_tokens, row.input_tokens + row.output_tokens + row.cache_read) for row in thinking_rows) == [(54, 1529), (84, 4589)], + json.dumps([row._asdict() for row in thinking_rows]), +) +# pi 1.4.0 counts a compaction's and a branch summary's usage, a usage entry +# (cache warming), and a tool's nested model work in the session's totals. +# Only the usage entry names its provider; the rest ran on the session's +# model, the last switch or reply before them, as pi wrote this compaction. +summaries = work / "summaries" / "sessions" +pi_switch = {"type": "model_change", "id": "mc1", "timestamp": stamp, "provider": "demo", "modelId": "demo-model"} +compaction = {"type": "compaction", "id": "c1", "parentId": "x", "timestamp": stamp, "summary": "Mock reply 210.", "firstKeptEntryId": "x", "tokensBefore": 44000, + "usage": {"input": 44102, "output": 521, "cacheRead": 561, "cacheWrite": 0, "reasoning": 427, "totalTokens": 45184}, "fromHook": False} +write_session(summaries / "pi.jsonl", [ + pi_switch, + compaction, + {"type": "usage", "id": "u1", "timestamp": stamp, "kind": "cache_warm", "provider": "demo", "model": "warm-model", "usage": {"input": 0, "output": 0, "cacheRead": 500, "cacheWrite": 0, "totalTokens": 500}}, + {"type": "usage", "id": "u2", "timestamp": stamp, "kind": "cache_warm", "provider": "other", "model": "warm-model", "usage": {"input": 9000}}, + {"type": "message", "id": "t1", "timestamp": stamp, "message": {"role": "toolResult", "toolCallId": "call_1", "toolName": "agent", "content": [], "usage": {"input": 30, "output": 3}, "isError": False}}, + {"type": "message", "id": "o1", "timestamp": stamp, "message": {"role": "assistant", "provider": "other", "model": "o", "usage": {"input": 9000}}}, + {**compaction, "id": "c2", "usage": {"input": 9000}}, +]) +# omp names the model as provider/model; a branch summary after it bills it, +# one before any switch or reply bills nobody this reader can name. +write_session(summaries / "omp.jsonl", [ + {"type": "branch_summary", "id": "b0", "timestamp": stamp, "fromId": "x", "summary": "s", "usage": {"input": 9000}}, + {"type": "model_change", "id": "mc2", "timestamp": stamp, "model": "demo/omp-model"}, + {"type": "branch_summary", "id": "b1", "timestamp": stamp, "fromId": "x", "summary": "s", "usage": {"input": 70, "output": 7}}, +]) +summary_rows, _ = collectors.scan_pi_sessions(r'"provider"\s*:\s*"(demo|other)"', lambda provider, model, entry: provider == "demo", roots=[summaries]) +check( + "summaries, usage entries, and nested tool usage count under the model the session was on, answering no prompt", + sorted((row.model, row.input_tokens + row.output_tokens + row.cache_read, row.prompts) for row in summary_rows) + == [("demo-model", 33, 0), ("demo-model", 45184, 0), ("omp-model", 77, 0), ("warm-model", 500, 0)], + json.dumps([row._asdict() for row in summary_rows]), +) +# A prompt is a typed user message, booked on the first reply answering it, +# as pi 1.4.0 wrote one that called a tool: the step after the tool result +# answers no new prompt. omp marks the user messages it writes itself (a +# handoff) as the agent's, and a prompt another provider answered is not +# this one's. +turns = work / "turns" / "sessions" +write_session(turns / "pi.jsonl", [ + {"type": "model_change", "id": "mc", "timestamp": stamp, "provider": "demo", "modelId": "m"}, + {"type": "message", "id": "sys", "timestamp": stamp, "message": {"role": "system", "content": "", "timestamp": 1}}, + {"type": "message", "id": "u1", "timestamp": stamp, "message": {"role": "user", "content": [{"type": "text", "text": "Use the tool"}], "timestamp": 1}}, + reply("demo", 0, "step1", usage={"input": 21401, "output": 47, "cacheRead": 274, "cacheWrite": 0, "totalTokens": 21722}), + {"type": "message", "id": "tr", "timestamp": stamp, "message": {"role": "toolResult", "toolCallId": "call_1", "toolName": "bash", "content": [{"type": "text", "text": "ok"}], "isError": False}}, + reply("demo", 0, "step2", usage={"input": 21501, "output": 47, "cacheRead": 275, "cacheWrite": 0, "totalTokens": 21823}), + {"type": "message", "id": "u2", "timestamp": stamp, "message": {"role": "user", "content": "Ask the other one", "timestamp": 2}}, + reply("other", 0, "o1", usage={"input": 9000}), + {"type": "message", "id": "u3", "timestamp": stamp, "message": {"role": "user", "content": [{"type": "text", "text": "handoff"}], "attribution": "agent"}}, + reply("demo", 0, "step3", usage={"input": 5}), +]) +turn_rows, _ = collectors.scan_pi_sessions(r'"provider"\s*:\s*"(demo|other)"', lambda provider, model, entry: provider == "demo", roots=[turns]) +check( + "a typed prompt books on its first reply, not on a tool step, a handoff, or another provider's reply", + sorted((row.input_tokens, row.prompts) for row in turn_rows) == [(5, 0), (21401, 1), (21501, 0)], + json.dumps([row._asdict() for row in turn_rows]), +) +check("a session without user messages keeps a prompt per reply", all(row.prompts == 1 for row in thinking_rows)) +check("a forked message counts toward the session it was forked from", next(row for row in rows if row.input_tokens == 100).session == os.path.realpath(parent)) +check("a fork that only copied its parent is no session of its own", os.path.realpath(sessions / "a" / "copy-only.jsonl") not in {row.session for row in rows}) +check("the provider filter is the caller's", all(row.provider == "demo" for row in rows)) + +# A writer that keeps the time on the message rather than the entry: a fork's +# copy still shares the message's (id, timestamp) and counts once. +inner = work / "inner" / "sessions" +inner_stamp = int(time.time() * 1000) +inner_line = {"type": "message", "id": "k1", "message": {"role": "assistant", "provider": "demo", "model": "demo", "timestamp": inner_stamp, "usage": {"input": 77}}} +write_session(inner / "parent.jsonl", [inner_line]) +write_session(inner / "fork.jsonl", [inner_line], header={"type": "session", "parentSession": str(inner / "parent.jsonl")}) +inner_rows, _ = collectors.scan_pi_sessions(r'"provider"\s*:\s*"demo"', lambda provider, model, entry: True, roots=[inner]) +check("a copy timed only on its message counts once", [row.input_tokens for row in inner_rows] == [77], json.dumps([row._asdict() for row in inner_rows])) + +# A header naming a FIFO is no parent to read: opening one would hang. +import signal # noqa: E402 + +piped = work / "piped" / "sessions" +fifo = work / "piped" / "parent.fifo" +fifo.parent.mkdir(parents=True) +os.mkfifo(fifo) +piped_line = {"type": "message", "id": "q1", "timestamp": inner_stamp, "message": {"role": "assistant", "provider": "demo", "model": "demo", "usage": {"input": 5}}} +write_session(piped / "one.jsonl", [piped_line], header={"type": "session", "parentSession": str(fifo)}) +write_session(piped / "two.jsonl", [piped_line]) + + +# Not an OSError, which the reader would swallow as an unreadable header. +class Hung(BaseException): + pass + + +def hung(*_): + raise Hung() + + +signal.signal(signal.SIGALRM, hung) +signal.alarm(5) +try: + piped_rows, _ = collectors.scan_pi_sessions(r'"provider"\s*:\s*"demo"', lambda provider, model, entry: True, roots=[piped]) +except Hung: + piped_rows = None +signal.alarm(0) +check("a parentSession naming a FIFO is never opened", piped_rows is not None and [row.input_tokens for row in piped_rows] == [5]) + +harness = work / "harness" / "sessions" +write_session(harness / "h.jsonl", [reply("demo", 300, "h1"), reply("demo", 100, "m1")]) +own_rows, own_complete = collectors.scan_pi_sessions(r'"provider"\s*:\s*"demo"', lambda provider, model, entry: True, roots=[harness]) +check("a collector can read a pi-format tree of its own instead", own_complete and sorted(row.input_tokens for row in own_rows) == [100, 300]) +(work / "alias").symlink_to(harness) +both_rows, _ = collectors.scan_pi_sessions(r'"provider"\s*:\s*"demo"', lambda provider, model, entry: provider == "demo", roots=[*collectors.pi_session_roots(), harness, work / "alias", work / "missing"]) +check( + "roots read beside pi's count a message copied between them once, an alias once, and a missing one as nothing", + sorted(row.input_tokens + row.output_tokens + row.cache_read + row.cache_write for row in both_rows) == [1, 1, 10, 16, 40, 100, 300, 1000], + json.dumps([row._asdict() for row in both_rows]), +) +check("canonical roots drop an alias and a root inside another", collectors.canonical_roots([harness, work / "alias", harness / "inner", work / "harness"]) == [(work / "harness").resolve()]) + +# A scan that couldn't read every session says so, and cached_scan above keeps +# such a scan out of the cache. +pattern = r'"provider"\s*:\s*"(demo|other)"' +real_path = os.environ["PATH"] +os.environ["PATH"] = str(work / "no-bin") +try: + # The warning it prints is the collector's to show, not this test's. + with contextlib.redirect_stderr(io.StringIO()): + rows, complete = collectors.scan_pi_sessions(pattern, lambda provider, model, entry: provider == "demo") +finally: + os.environ["PATH"] = real_path +check("without rg the sessions can't be read, and the scan says it is incomplete", rows == [] and not complete) +if os.geteuid() != 0: + # rg still prints the matches it could read, and only its exit status tells. + locked = sessions / "p" / "locked.jsonl" + write_session(locked, [reply("demo", 5000, "locked")]) + locked.chmod(0) + try: + rows, complete = collectors.scan_pi_sessions(pattern, lambda provider, model, entry: provider == "demo") + finally: + locked.chmod(0o600) + locked.unlink() + check( + "a session rg couldn't read makes the scan incomplete, the readable ones still counted", + not complete and sorted(row.input_tokens + row.output_tokens + row.cache_read + row.cache_write for row in rows) == by_tokens, + json.dumps([row._asdict() for row in rows]), + ) + +# --------------------------------------------------------------- opencode + +db = home / ".local" / "share" / "opencode" / "opencode.db" +db.parent.mkdir(parents=True) +now_ms = int(time.time() * 1000) - 1000 +fork_ms = now_ms - 60_000 + + +def v1(provider, tokens, created=now_ms): + return json.dumps({"role": "assistant", "providerID": provider, "modelID": "v1-model", "tokens": tokens, "time": {"created": created}}) + + +def v2(provider, tokens): + return json.dumps({"model": {"providerID": provider, "id": "v2-model"}, "tokens": tokens, "time": {"created": now_ms}}) + + +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text, time_created integer, time_updated integer, data text)") +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text, type text, time_created integer, data text)") +conn.execute("CREATE TABLE session_v2 (id text PRIMARY KEY, fork_session_id text, time_created integer)") +conn.execute("INSERT INTO session_v2 VALUES ('ses_fork', 'ses_v2', ?)", (fork_ms,)) +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ + ("msg_v1", "ses_v1", now_ms, now_ms, v1("demo", {"input": 10, "output": 5, "reasoning": 2, "cache": {"read": 3, "write": 1}})), + ("msg_both", "ses_v1", now_ms, now_ms, v1("demo", {"input": 40})), + ("msg_zero", "ses_v1", now_ms, now_ms, v1("demo", {"input": 8000}, created=0)), + ("msg_lookalike", "ses_v1", now_ms, now_ms, v1("demo-proxy", {"input": 9000})), + ("msg_other", "ses_v1", now_ms, now_ms, v1("other", {"input": 9000})), +]) +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, ?, ?)", [ + ("msg_compact", "ses_v2", "compaction", now_ms, v2("demo", {"input": 30})), + ("msg_both", "ses_v1", "assistant", now_ms, v2("demo", {"input": 40})), + ("msg_copy", "ses_fork", "assistant", fork_ms - 5_000, v2("demo", {"input": 7000})), + ("msg_forked", "ses_fork", "assistant", now_ms, v2("demo", {"input": 50})), + ("msg_bad", "ses_v2", "assistant", now_ms, '{"model":{"providerID":"demo"'), +]) +conn.commit() +conn.close() + +accept_demo = lambda provider, model, entry: provider == "demo" # noqa: E731 +rows, complete = collectors.scan_opencode(["demo"], accept_demo) +check( + "opencode rows come from both stores once each, compactions in, fork copies and untrusted times out", + complete and sorted(row.input_tokens + row.output_tokens for row in rows) == [17, 30, 40, 50], + json.dumps([row._asdict() for row in rows]), +) +check("the V2 store is read first", next(row for row in rows if row.input_tokens == 40).model == "v2-model") +check("opencode rows count toward their session", {row.session for row in rows} == {"opencode:ses_v1", "opencode:ses_v2", "opencode:ses_fork"}) +check("a closing quote keeps a look-alike provider out of the prefilter", collectors.scan_opencode(['demo"'], lambda *_: True)[0] == collectors.scan_opencode(["demo"], accept_demo)[0]) + +db.rename(db.parent / "elsewhere.db") +os.environ["OPENCODE_DB"] = "elsewhere.db" +check("OPENCODE_DB names the database, a relative name beside the default", len(collectors.scan_opencode(["demo"], accept_demo)[0]) == 4) +os.environ.pop("OPENCODE_DB") +(db.parent / "elsewhere.db").rename(db.parent / "opencode-v2.db") +check("opencode-v2.db is read when opencode.db is absent", len(collectors.scan_opencode(["demo"], accept_demo)[0]) == 4) + +before = (db.parent / "opencode-v2.db").read_bytes() +collectors.scan_opencode(["demo"], accept_demo) +check("the database is opened read-only", (db.parent / "opencode-v2.db").read_bytes() == before) + +(db.parent / "opencode-v2.db").write_bytes(b"not a database at all, just bytes" * 200) +rows, complete = collectors.scan_opencode(["demo"], accept_demo) +check("a database error marks the scan incomplete", not complete and rows == []) + +# opencode keeps its database open while it runs, so its newest rows sit in +# the write-ahead log; a read-only open still sees them. +wal_db = work / "wal" / "opencode.db" +wal_db.parent.mkdir() +WAL_WRITER = ''' +import os, sqlite3, sys +conn = sqlite3.connect(sys.argv[1]) +conn.execute("PRAGMA journal_mode=WAL") +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text, time_created integer, time_updated integer, data text)") +conn.execute("INSERT INTO message VALUES ('msg_wal', 'ses_wal', 0, 0, ?)", (sys.argv[2],)) +conn.commit() +# Killed with the connection open, so the rows stay in the log. +os.kill(os.getpid(), 9) +''' +subprocess.run([sys.executable, "-c", WAL_WRITER, str(wal_db), v1("demo", {"input": 61})]) +os.environ["OPENCODE_DB"] = str(wal_db) +wal_before = (wal_db.parent / "opencode.db-wal").is_file() +rows, complete = collectors.scan_opencode(["demo"], accept_demo) +check("a database left open with its rows in the write-ahead log still reads", wal_before and complete and [row.input_tokens for row in rows] == [61], + json.dumps({"wal": wal_before, "rows": [row._asdict() for row in rows]})) + +# A fresh V2 install has no `message` table at all, and some builds keep the +# provider and model at the top level of a session_message row, with the role +# in its type column. Either way the store reads on its own. +v2_only = work / "v2-only.db" +conn = sqlite3.connect(v2_only) +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL," + " seq integer NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") + + +def flat(provider, tokens): + return json.dumps({"providerID": provider, "modelID": "flat-model", "tokens": {"input": tokens}, "time": {"created": now_ms}}) + + +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, 1, ?, ?, ?)", [ + ("flat_reply", "ses_a", "assistant", now_ms, now_ms, flat("demo", 5)), + ("flat_prompt", "ses_a", "user", now_ms, now_ms, flat("demo", 900)), + ("flat_other", "ses_a", "assistant", now_ms, now_ms, flat("other", 900)), + ("nested_reply", "ses_b", "assistant", now_ms, now_ms, v2("demo", {"input": 7})), + ("broken", "ses_b", "assistant", now_ms, now_ms, "not json"), +]) +conn.commit() +conn.close() +os.environ["OPENCODE_DB"] = str(v2_only) +rows, complete = collectors.scan_opencode(["demo"], accept_demo) +os.environ.pop("OPENCODE_DB") +check( + "a database with only the V2 store reads on its own, flat rows and nested alike", + complete and sorted((row.session, row.model, row.input_tokens) for row in rows) == [("opencode:ses_a", "flat-model", 5), ("opencode:ses_b", "v2-model", 7)], + json.dumps([row._asdict() for row in rows]), +) + +# A prompt is a typed user message, so of the steps answering one, as +# opencode 1.19.3 writes them (the same parentID), the first books it. A +# compaction's summary answers the marker opencode adds, and a subagent's +# session holds the model's prompts, not typed ones. V2 keeps a prompt as a +# `user` row without a model, booked on the first reply after it. +turns_db = work / "turns.db" +conn = sqlite3.connect(turns_db) +conn.execute("CREATE TABLE session (id text PRIMARY KEY, parent_id text, time_created integer NOT NULL)") +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL," + " seq integer NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.executemany("INSERT INTO session VALUES (?, ?, ?)", [("ses_tool", None, fork_ms), ("ses_task", "ses_tool", fork_ms)]) + + +def step(parent, input_tokens, created, **extra): + return json.dumps({"parentID": parent, "role": "assistant", "mode": "build", "agent": "build", "cost": 0, "modelID": "oc-tooltest", "providerID": "demo", + "tokens": {"total": input_tokens, "input": input_tokens, "output": 0, "reasoning": 0, "cache": {"write": 0, "read": 0}}, + "time": {"created": created, "completed": created + 1}, **extra}) + + +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ + ("msg_u1", "ses_tool", now_ms, now_ms, json.dumps({"role": "user", "time": {"created": now_ms}, "agent": "build", "model": {"providerID": "demo", "modelID": "oc-tooltest"}})), + ("msg_a2", "ses_tool", now_ms + 2, now_ms + 2, step("msg_u1", 21501, now_ms + 2, finish="stop")), + ("msg_a1", "ses_tool", now_ms + 1, now_ms + 1, step("msg_u1", 21401, now_ms + 1, finish="tool-calls")), + ("msg_u2", "ses_tool", now_ms + 3, now_ms + 3, json.dumps({"role": "user", "time": {"created": now_ms + 3}, "model": {"providerID": "demo", "modelID": "oc-tooltest"}})), + ("msg_c1", "ses_tool", now_ms + 4, now_ms + 4, step("msg_u2", 23153, now_ms + 4, mode="compaction", agent="compaction", summary=True)), + ("msg_t1", "ses_task", now_ms + 5, now_ms + 5, step("msg_tu", 300, now_ms + 5)), +]) +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, ?, ?, ?, ?)", [ + ("v2_u1", "ses_v2", "user", 1, now_ms, now_ms, json.dumps({"text": "hi", "time": {"created": now_ms}})), + ("v2_a1", "ses_v2", "assistant", 2, now_ms, now_ms, v2("demo", {"input": 11})), + ("v2_a2", "ses_v2", "assistant", 3, now_ms, now_ms, v2("demo", {"input": 12})), + ("v2_c1", "ses_v2", "compaction", 4, now_ms, now_ms, v2("demo", {"input": 13})), + ("v2_lone", "ses_v2_lone", "assistant", 1, now_ms, now_ms, v2("demo", {"input": 14})), +]) +conn.commit() +conn.close() +os.environ["OPENCODE_DB"] = str(turns_db) +rows, complete = collectors.scan_opencode(["demo"], accept_demo) +os.environ.pop("OPENCODE_DB") +check( + "an opencode prompt books on its first step, never on a later one, a compaction, or a subagent's", + complete and sorted((row.input_tokens, row.prompts) for row in rows) == [(11, 1), (12, 0), (13, 0), (14, 1), (300, 0), (21401, 1), (21501, 0), (23153, 0)], + json.dumps([row._asdict() for row in rows]), +) + +# One prompt answered by steps from two providers, as a tool loop that +# switches models writes it, books once: on the first step, under the +# collector that takes it, in both stores. +switch_db = work / "switch.db" +conn = sqlite3.connect(switch_db) +conn.execute("CREATE TABLE session (id text PRIMARY KEY, parent_id text, time_created integer NOT NULL)") +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL," + " seq integer NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.execute("INSERT INTO session VALUES ('ses_switch', NULL, ?)", (fork_ms,)) +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ + ("msg_su", "ses_switch", now_ms, now_ms, json.dumps({"role": "user", "time": {"created": now_ms}, "agent": "build", "model": {"providerID": "anthropic", "modelID": "claude-sonnet-4-5"}})), + ("msg_sb", "ses_switch", now_ms + 2, now_ms + 2, json.dumps(json.loads(step("msg_su", 600, now_ms + 2, finish="stop")) | {"providerID": "openrouter", "modelID": "z-ai/glm-4.6"})), + ("msg_sa", "ses_switch", now_ms + 1, now_ms + 1, json.dumps(json.loads(step("msg_su", 500, now_ms + 1, finish="tool-calls")) | {"providerID": "anthropic", "modelID": "claude-sonnet-4-5"})), +]) +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, ?, ?, ?, ?)", [ + ("v2_su", "ses_switch_v2", "user", 1, now_ms, now_ms, json.dumps({"text": "hi", "time": {"created": now_ms}})), + ("v2_sa", "ses_switch_v2", "assistant", 2, now_ms, now_ms, v2("anthropic", {"input": 700})), + ("v2_sb", "ses_switch_v2", "assistant", 3, now_ms, now_ms, v2("openrouter", {"input": 800})), +]) +conn.commit() +conn.close() +os.environ["OPENCODE_DB"] = str(switch_db) +by_provider = { + provider: sorted((row.input_tokens, row.prompts) for row in collectors.scan_opencode([provider], lambda p, m, e, provider=provider: p == provider)[0]) + for provider in ("anthropic", "openrouter") +} +os.environ.pop("OPENCODE_DB") +check( + "a prompt answered by two providers' steps books once, on the first step's collector", + by_provider == {"anthropic": [(500, 1), (700, 1)], "openrouter": [(600, 0), (800, 0)]}, + json.dumps(by_provider), +) + +# A typed prompt no step answered, its request refused as opencode 1.19.3 +# records one (an assistant row with an APIError and no tokens), books +# itself without tokens on the model its user row names, so a session +# holding only that still counts. In V2 it books on the failed reply after +# it, and a message migrated into both stores books once. A compaction's +# marker, a subagent's prompt, a fork's copy and an untrusted time book none, +# and so does a prompt opencode 1.19.3 refused before sending it (an unknown +# model or provider): it writes the user row alone, with no step at all. +failed_db = work / "failed.db" +conn = sqlite3.connect(failed_db) +conn.execute("CREATE TABLE session (id text PRIMARY KEY, parent_id text, time_created integer NOT NULL)") +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL," + " seq integer NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.executemany("INSERT INTO session VALUES (?, ?, ?)", [("ses_fail", None, fork_ms), ("ses_only", None, fork_ms), ("ses_kid", "ses_fail", fork_ms), ("ses_copy", None, now_ms + 50), ("ses_unsent", None, fork_ms)]) + + +def asked(created, provider="demo", **extra): + return json.dumps({"role": "user", "time": {"created": created}, "agent": "build", "model": {"providerID": provider, "modelID": "oc-failoc"}, "summary": {"diffs": []}, **extra}) + + +def refused(parent, created, **extra): + return json.dumps(json.loads(step(parent, 0, created)) | {"error": {"name": "APIError", "data": {"message": "refused", "statusCode": 400}}, **extra}) + + +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ + # Refused, then `-c` and answered. + ("msg_fu1", "ses_fail", now_ms, now_ms, asked(now_ms)), + ("msg_fa1", "ses_fail", now_ms + 1, now_ms + 1, refused("msg_fu1", now_ms + 1)), + ("msg_fu2", "ses_fail", now_ms + 2, now_ms + 2, asked(now_ms + 2)), + ("msg_fa2", "ses_fail", now_ms + 3, now_ms + 3, step("msg_fu2", 900, now_ms + 3, finish="stop")), + # A compaction that failed answers its marker all the same. + ("msg_fu3", "ses_fail", now_ms + 4, now_ms + 4, asked(now_ms + 4)), + ("msg_fc3", "ses_fail", now_ms + 5, now_ms + 5, refused("msg_fu3", now_ms + 5, mode="compaction", agent="compaction", summary=True)), + # A session holding only a refused prompt, and one never sent. + ("msg_ou1", "ses_only", now_ms, now_ms, asked(now_ms)), + ("msg_oa1", "ses_only", now_ms + 1, now_ms + 1, refused("msg_ou1", now_ms + 1)), + ("msg_ou2", "ses_only", now_ms + 2, now_ms + 2, asked(now_ms + 2)), + # A session holding only a prompt never sent, as `opencode run -m` on an + # unknown model leaves it (no `summary` yet). + ("msg_nu1", "ses_unsent", now_ms, now_ms, json.dumps({k: v for k, v in json.loads(asked(now_ms)).items() if k != "summary"})), + # Another provider's, a subagent's, a fork's copy, and an untrusted time. + ("msg_xu1", "ses_only", now_ms + 3, now_ms + 3, asked(now_ms + 3, provider="other")), + ("msg_ku1", "ses_kid", now_ms, now_ms, asked(now_ms)), + ("msg_cu1", "ses_copy", now_ms, now_ms, asked(now_ms)), + ("msg_zu1", "ses_only", now_ms, now_ms, asked(0)), + # Migrated into V2 below under the same ids. + ("v2_mu", "ses_moved", now_ms, now_ms, asked(now_ms)), + ("v2_ma", "ses_moved", now_ms + 1, now_ms + 1, refused("v2_mu", now_ms + 1)), +]) +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, ?, ?, ?, ?)", [ + ("v2_mu", "ses_moved", "user", 1, now_ms, now_ms, json.dumps({"text": "hi", "time": {"created": now_ms}})), + ("v2_ma", "ses_moved", "assistant", 2, now_ms, now_ms, v2("demo", {"input": 0})), + ("v2_fu", "ses_v2_fail", "user", 1, now_ms, now_ms, json.dumps({"text": "hi", "time": {"created": now_ms}})), + ("v2_fa", "ses_v2_fail", "assistant", 2, now_ms, now_ms, v2("demo", {"input": 0})), +]) +conn.commit() +conn.close() +os.environ["OPENCODE_DB"] = str(failed_db) +rows, complete = collectors.scan_opencode(["demo"], accept_demo) +os.environ.pop("OPENCODE_DB") +booked = sorted((row.session, row.input_tokens, row.prompts) for row in rows if row.prompts or row.input_tokens) +check( + "an opencode prompt no step answered books once, on its user row's model, its session with it, and one never sent nowhere", + complete and booked == [("opencode:ses_fail", 0, 1), ("opencode:ses_fail", 900, 1), ("opencode:ses_moved", 0, 1), ("opencode:ses_only", 0, 1), ("opencode:ses_v2_fail", 0, 1)], + json.dumps([row._asdict() for row in rows]), +) + +# pi and omp write a refused request as a reply with its provider and no +# usage (pi 1.4.0's stopReason "error"), so the prompt books there, without +# tokens; one no reply came for at all books on the model switch before it. +failed_root = work / "failed-pi" +failed_root.mkdir() +failed_stamp = datetime.fromtimestamp(now_ms / 1000, timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") +zero_usage = {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0, "totalTokens": 0} + + +def pi_reply(n, provider, usage): + return {"type": "message", "id": f"a{n}", "timestamp": failed_stamp, "message": {"role": "assistant", "provider": provider, "model": "e2e-failpi", "usage": usage, + "stopReason": "error" if usage is zero_usage else "stop"}} + + +def pi_user(n): + return {"type": "message", "id": f"u{n}", "timestamp": failed_stamp, "message": {"role": "user", "content": [{"type": "text", "text": "go"}]}} + + +(failed_root / "failed.jsonl").write_text("".join(json.dumps(line) + "\n" for line in [ + {"type": "model_change", "id": "m1", "timestamp": failed_stamp, "provider": "demo", "modelId": "e2e-failpi"}, + pi_user(1), pi_reply(1, "other", zero_usage), + pi_user(2), pi_reply(2, "demo", zero_usage), + pi_user(3), pi_reply(3, "other", {"input": 40, "output": 1}), + {"type": "model_change", "id": "m2", "timestamp": failed_stamp, "provider": "demo", "modelId": "e2e-failpi"}, + pi_user(4), +])) +by_provider = { + provider: sorted((row.input_tokens, row.prompts) for row in collectors.scan_pi_sessions(r'"provider"\s*:\s*"' + provider + '"', lambda p, m, e, provider=provider: p == provider, [failed_root])[0]) + for provider in ("demo", "other") +} +check( + "a pi prompt no reply answered books on the refused reply, else on the model switch, once", + by_provider == {"demo": [(0, 1), (0, 1)], "other": [(0, 1), (40, 1)]}, + json.dumps(by_provider), +) + +# A damaged step the reader rejects (a boolean token count, an untrusted +# time) is never an opencode prompt's first step, so the valid step after it +# with the same parentID still books the prompt. +damaged_db = work / "damaged.db" +conn = sqlite3.connect(damaged_db) +conn.execute("CREATE TABLE session (id text PRIMARY KEY, parent_id text, time_created integer NOT NULL)") +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)") +# ses_when has no session row, so no fork time bounds its steps. +conn.execute("INSERT INTO session VALUES ('ses_bool', NULL, ?)", (fork_ms,)) +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ + ("msg_bu", "ses_bool", now_ms, now_ms, asked(now_ms)), + ("msg_b1", "ses_bool", now_ms + 1, now_ms + 1, json.dumps(json.loads(step("msg_bu", 0, now_ms + 1, finish="tool-calls")) | {"tokens": {"input": True, "output": 0}})), + ("msg_b2", "ses_bool", now_ms + 2, now_ms + 2, step("msg_bu", 610, now_ms + 2, finish="stop")), + ("msg_wu", "ses_when", now_ms, now_ms, asked(now_ms)), + ("msg_w1", "ses_when", now_ms + 1, now_ms + 1, json.dumps(json.loads(step("msg_wu", 50, now_ms + 1, finish="tool-calls")) | {"time": {"created": True}})), + ("msg_w2", "ses_when", now_ms + 2, now_ms + 2, step("msg_wu", 620, now_ms + 2, finish="stop")), +]) +conn.commit() +conn.close() +os.environ["OPENCODE_DB"] = str(damaged_db) +rows, complete = collectors.scan_opencode(["demo"], accept_demo) +os.environ.pop("OPENCODE_DB") +check( + "a step the reader rejects never takes an opencode prompt from the valid step after it", + complete and sorted((row.input_tokens, row.prompts) for row in rows) == [(610, 1), (620, 1)], + json.dumps([row._asdict() for row in rows]), +) + +# A prompt left unanswered in a pi session (its request failed) and answered +# in a fork that copied it, as pi 1.4.0 and omp 18.12.6 copy a message with +# its id and timestamp, is one prompt, counted once. +retried_root = work / "retried-pi" +retried_root.mkdir() +retried = [{"type": "model_change", "id": "m1", "timestamp": failed_stamp, "provider": "demo", "modelId": "e2e-failpi"}, pi_user(1), pi_reply(1, "demo", zero_usage)] +write_session(retried_root / "parent.jsonl", retried, header={"type": "session", "version": 3, "id": "s-parent", "timestamp": failed_stamp}) +write_session(retried_root / "fork.jsonl", retried + [pi_reply(2, "demo", {"input": 40, "output": 1})], + header={"type": "session", "version": 3, "id": "s-fork", "timestamp": failed_stamp, "parentSession": str(retried_root / "parent.jsonl")}) +rows, complete = collectors.scan_pi_sessions(r'"provider"\s*:\s*"demo"', lambda *_: True, [retried_root]) +check( + "a pi prompt unanswered in its session and answered in a fork counts once", + complete and sum(row.prompts for row in rows) == 1 and sum(row.input_tokens for row in rows) == 40, + json.dumps([row._asdict() for row in rows]), +) + +# A V2 prompt whose request failed books on its own user row's day, not the +# failed reply's, and one typed at an untrusted time counts nowhere. +late_db = work / "late.db" +yesterday_ms = now_ms - 86_400_000 +conn = sqlite3.connect(late_db) +conn.execute("CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL," + " seq integer NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +conn.executemany("INSERT INTO session_message VALUES (?, ?, ?, ?, ?, ?, ?)", [ + ("v2_lu", "ses_v2_late", "user", 1, yesterday_ms, yesterday_ms, json.dumps({"text": "hi", "time": {"created": yesterday_ms}})), + ("v2_la", "ses_v2_late", "assistant", 2, now_ms, now_ms, v2("demo", {"input": 0})), + ("v2_zu", "ses_v2_untimed", "user", 1, now_ms, now_ms, json.dumps({"text": "hi", "time": {"created": 0}})), + ("v2_za", "ses_v2_untimed", "assistant", 2, now_ms, now_ms, v2("demo", {"input": 0})), +]) +conn.commit() +conn.close() +os.environ["OPENCODE_DB"] = str(late_db) +rows, complete = collectors.scan_opencode(["demo"], accept_demo) +os.environ.pop("OPENCODE_DB") +check( + "a V2 prompt whose request failed books on the day it was typed, an untrusted one nowhere", + complete and sorted((row.session, row.day, row.prompts) for row in rows if row.prompts) == [("opencode:ses_v2_late", collectors.local_day(yesterday_ms), 1)], + json.dumps([row._asdict() for row in rows]), +) + +# ------------------------------------------------------------- endpoints + +# A base URL's host says who bills what it serves: a provider's own host or +# a subdomain of it, the user's own hardware, or nothing that can be placed. +for url, owner in ( + ("https://api.z.ai/api/anthropic", "zai"), + ("https://open.bigmodel.cn/api/anthropic", "zai"), + ("https://ollama.com", "ollama"), + ("https://api.minimax.io/anthropic", "minimax"), + ("https://api.minimaxi.com/anthropic", "minimax"), + ("https://api.deepseek.com/anthropic", "deepseek"), + ("https://api.kimi.com/coding/", "kimi"), + ("https://api.moonshot.ai/anthropic", "kimi"), + ("https://openrouter.ai/api", "openrouter"), + ("https://eu.openrouter.ai/api", "openrouter"), + ("http://localhost:11434", "local"), + ("http://127.0.0.1:4000", "local"), + ("http://192.168.1.20:8080/v1", "local"), + ("http://gpu-box:11434", "local"), + ("http://studio.lan:1234/v1", "local"), + ("https://llm-proxy.example.com/anthropic", ""), + ("https://notapi.z.ai.example.com", ""), + ("", ""), + ("http://[::1", ""), +): + check(f"the endpoint {url!r} is billed by {owner or 'nobody here'}", collectors.endpoint_owner(url) == owner, collectors.endpoint_owner(url)) + +# Claude Code's replies, by the model that answered and where Claude Code is +# sent: the one rule every collector reading its transcripts applies. +CLAUDE_CODE_OWNERS = { + "": {"claude-sonnet-4-5": "claude", None: "claude", "glm-4.6": "zai", "GLM-4.5-Air": "zai", "glm-4.6:cloud": "ollama", + "gpt-oss:120b-cloud": "ollama", "MiniMax-M2": "minimax", "deepseek-chat": "deepseek", "kimi-k2-0905": "kimi", + "moonshot-v1-8k": "kimi", "qwen3-coder": None, "": None}, + "https://api.anthropic.com": {"claude-sonnet-4-5": "claude", "glm-4.6": "zai"}, + "https://api.z.ai/api/anthropic": {"claude-sonnet-4-5": "zai", None: "zai", "glm-4.6": "zai", "glm-4.6:cloud": "zai"}, + "https://ollama.com": {"claude-sonnet-4-5": "ollama", "kimi-k2.7-code:cloud": "ollama", "glm-4.6": "ollama"}, + "https://openrouter.ai/api": {"claude-sonnet-4-5": "openrouter", "anthropic/claude-sonnet-4.5": "openrouter", None: "openrouter"}, + "https://api.minimax.io/anthropic": {"MiniMax-M2": "minimax", "claude-sonnet-4-5": "minimax"}, + "https://api.deepseek.com/anthropic": {"deepseek-chat": "deepseek"}, + "https://api.kimi.com/coding/": {"kimi-for-coding": "kimi", "claude-sonnet-4-5": "kimi"}, + "http://localhost:11434": {"qwen3-coder": "local", "glm-4.6": "local", "deepseek-r1:8b": "local", "glm-4.6:cloud": "ollama", + "claude-sonnet-4-5": "claude", None: "claude"}, + "https://llm-proxy.example.com/anthropic": {"claude-sonnet-4-5": "claude", None: "claude", "glm-4.6": None, "qwen3-coder": None, + "gpt-oss:120b-cloud": "ollama"}, +} +for url, models in CLAUDE_CODE_OWNERS.items(): + got = {model: collectors.claude_code_owner(model, url) for model in models} + check(f"Claude Code replies sent to {url or 'Anthropic'} go to the collector that bills them", got == models, json.dumps({str(k): v for k, v in got.items()})) + +# The shared transcript reader hands each owner its replies once, typed +# prompts apart from them, and the sessions in use; Stats books a row with no +# tokens as its prompts alone. +claude_home = home / "claude-code" +transcripts = claude_home / "projects" / "demo" +transcripts.mkdir(parents=True) + + +def claude_reply(message_id, model, input_tokens, session="s1"): + return {"type": "assistant", "timestamp": stamp, "sessionId": session, "requestId": "r-" + message_id, + "message": {"id": message_id, "role": "assistant", "model": model, "usage": {"input_tokens": input_tokens, "output_tokens": 1}}} + + +def claude_prompt(uuid, session="s1"): + return {"type": "user", "uuid": uuid, "timestamp": stamp, "sessionId": session, "message": {"role": "user", "content": "hi"}} + + +lines = [claude_prompt("p1"), claude_reply("m1", "claude-sonnet-4-5", 100), claude_prompt("p2"), claude_reply("m2", "glm-4.6", 200), + claude_reply("m3", "glm-4.6", 300)] +(transcripts / "s1.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) +# A fork copies the GLM reply into a session of its own: counted once. +(transcripts / "s2.jsonl").write_text(json.dumps({**claude_reply("m3", "glm-4.6", 300), "sessionId": "s2"}) + "\n") +by_owner = {} +for owner in ("claude", "zai", "ollama"): + owned = collectors.Stats() + for row in collectors.scan_claude_code(owner, claude_home / "projects", None, ""): + owned.add_row(row, row.model) + by_owner[owner] = owned.snapshot() +check( + "the Claude Code reader splits replies and prompts between their owners, a copied reply once", + [(by_owner[o]["todayTotalTokens"], by_owner[o]["todayPrompts"], by_owner[o]["todaySessions"]) for o in ("claude", "zai", "ollama")] + == [(101, 1, 1), (502, 1, 2), (0, 0, 0)], + json.dumps(by_owner), +) +check( + "a scan of Claude Code is matched to the kind of endpoint it was made under, never the URL itself", + collectors.claude_code_endpoint("https://user:secret@llm-proxy.example.com") == "other" + and collectors.claude_code_endpoint("https://api.anthropic.com") == "" + and collectors.claude_code_endpoint("http://127.0.0.1:4000") == "local", +) + +# MiniMax Code's routes: its own are MiniMax's, a custom provider its +# config.yaml baseURL's (read without a YAML library: block mappings, quoted +# and commented scalars), and one it no longer declares nobody's. +mcode_home = home / "mcode" +mcode_home.mkdir() +(mcode_home / "config.yaml").write_text("""# written by mcode +defaultModel: custom_provider:router/glm +custom_provider: + router: + models: + - glm + options: + apiKey: "sk: placeholder" + baseURL: 'https://openrouter.ai/api/v1' # preset + box: + options: + baseURL: http://127.0.0.1:8080/v1 + bare: {} +""") +routes = collectors.mcode_routes(mcode_home) +check("MiniMax Code's custom providers are read from config.yaml by their baseURL's owner", + routes == {"router": "openrouter", "box": "local", "bare": ""}, json.dumps(routes)) +check("a MiniMax Code reply goes to its route's owner", + [collectors.mcode_route_owner(provider, routes) for provider in (None, "minimax_oauth", "custom_provider:router", "custom_provider:box", "custom_provider:bare", "custom_provider:gone", "zai")] + == ["minimax", "minimax", "openrouter", "local", None, None, None]) + +# --------------------------------------------------------------- sign-ins + +import base64 # noqa: E402 + + +def jwt(claims): + body = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=") + return f"header.{body}.signature" + + +later = time.time() + 3600 +check("jwt_claims reads a JWT's claims", collectors.jwt_claims(jwt({"sub": "u1", "exp": later})) == {"sub": "u1", "exp": later}) +check( + "jwt_claims reads anything else as no claims", + [collectors.jwt_claims(value) for value in ("opaque-key", "a.b", "a.!!!.c", f"h.{base64.urlsafe_b64encode(b'[1]').decode()}.s", None, 5)] == [{}] * 6, +) +check("a token keys by its subject, an opaque key by itself", (collectors.account_key(jwt({"sub": "u1"})), collectors.account_key("opaque")) == ("u1", "opaque")) +check( + "a token lapses once its expiry passes; one with no expiry, or an opaque key, never does", + collectors.token_lapsed(jwt({"exp": time.time() - 5})) and not collectors.token_lapsed(jwt({"exp": later})) and not collectors.token_lapsed(jwt({})) and not collectors.token_lapsed("opaque"), +) +check("finite reads only finite numbers", [collectors.finite(value) for value in (True, "2.5", float("inf"), None, "x", 3)] == [0.0, 2.5, 0.0, 0.0, 0.0, 3.0]) + +config = work / "config" +os.environ["XDG_CONFIG_HOME"] = str(config) +for name in ("CURSOR_API_KEY", "CURSOR_AUTH_PATH", "CURSOR_AUTH_TOKEN"): + os.environ.pop(name, None) +check("no Cursor sign-in anywhere is no token", collectors.cursor_token() == "") +state_db = collectors.cursor_state_db() +state_db.parent.mkdir(parents=True) +conn = sqlite3.connect(state_db) +conn.execute("CREATE TABLE ItemTable (key TEXT, value BLOB)") +conn.execute("INSERT INTO ItemTable VALUES ('cursorAuth/accessToken', ?)", (jwt({"sub": "ide", "exp": later}).encode(),)) +conn.commit() +conn.close() +state_db_bytes = state_db.read_bytes() +check("the editor's sign-in stands when the CLI has none", collectors.account_key(collectors.cursor_token()) == "ide") +cli = config / "cursor" / "auth.json" +cli.parent.mkdir(parents=True) +cli.write_text(json.dumps({"accessToken": jwt({"sub": "cli", "exp": time.time() - 5})})) +check("a lapsed CLI sign-in gives way to the editor's", collectors.account_key(collectors.cursor_token()) == "ide") +cli.write_text(json.dumps({"accessToken": jwt({"sub": "cli", "exp": later})})) +check("the CLI's sign-in comes before the editor's", collectors.account_key(collectors.cursor_token()) == "cli") +moved_auth = work / "moved-auth.json" +moved_auth.write_text(json.dumps({"accessToken": "moved-key"})) +os.environ["CURSOR_AUTH_PATH"] = str(moved_auth) +check("CURSOR_AUTH_PATH moves the CLI's sign-in", collectors.cursor_auth_path() == moved_auth and collectors.cursor_token() == "moved-key") +check("an explicit auth path wins over CURSOR_AUTH_PATH", collectors.account_key(collectors.cursor_token(cli)) == "cli") +os.environ["CURSOR_API_KEY"] = " exported " +check("an exported API key is never sent as a bearer, and with no token traded for it there is none", + collectors.cursor_token() == "") +moved_auth.write_text(json.dumps({"accessToken": "traded-token", "apiKey": "exported"})) +check("an exported API key stands for the token the CLI traded it for", collectors.cursor_token() == "traded-token") +moved_auth.write_text(json.dumps({"accessToken": "traded-token", "apiKey": "another-key"})) +check("a token traded for another key never stands for this one", collectors.cursor_token() == "") +os.environ["CURSOR_AUTH_TOKEN"] = " direct " +check("CURSOR_AUTH_TOKEN is a bearer as it stands, ahead of everything", collectors.cursor_token() == "direct") +os.environ.pop("CURSOR_AUTH_TOKEN") +check( + "the editor's database is read without touching it", + state_db.read_bytes() == state_db_bytes and sorted(path.name for path in state_db.parent.iterdir()) == ["state.vscdb"], + json.dumps(sorted(path.name for path in state_db.parent.iterdir())), +) +for name in ("CURSOR_API_KEY", "CURSOR_AUTH_PATH"): + os.environ.pop(name) + +deep_header = work / "deep-header.jsonl" +deep_header.write_text("[" * 200_000 + "\n") +check("a session header nested deeper than the parser goes names no parent instead of stopping the scan", + collectors.pi_parent_session(str(deep_header)) is None) + +# --------------------------------------------------------------- accounts + +os.environ["XDG_STATE_HOME"] = str(work / "state") +registry = work / "state" / "omarchy" / "agents" / "accounts" / "demo.json" +registry.parent.mkdir(parents=True) +registry.write_text(json.dumps({"accounts": [{"id": "main", "primary": True}], "active": "main"})) +check("a registry with one account adds nothing to the record", collectors.registered_accounts("demo") == []) +check("no registry at all adds nothing either", collectors.registered_accounts("nothing") == []) +registry.write_text(json.dumps({"accounts": [{"id": "main", "primary": True}, {"id": "work", "label": "Work"}, {"label": "no id"}, "junk"], "active": "gone", "switch": "auto", "threshold": 80})) +accounts = collectors.registered_accounts("demo") +check( + "registered accounts carry which is active (the first when the registry names none it has) and the switch settings", + [(a["id"], a["active"], a["switch"]) for a in accounts] == [("main", True, {"mode": "auto", "threshold": 80}), ("work", False, {"mode": "auto", "threshold": 80})], + json.dumps(accounts), +) +for raw, held in ((20, 50), (180, 100), ("x", 95), (None, 95), (85.7, 85)): + registry.write_text(json.dumps({"accounts": [{"id": "main", "primary": True}, {"id": "work"}], "threshold": raw})) + check(f"a registry threshold of {raw!r} reads as {held}, as account state holds it", collectors.registered_accounts("demo")[0]["switch"]["threshold"] == held) +registry.write_text(json.dumps({"accounts": [{"id": "main", "primary": True}, {"id": "work", "label": "Work"}], "active": "gone", "switch": "auto", "threshold": 80})) +limits = {"limits": [{"label": "Weekly", "percent": 0.5, "resetsAt": ""}], "live": False, "fetchedAtMs": 5, "usageStatusText": "", "authHelpText": "help"} +entry = collectors.account_entry(accounts[1], limits, "", resetCredits=None) +check( + "an account entry names the account, its limits, and the collector's own fields last", + list(entry) == ["id", "label", "email", "plan", "active", "primary", "limits", "stale", "fetchedAt", "usageStatusText", "authHelpText", "retryAdvised", "resetCredits"] + and (entry["label"], entry["plan"], entry["stale"], entry["fetchedAt"], entry["retryAdvised"], entry["primary"]) == ("Work", "", True, 5, False, False) + and collectors.account_entry(dict(accounts[0], plan="Old"), dict(limits, live=True, retryAdvised=True), "Max")["plan"] == "Max" + and collectors.account_entry(dict(accounts[0], plan="Old"), limits, "")["plan"] == "Old", + json.dumps(entry), +) +check( + "an account whose check failed with nothing kept has no limits to call stale", + collectors.account_entry(accounts[1], dict(limits, limits=[]), "")["stale"] is False, +) +PY diff --git a/test/shell.d/agent-usage-synthetic-test.sh b/test/shell.d/agent-usage-synthetic-test.sh new file mode 100755 index 00000000000..60b9e1a210c --- /dev/null +++ b/test/shell.d/agent-usage-synthetic-test.sh @@ -0,0 +1,189 @@ +#!/bin/bash + +# The Synthetic collector's own sources: the API key from the environment, +# pi, omp, opencode or a named profile, Synthetic's quotas answer, and the +# Hugging Face model names it bills. The unused-machine record and counting +# pi, omp and opencode rows once are held centrally by the partition test, +# and the limits and scan caches by the contract checks, run here against a +# real stand-in for api.synthetic.new. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox +require_command rg + +python3 - <<'PY' +from agent_usage_fixtures import * + +MODERN = { + "subscription": {"requests": 0, "limit": 1000, "renewsAt": "2099-08-17T00:00:00Z"}, + "rollingFiveHourLimit": {"remaining": 750, "max": 1000, "limited": False, "nextTickAt": "2099-08-16T02:00:00Z"}, + "weeklyTokenLimit": {"percentRemaining": 64, "remainingCredits": "$31.05", "maxCredits": "$48.00", "nextRegenAt": "2099-08-16T04:00:00Z"}, + "search": {"hourly": {"requests": 25, "limit": 250, "renewsAt": "2099-08-16T03:00:00Z"}}, + "freeToolCalls": {"requests": "many", "limit": 100}, +} +state = {"mode": "up"} + + +# A stand-in for api.synthetic.new: /v2/quotas answers with MODERN while up, +# and with the mode's status otherwise. +def quotas(request): + # Something to do while the request is in flight, once. + if state.get("during"): + state.pop("during")() + if isinstance(state["mode"], int): + return state["mode"], {"error": "refused"} + return MODERN if request["path"] == "/v2/quotas" else {} + + +BASE_URL, reqs = stub_server(quotas) + + +# The stub's mode for the runs inside; an unreachable service is FakeHttp's, +# in front of the opener. +@contextlib.contextmanager +def serve(mode): + os.environ["SYNTHETIC_API_BASE_URL"] = BASE_URL + if mode in ("down", "timeout", "cut", "garbage"): + with FakeHttp(mode=mode): + yield + return + state["mode"] = mode + try: + yield + finally: + state["mode"] = "up" + + +def synthetic(*args): + with serve("up"): + return collect("synthetic", *(args or ("--force",))) + + +def sessions(): + return home() / ".pi" / "agent" / "sessions" + + +def key_used(**env): + for path in limits_cache_files("synthetic"): + path.unlink() + saved = {name: os.environ.get(name) for name in env} + os.environ.update(env) + try: + synthetic() + finally: + for name, value in saved.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + return reqs[-1]["headers"].get("Authorization") if reqs else None + + +# ---------------------------------------------------------------- the key + +# The partition test holds the rest of the unused-machine record; asking +# nothing without a key is Synthetic's own to say. +sandbox() +record = synthetic() +expect("Synthetic asks nothing without a key", [record["ready"], record["limits"], len(reqs)], [False, [], 0], record) + +pi_auth(home() / ".omp" / "profiles" / "work" / "agent", "synthetic", "profile-key") +expect("Synthetic reads a named profile's key", key_used(), "Bearer profile-key") +opencode_auth(None, "synthetic", "opencode-key") +expect("Synthetic prefers opencode's key to a profile's", key_used(), "Bearer opencode-key") +pi_auth(home() / ".omp" / "agent", "synthetic", "omp-key") +expect("Synthetic prefers omp's key to opencode's", key_used(), "Bearer omp-key") +pi_auth(None, "synthetic", "$PI_SYNTHETIC") +expect("Synthetic resolves pi's $ENV key", key_used(PI_SYNTHETIC="pi-key"), "Bearer pi-key") +write_json(home() / "moved" / "auth.json", {"synthetic": {"key": "moved-key"}}) +expect("Synthetic follows PI_CODING_AGENT_DIR to pi's key", key_used(PI_CODING_AGENT_DIR=str(home() / "moved")), "Bearer moved-key") +expect("Synthetic prefers SYNTHETIC_API_KEY to every stored key", key_used(SYNTHETIC_API_KEY="env-key"), "Bearer env-key") + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_usage(): + add_usage.count = getattr(add_usage, "count", 0) + 1 + write_pi_session(sessions() / "grown" / f"{add_usage.count}.jsonl", [pi_reply("synthetic", 10, id=f"grown-{add_usage.count}")]) + + +adapter = Adapter("synthetic", plant=lambda: os.environ.update(SYNTHETIC_API_KEY="secret"), serve=serve, count=lambda: len(reqs), + grow=add_usage, pi_root=sessions) +fresh = check_limits_contract(adapter) + +# The rolling limit and the weekly credits regenerate a little at a time, so +# their next tick is no reset time; search renews on the hour. +expect("Synthetic maps the quota windows, leaving regeneration ticks out of reset times", + [[entry["label"], round(entry["percent"], 3), entry["resetsAt"], entry.get("title", "")] for entry in fresh["limits"]], + [["Rolling (5h)", 0.25, "", ""], ["Weekly credits", 0.36, "", ""], ["Search (hourly)", 0.1, "2099-08-16T03:00:00Z", "Search"]], fresh) +expect("Synthetic's windows carry only contract fields and no plan", + [fresh["tierLabel"], sorted({key for entry in fresh["limits"] for key in entry})], ["", ["label", "percent", "resetsAt", "title"]], fresh) + +with serve(401): + record = collect("synthetic", "--force") +expect("a refused key keeps the last quotas beside the sign-in it asks for", + [record["usageStatusText"], record["limitsStale"], len(record["limits"])], ["Sign-in expired", True, 3], record) + +for path in limits_cache_files("synthetic"): + path.unlink() +with serve(500): + record = collect("synthetic", "--force") +expect("a failed check with nothing kept says the limits are unavailable", + [record["usageStatusText"], record["limits"], record["limitsFetchedAt"]], ["Synthetic limits unavailable", [], 0], record) +check("a failed check says what failed", "status 500" in record["authHelpText"], record) + +# An answer never rolls back one another run stamped while this one waited +# on the network. +synthetic() +state["during"] = lambda: edit_limits_cache("synthetic", fetched_ago=-60) +synthetic() +expect("Synthetic never overwrites a newer cached answer with its own", + [json.loads(path.read_text())["fetchedAtMs"] > ms_in(30) for path in limits_cache_files("synthetic")], [True]) +for path in limits_cache_files("synthetic"): + path.unlink() + +check_scan_cache(adapter) + +# ------------------------------------------------------------- the answer + +parse_quotas = load_collector("synthetic").parse_quotas +expect("Synthetic reads the documented subscription bucket", + [[entry["label"], entry["percent"]] for entry in parse_quotas({"subscription": {"requests": 27, "limit": 135, "renewsAt": "2099-08-17T00:00:00Z"}})], + [["Subscription requests", 0.2]]) +expect("Synthetic fills a limited rolling window and skips an unreadable one", + [entry["percent"] for entry in parse_quotas({"rollingFiveHourLimit": {"remaining": 900, "max": 1000, "limited": True}, "weeklyTokenLimit": {"percentRemaining": True}})], + [1.0]) + +# SYNTHETIC_API_BASE_URL moves the endpoint, never the key onto plain HTTP +# to another machine, nor into a URL that isn't one. +sandbox() +os.environ["SYNTHETIC_API_KEY"] = "secret" +for url in ("http://evil.example", "not a url"): + os.environ["SYNTHETIC_API_BASE_URL"] = url + with FakeHttp({"/v2/quotas": MODERN}) as http: + record = collect("synthetic", "--force") + expect(f"Synthetic sends its key to no untrusted endpoint ({url})", + [http.count, record["usageStatusText"], "HTTPS" in record["authHelpText"]], [0, "Synthetic limits unavailable", True], record) + expect(f"a Synthetic key with no sessions here has no local stats ({url})", + [record["ready"], record["totalSessions"], record["hasLocalStats"]], [True, 0, False], record) +del os.environ["SYNTHETIC_API_KEY"], os.environ["SYNTHETIC_API_BASE_URL"] + +# ------------------------------------------------------------ local stats + +# Synthetic serves Hugging Face models as hf:org/model; the model is the last +# part, from pi and opencode alike. +sandbox() +write_pi_session(sessions() / "work" / "one.jsonl", [pi_reply("synthetic", 100, id="a", model="hf:moonshotai/Kimi-K2")]) +opencode_db(message=[("m1", "s1", oc_v1("synthetic", 40, model="hf:zai-org/GLM-4.6"))]) +record = synthetic() +expect("Synthetic names its Hugging Face models by their last part", sorted(record["modelUsage"]), ["GLM-4.6", "Kimi-K2"], record) +expect("Synthetic in use without a key waits for auth", + [record["ready"], record["usageStatusText"], record["limits"]], [True, "Waiting for auth", []], record) +expect("Synthetic's sessions here are local stats", record["hasLocalStats"], True, record) +PY diff --git a/test/shell.d/agent-usage-vibe-test.sh b/test/shell.d/agent-usage-vibe-test.sh new file mode 100755 index 00000000000..9b2e751214f --- /dev/null +++ b/test/shell.d/agent-usage-vibe-test.sh @@ -0,0 +1,367 @@ +#!/bin/bash + +# The Mistral Vibe collector's own sources: Vibe's Unified Harness store, the +# session logs older releases wrote and their per-session index, the model +# and provider each session resolves to, the plan Vibe caches, VIBE_HOME and +# its session_logging settings, and the model names pi, omp and opencode give +# Mistral's models. How pi and opencode +# stores are read (roots, forks, both opencode stores, PI_CODING_AGENT_DIR, +# untrusted values) is held by the shared and partition tests, as is the +# unused-machine record; the scan cache goes through the contract check. +# Mistral publishes no limits, so there is no limits contract to hold. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +MISTRAL_CONFIG = { + "active_model": "mistral-medium-3.5", + "models": { + "mistral-medium-3.5": {"name": "mistral-vibe-cli-latest", "provider": "mistral", "alias": "mistral-medium-3.5"}, + "local": {"name": "qwen", "provider": "llamacpp", "alias": "local"}, + }, + "providers": [ + {"name": "mistral", "api_base": "https://api.mistral.ai/v1", "backend": "mistral"}, + {"name": "llamacpp", "api_base": "http://127.0.0.1:8080/v1", "backend": "generic"}, + ], +} +LOCAL_CONFIG = {**MISTRAL_CONFIG, "active_model": "local"} +NOW = iso_in(0) + + +def vibe(*args): + return collect("vibe", *(args or ("--force",))) + + +def session_root(): + return home() / ".vibe" / "logs" / "session" + + +# A Vibe session's meta.json: its running totals and the config it ran with. +def write_meta(directory, start, prompt, cached, completion, config=MISTRAL_CONFIG): + stats = {"session_prompt_tokens": prompt, "session_cached_tokens": cached, "session_completion_tokens": completion} + return write_json(Path(directory) / "meta.json", {"session_id": Path(directory).name, "start_time": start, "stats": stats, "config": config}) + + +def write_messages(directory, lines): + (Path(directory) / "messages.jsonl").write_text("".join(json.dumps(line) + "\n" for line in lines)) + + +# ------------------------------------------------------------ Vibe sessions + +sandbox() +main_session = session_root() / "session_20260101_120000_aaaa" +write_meta(main_session, NOW, 1000, 400, 200) +write_messages(main_session, [ + {"role": "user", "content": "Fix the bug"}, + {"role": "assistant", "content": "Done"}, + {"role": "user", "content": "", "injected": True}, + {"role": "user", "content": "summary", "context_boundary": "compaction"}, + {"role": "tool", "content": "ok"}, + {"role": "user", "content": "Thanks"}, +]) +# A subagent's totals belong to the session that spawned it; its task is not +# a prompt anyone typed. +subagent = main_session / "agents" / "explore_20260101_120100_bbbb" +write_meta(subagent, NOW, 300, 0, 50) +write_messages(subagent, [{"role": "user", "content": "Look around"}]) +# A session on the local llama.cpp provider is not Mistral usage, but the +# local collector's. +write_meta(session_root() / "session_20260101_130000_cccc", NOW, 9000, 0, 9000, LOCAL_CONFIG) +write_messages(session_root() / "session_20260101_130000_cccc", [{"role": "user", "content": "local"}]) +# Running totals are booked on the day the session started. +write_meta(session_root() / "session_20260101_140000_dddd", iso_in(-2 * 86400), 100, 0, 10) + +record = vibe() +two_days_ago = (datetime.now() - timedelta(days=2)).strftime("%Y-%m-%d") +expect("Vibe counts Mistral sessions, their subagents, and typed prompts only, booked on the day each started", { + "name": record["name"], + "counts": {key: record[key] for key in ("ready", "todayPrompts", "todaySessions", "todayTotalTokens", "totalPrompts", "totalSessions", "activeDays")}, + "models": record["modelUsage"], + "two days ago": next(day["messageCount"] for day in record["recentDays"] if day["date"] == two_days_ago), +}, { + "name": "Mistral Vibe", + "counts": {"ready": True, "todayPrompts": 2, "todaySessions": 1, "todayTotalTokens": 1550, "totalPrompts": 2, "totalSessions": 2, "activeDays": 2}, + "models": {"mistral-medium-3.5": {"inputTokens": 1000, "outputTokens": 260, "cacheReadInputTokens": 400, "cacheCreationInputTokens": 0}}, + "two days ago": 110, +}, record) +local = collect("local", "--force") +expect("Local counts the Vibe session on the llama.cpp server, with its prompt", + [local["todayTokensByModel"], local["todayPrompts"], local["todaySessions"]], [{"local": 18000}, 1, 1], local) + +# A zoneless start at the edge of the calendar, in a zone off UTC, is no +# start to book; the scan still completes, and is cached. +sandbox() +saved_tz = os.environ.get("TZ") +os.environ["TZ"] = "Asia/Tokyo" +time.tzset() +write_meta(session_root() / "session_20260101_150000_eeee", "0001-01-01T00:00:00", 100, 0, 10) +write_meta(session_root() / "session_20260101_160000_ffff", NOW, 100, 0, 10) +try: + record = vibe() +finally: + if saved_tz is None: + os.environ.pop("TZ") + else: + os.environ["TZ"] = saved_tz + time.tzset() +expect("a zoneless start at year 1 off UTC skips that session, and the scan is cached", + [record["totalSessions"], bool(scan_cache_files("vibe"))], [1, True], record) + +# --------------------------------------------------- Unified Harness store + +import hashlib + + +# One session in the Unified Harness store, as Vibe 3 writes it: CURRENT names +# a generation whose projection-state.json holds the camelCase public state, +# its history kept in the chunk pool when chunked, and the journal it names +# holds the records since. meta.json names only the model. +def write_unified(session_id, usage, entries=(), journal=(), alias="mistral-medium-3.5", chunked=False, metadata=None, torn=b""): + root = session_root() / "unified" / session_id + generation = "0000000000000003-abcd" + gen = root / "generations" / generation + gen.mkdir(parents=True) + def doc(path, value): + path.write_text(json.dumps(value, sort_keys=True, separators=(",", ":")) + "\n") + session = {"id": session_id, "rootSessionId": (metadata or {}).get("root_session_id", session_id), "status": {"type": "idle"}, + "createdAt": ms_in(0), "updatedAt": ms_in(0), "tokenUsage": usage} + history = {"range": "latest", "entries": [] if chunked else list(entries), "cursor": {}} + projection = {"projection_state_version": 1, "session_id": session_id, "snapshot_sequence": 3, "watermark": 9, + "snapshot": {"format": "harness.public-session-state/v1", "session": session, "history": history, "turnQueue": {}}} + described = {"path": "projection-state.json", "sha256": "0" * 64} + if chunked: + body = json.dumps(list(entries), sort_keys=True, separators=(",", ":")).encode() + digest = hashlib.sha256(body).hexdigest() + (root / "chunks").mkdir() + (root / "chunks" / f"{digest}.json").write_bytes(body) + described["chunks"] = [digest] + doc(gen / "projection-state.json", projection) + doc(gen / "runtime-state.json", {"session_id": session_id, "snapshot_sequence": 3, + "session_metadata": {"cwd": "/tmp", "root_session_id": session_id, **(metadata or {})}}) + doc(gen / "manifest.json", {"manifest_version": 1, "session_id": session_id, "generation": generation, "snapshot_sequence": 3, + "execution_state": "quiescent", "projection_state": described, + "runtime_state": {"path": "runtime-state.json", "sha256": "0" * 64}, + "checkpoint": {"path": "checkpoint.json", "sha256": "0" * 64}, + "recovery_journal_segment": {"path": "journal/0000000000000004.jsonl", "first_sequence": 4}}) + (root / "journal").mkdir() + lines = [json.dumps({"recovery_journal_record_version": 1, "sequence": 4 + index, "type": kind, "payload": payload}) + "\n" + for index, (kind, payload) in enumerate(journal)] + (root / "journal" / "0000000000000004.jsonl").write_bytes("".join(lines).encode() + torn) + doc(root / "CURRENT", {"store_format": "mistral.vibe.unified-session-store/v1", "store_format_minor": 3, "session_id": session_id, + "generation": generation, "snapshot_sequence": 3, "manifest_sha256": "0" * 64}) + config = {"active_model": alias, "model_name": "mistral-vibe-cli-latest"} if alias else None + write_json(root / "meta.json", {"session_id": session_id, "start_time": NOW, "config": config}) + return root + + +def usage(input_tokens, output_tokens, cached=0): + return {"inputTokens": input_tokens, "outputTokens": output_tokens, "totalTokens": input_tokens + output_tokens, "cachedInputTokens": cached} + + +def message(entry_id, role, source): + return {"type": "message", "id": entry_id, "role": role, "source": source, "content": [{"type": "text", "text": entry_id}]} + + +sandbox() +# The journal's newest envelope carries the totals since the generation was +# published, and a record still being written counts for nothing. +newer = {"format": "harness.public-session-state/v1", "history": {"entries": []}, + "session": {"id": "ses_main", "createdAt": ms_in(0), "tokenUsage": usage(5000, 700, 2000)}} +write_unified("ses_main", usage(1000, 100, 400), chunked=True, entries=[ + message("u1", "user", "turn_start"), + message("a1", "assistant", "turn_start"), + message("r1", "user", "harness"), +], journal=[ + ("projection_delta", {"watermark": 10, "delta": [{"op": "append_entry", "entry": message("u2", "user", "turn_steer")}, {"op": "set_envelope", "state": newer}]}), +], torn=b'{"type":"projection_delta","payload":{"delta":[{"op":"set_envelope","state":{"session":{"tokenUsage":{"inputTokens":999999') +# A subagent's totals belong to the session that spawned it; its task is not a +# prompt anyone typed. +write_unified("ses_child", usage(300, 50), entries=[message("t1", "user", "turn_start")], + metadata={"root_session_id": "ses_main", "parent_session_id": "ses_main", "subagent_spawn_key": "spawn-1"}) +# A fork starts from no usage of its own, and what it copied is not its prompts. +write_unified("ses_fork", usage(40, 10), entries=[message("c1", "user", "harness"), message("f1", "user", "turn_start")], + metadata={"root_session_id": "ses_main", "parent_session_id": "ses_main"}) +# A session on the local llama.cpp model config.toml adds is not Mistral +# usage, but the local collector's. +(home() / ".vibe").mkdir(parents=True, exist_ok=True) +(home() / ".vibe" / "config.toml").write_text('[[models]]\nname = "qwen"\nprovider = "llamacpp"\nalias = "local"\n\n[[providers]]\nname = "llamacpp"\napi_base = "http://127.0.0.1:8080/v1"\n') +write_unified("ses_local", usage(9000, 9000), entries=[message("l1", "user", "turn_start")], alias="local") +with contextlib.redirect_stderr(io.StringIO()): + record = vibe() +expect("Vibe counts Unified Harness sessions from the store's newest state, folding subagents into their session", { + key: record[key] for key in ("ready", "todayPrompts", "todaySessions", "todayTotalTokens", "totalSessions", "modelUsage") +}, { + "ready": True, "todayPrompts": 3, "todaySessions": 2, "todayTotalTokens": 6100, "totalSessions": 2, + "modelUsage": {"mistral-medium-3.5": {"inputTokens": 3340, "outputTokens": 760, "cacheReadInputTokens": 2000, "cacheCreationInputTokens": 0}}, +}, record) +check("Vibe caches a Unified Harness scan whose torn last journal record is still being written", bool(scan_cache_files("vibe"))) +with contextlib.redirect_stderr(io.StringIO()): + local = collect("local", "--force") +expect("Local counts the Unified Harness session on the llama.cpp server, with its prompt", + [local["todayTokensByModel"], local["todayPrompts"]], [{"local": 18000}, 1], local) + +# A store whose CURRENT names a generation that isn't there leaves the scan +# incomplete, so it isn't cached. +sandbox() +write_unified("ses_ok", usage(10, 5)) +write_unified("ses_broken", usage(10, 5)) +write_json(session_root() / "unified" / "ses_broken" / "CURRENT", {"generation": "missing"}) +with contextlib.redirect_stderr(io.StringIO()): + record = vibe() +expect("Vibe skips a Unified Harness session it cannot read without caching the scan", + [record["todayTotalTokens"], bool(scan_cache_files("vibe"))], [15, False], record) + +# A Unified Harness session goes where config.toml routes its alias now, so a +# route change reaches the cached scan at once, even on the panel's +# --limits-only refresh that otherwise reuses a scan for minutes. +sandbox() +MISTRAL_TOML = '[[providers]]\nname = "mistral"\napi_base = "{}"\napi_key_env_var = "MISTRAL_API_KEY"\nbackend = "mistral"\n\n[[models]]\nname = "devstral-2"\nprovider = "mistral"\nalias = "devstral-mock"\n' +(home() / ".vibe").mkdir(parents=True, exist_ok=True) +(home() / ".vibe" / "config.toml").write_text(MISTRAL_TOML.format("https://api.mistral.ai/v1")) +write_unified("ses_routed", usage(67000, 896), entries=[message("u1", "user", "turn_start")], alias="devstral-mock") +before = vibe("--limits-only") +(home() / ".vibe" / "config.toml").write_text(MISTRAL_TOML.format("http://127.0.0.1:18083/v1")) +after = vibe("--limits-only") +expect("a Vibe route change reaches Vibe's cached scan at once", + [before["todayTokensByModel"], after["todayTokensByModel"], after["todayPrompts"]], [{"devstral-mock": 67896}, {}, 0], after) + +# A session whose alias config.toml no longer names isn't Mistral's because +# Vibe would now fall back to its default: no one can say who served it. +sandbox() +(home() / ".vibe").mkdir(parents=True, exist_ok=True) +(home() / ".vibe" / "config.toml").write_text('[[models]]\nname = "qwen"\nprovider = "llamacpp"\nalias = "local"\n') +write_unified("ses_dropped", usage(5000, 5000), entries=[message("d1", "user", "turn_start")], alias="removed-local") +with contextlib.redirect_stderr(io.StringIO()): + record = vibe("--force") + local = collect("local", "--force") +expect("a Unified Harness session on an alias config.toml dropped is billed to no one", + [record["todayTokensByModel"], local["todayTokensByModel"]], [{}, {}], [record, local]) + +# ------------------------------------------------------- model resolution + +sandbox() +# Unpinned: Vibe falls back to its default alias. +write_meta(session_root() / "session_1", NOW, 10, 0, 0, { + "active_model": "", + "models": {"local": {"name": "q", "provider": "llamacpp", "alias": "local"}, "mistral-medium-3.5": {"name": "m", "provider": "mistral", "alias": "mistral-medium-3.5"}}, + "providers": [{"name": "mistral", "backend": "mistral"}, {"name": "llamacpp", "backend": "generic"}], +}) +# Older releases dump the models as a list. +write_meta(session_root() / "session_2", NOW, 20, 0, 0, { + "active_model": "devstral", "models": [{"name": "devstral-2", "provider": "mistral", "alias": "devstral"}], "providers": [{"name": "mistral", "backend": "mistral"}], +}) +# A generic provider aimed at Mistral's API is still Mistral. +write_meta(session_root() / "session_3", NOW, 40, 0, 0, { + "active_model": "codestral", "models": {"codestral": {"name": "codestral-latest", "provider": "la-plateforme", "alias": "codestral"}}, + "providers": [{"name": "la-plateforme", "api_base": "https://codestral.mistral.ai/v1", "backend": "generic"}], +}) +# A lookalike host is not. +write_meta(session_root() / "session_4", NOW, 80, 0, 0, { + "active_model": "x", "models": {"x": {"name": "x", "provider": "proxy", "alias": "x"}}, + "providers": [{"name": "proxy", "api_base": "https://mistral.ai.example.com/v1", "backend": "generic"}], +}) +# The mistral backend sends to its api_base too: one aimed at a server you +# run is the local collector's, and one at another provider's host is not +# Mistral's, as Vibe 3.0.1 wrote them against a local mock. +for n, api_base in ((5, "http://127.0.0.1:18083/v1"), (6, "https://api.deepseek.com/v1")): + write_meta(session_root() / f"session_{n}", NOW, 1000 * n, 0, 0, { + "active_model": "devstral-mock", "models": [{"name": "devstral-2", "provider": "mistral", "alias": "devstral-mock"}], + "providers": [{"name": "mistral", "api_base": api_base, "api_key_env_var": "MISTRAL_API_KEY", "backend": "mistral"}], + }) +record = vibe() +expect("Vibe resolves the session model and its provider as Vibe does", + {model: usage["inputTokens"] for model, usage in record["modelUsage"].items()}, {"mistral-medium-3.5": 10, "devstral": 20, "codestral": 40}, record) +local = collect("local", "--force") +expect("Local counts a mistral-backend session on a server you run", + local["todayTokensByModel"], {"devstral-mock": 5000}, local) + +# --------------------------------------------------------- untrusted fields + +sandbox() +write_meta(session_root() / "session_bad_time", "not a time", 500, 0, 500) +write_meta(session_root() / "session_future", "2999-01-01T00:00:00+00:00", 500, 0, 500) +write_meta(session_root() / "session_odd_tokens", NOW, "700", True, float("nan")) +write_meta(session_root() / "session_negative", NOW, -5, 0, 7) +(session_root() / "session_not_json").mkdir() +(session_root() / "session_not_json" / "meta.json").write_text("[1, 2") +with contextlib.redirect_stderr(io.StringIO()): + record = vibe() +expect("Vibe skips bad times and bad token fields without losing the record", + [record["ready"], record["todayTotalTokens"], record["totalSessions"]], [True, 7, 1], record) + +# A meta.json that fails to parse leaves the scan incomplete, so it isn't cached. +check("Vibe does not cache a scan that could not read every session", not scan_cache_files("vibe"), + ", ".join(path.name for path in scan_cache_files("vibe"))) + +# --------------------------------------------------------------------- plan + +sandbox() +write_json(home() / ".vibe" / "whoami_cache.json", { + "old": {"stored_at_timestamp": 100, "payload": {"plan_type": "api", "plan_name": "FREE_TRIAL"}}, + "new": {"stored_at_timestamp": 200, "payload": {"plan_type": "chat", "plan_name": "individual"}}, +}) +record = vibe() +expect("Vibe names the newest cached plan as Vibe does, with no local stats while no session has run", + [record["ready"], record["tierLabel"], record["hasLocalStats"]], [True, "Pro", False], record) + +write_json(home() / ".vibe" / "whoami_cache.json", {"a": {"stored_at_timestamp": 1, "payload": {"plan_type": "chat", "plan_name": "MYSTERY"}}}) +record = vibe() +expect("Vibe leaves a plan Vibe cannot name unnamed", [record["ready"], record["tierLabel"]], [False, ""], record) + +# ---------------------------------------------------- VIBE_HOME and save_dir + +sandbox() +(home() / "vibe-home").mkdir() +(home() / "vibe-home" / "config.toml").write_text(f'[session_logging]\nsave_dir = "{home() / "elsewhere"}"\nsession_prefix = "run"\n') +write_meta(home() / "elsewhere" / "run_1", NOW, 30, 0, 3) +write_meta(home() / "elsewhere" / "session_1", NOW, 999, 0, 999) +os.environ["VIBE_HOME"] = str(home() / "vibe-home") +record = vibe() +del os.environ["VIBE_HOME"] +expect("Vibe honors VIBE_HOME and session_logging in config.toml", record["todayTotalTokens"], 33, record) + +# ------------------------------------------------------ pi, omp, opencode + +# Only the mistral provider counts, and Mistral's models keep their own names: +# a thinking suffix and a provider prefix come off. +sandbox() +write_pi_session(home() / ".pi" / "agent" / "sessions" / "p" / "a.jsonl", [ + pi_reply("mistral", id="m1", model="mistral-large:high", usage={"input": 6, "output": 1}), + pi_reply("openrouter", id="m2", model="mistralai/devstral", usage={"input": 999, "output": 999}), +]) +opencode_db(tables=("message",), message=[ + ("o1", "ses_1", oc_v1("mistral", {"input": 100, "output": 20, "reasoning": 5, "cache": {"read": 30, "write": 0}}, model="mistral/codestral-latest")), + ("o2", "ses_1", oc_v1("mistral-proxy", {"input": 999, "output": 999}, model="mistral/codestral-latest")), +]) +record = vibe() +expect("Vibe counts pi and opencode messages on the mistral provider under Mistral's model names", record["modelUsage"], { + "mistral-large": {"inputTokens": 6, "outputTokens": 1, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, + "codestral-latest": {"inputTokens": 100, "outputTokens": 25, "cacheReadInputTokens": 30, "cacheCreationInputTokens": 0}, +}, record) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_session(): + add_session.count = getattr(add_session, "count", 0) + 1 + write_meta(session_root() / f"session_grown_{add_session.count}", iso_in(0), 5, 0, 0) + + +check_scan_cache(Adapter("vibe", grow=add_session, pi_root=lambda: home() / ".pi" / "agent" / "sessions")) + +# The per-session index is only an optimization: a session whose meta.json +# changed is read again. +before = vibe()["todayTotalTokens"] +write_meta(session_root() / "session_grown_1", iso_in(0), 50, 0, 0) +record = vibe() +expect("Vibe rereads a session whose meta.json changed", record["todayTotalTokens"] - before, 45, record) +PY 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..a10e057763f --- /dev/null +++ b/test/shell.d/agent-usage-zai-limits-test.sh @@ -0,0 +1,188 @@ +#!/bin/bash + +# The Z.ai collector's limits: how it reads the monitor endpoint's windows and +# plan, a key without a Coding Plan, a rejected key, the keys pi saves, and its +# adapter for the contract checks. The unused-machine record is held centrally +# by the partition test, the limits cache's mechanics by the shared test, and +# the scan cache by the scanner test. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + +QUOTA = "api.z.ai/api/monitor/usage/quota/limit" +ANSWER = {"data": {"level": "max", "limits": [ + {"type": "CREDIT_LIMIT", "unit": 3, "number": 5, "percentage": 44.0, "nextResetTime": ms_in(3 * 3600)}, + {"type": "CREDIT_LIMIT", "unit": 6, "number": 1, "percentage": 10, "nextResetTime": ms_in(5 * 86400)}, +]}} + + +def keyed(key="token"): + os.environ["ZAI_API_KEY"] = key + + +def zai(answer=None, mode="up", args=("--force",)): + with FakeHttp({QUOTA: answer}, mode=mode) as http: + record = collect("zai", *args) + return record, http + + +# ---------------------------------------------------------------- contract + +sandbox() +fresh = check_limits_contract(Adapter("zai", plant=keyed, answer={QUOTA: ANSWER})) +expect("Z.ai's plan and windows come from the monitor endpoint", + [fresh["tierLabel"], [(entry["label"], entry["percent"]) for entry in fresh["limits"]]], + ["Max", [("Session (5-hour)", 0.44), ("Weekly (7-day)", 0.1)]], fresh) + +# A failed check keeps the plan with the windows. +record, _ = zai(mode="down") +expect("Z.ai collector keeps the cached plan past a failed check", [record["tierLabel"], record["limitsStale"]], ["Max", True], record) + +# A rejected key is a sign-in to redo, and only then is the key blamed. +record, _ = zai(mode=401) +check("Z.ai collector blames the key only when the key is rejected", + record["usageStatusText"] == "Sign-in expired" and "rejected the saved API key" in record["authHelpText"] and len(record["limits"]) == 2, record) + +# Z.ai names a refused credential in the body under HTTP 200, by its +# documented authentication codes; this is its live answer to no key at all. +for code in (1001, "1000"): + record, _ = zai({"code": code, "msg": "Authentication parameter not received in Header, unable to authenticate", "success": False}) + expect(f"Z.ai collector reads a body-level authentication failure ({code}) as a refused key", + [record["usageStatusText"], "rejected the saved API key" in record["authHelpText"], len(record["limits"])], + ["Sign-in expired", True, 2], record) + +# Another key never shows this one's numbers. +keyed("other-key") +record, _ = zai(mode="down") +expect("Z.ai collector keeps each key's limits apart", [record["limits"], record["usageStatusText"]], [[], "Z.ai limits unavailable"], record) + +# ------------------------------------------------------------- the answer + +# The mapped session and weekly windows, then the entries that must drop: a +# window tagged with an unmapped (unit, number), one tagged with a list, one +# whose percentage will not parse, and one whose percentage is a boolean. +# nextResetTime is milliseconds for the first and plain seconds for the +# second, to prove a seconds epoch is not rendered 1000x in the future. +sandbox() +keyed() +week = int(time.time() + 5 * 86400) +with contextlib.redirect_stderr(io.StringIO()): + record, _ = zai({"data": {"level": "pro", "limits": [ + {"type": "CREDIT_LIMIT", "unit": 3, "number": 5, "percentage": 42.0, "nextResetTime": ms_in(3 * 3600)}, + {"type": "TOKENS_LIMIT", "unit": 6, "number": 1, "percentage": 10, "nextResetTime": week}, + {"type": "CREDIT_LIMIT", "unit": 9, "number": 9, "percentage": 5, "nextResetTime": ms_in(3 * 3600)}, + {"type": "CREDIT_LIMIT", "unit": [3], "number": 5, "percentage": 7, "nextResetTime": ms_in(3 * 3600)}, + {"type": "CREDIT_LIMIT", "unit": 3, "number": 5, "percentage": None, "nextResetTime": ms_in(3 * 3600)}, + {"type": "CREDIT_LIMIT", "unit": 3, "number": 5, "percentage": True, "nextResetTime": ms_in(3 * 3600)}, + ]}}) +expect("Z.ai collector reads windows, tier, seconds/ms resets, and drops bad rows", + [[(entry["label"], entry["percent"]) for entry in record["limits"]], record["tierLabel"], record["limits"][1]["resetsAt"][:10]], + [[("Session (5-hour)", 0.42), ("Weekly (7-day)", 0.1)], "Pro", datetime.fromtimestamp(week, timezone.utc).date().isoformat()], record) + +# A key with no plan is a pay-as-you-go account: no trouble, no blame on the +# key, and no limits kept from a plan that lapsed. +record, _ = zai({"success": False, "msg": "This API key has no coding plan"}) +expect("Z.ai collector treats a no-plan key as healthy, without limits", + [record["usageStatusText"], record["authHelpText"], bool(record.get("retryAdvised")), record["limits"], record["limitsStale"]], + ["", "", False, [], False], record) + +# A confirmed no-plan answer replaces the lapsed plan's windows, so an outage +# right after it can't bring the cancelled allowance or its tier back, and +# with no limits to lose, the failed check is no news either. +sandbox() +keyed() +zai(ANSWER) +zai({"success": False, "msg": "This API key has no coding plan"}) +record, _ = zai(mode="down") +expect("Z.ai collector does not restore a cancelled plan's windows after a failed check, nor call it trouble", + [record["limits"], record["tierLabel"], record["limitsStale"], record["usageStatusText"], record["authHelpText"]], + [[], "", False, "", ""], record) + +# The same answer comes every time, so a no-plan key is asked again hourly, +# not every run. +zai({"success": False, "msg": "This API key has no coding plan"}) +edit_limits_cache("zai", fetched_ago=1800) +record, http = zai(mode="down", args=()) +expect("Z.ai collector reuses a no-plan answer for an hour", + [len(http.requests), record["usageStatusText"], record["limits"]], [0, "", []], record) +edit_limits_cache("zai", fetched_ago=3601) +record, http = zai(mode="down", args=()) +expect("Z.ai collector asks a no-plan key again after an hour, and a failure then is still no trouble", + [len(http.requests), record["usageStatusText"], record["authHelpText"], record["limits"]], [1, "", "", []], record) + +# With nothing kept, a failed check says the limits are unavailable; only a +# transport failure asks for an early retry. +for path in limits_cache_files("zai"): + path.unlink() +record, _ = zai(mode="down") +check("Z.ai collector reports unavailable limits when offline with nothing kept, and advises a retry", + record["usageStatusText"] == "Z.ai limits unavailable" and record.get("retryAdvised") is True and record["limits"] == [] + and "reach" in record["authHelpText"], record) +record, _ = zai(mode=500) +expect("Z.ai collector does not advise a retry for a server's error status", + [record["usageStatusText"], bool(record.get("retryAdvised"))], ["Z.ai limits unavailable", False], record) + +# A predictable lock path must not follow and truncate a planted symlink. +lock = load_collector("zai").limits_cache_paths("token", "zai")[1] +sentinel = home() / "sentinel" +sentinel.write_text("intact") +lock.parent.mkdir(parents=True, exist_ok=True) +lock.unlink(missing_ok=True) +lock.symlink_to(sentinel) +zai(ANSWER) +expect("Z.ai collector does not follow a planted lock symlink", sentinel.read_text(), "intact") + +# ------------------------------------------------------------ credentials + +# A key pi saved with /login counts as a credential. +sandbox() +pi_auth(None, "zai", "pi-saved") +expect("Z.ai collector reads the key pi saved in auth.json", list(load_collector("zai").credentials()), ["pi-saved", "zai"]) +pi_auth(home() / ".omp" / "agent", "zai", "omp-saved") +expect("Z.ai collector prefers pi's auth.json to omp's", list(load_collector("zai").credentials()), ["pi-saved", "zai"]) + +# A pi `$VAR` or `!command` entry is pi's to resolve, never a key to send. +pi_auth(None, "zai", "$ZAI_KEY") +pi_auth(home() / ".omp" / "agent", "zai", "!pass show zai") +expect("Z.ai collector never sends a pi $VAR or !command entry as a key", list(load_collector("zai").credentials()), ["", "zai"]) + +# Keys in omp's and pi's .env files, omp's first, each naming its platform. +sandbox() +(home() / ".pi" / "agent").mkdir(parents=True) +(home() / ".pi" / "agent" / ".env").write_text("# pi\nZHIPU_API_KEY=pi-env\n") +expect("Z.ai collector reads a Zhipu key from pi's .env", list(load_collector("zai").credentials()), ["pi-env", "zhipu"]) +(home() / ".omp" / "agent").mkdir(parents=True) +(home() / ".omp" / "agent" / ".env").write_text("export ZAI_API_KEY='omp-env'\n") +expect("Z.ai collector prefers omp's .env to pi's", list(load_collector("zai").credentials()), ["omp-env", "zai"]) +(home() / ".omp" / "agent" / ".env").write_bytes(b"ZAI_API_KEY=\xff\xfe\n") +expect("Z.ai collector skips a .env that isn't UTF-8", list(load_collector("zai").credentials()), ["pi-env", "zhipu"]) + +# A Zhipu key asks BigModel's host, with the raw key as its Authorization: +# the endpoint takes no Bearer prefix. +sandbox() +os.environ["ZHIPU_API_KEY"] = "zhipu-key" +BIGMODEL = "open.bigmodel.cn" + load_collector("zai").QUOTA_PATH +with FakeHttp({BIGMODEL: ANSWER}) as http: + record = collect("zai", "--force") +expect("a Zhipu key asks open.bigmodel.cn with the raw key, no Bearer", + [[request["url"] for request in http.requests], http.requests[0]["headers"].get("authorization"), len(record["limits"])], + [["https://" + BIGMODEL], "zhipu-key", 2], record) + +# The same key on the other platform is another account: its limits are kept +# apart, per platform. +write_json(config_home() / "omarchy" / "agents" / "zai.json", {"apiKey": "zhipu-key", "platform": "zai"}) +with FakeHttp(mode="down"): + other = collect("zai", "--force") +(config_home() / "omarchy" / "agents" / "zai.json").unlink() +with FakeHttp(mode="down"): + kept = collect("zai", "--force") +expect("Z.ai collector keeps one key's limits per platform", + [other["limits"], other["usageStatusText"], len(kept["limits"]), kept["limitsStale"]], [[], "Z.ai limits unavailable", 2, True], kept) +PY 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..f8d3e6aa5f9 --- /dev/null +++ b/test/shell.d/agent-usage-zai-scanner-test.sh @@ -0,0 +1,108 @@ +#!/bin/bash + +# The Z.ai collector's local stats: pi and omp turns on the zai and zhipu +# providers, dated as Z.ai's harnesses date them, and the keyless wait for +# auth. How pi and opencode stores are read (roots, forks, repeated ids, both +# opencode stores, OPENCODE_DB, opencode-v2.db, untrusted values) is held +# centrally by the shared and partition tests, the unused-machine record by +# the partition test, and the scan cache by the contract check. The limits +# are agent-usage-zai-limits-test.sh's. + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +source "$SHELL_TEST_DIR/agent-usage-helpers.sh" + +agent_sandbox + +python3 - <<'PY' +from agent_usage_fixtures import * + + +def zai(*args): + with FakeHttp(): + return collect("zai", *(args or ("--force",))) + + +def sessions(): + return home() / ".pi" / "agent" / "sessions" + + +# One session file exercises the scan: +# 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 +sandbox() +stamp = today_stamp() +write_pi_session(sessions() / "project" / "pi.jsonl", [ + {"type": "message", "id": "u1", "timestamp": stamp, "message": {"role": "user", "content": "hi"}}, + pi_reply("zai", id="a", model="glm-5.3", timestamp=stamp, usage={"input": 10, "output": 10, "reasoning": 6, "cacheRead": 3, "cacheWrite": 2}), + pi_reply("zhipu", id="b", model="glm-4.7", timestamp=stamp, usage={"input": 20, "output": 5}), + pi_reply("openai", id="c", model="gpt", timestamp=stamp, usage={"input": 999, "output": 999}), + {"type": "message", "id": "d", "timestamp": stamp, "message": {"role": "user", "provider": "zai", "model": "glm-5.3", "usage": {"input": 500, "output": 500}}}, + pi_reply("zai", id="a", model="glm-5.3", timestamp=stamp, usage={"input": 77, "output": 77}), + {"type": "message", "id": "u2", "timestamp": stamp, "message": {"role": "user", "content": "again"}}, + {"type": "message", "id": "f", "message": {"role": "assistant", "provider": "zai", "model": "glm-4.6", "timestamp": "2020-01-01T00:00:00Z", "usage": {"input": 7}}}, + '{"type":"message","id":"e","message":{"role":"assistant","provider":"zai","usage":{ broken', +]) +record = zai() + +# In use here, but no key to ask for its limits. +expect("Z.ai is ready on local stats alone and waits for a key, naming ZAI_API_KEY", + [record["name"], record["ready"], record["limits"], record["usageStatusText"], "ZAI_API_KEY" in record["authHelpText"]], + ["Z.ai", True, [], "Waiting for auth", True], record) + +# Reasoning tokens fold into output and cache tokens split out; openai, the +# user turn and the malformed line drop out. +expect("Z.ai counts zai and zhipu assistant turns, reasoning already in output", { + "models": sorted(record["modelUsage"]), "glm-5.3": record["modelUsage"]["glm-5.3"], +}, { + "models": ["glm-4.6", "glm-4.7", "glm-5.3"], + "glm-5.3": {"inputTokens": 10, "outputTokens": 10, "cacheReadInputTokens": 3, "cacheCreationInputTokens": 2}, +}, record) + +# a + b land today (25 + 25); f is dated by its message's own time, an old +# day, so it counts in the totals and never in today's. A prompt books on the +# first reply after it, a's and f's, b being a further step, and counts on +# the day it was typed: both today. +expect("Z.ai dates a turn by message.timestamp, keeping an old one out of today", + [record["todayTotalTokens"], record["todayPrompts"], record["totalPrompts"], record["activeDays"]], [50, 2, 2, 2], record) + +# Claude Code pointed at Z.ai (its settings' env, as Z.ai's guide sets it) is +# Z.ai's whatever model answered; with no endpoint, only a GLM reply is. A +# Claude reply sent to Anthropic stays Claude's. +sandbox() +claude_code_session("glm", [("glm-4.6", 500), ("claude-sonnet-4-5", 40)]) +record = zai() +expect("Z.ai counts Claude Code's GLM replies and their prompts, and no Claude one", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"], record["ready"]], [{"glm-4.6": 500}, 1, 1, True], record) +claude_code_session("glm", [("glm-4.6", 500), ("claude-sonnet-4-5", 40)], base_url="https://api.z.ai/api/anthropic") +record = zai() +expect("Z.ai counts every Claude Code reply while Claude Code is pointed at Z.ai", + [record["todayTokensByModel"], record["todayPrompts"]], [{"glm-4.6": 500, "claude-sonnet-4-5": 40}, 2], record) + +# MiniMax Code's Z.AI preset is a custom provider at api.z.ai: Z.ai's, with +# the prompt it answered. One at another host isn't. +sandbox() +mcode_routed_session("zed", "custom_provider:zai-coding-plan", "glm-4.7", 2358, base_url="https://api.z.ai/api/coding/paas/v4") +mcode_routed_session("other", "custom_provider:acme", "glm-elsewhere", 40, base_url="https://llm.acme.example/v1") +record = zai() +expect("Z.ai counts MiniMax Code replies on a custom provider at api.z.ai", + [record["todayTokensByModel"], record["todayPrompts"], record["todaySessions"]], [{"glm-4.7": 2358}, 1, 1], record) + +# ---------------------------------------------------------------- contract + +sandbox() + + +def add_reply(): + add_reply.count = getattr(add_reply, "count", 0) + 1 + write_pi_session(sessions() / "grown" / f"{add_reply.count}.jsonl", [pi_reply("zai", 5, id=f"grown-{add_reply.count}")]) + + +check_scan_cache(Adapter("zai", grow=add_reply, pi_root=sessions)) +PY diff --git a/test/shell.d/agent_usage_fixtures.py b/test/shell.d/agent_usage_fixtures.py new file mode 100644 index 00000000000..210bc3276be --- /dev/null +++ b/test/shell.d/agent_usage_fixtures.py @@ -0,0 +1,954 @@ +"""Fixtures the agent usage collector tests share. + +Imported from a test's Python heredoc once `agent_sandbox` (in +agent-usage-helpers.sh) has put this directory on PYTHONPATH: + + from agent_usage_fixtures import * + +Everything here builds a machine for a collector to read, runs the collector, +or holds its record to the collector contract (docs/agent-usage-collectors.md). +The contract checks drive a collector through main() and assert on the record +it prints, so every collector is held to the same field names. + +Files are found through ROOT and OMARCHY_PATH, which base-test.sh exports; +never through this file's own location or HOME. +""" + +from __future__ import annotations + +import base64 +import contextlib +import email.message +import glob +import http.client +import importlib.machinery +import importlib.util +import io +import json +import os +import sqlite3 +import sys +import tempfile +import threading +import time +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any, Callable + +sys.dont_write_bytecode = True + +ROOT = Path(os.environ["ROOT"]) +OMARCHY_PATH = Path(os.environ["OMARCHY_PATH"]) +sys.path.insert(0, str(OMARCHY_PATH / "shell" / "plugins" / "agents")) + +# Variables a developer's own shell may carry that point a collector at a +# real credential or a real store. The sandbox drops every one of them, and +# any *_API_KEY, so a test only ever sees what it planted. +DEVELOPER_PREFIXES = ( + "AGY_", "ANTHROPIC_", "CLAUDE_", "CLOUD_CODE_", "CODEX_", "COPILOT_", "CURSOR_", "DEEPSEEK_", + "DEVIN_", "DSH_", "FIREWORKS_", "GH_", "GITHUB_", "GROK_", "HERMES_", "KIMCHI_", "KIMI_", + "MAVIS_", "MINIMAX_", "MISE_", "MMX_", "MUSE_", "OLLAMA_", "OPENAI_", "OPENCLAW_", "OPENCODE_", "OPENROUTER_", + "PI_CODING_AGENT_", "QWEN_", "SYNTHETIC_", "VIBE_", "ZAI_", "ZHIPU", +) + +# A request that slips past FakeHttp or a stub server goes to a port nothing +# listens on, never to a real service. Loopback stays direct for stub servers. +DEAD_PROXY = "http://127.0.0.1:9" + + +# ------------------------------------------------------------------ output + + +def check(description: str, condition: Any, detail: Any = "") -> None: + """A TAP line: `ok - description`, or `not ok` and the detail on stderr, + ending the test, as check() in agent-usage-shared-test.sh does.""" + if condition: + print("ok - " + description, flush=True) + return + print("not ok - " + description, file=sys.stderr) + if detail: + print(detail if isinstance(detail, str) else json.dumps(detail, default=str), file=sys.stderr) + sys.exit(1) + + +def expect(description: str, actual: Any, expected: Any, record: Any = None) -> None: + """check() that actual == expected, printing both (and the record) when not.""" + detail = f"expected: {json.dumps(expected, default=str)}\nactual: {json.dumps(actual, default=str)}" + if record is not None: + detail += "\nrecord: " + json.dumps(record, default=str) + check(description, actual == expected, detail) + + +# ----------------------------------------------------------------- sandbox + + +def sandbox(root: str | Path | None = None) -> Path: + """Point HOME and every XDG directory under `root` (a new temporary + directory when omitted) and drop the developer's agent variables. Returns + the home, created with the usual directories in it.""" + if root is None: + root = tempfile.mkdtemp(prefix="agent-usage.", dir=os.environ.get("SANDBOX") or None) + root = Path(root) + home = root / "home" + for name in list(os.environ): + if name.endswith("_API_KEY") or name.startswith(DEVELOPER_PREFIXES): + del os.environ[name] + os.environ.update( + HOME=str(home), + XDG_CONFIG_HOME=str(home / ".config"), + XDG_DATA_HOME=str(home / ".local" / "share"), + XDG_STATE_HOME=str(home / ".local" / "state"), + XDG_CACHE_HOME=str(home / ".cache"), + XDG_RUNTIME_DIR=str(home / ".run"), + http_proxy=DEAD_PROXY, https_proxy=DEAD_PROXY, HTTP_PROXY=DEAD_PROXY, HTTPS_PROXY=DEAD_PROXY, + no_proxy="127.0.0.1,localhost", NO_PROXY="127.0.0.1,localhost", + ) + for path in (home / ".config", home / ".local" / "share", home / ".local" / "state", home / ".cache", home / ".run"): + path.mkdir(parents=True, exist_ok=True) + return home + + +def home() -> Path: + return Path(os.environ["HOME"]) + + +def data_home() -> Path: + return Path(os.environ.get("XDG_DATA_HOME") or home() / ".local" / "share") + + +def config_home() -> Path: + return Path(os.environ.get("XDG_CONFIG_HOME") or home() / ".config") + + +def state_home() -> Path: + return Path(os.environ.get("XDG_STATE_HOME") or home() / ".local" / "state") + + +def cache_root() -> Path: + """$XDG_CACHE_HOME/omarchy/agent-usage, where every collector caches.""" + import collectors + return collectors.cache_root() + + +# -------------------------------------------------------------- collectors + +_loaded = 0 + + +def load_collector(agent_id: str) -> Any: + """bin/omarchy-agent-usage- as a fresh module. The shared module stays + imported once, so anything patched on urllib reaches every load.""" + global _loaded + _loaded += 1 + path = ROOT / "bin" / f"omarchy-agent-usage-{agent_id}" + loader = importlib.machinery.SourceFileLoader(f"agent_usage_{agent_id.replace('-', '_')}_{_loaded}", str(path)) + spec = importlib.util.spec_from_loader(loader.name, loader) + module = importlib.util.module_from_spec(spec) + loader.exec_module(module) + module.FIXTURE_ID = agent_id + return module + + +def run_main(module: Any, *args: str) -> dict: + """Run a loaded collector's main() with `args` and return the record it + printed. A collector that exits non-zero or prints no record fails the test.""" + saved = sys.argv + sys.argv = [f"omarchy-agent-usage-{module.FIXTURE_ID}", *args] + output = io.StringIO() + try: + with contextlib.redirect_stdout(output): + try: + code = module.main() + except SystemExit as exit_: + code = exit_.code + finally: + sys.argv = saved + lines = [line for line in output.getvalue().splitlines() if line.strip()] + if code not in (None, 0) or not lines: + check(f"{module.FIXTURE_ID} {' '.join(args)} prints a record", False, f"exit {code!r}, stdout {output.getvalue()!r}") + return json.loads(lines[-1]) + + +def collect(agent_id: str, *args: str) -> dict: + """load_collector() and run_main() in one: a fresh module per run.""" + return run_main(load_collector(agent_id), *args) + + +# -------------------------------------------------------------------- http + + +class FakeResponse(io.BytesIO): + def __init__(self, body: bytes, status: int = 200, url: str = "") -> None: + super().__init__(body) + self.status = self.code = status + self.url = url + self.headers = self.msg = email.message.Message() + self.headers["Content-Type"] = "application/json" + + def getcode(self) -> int: + return self.status + + def geturl(self) -> str: + return self.url + + def info(self) -> email.message.Message: + return self.headers + + +class _CutShort(FakeResponse): + def read(self, *args: Any) -> bytes: + raise http.client.IncompleteRead(b'{"cut', 64) + + +def _body(payload: Any) -> bytes: + if isinstance(payload, bytes): + return payload + if isinstance(payload, str): + return payload.encode() + return json.dumps(payload).encode() + + +class FakeHttp: + """Stands in for every urllib.request.urlopen call, the one door every + collector request goes through (collectors.open_json included). + + `routes` maps a URL substring to the answer: a JSON-able payload, bytes or + a str sent as is, an int answered as that HTTP status, an exception + raised, or a callable taking the Request and returning any of those. The + first route whose key is in the URL answers; one matching none is logged + in `unrouted` and refused with 404. + + `mode` decides the whole service at once, ahead of the routes: + "up" (the routes answer), "down" (no DNS: URLError), "timeout", + "reset" (a connection reset), "cut" (a reply cut short), "garbage" (a body + that isn't JSON), or an int status every request is refused with. Change + it between runs with `http.mode = 500`. + + Every request lands in `requests` as {url, method, headers, data}. + """ + + def __init__(self, routes: dict | None = None, mode: Any = "up") -> None: + self.routes = dict(routes or {}) + self.mode = mode + self.requests: list[dict] = [] + self.unrouted: list[str] = [] + self._saved = None + + def __enter__(self) -> FakeHttp: + return self.start() + + def __exit__(self, *exc: Any) -> None: + self.stop() + + def start(self) -> FakeHttp: + self._saved = urllib.request.urlopen + urllib.request.urlopen = self.urlopen + return self + + def stop(self) -> None: + if self._saved is not None: + urllib.request.urlopen = self._saved + self._saved = None + + @property + def count(self) -> int: + return len(self.requests) + + @staticmethod + def refusal(url: str, status: int, body: Any = b"", headers: dict | None = None) -> urllib.error.HTTPError: + message = email.message.Message() + for key, value in (headers or {}).items(): + message[key] = value + return urllib.error.HTTPError(url, status, f"status {status}", message, io.BytesIO(_body(body))) + + def urlopen(self, request: Any, data: Any = None, timeout: Any = None, **_: Any) -> Any: + if isinstance(request, str): + request = urllib.request.Request(request, data=data) + url = request.full_url + self.requests.append({ + "url": url, + "method": request.get_method(), + "headers": {key.lower(): value for key, value in request.header_items()}, + "data": request.data, + }) + mode = self.mode + if mode == "down": + raise urllib.error.URLError(OSError("Name or service not known")) + if mode == "timeout": + raise TimeoutError("timed out") + if mode == "reset": + raise ConnectionResetError("connection reset by peer") + if mode == "cut": + return _CutShort(b"", url=url) + if mode == "garbage": + return FakeResponse(b"not json", url=url) + if isinstance(mode, int) and not isinstance(mode, bool): + raise self.refusal(url, mode, {"error": "refused"}) + for key, answer in self.routes.items(): + if key in url: + if callable(answer) and not isinstance(answer, type): + answer = answer(request) + if isinstance(answer, BaseException): + raise answer + if isinstance(answer, int) and not isinstance(answer, bool): + raise self.refusal(url, answer, {"error": "refused"}) + return FakeResponse(_body(answer), url=url) + self.unrouted.append(url) + raise self.refusal(url, 404, {"error": "no route in the test"}) + + +def stub_server(handler: Any, tls: bool = False) -> tuple[str, list]: + """A real HTTP server on loopback, for a collector that must be driven + over a socket: (base_url, requests). + + `handler` is a BaseHTTPRequestHandler subclass, or a function taking + {method, path, headers, body} and returning a payload, (status, payload) + or (status, payload, headers). Every request lands in `requests`. + + With `tls`, every https:// request any opener makes is carried to this + server instead, so a collector keeps its own https URLs, redirect refusal + and error handling, and only the transport is the fixture's. + """ + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + requests: list[dict] = [] + + if isinstance(handler, type): + base = handler + + class Logged(base): # type: ignore[valid-type, misc] + def log_message(self, *args: Any) -> None: + pass + + def parse_request(self) -> bool: + parsed = super().parse_request() + if parsed: + requests.append({"method": self.command, "path": self.path, "headers": dict(self.headers), "body": None}) + return parsed + else: + answer = handler + + class Logged(BaseHTTPRequestHandler): # type: ignore[no-redef] + def log_message(self, *args: Any) -> None: + pass + + def _serve(self) -> None: + length = int(self.headers.get("Content-Length") or 0) + body = self.rfile.read(length) if length else b"" + entry = {"method": self.command, "path": self.path, "headers": dict(self.headers), "body": body} + requests.append(entry) + result = answer(entry) + status, payload, headers = 200, result, {} + if isinstance(result, tuple): + status, payload, *rest = result + headers = rest[0] if rest else {} + data = _body(payload) + self.send_response(status) + for key, value in headers.items(): + self.send_header(key, value) + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + do_GET = do_POST = do_PUT = do_DELETE = _serve + + server = ThreadingHTTPServer(("127.0.0.1", 0), Logged) + threading.Thread(target=server.serve_forever, daemon=True).start() + base_url = f"http://127.0.0.1:{server.server_port}" + if tls: + _route_tls_to(base_url) + return base_url, requests + + +def _route_tls_to(base_url: str) -> None: + real_opener = getattr(urllib.request, "_fixture_real_build_opener", urllib.request.build_opener) + urllib.request._fixture_real_build_opener = real_opener + + # The fixture's answer goes back unprocessed, so a redirect or an error + # status is handled by the collector's own opener, not this inner one. + class PassThrough(urllib.request.HTTPErrorProcessor): + def http_response(self, request: Any, response: Any) -> Any: + return response + + class LocalTLS(urllib.request.HTTPSHandler): + def https_open(self, request: Any) -> Any: + split = urllib.parse.urlsplit(request.full_url) + path = split.path + (f"?{split.query}" if split.query else "") + mapped = urllib.request.Request(base_url + path, data=request.data, headers=dict(request.header_items()), method=request.get_method()) + return real_opener(PassThrough()).open(mapped, timeout=10) + + urllib.request.build_opener = lambda *handlers: real_opener(LocalTLS(), *handlers) + # The shared module installs its redirect-refusing opener when it is first + # imported; install it again so it carries the fixture transport. + import collectors + urllib.request.install_opener(urllib.request.build_opener(collectors._RefuseRedirects())) + + +# -------------------------------------------------------------------- time + +_frozen_day = None + + +class _DayClock(datetime): + """datetime whose now() runs as usual until local midnight past the frozen + day, then stops at that day's last microsecond. freeze_today() hands it to + the shared module.""" + + @classmethod + def now(cls, tz: Any = None) -> datetime: + real = datetime.now(tz) + if _frozen_day is None or datetime.now().date() <= _frozen_day: + return real + end = datetime.combine(_frozen_day, datetime.max.time()).astimezone() + return end.astimezone(tz) if tz else end.replace(tzinfo=None) + + +def freeze_today() -> None: + """Hold "today" to the local day it is now, for the rest of the process: + the shared module's clock, and the rows the fixtures stamp "now" + (pi_reply(), opencode rows, today_stamp()), run on as usual until local + midnight and then stop at that day's last microsecond. A test that + straddles midnight then dates every row and scan by the day it began; + time.time(), and so every cache age, runs on. check_scan_cache freezes + today itself; a test comparing today's totals across runs calls this first.""" + global _frozen_day + import collectors + if _frozen_day is None: + _frozen_day = datetime.now().date() + collectors.datetime = _DayClock + + +def fixture_now() -> datetime: + """Now in UTC, held to the frozen day once freeze_today() has run.""" + return _DayClock.now(timezone.utc) + + +def iso_in(seconds: float) -> str: + """An ISO UTC time `seconds` from now (negative for the past).""" + return (datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat() + + +def ms_in(seconds: float) -> int: + """An epoch-millisecond stamp `seconds` from now.""" + return int((time.time() + seconds) * 1000) + + +def today_stamp(hour: int = 12) -> str: + """An ISO time at `hour` today in local time, with its offset, so it lands + on today whatever the machine's zone.""" + return fixture_now().astimezone().replace(hour=hour, minute=0, second=0, microsecond=0).isoformat() + + +def jwt(claims: dict) -> str: + """An unsigned JWT carrying `claims`, as collectors read them.""" + body = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=") + return f"header.{body}.signature" + + +# ------------------------------------------------------- pi/omp and opencode + + +def pi_reply(provider: str, tokens: int = 1, id: str | None = None, model: str = "m", + timestamp: Any = "now", usage: dict | None = None, api: str | None = None) -> dict: + """One pi/omp assistant line. `timestamp` defaults to now; pass None to + leave it out.""" + message = {"role": "assistant", "provider": provider, "model": model, "usage": usage or {"input": tokens}} + if api: + message["api"] = api + line: dict[str, Any] = {"type": "message", "message": message} + if timestamp == "now": + timestamp = fixture_now().isoformat(timespec="milliseconds").replace("+00:00", "Z") + if timestamp is not None: + line["timestamp"] = timestamp + if id: + line["id"] = id + return line + + +def write_pi_session(path: str | Path, lines: list, parent: str | Path | None = None) -> Path: + """A pi/omp session file; `parent` makes it a fork naming that session.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w") as handle: + if parent is not None: + handle.write(json.dumps({"type": "session", "parentSession": str(parent)}) + "\n") + for line in lines: + handle.write((line if isinstance(line, str) else json.dumps(line)) + "\n") + return path + + +def _oc_now() -> int: + """An opencode row's default time, in ms on the fixtures' clock: a second + ago, but never before today's local midnight.""" + now = fixture_now() + midnight = now.astimezone().replace(hour=0, minute=0, second=0, microsecond=0) + return int(max(now - timedelta(seconds=1), midnight).timestamp() * 1000) + + +def _oc_tokens(tokens: Any) -> dict: + if isinstance(tokens, dict): + return tokens + return {"input": tokens, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}} + + +def oc_v1(provider: str, tokens: Any, model: str = "m", created: int | None = None, role: str = "assistant") -> str: + """An opencode `message` row's data: provider and model at the top level.""" + return json.dumps({"role": role, "providerID": provider, "modelID": model, "tokens": _oc_tokens(tokens), + "time": {"created": _oc_now() if created is None else created}}) + + +def oc_v2(provider: str, tokens: Any, model: str = "m", created: int | None = None) -> str: + """An opencode V2 `session_message` row's data: the model nested.""" + return json.dumps({"model": {"providerID": provider, "id": model}, "tokens": _oc_tokens(tokens), + "time": {"created": _oc_now() if created is None else created}}) + + +OC_TABLES = { + "message": "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)", + "session_message": "CREATE TABLE session_message (id text PRIMARY KEY, session_id text NOT NULL, type text NOT NULL, time_created integer NOT NULL, data text NOT NULL)", + "session_v2": "CREATE TABLE session_v2 (id text PRIMARY KEY, fork_session_id text, time_created integer NOT NULL)", +} + + +def opencode_db(path: str | Path | None = None, message=(), session_message=(), forks=(), + tables=("message", "session_message", "session_v2")) -> Path: + """opencode's database (the default one unless `path`), with the stores in + `tables` and the rows oc_add() takes.""" + path = Path(path) if path else data_home() / "opencode" / "opencode.db" + path.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(path) + for table in tables: + conn.execute(OC_TABLES[table]) + conn.commit() + conn.close() + oc_add(path, message=message, session_message=session_message, forks=forks) + return path + + +def oc_add(path: str | Path, message=(), session_message=(), forks=()) -> None: + """Rows into an opencode database: `message` rows (id, session, data[, + created_ms]), `session_message` rows (id, session, type, data[, + created_ms]), and `forks` (session, source session, created_ms).""" + conn = sqlite3.connect(path) + now = _oc_now() + for row in message: + message_id, session, data, *created = row + stamp = created[0] if created else now + conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", (message_id, session, stamp, stamp, data)) + for row in session_message: + message_id, session, kind, data, *created = row + conn.execute("INSERT INTO session_message VALUES (?, ?, ?, ?, ?)", (message_id, session, kind, created[0] if created else now, data)) + for session, source, created in forks: + conn.execute("INSERT INTO session_v2 VALUES (?, ?, ?)", (session, source, created)) + conn.commit() + conn.close() + + +# ------------------------------------------------------------ Claude Code + + +def claude_code_session(name: str, replies: list, base_url: str | None = None) -> Path: + """A Claude Code transcript under ~/.claude/projects: a typed prompt and its + reply for each (model, input tokens) in `replies`, as Claude Code 2.1 + writes them. `base_url` sets ANTHROPIC_BASE_URL in Claude Code's settings, + where pointing it at another provider is documented.""" + stamp = fixture_now().isoformat(timespec="milliseconds").replace("+00:00", "Z") + lines = [] + for n, (model, tokens) in enumerate(replies): + lines.append({"type": "user", "uuid": f"{name}-prompt-{n}", "sessionId": name, "timestamp": stamp, "message": {"role": "user", "content": "go"}}) + lines.append({"type": "assistant", "sessionId": name, "timestamp": stamp, "requestId": f"{name}-req-{n}", + "message": {"id": f"{name}-msg-{n}", "role": "assistant", "model": model, "usage": {"input_tokens": tokens, "output_tokens": 0}}}) + path = home() / ".claude" / "projects" / "project" / f"{name}.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(line) + "\n" for line in lines)) + if base_url is not None: + _merge(home() / ".claude" / "settings.json", "env", {"ANTHROPIC_BASE_URL": base_url}) + return path + + +def mcode_routed_session(name, provider, model, tokens, base_url=None): + """A MiniMax Code session with one prompt and its reply on `provider`, and, + for a `custom_provider:`, that provider declared at `base_url` in + config.yaml, as `mcode provider add` writes it.""" + root = home() / ".minimax" + session = root / "v2" / "sessions" / "2026" / "10" / "22" / name / "messages.jsonl" + session.parent.mkdir(parents=True, exist_ok=True) + session.write_text("".join(json.dumps(line) + "\n" for line in ( + {"message": {"role": "user", "timestamp": ms_in(0), "content": "go"}}, + {"message": {"role": "assistant", "timestamp": ms_in(0), "provider": provider, "model": model, "usage": {"input": tokens}}}, + ))) + if base_url is not None: + config = root / "config.yaml" + text = config.read_text() if config.exists() else "custom_provider:\n" + pid = provider.partition(":")[2] + config.write_text(text + f" {pid}:\n name: {pid}\n kind: custom\n options:\n apiKey: sk-placeholder\n baseURL: {base_url}\n") + return session + + +# ------------------------------------------------------ credentials and state + + +def write_json(path: str | Path, data: Any) -> Path: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data)) + return path + + +def _merge(path: Path, key: str, value: Any) -> Path: + try: + data = json.loads(path.read_text()) + except (OSError, ValueError): + data = {} + data[key] = value + return write_json(path, data) + + +def cursor_state_db(path: str | Path | None, token: str) -> Path: + """The Cursor editor's state.vscdb holding `token` as its sign-in (the + default place under XDG_CONFIG_HOME when `path` is None).""" + path = Path(path) if path else config_home() / "Cursor" / "User" / "globalStorage" / "state.vscdb" + path.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(path) + conn.execute("CREATE TABLE IF NOT EXISTS ItemTable (key TEXT UNIQUE ON CONFLICT REPLACE, value BLOB)") + conn.execute("INSERT INTO ItemTable VALUES ('cursorAuth/accessToken', ?)", (token,)) + conn.commit() + conn.close() + return path + + +def pi_auth(directory: str | Path | None, provider: str, key: str) -> Path: + """A key pi (or omp) keeps from /login: /auth.json, ~/.pi/agent by + default, merged with what is there.""" + directory = Path(directory) if directory else home() / ".pi" / "agent" + return _merge(directory / "auth.json", provider, {"type": "api_key", "key": key}) + + +def opencode_auth(path: str | Path | None, provider: str, key: str) -> Path: + """A key opencode keeps: its auth.json under XDG_DATA_HOME by default.""" + path = Path(path) if path else data_home() / "opencode" / "auth.json" + return _merge(path, provider, {"type": "api", "key": key}) + + +def registry(agent_id: str, accounts: list, active: str | None = None, **switch: Any) -> Path: + """The accounts `omarchy agent account` registered: each account a dict or + an id; `switch` adds the registry's own settings (switch="auto", + threshold=90).""" + entries = [{"id": account} if isinstance(account, str) else account for account in accounts] + data: dict[str, Any] = {"accounts": entries, **switch} + if active is not None: + data["active"] = active + return write_json(state_home() / "omarchy" / "agents" / "accounts" / f"{agent_id}.json", data) + + +# ------------------------------------------------------------------ caches + + +def limits_cache_files(agent_id: str) -> list[Path]: + """Every limits cache a collector keeps: -limits*.json.""" + return sorted(Path(path) for path in glob.glob(str(cache_root() / f"{glob.escape(agent_id)}-limits*.json"))) + + +def scan_cache_files(agent_id: str) -> list[Path]: + return sorted(Path(path) for path in glob.glob(str(cache_root() / f"{glob.escape(agent_id)}-scan-*.json"))) + + +def edit_limits_cache(agent_id: str, fetched_ago: float | None = None, resets_past: bool = False) -> list[Path]: + """Rewrite every limits cache of `agent_id` in place: stamped `fetched_ago` + seconds back, and with `resets_past` every window's reset an hour gone, as + if the window elapsed while nothing ran. No sleeping.""" + files = limits_cache_files(agent_id) + if not files: + check(f"{agent_id} keeps a limits cache to edit", False, f"nothing under {cache_root()}") + for path in files: + data = json.loads(path.read_text()) + if fetched_ago is not None: + data["fetchedAtMs"] = ms_in(-fetched_ago) + if resets_past: + for entry in data.get("limits") or []: + if isinstance(entry, dict): + entry["resetsAt"] = iso_in(-3600) + path.write_text(json.dumps(data)) + return files + + +def age_scan_cache(target: str | Path, seconds: float) -> None: + """Move a scan cache's stamp `seconds` into the past (negative: the future, + a clock that went back). `target` is the file, or an agent id for all of + its scans. Set `scanDate` with edit_scan_cache instead for another day.""" + edit_scan_cache(target, lambda data: data.update(scannedAtMs=int(data.get("scannedAtMs") or 0) - int(seconds * 1000))) + + +def edit_scan_cache(target: str | Path, change: Callable[[dict], Any]) -> None: + files = [Path(target)] if Path(str(target)).is_file() else scan_cache_files(str(target)) + if not files: + check(f"{target} has a scan cache to edit", False, f"nothing under {cache_root()}") + for path in files: + data = json.loads(path.read_text()) + change(data) + path.write_text(json.dumps(data)) + + +# -------------------------------------------------------- contract checks + + +def _default_measure(record: dict) -> tuple: + return (record.get("totalPrompts"), record.get("totalSessions"), record.get("todayTotalTokens")) + + +@dataclass +class Adapter: + """What the contract checks need to know about one collector. + + - `id`: the collector, bin/omarchy-agent-usage-. + - `plant()`: sign the sandbox in (a credential on disk or in os.environ), + so a healthy service gives limits. + - `answer`: FakeHttp routes for the healthy service. Its windows' resets + must lie ahead, so a fresh answer reads as it was given. + - `requests_per_probe`: requests one probe makes (credits and limits are 2). + - `refused_status`: usageStatusText for a refused credential (401); None + when the collector has no such answer. + - `serve(mode)`: instead of FakeHttp, a context manager making the service + behave as `mode` for the runs inside it: "up", 500 (a failed check), + "cut" and "garbage" (an answer cut short, or not JSON), "down" and + "timeout" (unreachable), 401 (refused). Pair it with + `count()`, the requests (or CLI starts) made so far, or None when there + is nothing to count. + - `retry_modes`: the unreachable modes the service can take; () when it has + no network at all. + - `skip`: contract checks that can't apply, each with its reason beside the + adapter: "reuse", "force", "stale", "cut", "retry", "elapsed", "refused". + - `grow()`: one more unit of local usage on disk, for check_scan_cache; + `measure(record)` what it changes (prompts, sessions, today's tokens). + - `pi_root()`: for a collector reading pi-format sessions, the root it + reads them from (~/.pi/agent/sessions, or a harness's own), so + check_scan_cache can hold it to never caching a scan rg couldn't run. + """ + + id: str + plant: Callable[[], Any] = lambda: None + answer: Any = None + requests_per_probe: int = 1 + refused_status: str | None = "Sign-in expired" + serve: Callable[[Any], Any] | None = None + count: Callable[[], int] | None = None + retry_modes: tuple = ("down", "timeout") + skip: tuple = () + grow: Callable[[], Any] | None = None + measure: Callable[[dict], Any] = _default_measure + pi_root: Callable[[], Path] | None = None + + +class _Service: + """The adapter's service, FakeHttp unless it brings its own.""" + + def __init__(self, adapter: Adapter) -> None: + self.adapter = adapter + self.http = None if adapter.serve else FakeHttp(adapter.answer or {}) + + @contextlib.contextmanager + def mode(self, mode: Any): + if self.adapter.serve: + with self.adapter.serve(mode): + yield + else: + self.http.mode = mode + with self.http: + yield + + def count(self) -> int | None: + if self.adapter.serve: + return self.adapter.count() if self.adapter.count else None + return self.http.count + + def run(self, mode: Any, *args: str) -> tuple[dict, int | None]: + """The record a run prints under `mode`, and the requests it made.""" + with self.mode(mode): + before = self.count() + record = collect(self.adapter.id, *args) + after = self.count() + return record, None if before is None else after - before + + +def _windows(record: dict) -> list: + return [(entry.get("percent"), entry.get("resetsAt")) for entry in record.get("limits") or []] + + +def check_limits_contract(adapter: Adapter) -> dict: + """Hold a collector's limits to the contract, through main() and the record + it prints: + + - fresh: --force asks the service once per probe, and prints live limits + stamped with when they were measured, no status and no help. + - reuse: a normal run and --limits-only right after ask nothing and print + the same limits, live. + - --force asks again. + - stale: a failed check (500) keeps the last limits, stale, with their + limitsFetchedAt, no status, and no retryAdvised. + - cut: a reply cut short and a body that isn't JSON are answers, not + transport: they keep the last limits, stale, and advise no retry. + - retry: an unreachable service (down, timeout) keeps them too, and sets + retryAdvised. + - elapsed: kept windows whose reset passed (a cache rewrite, no sleeping) + come back at 0 with no reset, reused or kept past a failed check. + - refused: a refused credential (401) says `refused_status`, never presents + the limits as live, and advises no retry. + + Returns the fresh record, for the collector's own reading of its answer. + """ + name, skip = adapter.id, set(adapter.skip) + for path in limits_cache_files(name): + path.unlink() + adapter.plant() + service = _Service(adapter) + per_probe = adapter.requests_per_probe + + fresh, asked = service.run("up", "--force") + check( + f"{name}: fresh limits are live, stamped, healthy, and asked for once per probe", + fresh.get("ready") is True and fresh.get("limits") and fresh.get("limitsStale") is False and fresh.get("limitsFetchedAt", 0) > 0 + and fresh.get("usageStatusText") == "" and fresh.get("authHelpText") == "" and not fresh.get("retryAdvised") + and asked in (None, per_probe), + f"requests: {asked} (expected {per_probe})\nrecord: {json.dumps(fresh)}", + ) + fetched_at = fresh["limitsFetchedAt"] + + if "reuse" not in skip: + for args in ((), ("--limits-only",)): + reused, asked = service.run(500, *args) + check( + f"{name}: a {'--limits-only' if args else 'normal'} run inside the probe interval reuses the answer, live", + asked in (None, 0) and reused.get("limits") == fresh["limits"] and reused.get("limitsFetchedAt") == fetched_at + and reused.get("limitsStale") is False and reused.get("usageStatusText") == "" and reused.get("authHelpText") == "", + f"requests: {asked}\nfresh: {json.dumps(fresh)}\nreused: {json.dumps(reused)}", + ) + + if "force" not in skip: + forced, asked = service.run("up", "--force") + check(f"{name}: --force asks again inside the probe interval", asked in (None, per_probe) and forced.get("limitsStale") is False, + f"requests: {asked} (expected {per_probe})\nrecord: {json.dumps(forced)}") + fetched_at = forced["limitsFetchedAt"] + fresh = forced + + if "stale" not in skip: + kept, _ = service.run(500, "--force") + check( + f"{name}: a failed check keeps the last limits, stale, with when they were measured", + kept.get("limits") == fresh["limits"] and kept.get("limitsStale") is True and kept.get("limitsFetchedAt") == fetched_at + and kept.get("usageStatusText") == "" and not kept.get("retryAdvised"), + f"fresh: {json.dumps(fresh)}\nkept: {json.dumps(kept)}", + ) + + if "cut" not in skip: + for mode in ("cut", "garbage"): + answered, _ = service.run(mode, "--force") + check( + f"{name}: an answer {'cut short' if mode == 'cut' else 'that is not JSON'} keeps the last limits, stale, and advises no retry", + not answered.get("retryAdvised") and answered.get("limits") == fresh["limits"] + and answered.get("limitsStale") is True and answered.get("limitsFetchedAt") == fetched_at + and answered.get("usageStatusText") == "", + f"fresh: {json.dumps(fresh)}\nanswered: {json.dumps(answered)}", + ) + + if "retry" not in skip: + for mode in adapter.retry_modes: + unreachable, _ = service.run(mode, "--force") + check( + f"{name}: an unreachable service ({mode}) advises a retry and keeps the last limits, stale", + unreachable.get("retryAdvised") is True and unreachable.get("limits") == fresh["limits"] + and unreachable.get("limitsStale") is True and unreachable.get("limitsFetchedAt") == fetched_at + and unreachable.get("usageStatusText") == "", + json.dumps(unreachable), + ) + + if "elapsed" not in skip: + edit_limits_cache(name, resets_past=True) + zeroed = [(0.0, "")] * len(fresh["limits"]) + runs = [service.run(500, "--force")[0]] + if "reuse" not in skip: + runs.append(service.run(500)[0]) + check( + f"{name}: kept windows whose reset passed read 0 with no reset, past a failed check{' and reused' if len(runs) > 1 else ''}", + all(_windows(record) == zeroed for record in runs) and runs[0].get("limitsStale") is True, + "\n".join(json.dumps(record) for record in runs), + ) + service.run("up", "--force") + + if adapter.refused_status is not None and "refused" not in skip: + refused, _ = service.run(401, "--force") + check( + f"{name}: a refused credential says {adapter.refused_status!r}, advises no retry, and shows nothing as live", + refused.get("usageStatusText") == adapter.refused_status and not refused.get("retryAdvised") + and (refused.get("limitsStale") is True or not refused.get("limits")), + json.dumps(refused), + ) + + return fresh + + +def check_scan_cache(adapter: Adapter) -> None: + """Hold a collector's local scan cache to the contract, by growing its + usage on disk between runs and rewriting the cache's stamps (no sleeping): + + - a normal run reuses a scan for 20 seconds, --limits-only for 900, and + --force never; + - a scan from another day, or stamped in the future, is a miss; + - with `pi_root`, a scan that couldn't read the pi sessions (no rg) serves + its run but is never cached. + """ + name = adapter.id + if adapter.grow is None: + check(f"{name}: the scan adapter says how to add usage (grow)", False) + # Every run below must date its scan by one day, or a run straddling local + # midnight misses a cache it should reuse. + freeze_today() + for path in scan_cache_files(name): + path.unlink() + adapter.plant() + adapter.grow() + service = _Service(adapter) + + def seen(*args: str) -> Any: + return adapter.measure(service.run("up", *args)[0]) + + base = seen("--force") + adapter.grow() + check(f"{name}: a normal run reuses a scan this recent", seen() == base) + check(f"{name}: a --limits-only run reuses a scan this recent", seen("--limits-only") == base) + forced = seen("--force") + check(f"{name}: --force scans again", forced != base, f"{base} -> {forced}") + + adapter.grow() + age_scan_cache(name, 60) + limits_only = seen("--limits-only") + normal = seen() + check(f"{name}: a minute-old scan serves --limits-only but not a normal run", limits_only == forced and normal != forced, + f"forced {forced}, --limits-only {limits_only}, normal {normal}") + + adapter.grow() + edit_scan_cache(name, lambda data: data.update(scanDate="2000-01-01")) + other_day = seen("--limits-only") + check(f"{name}: a scan from another day is made again", other_day != normal, f"{normal} -> {other_day}") + + adapter.grow() + age_scan_cache(name, -3600) + future = seen("--limits-only") + check(f"{name}: a scan stamped in the future is made again", future != other_day, f"{other_day} -> {future}") + + if adapter.pi_root is not None: + for path in scan_cache_files(name): + path.unlink() + write_pi_session(Path(adapter.pi_root()) / "rg-check" / "session.jsonl", [pi_reply("rg-check", 1, id="rg-check")]) + real_path = os.environ.get("PATH", "") + os.environ["PATH"] = str(home() / "no-rg") + try: + with contextlib.redirect_stderr(io.StringIO()): + service.run("up", "--force") + finally: + os.environ["PATH"] = real_path + check(f"{name}: a scan that couldn't read the pi sessions (no rg) is never cached", not scan_cache_files(name), + ", ".join(path.name for path in scan_cache_files(name))) diff --git a/test/shell.d/agents-panel-test.sh b/test/shell.d/agents-panel-test.sh index d82cc494e79..a23a95c701b 100755 --- a/test/shell.d/agents-panel-test.sh +++ b/test/shell.d/agents-panel-test.sh @@ -8,6 +8,20 @@ run_node_test <<'JS' const fs = require('fs') const panelSource = fs.readFileSync(root + '/shell/plugins/agents/Panel.qml', 'utf8') +const signInFunctions = ['supportsSignIn', 'needsSignIn', 'otherTrouble'].map(name => { + const match = panelSource.match(new RegExp('function ' + name + '\\([^]*?\\n \\}')) + return match ? match[0] : '' +}).join('\n') +const signIn = new Function('addProviders', signInFunctions + '\nreturn {needsSignIn, otherTrouble}')( + [{providerId: 'claude'}, {providerId: 'codex'}, {providerId: 'grok'}] +) +for (const id of ['claude', 'codex', 'grok']) + assert(signIn.needsSignIn({providerId: id, usageStatusText: 'Sign-in expired'}), id + ' keeps its sign-in link') +assert(signIn.needsSignIn({usageStatusText: 'Sign-in expired'}), 'account rows keep their sign-in link') +for (const status of ['Sign-in expired', 'Waiting for auth']) { + assert(!signIn.needsSignIn({providerId: 'agy', usageStatusText: status}), 'an agent the panel cannot sign in gets no link for ' + status) + assert(signIn.otherTrouble({providerId: 'agy', usageStatusText: status}) === status, 'an agent the panel cannot sign in still shows ' + status) +} assert(/function launchAgent\(\)/.test(panelSource), 'agents panel launches the default agent') assert(/root\.bar\.run\("omarchy-agent --pick"\)/.test(panelSource), 'agents panel uses the desktop agent launcher') assert(/if \(buttonCode === Qt\.RightButton\) root\.launchAgent\(\)/.test(panelSource), 'agents right click launches the agent') @@ -50,3 +64,66 @@ const catcherSource = fs.readFileSync(root + '/shell/Ui/PanelKeyCatcher.qml', 'u assert(/if \(reorderable && \(event\.modifiers & Qt\.ControlModifier\)\)/.test(catcherSource), 'only panels that ask for it turn Ctrl+Up/Down into a reorder') assert(/if \(\["account", "autoswitch", "signin"\]\.indexOf\(target\.kind\) < 0\) return -1/.test(panelSource), 'Ctrl+Up/Down does nothing outside an agent') JS + +# The panel reads a window out of a label, and that guess cannot survive a +# model name — "Opus 5 (1M context)" parses as a one-minute window. A collector +# that states the title outright is taken at its word. +run_node_test <<'JS' +const fs = require('fs') +const source = fs.readFileSync(root + '/shell/plugins/agents/Panel.qml', 'utf8') +const panelRoot = { nowMs: Date.now() } +const start = source.indexOf('function windowIsLong') +const end = source.indexOf('// The window that decides') +assert(start > 0 && end > start, 'agents panel exposes its limit-window helpers') +eval(source.slice(start, end).replaceAll('root.nowMs', 'panelRoot.nowMs')) +eval(source.slice(source.indexOf('function needsSignIn'), source.indexOf('// Sign an account'))) + +const paused = { usageStatusText: 'Limits paused', limits: [] } +assertEqual(pausedWithoutLimits(paused), true, 'paused usage without a cache shows a recovery placeholder') +assertEqual(otherTrouble(paused), '', 'paused usage never puts the confusing status in the header') +assertEqual(pausedWithoutLimits({ ...paused, limits: [{ label: 'Session (5-hour)', percent: 0, resetsAt: '' }] }), false, 'cached usage replaces the recovery placeholder') +assertEqual(pausedWithoutLimits({ usageStatusText: 'Sign-in expired', limits: [] }), false, 'expired sign-ins retain their sign-in action') +assertEqual(pausedWithoutLimits(null), false, 'missing records do not show a paused-usage placeholder') + +assertDeepEqual( + limitWindows({ limits: [ + { label: 'Session (5-hour)', percent: 0.78, resetsAt: '' }, + { label: 'Opus 5 (1M context) Weekly', title: 'Opus 5 (1M context) Weekly', percent: 0.42, resetsAt: '' } + ] }), + [ + { title: 'Session', percent: 0.78, resetAt: '' }, + { title: 'Opus 5 (1M context) Weekly', percent: 0.42, resetAt: '' } + ], + 'agents panel titles a limit off the collector when it states one' +) + +assertDeepEqual( + limitWindows({ limits: [{ label: 'Weekly (7-day)', percent: 0.12, resetsAt: '' }] }), + [{ title: 'Weekly', percent: 0.12, resetAt: '' }], + 'agents panel still reads a window out of a label that carries no title' +) + +const resetAt = new Date(panelRoot.nowMs + 1000).toISOString() +const measured = { limits: [{ label: 'Session (5-hour)', percent: 0.78, resetsAt: resetAt }] } +assertEqual(limitWindows(measured)[0].percent, 0.78, 'usage remains until the reset time') +panelRoot.nowMs += 1000 +assertDeepEqual(limitWindows(measured), [{ title: 'Session', percent: 0, resetAt: '' }], 'usage resets at the deadline while the panel stays open') +assertEqual(measured.limits[0].percent, 0.78, 'display reset leaves the measured record intact') + +// A model's weekly allowance rides under the Weekly row; one whose window has +// no row of its own keeps a row. +assertDeepEqual( + displayWindows({ limits: [ + { label: 'Session (5-hour)', percent: 0.94, resetsAt: '' }, + { label: 'Weekly (7-day)', percent: 0.25, resetsAt: 'w' }, + { label: 'Fable Weekly', title: 'Fable Weekly', percent: 0.09, resetsAt: 'w' }, + { label: 'Opus Monthly', title: 'Opus Monthly', percent: 0.5, resetsAt: '' } + ] }), + [ + { title: 'Session', percent: 0.94, resetAt: '', scoped: [] }, + { title: 'Weekly', percent: 0.25, resetAt: 'w', scoped: [{ title: 'Fable', percent: 0.09, resetAt: 'w' }] }, + { title: 'Opus Monthly', percent: 0.5, resetAt: '', scoped: [] } + ], + 'agents panel shows a model-scoped window under the window it runs on' +) +JS diff --git a/test/shell.d/base-test.sh b/test/shell.d/base-test.sh index 4522d9e011a..c772e943e9c 100644 --- a/test/shell.d/base-test.sh +++ b/test/shell.d/base-test.sh @@ -10,6 +10,11 @@ SHELL_TEST_DIR="$ROOT/test/shell.d" export ROOT +# Commands under test belong to this checkout, never to the one the running +# session points OMARCHY_PATH at; the agent usage collectors load their +# shared module from it. +export OMARCHY_PATH=$ROOT + # Commands print user-facing text in the language picked under Setup > # Language, so pin tests to English (a missing file means English) rather than # whatever the machine running them is set to. Tests of other languages point diff --git a/test/shell.d/default-agent-test.sh b/test/shell.d/default-agent-test.sh index 6bfeb171144..ab07fcfeedc 100644 --- a/test/shell.d/default-agent-test.sh +++ b/test/shell.d/default-agent-test.sh @@ -20,6 +20,7 @@ stub_log="$test_tmp/stubs" terminal_log="$test_tmp/terminal" menu_log="$test_tmp/menu" muse_login_log="$test_tmp/muse-login" +clipboard_log="$test_tmp/clipboard" mkdir -p "$mock_bin" "$test_home" cat >"$mock_bin/omarchy-install-chromium-claude" <<'SH' @@ -36,6 +37,11 @@ cat >"$mock_bin/omarchy-notification-send" <<'SH' printf '%s\0' "$@" >>"$OMARCHY_TEST_NOTIFICATION_HISTORY" SH +cat >"$mock_bin/wl-copy" <<'SH' +#!/bin/bash +cat >"$OMARCHY_TEST_CLIPBOARD" +SH + cat >"$mock_bin/omarchy-cmd-missing" <<'SH' #!/bin/bash [[ $1 == ${OMARCHY_TEST_MISSING_COMMAND:-} ]] @@ -65,8 +71,21 @@ cat >"$mock_bin/mise" <<'SH' #!/bin/bash printf '%s\0' "$@" >"$OMARCHY_TEST_MISE_LOG" printf '%s\n' "$*" >>"$OMARCHY_TEST_MISE_HISTORY" +if [[ -n ${MISE_PRERELEASES:-} ]]; then + printf 'prereleases=%s %s\n' "$MISE_PRERELEASES" "$*" >>"$OMARCHY_TEST_MISE_HISTORY" +fi -if [[ $1 == "where" ]]; then +if [[ $1 == "tool" && $3 == "--backend" ]]; then + case "$2" in + npm:*) printf '%s\n' "$2" ;; + amp) printf '%s\n' "npm:@ampcode/cli" ;; + qwen) printf '%s\n' "npm:@qwen-code/qwen-code" ;; + *) printf '%s\n' "aqua:example/$2" ;; + esac + exit 0 +fi + +if [[ $1 == "where" || $1 == "which" ]]; then [[ ${OMARCHY_TEST_AGENT_INSTALLED:-false} == "true" ]] exit fi @@ -128,6 +147,7 @@ export OMARCHY_TEST_STUB_LOG="$stub_log" export OMARCHY_TEST_AGENT_TERMINAL_LOG="$terminal_log" export OMARCHY_TEST_AGENT_MENU_LOG="$menu_log" export OMARCHY_TEST_MUSE_LOGIN_LOG="$muse_login_log" +export OMARCHY_TEST_CLIPBOARD="$clipboard_log" export OMARCHY_PATH="$ROOT" export OMARCHY_MISE_CONFIG_PATH="$test_tmp/etc/mise/config.toml" @@ -140,6 +160,20 @@ agy_package="antigravity-cli" ori_package="ori" legacy_ori_package="github:OpenRouterLabs/ori-releases" kimi_package="npm:@moonshot-ai/kimi-code" +qwen_package="qwen" +dsh_package="npm:@deepseek-ai/dsh" +# Omarchy declares a tool with backend options under its bare name; the +# bracketed package is what a user without the declarations gets. +mcode_tool="npm:@minimax-ai/code" +mcode_package="$mcode_tool[allow_builds=better-sqlite3]" +kiro_cli_tool="http:kiro-cli" +kiro_cli_package="$kiro_cli_tool[platforms_linux_x64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-x86_64-linux.tar.gz,platforms_linux_arm64_url=https://prod.download.cli.kiro.dev/stable/{{version}}/kirocli-aarch64-linux.tar.gz,version_list_url=https://prod.download.cli.kiro.dev/stable/latest/manifest.json,version_json_path=.version]" +freebuff_package="npm:freebuff" +kilo_package="npm:@kilocode/cli" +devin_tool="http:devin" +devin_package="$devin_tool[platforms_linux_x64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-x86_64-unknown-linux.tar.gz,platforms_linux_arm64_url=https://static.devin.ai/cli/{{ version }}/devin-{{ version }}-aarch64-unknown-linux.tar.gz,bin_path=bin,version_list_url=https://static.devin.ai/cli/current/manifest.json,version_json_path=.version]" +amp_package="amp" +commandcode_package="npm:command-code" cursor_agent_package="cursor-agent" muse_package="muse" legacy_muse_package="http:muse[url=https://api.meta.ai/muse-launcher.sh,bin=muse,version_list_url=https://api.meta.ai/muse-code/channels/muse-stable,version_json_path=.version]" @@ -183,12 +217,46 @@ for tool in codex claude copilot opencode pi "\"$kimi_package\""; do grep -Eq "^$tool = \\{ version = \"latest\", lazy = true, auto_update = \"4h\", " "$lazy_config" || fail "user setup declares $tool as a self-updating native lazy tool" done -[[ $(grep -c 'auto_update = "4h"' "$lazy_config") == 12 ]] || fail "only the agent CLIs update themselves" +grep -Eq "^$qwen_package = \\{ version = \"latest\", lazy = true, auto_update = \"4h\", minimum_release_age = \"0s\" \\}$" "$lazy_config" || + fail "user setup declares Qwen Code as a self-updating native lazy tool" +# A tool without registry bin metadata names the commands mise shims for it. +declare -A declared_bins=( + [$amp_package]='"amp"' + ["\"$dsh_package\""]='"dsh"' + ["\"$freebuff_package\""]='"freebuff"' + ["\"$kilo_package\""]='"kilo"' + ["\"$commandcode_package\""]='"commandcode"' +) +for tool in "${!declared_bins[@]}"; do + grep -Fx "$tool = { version = \"latest\", lazy = true, auto_update = \"4h\", lazy_bins = [${declared_bins[$tool]}], minimum_release_age = \"0s\" }" "$lazy_config" >/dev/null || + fail "user setup declares $tool as a self-updating native lazy tool with its commands" +done +# mise drops bracket options from a table-valued key, so a tool with backend +# options is declared by its bare name with the options inside the table: the +# same options the launcher hands a user without the declarations. +declare -A declared_option_bins=( + [mcode]='"mcode"' + [kiro_cli]='"kiro-cli"' + [devin]='"devin"' +) +for agent in "${!declared_option_bins[@]}"; do + tool_var="${agent}_tool" + package_var="${agent}_package" + declaration=$(grep -F "\"${!tool_var}\" = { version = \"latest\", lazy = true, auto_update = \"4h\", lazy_bins = [${declared_option_bins[$agent]}], minimum_release_age = \"0s\", " "$lazy_config") || + fail "user setup declares ${!tool_var} as a self-updating native lazy tool with its commands" + options=${!package_var#*[} + IFS=, read -ra options <<<"${options%]}" + for option in "${options[@]}"; do + [[ $declaration == *" ${option%%=*} = \"${option#*=}\""* ]] || + fail "user setup declares ${!tool_var} with ${option%%=*} inside its table" + done +done +[[ $(grep -c 'auto_update = "4h"' "$lazy_config") == 21 ]] || fail "only the agent CLIs update themselves" for tool in node uv; do grep -Eq "^$tool = \\{ version = \"latest\", lazy = true, minimum_release_age = \"0s\" \\}$" "$lazy_config" || fail "user setup declares $tool as a native lazy tool" done -[[ $(grep -c 'lazy = true' "$lazy_config") == 22 ]] || fail "user setup declares every default mise tool as lazy" +[[ $(grep -c 'lazy = true' "$lazy_config") == 31 ]] || fail "user setup declares every default mise tool as lazy" grep -Fx "reshim --system" "$mise_history" >/dev/null || fail "user setup builds the system lazy bootstrap shims" pass "user setup configures system lazy tools from registry shorthands" : >"$stub_log" @@ -535,6 +603,27 @@ declare -A expected_agents=( [github-copilot]="copilot" [kimi]="kimi" [kimi-code]="kimi" + [qwen]="qwen" + [qwen-code]="qwen" + [dsh]="dsh" + [deepseek]="dsh" + [deepseek-harness]="dsh" + [mcode]="mcode" + [minimax-code]="mcode" + [minimax]="mcode" + [kiro-cli]="kiro-cli" + [kiro]="kiro-cli" + [freebuff]="freebuff" + [kilo]="kilo" + [kilocode]="kilo" + [kilo-code]="kilo" + [devin]="devin" + [devin-cli]="devin" + [amp]="amp" + [ampcode]="amp" + [commandcode]="commandcode" + [command-code]="commandcode" + [cmdc]="commandcode" [cursor]="cursor-agent" [cursor-agent]="cursor-agent" [muse]="muse" @@ -554,6 +643,15 @@ declare -A expected_packages=( [agy]="$agy_package" [copilot]="copilot" [kimi]="$kimi_package" + [qwen]="$qwen_package" + [dsh]="$dsh_package" + [mcode]="$mcode_package" + [kiro-cli]="$kiro_cli_package" + [freebuff]="$freebuff_package" + [kilo]="$kilo_package" + [devin]="$devin_package" + [amp]="$amp_package" + [commandcode]="$commandcode_package" [cursor-agent]="$cursor_agent_package" [muse]="$muse_package" ) @@ -736,6 +834,23 @@ printf 'ignored_config_paths = ["%s"]\n' "$OMARCHY_MISE_CONFIG_PATH" >"$test_hom omarchy-default-agent --install muse >/dev/null grep -Fx "use -g $muse_package" "$mise_history" >/dev/null || fail "Muse selection adds Muse to the user config after Remove Preinstalls" +# Every DSH release is a prerelease, which `mise use` skips unless told otherwise. +: >"$mise_history" +omarchy-default-agent --install dsh >/dev/null +grep -Fx "prereleases=true use -g $dsh_package" "$mise_history" >/dev/null || + fail "DSH selection lets mise pick a prerelease after Remove Preinstalls" +: >"$mise_history" +omarchy-default-agent --install kilo >/dev/null +grep -q '^prereleases=' "$mise_history" && fail "other agents install only stable releases after Remove Preinstalls" +# An npm-backed agent brings node along here too, since node may still be +# declared without being installed; Muse isn't npm-backed and doesn't. +node_line=$(grep -nx "install node" "$mise_history" | cut -d: -f1) +kilo_line=$(grep -nFx "use -g $kilo_package" "$mise_history" | cut -d: -f1) +[[ -n $node_line && -n $kilo_line ]] && (( node_line < kilo_line )) || + fail "an npm-backed agent installs node before it after Remove Preinstalls" +: >"$mise_history" +omarchy-default-agent --install muse >/dev/null +grep -qx "install node" "$mise_history" && fail "an agent that isn't npm-backed installs no node after Remove Preinstalls" rm -f "$agent_file" "$test_home/.config/mise/miserc.toml" # With Omarchy's system config, installing is enough; the user config stays alone. @@ -752,6 +867,28 @@ mapfile -d '' -t agent_open_args <"$agent_open_log" fail "newly installed Muse opens in the installation terminal" pass "Muse installs visibly through mise and opens directly" +saved_agent=$(<"$agent_file") +for agent in qwen dsh mcode kiro-cli freebuff kilo devin amp commandcode; do + tool_var="${agent//-/_}_tool" + package_var="${agent//-/_}_package" + declared_tool=${!tool_var:-${!package_var}} + : >"$mise_history" + omarchy-default-agent --install "$agent" >/dev/null + # npm-backed agents bring node along, which a fresh machine doesn't have yet. + case "$agent" in + kiro-cli | devin) expected_install="install $declared_tool" ;; + *) expected_install="install node $declared_tool" ;; + esac + grep -Fx "$expected_install" "$mise_history" >/dev/null || fail "visible $agent installation installs Omarchy's declaration (expected: $expected_install)" + : >"$mise_history" + OMARCHY_TEST_AGENT_INSTALLED=true omarchy-default-agent "$agent" >/dev/null + grep -Fx "where $declared_tool" "$mise_history" >/dev/null || fail "$agent selection looks for Omarchy's declaration" + grep -q '^use ' "$mise_history" && fail "$agent selection leaves the user config alone when the system config declares it" + [[ $(omarchy-default-agent) == "$agent" ]] || fail "visible $agent installation changes the selection" +done +printf '%s\n' "$saved_agent" >"$agent_file" +pass "the added agents install through Omarchy's mise declarations" + : >"$terminal_log" : >"$muse_login_log" : >"$agent_open_log" @@ -845,6 +982,30 @@ assert_launch hermes env -u HERMES_SESSION_SOURCE hermes chat --yolo --tui "--qu assert_launch agy agy --dangerously-skip-permissions --prompt-interactive "Review this project" assert_launch copilot copilot --allow-all --interactive "Review this project" assert_launch kimi kimi --prompt "Review this project" +assert_launch qwen qwen --yolo "--prompt-interactive=Review this project" +printf '%s\n' qwen >"$agent_file" +omarchy-agent-prompt '-n --help; echo literal' +assert_launched qwen "keeps an option-like prompt literal" qwen --yolo "--prompt-interactive=-n --help; echo literal" +# DSH's browser UI takes no prompt, and its headless profile would run one with +# full access in a window that closes with the answer. +: >"$notification_history" +printf '%s\n' dsh >"$agent_file" +omarchy-agent-prompt '-n --help; echo literal' +assert_launched dsh "opens its browser UI for a prompt" env DSH_PERMISSION_MODE=danger-full-access dsh web +[[ $(<"$clipboard_log") == '-n --help; echo literal' ]] || fail "DSH launcher puts the prompt on the clipboard" +mapfile -d '' -t notification_args <"$notification_history" +[[ ${notification_args[*]} == *"DeepSeek Harness takes no starting prompt"* ]] || fail "DSH launcher says where the prompt went" +assert_launch mcode mcode -- "Review this project" +assert_launch kiro-cli kiro-cli chat --trust-all-tools -- "Review this project" +: >"$notification_history" +assert_launch freebuff freebuff +assert_launch kilo kilo --auto "--prompt=Review this project" +assert_launch devin devin --permission-mode dangerous --respect-workspace-trust false -- "Review this project" +assert_launch amp bash -c 'printf "%s\n" "$1" | exec "$0" --dangerously-allow-all' amp "Review this project" +assert_launch commandcode commandcode --yolo -- "Review this project" +[[ $(<"$clipboard_log") == "Review this project" ]] || fail "Freebuff launcher puts the prompt on the clipboard" +mapfile -d '' -t notification_args <"$notification_history" +[[ ${notification_args[*]} == *"Freebuff takes no starting prompt"* ]] || fail "Freebuff launcher says where the prompt went" pass "agent launcher adapts initial prompts for every supported agent" literal_muse_prompt=$'--disable-sandbox !Crash {$(touch must-not-run)}\ntrailing\\ ' @@ -874,6 +1035,18 @@ assert_bypass hermes hermes --yolo assert_bypass agy agy --dangerously-skip-permissions assert_bypass copilot copilot --allow-all assert_bypass kimi kimi --yolo +assert_bypass qwen qwen --yolo +assert_bypass dsh env DSH_PERMISSION_MODE=danger-full-access dsh web +assert_bypass mcode mcode +assert_bypass kiro-cli kiro-cli chat --trust-all-tools +: >"$clipboard_log" +: >"$notification_history" +assert_bypass freebuff freebuff +assert_bypass kilo kilo --auto +assert_bypass devin devin --permission-mode dangerous --respect-workspace-trust false +assert_bypass amp amp --dangerously-allow-all +assert_bypass commandcode commandcode --yolo +[[ ! -s $clipboard_log && ! -s $notification_history ]] || fail "Freebuff launcher leaves the clipboard alone without a prompt" pass "agent launcher skips permission prompts for every supported agent" printf '%s\n' "opencode" >"$agent_file" @@ -1017,13 +1190,28 @@ fi pass "OpenClaw's agent hands a missing command to its launcher only when the package is there" # Native lazy shims can precede ~/.local/bin. The agent launcher must still run -# the official command when the default-agent selector recognized one. -for agent in cursor-agent muse; do +# the official command when the default-agent selector recognized one, for +# every agent and whether or not it is wrapped in env or bash. +lazy_shims="$test_tmp/lazy-shims" +mkdir -p "$lazy_shims" +for agent in cursor-agent muse claude kiro-cli devin amp dsh hermes; do printf '#!/bin/bash\nprintf "official %s\\n"\n' "$agent" >"$test_home/.local/bin/$agent" - chmod +x "$test_home/.local/bin/$agent" + printf '#!/bin/bash\nprintf "lazy shim %s\\n"\n' "$agent" >"$lazy_shims/$agent" + chmod +x "$test_home/.local/bin/$agent" "$lazy_shims/$agent" printf '%s\n' "$agent" >"$agent_file" - output=$("$ROOT/bin/omarchy-agent" --inline) + output=$(PATH="$lazy_shims:$PATH" "$ROOT/bin/omarchy-agent" --inline) [[ $output == "official $agent" ]] || fail "agent launcher preserves the user-installed $agent ahead of a mise shim" - rm "$test_home/.local/bin/$agent" + output=$(PATH="$lazy_shims:$PATH" omarchy-agent-prompt --inline "Review this project") + [[ $output == "official $agent" ]] || fail "prompted agent launcher preserves the user-installed $agent ahead of a mise shim" + # Once mise has its own copy installed, a terminal runs that one, so the + # launcher does too rather than a leftover install in ~/.local/bin. + output=$(OMARCHY_TEST_AGENT_INSTALLED=true PATH="$lazy_shims:$PATH" "$ROOT/bin/omarchy-agent" --inline) + [[ $output == "lazy shim $agent" ]] || fail "agent launcher runs mise's installed $agent over a leftover ~/.local/bin copy" + # An older wrapper set its variable on the install line; it is still a + # wrapper, not the user's own install. + printf '#!/bin/bash\nMISE_MINIMUM_RELEASE_AGE=0 mise use -g --quiet "%s" || exit 1\nprintf "old wrapper %s\\n"\n' "$agent" "$agent" >"$test_home/.local/bin/$agent" + output=$(PATH="$lazy_shims:$PATH" "$ROOT/bin/omarchy-agent" --inline) + [[ $output == "lazy shim $agent" ]] || fail "agent launcher treats an older one-line mise wrapper for $agent as a wrapper" + rm "$test_home/.local/bin/$agent" "$lazy_shims/$agent" done -pass "agent launcher runs user-installed Cursor and Muse ahead of native lazy shims" +pass "agent launcher runs user-installed agents ahead of native lazy shims" diff --git a/test/shell.d/menu-test.sh b/test/shell.d/menu-test.sh index af3c01c2377..1179a2666a3 100644 --- a/test/shell.d/menu-test.sh +++ b/test/shell.d/menu-test.sh @@ -257,22 +257,30 @@ assertEqual( 'menu lists Reset Computer last under Setup' ) const expectedAgents = { - agy: { icon: '󰫢', label: 'Antigravity' }, + agy: { icon: '\ue91c', iconFont: 'omarchy', label: 'Antigravity' }, pi: { icon: '\ue901', iconFont: 'omarchy', label: 'Pi' }, omp: { icon: '\ue903', iconFont: 'omarchy', label: 'omp' }, opencode: { icon: '\ue902', iconFont: 'omarchy', label: 'OpenCode' }, ori: { icon: '\ue909', iconFont: 'omarchy', label: 'Ori' }, - claude: { icon: '󰛄', label: 'Claude' }, + claude: { icon: '\ue90e', iconFont: 'omarchy', label: 'Claude' }, codex: { icon: '\ue905', iconFont: 'omarchy', label: 'Codex' }, grok: { icon: '\ue904', iconFont: 'omarchy', label: 'Grok' }, - kimi: { icon: '\u{f06a9}', label: 'Kimi' }, + kimi: { icon: '\ue91d', iconFont: 'omarchy', label: 'Kimi' }, + qwen: { icon: '\ue910', iconFont: 'omarchy', label: 'Qwen Code' }, + dsh: { icon: '\ue911', iconFont: 'omarchy', label: 'DeepSeek Harness' }, + mcode: { icon: '\ue912', iconFont: 'omarchy', label: 'MiniMax Code' }, + 'kiro-cli': { icon: '\ue914', iconFont: 'omarchy', label: 'Kiro CLI' }, + freebuff: { icon: '\ue916', iconFont: 'omarchy', label: 'Freebuff' }, + kilo: { icon: '\ue917', iconFont: 'omarchy', label: 'Kilo' }, + devin: { icon: '\ue918', iconFont: 'omarchy', label: 'Devin CLI' }, + amp: { icon: '\ue919', iconFont: 'omarchy', label: 'Amp' }, + commandcode: { icon: '\ue91a', iconFont: 'omarchy', label: 'Command Code' }, hermes: { icon: '\ue90a', iconFont: 'omarchy', label: 'Hermes' }, openclaw: { icon: '\ue90c', iconFont: 'omarchy', label: 'OpenClaw' }, copilot: { icon: '', label: 'Copilot' }, - crush: { icon: '󰋑', label: 'Crush' }, - muse: { icon: '󰛤', label: 'Muse Code' }, + crush: { icon: '\ue91e', iconFont: 'omarchy', label: 'Crush' }, + muse: { icon: '\ue91f', iconFont: 'omarchy', label: 'Muse Code' }, 'cursor-agent': { icon: '\ue90d', iconFont: 'omarchy', label: 'Cursor CLI' }, - } assert( Object.entries(expectedAgents).every(([agent, expected]) => { @@ -291,7 +299,7 @@ assertDeepEqual( defaultItems .filter(item => item.parent === 'setup.default.agent') .map(item => item.label), - ['Antigravity', 'Claude', 'Codex', 'Copilot', 'Crush', 'Cursor CLI', 'Grok', 'Hermes', 'Kimi', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi'], + ['Amp', 'Antigravity', 'Claude', 'Codex', 'Command Code', 'Copilot', 'Crush', 'Cursor CLI', 'DeepSeek Harness', 'Devin CLI', 'Freebuff', 'Grok', 'Hermes', 'Kilo', 'Kimi', 'Kiro CLI', 'MiniMax Code', 'Muse Code', 'omp', 'OpenClaw', 'OpenCode', 'Ori', 'Pi', 'Qwen Code'], 'menu sorts coding agents alphabetically' ) const expectedDefaults = { @@ -943,5 +951,5 @@ assert( JS font_charset=$(fc-query --format='%{charset}' "$ROOT/default/fonts/omarchy/omarchy.ttf") -[[ $font_charset == *"e900-e90f"* ]] || fail "Omarchy icon font includes every custom menu glyph" +[[ $font_charset == *"e900-e91f"* ]] || fail "Omarchy icon font includes every custom menu glyph" pass "Omarchy icon font includes the official agent marks" diff --git a/test/shell.d/mise-agent-shims-migration-test.sh b/test/shell.d/mise-agent-shims-migration-test.sh new file mode 100644 index 00000000000..2a3a3e636f2 --- /dev/null +++ b/test/shell.d/mise-agent-shims-migration-test.sh @@ -0,0 +1,65 @@ +#!/bin/bash + +set -euo pipefail + +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" + +require_command mise + +test_tmp=$(mktemp -d) +trap 'rm -rf "$test_tmp"' EXIT +migration="$ROOT/migrations/1792504295.sh" +agents=(qwen dsh mcode kiro-cli freebuff kilo devin amp commandcode) + +# The settings package ships the declarations before migrations run. +system_dir="$test_tmp/etc/mise" +mkdir -p "$system_dir" +cp -r "$ROOT/etc/mise/conf.d" "$system_dir/" + +run_migration() { + local test_home="$1" tools_config="${2:-$system_dir/conf.d/omarchy-tools.toml}" + + ( + cd "$test_home" + env -i \ + HOME="$test_home" \ + XDG_CACHE_HOME="$test_home/.cache" \ + XDG_CONFIG_HOME="$test_home/.config" \ + XDG_DATA_HOME="$test_home/.local/share" \ + XDG_STATE_HOME="$test_home/.local/state" \ + MISE_SYSTEM_DIR="$system_dir" \ + OMARCHY_MISE_CONFIG_PATH="$tools_config" \ + PATH="$ROOT/bin:/usr/bin" \ + bash -euo pipefail "$migration" + ) +} + +new_home() { + local test_home="$test_tmp/$1" + mkdir -p "$test_home/.config/mise" + echo "$test_home" +} + +home=$(new_home declared) +run_migration "$home" >/dev/null 2>&1 +for agent in "${agents[@]}"; do + [[ -x $home/.local/share/mise/shims/$agent ]] || fail "migration builds the lazy $agent shim" +done +[[ ! -e $home/.local/share/mise/installs ]] || fail "migration installs nothing before first launch" +run_migration "$home" >/dev/null 2>&1 || fail "migration is idempotent" +pass "migration builds the lazy shims for the newly declared agents" + +# Until the settings package declares the agents, the migration stays pending. +waiting_home=$(new_home waiting) +if run_migration "$waiting_home" "$test_tmp/missing/omarchy-tools.toml" >/dev/null 2>&1; then + fail "migration waits for the settings package's declarations" +fi +[[ ! -e $waiting_home/.local/share/mise/shims/qwen ]] || fail "migration builds nothing before the declarations arrive" +pass "migration stays pending until the declarations arrive" + +# Remove Preinstalls opts the user out of Omarchy's declarations. +opted_out_home=$(new_home opted-out) +printf 'ignored_config_paths = ["%s"]\n' "$system_dir/conf.d/omarchy-tools.toml" >"$opted_out_home/.config/mise/miserc.toml" +run_migration "$opted_out_home" >/dev/null 2>&1 || fail "migration completes after Remove Preinstalls" +[[ ! -e $opted_out_home/.local/share/mise/shims/qwen ]] || fail "migration builds no shims after Remove Preinstalls" +pass "migration leaves a user who removed the preinstalls alone" diff --git a/test/shell.d/mise-provision-test.sh b/test/shell.d/mise-provision-test.sh index c2ccd4f1dbc..97c3fbdc7cf 100644 --- a/test/shell.d/mise-provision-test.sh +++ b/test/shell.d/mise-provision-test.sh @@ -17,15 +17,65 @@ SH chmod +x "$mock_bin/mise" source "$ROOT/install/user/mise.sh" -for agent in kimi cursor-agent muse; do +for agent in kimi cursor-agent muse qwen dsh mcode kiro-cli freebuff kilo devin amp commandcode; do [[ ! -e $HOME/.local/bin/$agent ]] || fail "provisioning must not create a handwritten $agent wrapper" done mapfile -d '' -t args <"$OMARCHY_TEST_MISE_LOG" [[ ${args[*]} == "reshim --system reshim" ]] || fail "provisioning builds the lazy shims without installing packages or writing the user's config" -pass "provisioning leaves lazy Kimi, Cursor, and Muse shims to mise" +pass "provisioning leaves the lazy agent shims to mise" # The settings package ships this, so the user's config stays theirs. grep -Fx 'upgrade.auto_prune = false' "$ROOT/etc/mise/conf.d/omarchy.toml" >/dev/null || fail "Omarchy's mise config keeps upgrades from pruning versions still in use" pass "upgrade pruning is off in the package-owned mise config" + +# mise drops "tool[...]" bracket options from a key whose value is a table, so +# such a declaration installs without its options: Kiro and Devin found no +# versions at all. Backend options belong inside the table. +tools_config="$ROOT/etc/mise/conf.d/omarchy-tools.toml" +bracketed=$(awk '/^\[/ { in_tools = ($0 == "[tools]"); next } in_tools && /^"[^"]*\[[^"]*" = \{/' "$tools_config") +[[ -z $bracketed ]] || fail "no table-valued mise declaration carries bracket options" "$bracketed" +pass "mise declarations keep their backend options inside the table" + +# Ask the real mise which backend options it resolves for each declaration that +# has some, in a throwaway home whose only config is that declaration. mise +# reports a tool's options only once it resolves a version, which takes the +# network, so the check is skipped without it. +real_mise=$(PATH=/usr/local/bin:/usr/bin command -v mise || true) +if [[ -z $real_mise ]]; then + skip "mise is not installed; skipping the backend option resolution check" +elif ! curl -sfI --max-time 5 https://registry.npmjs.org/ >/dev/null 2>&1; then + skip "no network; skipping the backend option resolution check" +else + # Prints each declaration's key, the tool it names, and its backend options: + # those in brackets on the key and every table key other than the ones that + # say how mise manages the tool. + declared_options=$(python3 -I -c ' +import sys, tomllib +managed = {"version", "lazy", "lazy_bins", "auto_update", "minimum_release_age"} +for key, spec in tomllib.load(open(sys.argv[1], "rb"))["tools"].items(): + tool, _, bracket = key.partition("[") + options = [option.split("=")[0] for option in bracket.rstrip("]").split(",") if option] + if isinstance(spec, dict): + options += sorted(set(spec) - managed) + if options: + print(key, tool, *options, sep="\t") +' "$tools_config") + [[ -n $declared_options ]] || fail "Omarchy declares tools with backend options" + + while IFS=$'\t' read -r key tool options; do + mise_root="$test_tmp/real-mise/${tool//[^a-z]/-}" + mkdir -p "$mise_root/etc/mise/conf.d" "$mise_root/home" + { echo "[tools]"; grep -F "\"$key\" = " "$tools_config"; } >"$mise_root/etc/mise/conf.d/tools.toml" + resolved=$(cd "$mise_root/home" && env -i HOME="$mise_root/home" PATH=/usr/local/bin:/usr/bin \ + XDG_CONFIG_HOME="$mise_root/home/.config" XDG_CACHE_HOME="$mise_root/home/.cache" \ + XDG_DATA_HOME="$mise_root/home/.local/share" XDG_STATE_HOME="$mise_root/home/.local/state" \ + MISE_SYSTEM_DIR="$mise_root/etc/mise" timeout 60 "$real_mise" tool --tool-options "$tool" 2>/dev/null || true) + for option in $options; do + [[ $resolved == *"$option="* ]] || + fail "mise resolves $option for $tool from Omarchy's declaration" "mise resolved: ${resolved:-nothing}" + done + done <<<"$declared_options" + pass "mise resolves the backend options of Omarchy's declarations" +fi