diff --git a/bin/omarchy-agent-usage-opencodego b/bin/omarchy-agent-usage-opencodego new file mode 100755 index 00000000000..4cd6c1c8372 --- /dev/null +++ b/bin/omarchy-agent-usage-opencodego @@ -0,0 +1,628 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the OpenCode Go usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect OpenCode Go subscription usage into one display-ready JSON record. + +Limits come from OpenCode's Go usage endpoint (the rolling 5-hour, weekly, +and monthly windows the console dashboard shows). Machine-local stats come +from OpenCode v1/v2 databases and pi/omp sessions on the exact opencode-go +provider, grouped by day and model. The agents panel only reads the JSON +this prints. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import hashlib +import http.client +import json +import math +import os +import sqlite3 +import stat +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + +AGENT_ID = "opencodego" +AGENT_NAME = "OpenCode Go" +AUTH_HELP = "Run /connect in opencode and choose OpenCode Go, or set OPENCODE_GO_API_KEY." +USAGE_URL = "https://opencode.ai/zen/go/v1/usage" +PROVIDER_ID = "opencode-go" +DEFAULT_CACHE_SECONDS = 900 + + +class LocalScanError(Exception): + """A local history could not be read completely; do not publish a record.""" + + +def source_exists(path: Path) -> bool: + try: + path.stat() + return True + except FileNotFoundError: + return False + except OSError as exc: + raise LocalScanError from exc + + +def session_files(root: Path): + if not source_exists(root): + return + + def fail_walk(error): + raise LocalScanError from error + + # pathlib globbing can suppress traversal errors, yielding partial histories. + for directory, _, filenames in os.walk(root, onerror=fail_walk): + for name in filenames: + if name.endswith(".jsonl"): + yield Path(directory) / name + + +def number(value: Any) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int): + return max(0, value) + if isinstance(value, float) and math.isfinite(value): + return max(0, int(value)) + return 0 + + +def empty_bucket() -> dict[str, int]: + return { + "inputTokens": 0, + "outputTokens": 0, + "cacheReadInputTokens": 0, + "cacheCreationInputTokens": 0, + } + + +def cache_root() -> Path: + root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "omarchy" / "agent-usage" + # Caching must not make a readable local history unavailable. + try: + root.mkdir(parents=True, exist_ok=True) + except OSError: + pass + return root + + +def date_string(value: dt.date) -> str: + return value.strftime("%Y-%m-%d") + + +def local_date_string() -> str: + return date_string(dt.date.today()) + + +def recent_date_strings() -> list[str]: + today = dt.date.today() + return [date_string(today - dt.timedelta(days=offset)) for offset in range(6, -1, -1)] + + +def data_home() -> Path: + return Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) + + +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() + + # Session 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 number(round(float(value or 0))) + except Exception: + return 0 + + +def merge_stats(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: + merged = dict(base) + for key in ("todayPrompts", "todaySessions", "todayTotalTokens", "totalPrompts", "totalSessions"): + merged[key] = number(base.get(key)) + number(extra.get(key)) + + combined = dict(base.get("todayTokensByModel") or {}) + for model, count in (extra.get("todayTokensByModel") or {}).items(): + combined[model] = number(combined.get(model)) + number(count) + merged["todayTokensByModel"] = combined + + usage = {model: dict(bucket) for model, bucket in (base.get("modelUsage") or {}).items()} + for model, bucket in (extra.get("modelUsage") or {}).items(): + target = usage.setdefault(model, empty_bucket()) + for field, count in (bucket or {}).items(): + target[field] = number(target.get(field)) + number(count) + merged["modelUsage"] = usage + + by_date: dict[str, int] = {} + for source in (base.get("recentDays") or [], extra.get("recentDays") or []): + for day in source: + date = str((day or {}).get("date") or "") + if date: + by_date[date] = by_date.get(date, 0) + number((day or {}).get("messageCount")) + merged["recentDays"] = [{"date": date, "messageCount": by_date[date]} for date in sorted(by_date)] + + # Sources overlap in time, so union dates rather than summing counts. + dates = set(base.get("activeDates") or []) | set(extra.get("activeDates") or []) + merged["activeDates"] = sorted(dates) + merged["activeDays"] = max(len(dates), number(base.get("activeDays")), number(extra.get("activeDays"))) + return merged + + +def cache_identity() -> dict[str, Any]: + # Offset alone cannot distinguish zones with different daylight-saving rules. + return { + "version": 1, + "localDay": local_date_string(), + "timezone": [os.environ.get("TZ", ""), list(time.tzname), time.timezone, + dt.datetime.now().astimezone().strftime("%z")], + } + + +def valid_stats(stats: Any) -> bool: + if not isinstance(stats, dict): + return False + + def count(value: Any) -> bool: + return type(value) is int and value >= 0 + + def date(value: Any) -> bool: + try: + return isinstance(value, str) and dt.date.fromisoformat(value).isoformat() == value + except ValueError: + return False + + def counts(value: Any) -> bool: + return isinstance(value, dict) and all(isinstance(key, str) and count(n) for key, n in value.items()) + + if not all(count(stats.get(key)) for key in ( + "todayPrompts", "todaySessions", "todayTotalTokens", "totalPrompts", "totalSessions", "activeDays")): + return False + if not counts(stats.get("todayTokensByModel")): + return False + days = stats.get("recentDays") + if not isinstance(days, list) or len(days) != 7: + return False + if not all(isinstance(day, dict) and date(day.get("date")) and count(day.get("messageCount")) for day in days): + return False + if [day["date"] for day in days] != recent_date_strings(): + return False + dates = stats.get("activeDates") + if not isinstance(dates, list) or not all(date(day) for day in dates): + return False + if dates != sorted(set(dates)) or stats["activeDays"] != len(dates): + return False + models = stats.get("modelUsage") + fields = set(empty_bucket()) + return isinstance(models, dict) and all( + isinstance(model, str) and counts(bucket) and set(bucket) == fields + for model, bucket in models.items()) + + +def read_fresh_json(path: Path, max_age_seconds: float) -> dict[str, Any] | None: + try: + age = dt.datetime.now().timestamp() - path.stat().st_mtime + if max_age_seconds <= 0 or not 0 <= age <= max_age_seconds: + return None + cached = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(cached, dict) or any(cached.get(key) != value for key, value in cache_identity().items()): + return None + return cached if valid_stats(cached.get("stats")) else None + except (OSError, ValueError): + return None + + +def write_json_atomic(path: Path, payload: dict[str, Any]) -> None: + tmp = None + try: + # A unique sibling supports overlapping collector invocations; replacement + # is atomic and a failed write leaves the previous successful scan intact. + with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=path.parent, + prefix=f".{path.name}.", suffix=".tmp", delete=False) as handle: + tmp = Path(handle.name) + json.dump({**cache_identity(), **payload}, handle, allow_nan=False) + os.replace(tmp, path) + except (OSError, ValueError): + pass + finally: + if tmp is not None: + try: + tmp.unlink(missing_ok=True) + except OSError: + pass + + +# ------------------------------------------------------------------- limits +# +# OpenCode's Go usage endpoint answers with the same three windows the +# console dashboard draws: a rolling 5-hour session, a weekly, and a monthly +# allowance, each as a percent of the plan used plus a reset timestamp. + + +def api_key() -> str: + env_key = os.environ.get("OPENCODE_GO_API_KEY") + if env_key: + return env_key + auth_file = data_home() / "opencode" / "auth.json" + try: + auth = json.loads(auth_file.read_text(encoding="utf-8")) + except (OSError, ValueError): + return "" + entry = auth.get(PROVIDER_ID) if isinstance(auth, dict) else None + if isinstance(entry, str): + return entry + if isinstance(entry, dict): + key = entry.get("key") + return key if isinstance(key, str) else "" + return "" + + +def probe_limits(key: str) -> dict[str, Any]: + if not key: + return { + "ok": False, + "limits": [], + "usageStatusText": "OpenCode Go unavailable", + "authHelpText": AUTH_HELP, + } + # The endpoint rejects the default Python-urllib user agent, so identify + # like the agent clients the service expects. + request = urllib.request.Request( + USAGE_URL, + headers={ + "Authorization": f"Bearer {key}", + "User-Agent": "omarchy-agent-usage-opencodego/1.0", + }, + ) + try: + with urllib.request.urlopen(request, timeout=10) as response: + payload = json.loads(response.read().decode("utf-8")) + except (urllib.error.URLError, OSError, ValueError, http.client.HTTPException) as exc: + # Exception strings may contain request details or credentials. Only + # expose a numeric HTTP status, never the remote/network error text. + code = getattr(exc, "code", None) + reason = f"HTTP {code}" if type(code) is int else "request failed" + return { + "ok": False, + "limits": [], + "usageStatusText": "OpenCode Go unavailable", + "authHelpText": f"OpenCode Go usage endpoint error: {reason}", + "retryAdvised": code not in (401, 403), + } + + windows = { + "rolling": ("Session (5-hour)", "Rolling 5-hour window"), + "weekly": ("Weekly (7-day)", "Weekly window"), + "monthly": ("Monthly", "Monthly window"), + } + usage = payload.get("usage") if isinstance(payload, dict) else None + usage = usage if isinstance(usage, dict) else {} + limits: list[dict[str, Any]] = [] + for window, (label, title) in windows.items(): + entry = usage.get(window) + if not isinstance(entry, dict): + continue + percent = entry.get("percent") + if isinstance(percent, bool) or not isinstance(percent, (int, float)) or not 0 <= percent <= 100: + continue + reset = entry.get("resetsAt") + if not isinstance(reset, str) or "T" not in reset: + continue + try: + dt.datetime.fromisoformat(reset.replace("Z", "+00:00")) + except ValueError: + continue + limits.append({ + "label": label, + "title": title, + "percent": float(percent) / 100.0, + "resetsAt": reset, + }) + if not limits: + return { + "ok": False, + "limits": [], + "usageStatusText": "OpenCode Go unavailable", + "authHelpText": "OpenCode Go usage endpoint returned no limits.", + "retryAdvised": True, + } + return {"ok": True, "limits": limits, "usageStatusText": "", "authHelpText": ""} + + +# -------------------------------------------------------------------- stats +# +# opencode records per-message provider, model, and token usage in its own +# database. Scan it for opencode-go messages and aggregate them by day and +# model. A big history walked on every panel open would be wasteful, so the +# scan caches its result with the same freshness contract as the claude +# collector's scans. + + +def scan_opencode_stats(max_age_seconds: float) -> dict[str, Any]: + 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 + + db = data_home() / "opencode" / "opencode.db" + cache_file = None + scan_ok = False + if source_exists(db): + cache_file = cache_root() / f"opencodego-{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") or {} + + conn = None + try: + # Read-only: opencode may be writing right now. + conn = sqlite3.connect(db.resolve().as_uri() + "?mode=ro", uri=True, timeout=2) + conn.execute("PRAGMA query_only = ON") + tables = {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = 'table'")} + 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 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 + # v2 moves role to the type column, and provider/model into data.model. + if row_type is None: + role, provider, model_id = entry.get("role"), entry.get("providerID"), entry.get("modelID") + else: + model_info = entry.get("model") or {} + role, provider, model_id = row_type, model_info.get("providerID"), model_info.get("id") + # Exact match: a custom gateway is not this subscription. + if role != "assistant" or provider != PROVIDER_ID: + 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 PROVIDER_ID).rstrip("/").split("/")[-1] + except Exception: + continue + session_key = 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 + scan_ok = bool(queries) + except (sqlite3.Error, OSError) as exc: + raise LocalScanError from exc + finally: + if conn is not None: + conn.close() + + stats = { + "todayPrompts": today_prompt_count, + "todaySessions": len(today_sessions), + "todayTotalTokens": today_token_total, + "todayTokensByModel": today_tokens, + "recentDays": [recent[day] for day in recent_dates], + "totalPrompts": prompts, + "totalSessions": len(sessions), + "activeDays": len(active_days), + "activeDates": sorted(active_days), + "modelUsage": usage_by_model, + } + if cache_file is not None and scan_ok: + write_json_atomic(cache_file, {"stats": stats}) + return stats + + +# ------------------------------------------------------- pi and omp sessions +# +# pi (and omp) can burn an OpenCode Go subscription without ever touching +# the opencode database. Their compatible JSONL session formats carry the +# provider, model, and token usage on every assistant message. + + +def scan_pi_usage(max_age_seconds: float) -> dict[str, Any] | None: + agent_dir = Path(os.environ.get("PI_CODING_AGENT_DIR") or (Path.home() / ".pi" / "agent")) + roots = [ + agent_dir / "sessions", + Path.home() / ".omp" / "agent" / "sessions", + ] + profiles = Path.home() / ".omp" / "profiles" + if source_exists(profiles): + try: + roots.extend(profile / "agent" / "sessions" for profile in profiles.iterdir() + if stat.S_ISDIR(profile.stat().st_mode)) + except OSError as exc: + raise LocalScanError from exc + root_identity = sorted(str(root.absolute()) for root in roots) + cache_file = cache_root() / "opencodego-pi-sessions.json" + cached = read_fresh_json(cache_file, max_age_seconds) + if cached is not None and cached.get("sessionRoots") == root_identity: + 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: + for path in session_files(root): + 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 != PROVIDER_ID: + 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 PROVIDER_ID) + 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 as exc: + raise LocalScanError from exc + + 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), + } + if stats is not None: + write_json_atomic(cache_file, {"stats": stats, "sessionRoots": root_identity}) + return stats + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--force", action="store_true", help="rescan the database and re-probe limits, ignoring caches") + parser.add_argument("--limits-only", action="store_true", help="reuse any recent database scan; only the limits probe must be fresh") + args = parser.parse_args() + + scan_age = 0 if args.force else DEFAULT_CACHE_SECONDS + try: + stats = scan_opencode_stats(scan_age) + pi_usage = scan_pi_usage(scan_age) + except LocalScanError: + print("OpenCode Go local usage scan failed.", file=sys.stderr) + return 1 + if pi_usage is not None: + stats = merge_stats(stats, pi_usage) + limits = probe_limits(api_key()) + + 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, + # The endpoint reports allowances, not the Go / Go Plus plan name. + "tierLabel": "", + "usageStatusText": limits["usageStatusText"], + "authHelpText": limits["authHelpText"], + "limits": limits["limits"], + } + if limits.get("retryAdvised"): + record["retryAdvised"] = True + record.update(stats) + print(json.dumps(record, separators=(",", ":"), sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index d05eb8aec1f..96b8a92adfc 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -72,6 +72,7 @@ Adding an agent needs no QML changes: ship a collector that prints the record co | `grok` | The credits endpoint behind Grok's `/usage` view (the billing period's included usage) | Each session's `usage.json` (the ledger `grok usage` prints: tokens by model per finished turn), plus `summary.json` for sessions | | `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 | +| `opencodego` | OpenCode Go usage endpoint (rolling 5-hour, weekly, monthly) | OpenCode v1/v2 databases and pi/omp sessions, including omp profiles, on the exact `opencode-go` provider | When `~/.local/state/omarchy/agents/accounts/.json` registers more than one account, the `claude`, `codex`, and `grok` records @@ -108,6 +109,10 @@ Expired tokens are refreshed under the current CLI's renewable directory lock at 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. +### OpenCode Go + +The collector reads `OPENCODE_GO_API_KEY` first, then the `opencode-go` key in `$XDG_DATA_HOME/opencode/auth.json` (default `~/.local/share/opencode/auth.json`, written by `/connect`). Limits are account-global; token stats come only from this machine's OpenCode database and pi/omp session files. The endpoint does not identify Go versus Go Plus, so no plan label is inferred. Missing credentials or an unavailable endpoint leave local stats usable; missing local sources produce an empty record. Local scan caches are best-effort and expire after 15 minutes or a local day/timezone change. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/assets/opencodego.svg b/shell/plugins/agents/assets/opencodego.svg new file mode 100644 index 00000000000..33a257ad482 --- /dev/null +++ b/shell/plugins/agents/assets/opencodego.svg @@ -0,0 +1,8 @@ + + + + + + diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 21dc959dcf2..932e12147c8 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": "Claude Code, Codex, Grok, Fireworks, Kimi, and OpenCode Go usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -23,7 +23,8 @@ "codex": { "enabled": true }, "grok": { "enabled": true }, "fireworks": { "enabled": true }, - "kimi": { "enabled": true } + "kimi": { "enabled": true }, + "opencodego": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-opencodego-limits-test.sh b/test/shell.d/agent-usage-opencodego-limits-test.sh new file mode 100644 index 00000000000..8af584855cd --- /dev/null +++ b/test/shell.d/agent-usage-opencodego-limits-test.sh @@ -0,0 +1,155 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +require_command python3 + +PYTHONDONTWRITEBYTECODE=1 python3 - "$ROOT/bin/omarchy-agent-usage-opencodego" <<'PY' +import contextlib +import http.client +import importlib.machinery +import importlib.util +import io +import json +import os +import sys +import tempfile +import unittest +import urllib.error +from pathlib import Path +from unittest.mock import patch + +loader = importlib.machinery.SourceFileLoader("opencodego", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +collector = importlib.util.module_from_spec(spec) +loader.exec_module(collector) +RESET = "2026-10-13T12:00:00Z" +KEY = "synthetic-secret-key" + + +class LimitsTests(unittest.TestCase): + def setUp(self): + self.scratch = tempfile.TemporaryDirectory() + self.addCleanup(self.scratch.cleanup) + self.home = Path(self.scratch.name) + env = patch.dict(os.environ, { + "HOME": str(self.home), "XDG_DATA_HOME": str(self.home / "data"), + "XDG_CACHE_HOME": str(self.home / "cache"), "XDG_CONFIG_HOME": str(self.home / "config"), + "OPENCODE_GO_API_KEY": "", "PI_CODING_AGENT_DIR": "", + }) + env.start() + self.addCleanup(env.stop) + network = patch.object(collector.urllib.request, "urlopen", side_effect=AssertionError("unmocked network")) + network.start() + self.addCleanup(network.stop) + + def probe(self, payload): + def respond(request, timeout): + self.assertEqual(request.full_url, collector.USAGE_URL) + self.assertEqual(request.get_header("Authorization"), f"Bearer {KEY}") + self.assertEqual(request.get_header("User-agent"), "omarchy-agent-usage-opencodego/1.0") + self.assertEqual(timeout, 10) + return contextlib.closing(io.BytesIO(json.dumps(payload).encode())) + with patch.object(collector.urllib.request, "urlopen", side_effect=respond): + result = collector.probe_limits(KEY) + self.assertNotIn(KEY, json.dumps(result)) + return result + + def test_success_boundaries_partial_windows(self): + result = self.probe({"usage": { + "rolling": {"percent": 0, "resetsAt": RESET}, + "weekly": {"percent": 100, "resetsAt": "2026-10-14T12:30:00+02:00"}, + "monthly": {"percent": 45.5, "resetsAt": RESET}, + }}) + self.assertTrue(result["ok"]) + self.assertEqual([limit["percent"] for limit in result["limits"]], [0, 1, .455]) + partial = self.probe({"usage": {"rolling": {"percent": 20, "resetsAt": RESET}, "weekly": []}}) + self.assertEqual(len(partial["limits"]), 1) + + def test_invalid_percentages(self): + for value in (True, False, "50", None, -1, 101, float("nan"), float("inf"), -float("inf")): + with self.subTest(percent=value): + result = self.probe({"usage": {"rolling": {"percent": value, "resetsAt": RESET}}}) + self.assertFalse(result["ok"]) + self.assertEqual(result["limits"], []) + + def test_invalid_reset_timestamps(self): + for value in (None, "", " ", "not-a-date", "2026-99-99T12:00:00Z", 123, [], True, "2026-10-13"): + with self.subTest(reset=value): + result = self.probe({"usage": {"rolling": {"percent": 20, "resetsAt": value}}}) + self.assertFalse(result["ok"]) + + def test_malformed_payload_shapes(self): + for payload in (None, [], 42, "text", {}, {"usage": []}, {"usage": {"rolling": None}}): + with self.subTest(payload=payload): + result = self.probe(payload) + self.assertFalse(result["ok"]) + self.assertEqual(result["limits"], []) + + def test_http_auth_network_and_json_errors_do_not_leak_keys(self): + for error in ( + urllib.error.HTTPError(collector.USAGE_URL, 401, KEY, {}, None), + urllib.error.HTTPError(collector.USAGE_URL, 403, KEY, {}, None), + urllib.error.HTTPError(collector.USAGE_URL, 500, KEY, {}, None), + urllib.error.URLError(f"failed request {KEY}"), OSError(KEY), + ): + if isinstance(error, urllib.error.HTTPError): + self.addCleanup(error.close) + with self.subTest(error=type(error).__name__), \ + patch.object(collector.urllib.request, "urlopen", side_effect=error): + result = collector.probe_limits(KEY) + self.assertFalse(result["ok"]) + self.assertEqual(bool(result.get("retryAdvised")), getattr(error, "code", None) not in (401, 403)) + self.assertNotIn(KEY, json.dumps(result)) + with patch.object(collector.urllib.request, "urlopen", return_value=contextlib.closing(io.BytesIO(b"not json"))): + self.assertFalse(collector.probe_limits(KEY)["ok"]) + + def test_incomplete_response_retains_local_stats_without_leaking(self): + class BrokenResponse(io.BytesIO): + def read(self): + raise http.client.IncompleteRead(b"synthetic-secret-response " + KEY.encode(), 100) + output, errors = io.StringIO(), io.StringIO() + stats = {"totalPrompts": 7, "todayTotalTokens": 123} + with patch.object(collector.urllib.request, "urlopen", return_value=BrokenResponse()), \ + patch.object(collector, "scan_opencode_stats", return_value=stats), \ + patch.object(collector, "scan_pi_usage", return_value=None), \ + patch.dict(os.environ, {"OPENCODE_GO_API_KEY": KEY}), \ + patch.object(sys, "argv", ["collector", "--force"]), \ + contextlib.redirect_stdout(output), contextlib.redirect_stderr(errors): + self.assertEqual(collector.main(), 0) + record = json.loads(output.getvalue()) + self.assertEqual(record["totalPrompts"], 7) + self.assertEqual(record["todayTotalTokens"], 123) + self.assertEqual(record["limits"], []) + self.assertTrue(record["retryAdvised"]) + self.assertEqual(record["authHelpText"], "OpenCode Go usage endpoint error: request failed") + self.assertNotIn(KEY, output.getvalue() + errors.getvalue()) + self.assertNotIn("synthetic-secret-response", output.getvalue() + errors.getvalue()) + self.assertEqual(errors.getvalue(), "") + + def test_key_sources_missing_auth_and_neutral_tier(self): + self.assertEqual(collector.api_key(), "") + self.assertEqual(collector.probe_limits("")["limits"], []) + auth = self.home / "data/opencode/auth.json" + auth.parent.mkdir(parents=True) + for value in ([], None, {"opencode-go": {"key": 123}}): + auth.write_text(json.dumps(value)) + self.assertEqual(collector.api_key(), "") + auth.write_text(json.dumps({"opencode-go": {"key": KEY}})) + self.assertEqual(collector.api_key(), KEY) + with patch.dict(os.environ, {"OPENCODE_GO_API_KEY": "environment-test-key"}): + self.assertEqual(collector.api_key(), "environment-test-key") + output = io.StringIO() + with patch.object(sys, "argv", ["collector", "--limits-only"]), \ + patch.object(collector, "probe_limits", return_value={"ok": True, "limits": [], "usageStatusText": "", "authHelpText": ""}), \ + contextlib.redirect_stdout(output): + collector.main() + record = json.loads(output.getvalue()) + self.assertEqual(record["tierLabel"], "") # Endpoint does not identify Go vs Go Plus. + self.assertNotEqual(record.get("scope"), "account") # Token stats are machine-local. + self.assertNotIn(KEY, output.getvalue()) + + +unittest.main(argv=[sys.argv[0]], verbosity=2) +PY +pass "OpenCode Go validates mocked account limits without disclosing credentials" diff --git a/test/shell.d/agent-usage-opencodego-scanner-test.sh b/test/shell.d/agent-usage-opencodego-scanner-test.sh new file mode 100644 index 00000000000..b18a882d323 --- /dev/null +++ b/test/shell.d/agent-usage-opencodego-scanner-test.sh @@ -0,0 +1,370 @@ +#!/bin/bash + +set -euo pipefail +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh" +require_command python3 + +PYTHONDONTWRITEBYTECODE=1 python3 - "$ROOT/bin/omarchy-agent-usage-opencodego" <<'PY' +import contextlib +import datetime as dt +import hashlib +import importlib.machinery +import importlib.util +import io +import json +import os +import sqlite3 +import sys +import tempfile +import time +import unittest +from pathlib import Path +from unittest.mock import patch + +loader = importlib.machinery.SourceFileLoader("opencodego", sys.argv[1]) +spec = importlib.util.spec_from_loader(loader.name, loader) +scanner = importlib.util.module_from_spec(spec) +loader.exec_module(scanner) + + +class ScannerTests(unittest.TestCase): + def setUp(self): + self.scratch = tempfile.TemporaryDirectory() + self.addCleanup(self.scratch.cleanup) + self.home = Path(self.scratch.name) + self.env = patch.dict(os.environ, { + "HOME": str(self.home), "XDG_DATA_HOME": str(self.home / "data"), + "XDG_CACHE_HOME": str(self.home / "cache"), "XDG_CONFIG_HOME": str(self.home / "config"), + "OPENCODE_GO_API_KEY": "", "PI_CODING_AGENT_DIR": "", "TZ": "UTC", + }) + self.env.start() + self.addCleanup(self.env.stop) + time.tzset() + self.addCleanup(time.tzset) + # Any accidental network use fails, including in the full main() path. + self.network = patch.object(scanner.urllib.request, "urlopen", side_effect=AssertionError("unexpected network")) + self.network.start() + self.addCleanup(self.network.stop) + self.db = self.home / "data/opencode/opencode.db" + self.now = int(time.time() * 1000) + + def database(self, versions=(1,), providers=True): + self.db.parent.mkdir(parents=True, exist_ok=True) + with contextlib.closing(sqlite3.connect(self.db)) as conn, conn: + for version in versions: + table = "message" if version == 1 else "session_message" + columns = "session_id text, data text" + (", type text" if version == 2 else "") + conn.execute(f"CREATE TABLE IF NOT EXISTS {table} ({columns})") + def insert(provider, role="assistant", raw=None): + entry = { + "tokens": {"input": 10, "output": 4, "reasoning": 3, "cache": {"read": 2, "write": 1}}, + "time": {"created": self.now}, + } + if version == 1: + entry.update(role=role, providerID=provider, modelID="models/shared") + else: + entry["model"] = {"providerID": provider, "id": "models/shared"} + values = ("same-session", json.dumps(entry) if raw is None else raw) + if version == 2: + values += (role,) + conn.execute(f"INSERT INTO {table} VALUES ({','.join('?' for _ in values)})", values) + insert("opencode-go") + if providers: + for provider in ("opencode-go-proxy", "opencode", "opencode-go-plus", "openai"): + insert(provider) + insert("opencode-go", role="user") + for raw in ('not json', '[]', 'null', '{"role":"assistant","providerID":"opencode-go","tokens":[]}'): + insert("opencode-go", raw=raw) + + def db_cache(self): + digest = hashlib.sha1(str(self.db).encode()).hexdigest()[:16] + return self.home / f"cache/omarchy/agent-usage/opencodego-{digest}.json" + + def assert_stats(self, stats, prompts=1, tokens=20, sessions=1): + self.assertEqual(stats["totalPrompts"], prompts) + self.assertEqual(stats["todayPrompts"], prompts) + self.assertEqual(stats["todayTotalTokens"], tokens) + self.assertEqual(stats["todaySessions"], sessions) + self.assertEqual(stats["totalSessions"], sessions) + self.assertEqual(len(stats["recentDays"]), 7) + self.assertEqual(sum(day["messageCount"] for day in stats["recentDays"]), tokens) + + def test_v1_exact_provider_malformed_rows_and_read_only(self): + self.database() + before = self.db.read_bytes() + stats = scanner.scan_opencode_stats(0) + self.assert_stats(stats) + self.assertEqual(stats["modelUsage"]["shared"], { + "inputTokens": 10, "outputTokens": 7, + "cacheReadInputTokens": 2, "cacheCreationInputTokens": 1, + }) + self.assertEqual(self.db.read_bytes(), before) + + def test_v2_only(self): + self.database((2,)) + self.assert_stats(scanner.scan_opencode_stats(0)) + + def test_mixed_database(self): + self.database((1, 2)) + self.assert_stats(scanner.scan_opencode_stats(0), prompts=2, tokens=40, sessions=2) + + def test_missing_database_and_no_auth_record(self): + output = io.StringIO() + with patch.object(sys, "argv", ["collector", "--force"]), contextlib.redirect_stdout(output): + self.assertEqual(scanner.main(), 0) + record = json.loads(output.getvalue()) + self.assertFalse(record["ready"]) + self.assertEqual(record["limits"], []) + self.assert_stats(record, prompts=0, tokens=0, sessions=0) + self.assertFalse(self.db.exists()) + + def test_failed_scan_preserves_cache_and_recovers(self): + self.database() + scanner.scan_opencode_stats(0) + cache = self.db_cache() + good = cache.read_bytes() + # A readable file that is temporarily not a database must not replace + # the good scan with a successful-looking empty cache. + self.db.write_bytes(b"not sqlite") + with self.assertRaises(scanner.LocalScanError): + scanner.scan_opencode_stats(0) + self.assertEqual(cache.read_bytes(), good) + self.db.write_bytes(b"") + self.database((2,)) + self.assert_stats(scanner.scan_opencode_stats(0)) + + def test_corrupt_cache_envelopes_and_stats(self): + self.database() + good = scanner.scan_opencode_stats(0) + invalid = [None, [], 42, "scalar", {}, {"stats": []}, {"stats": {}}] + for field, value in ( + ("todayPrompts", True), ("todayTotalTokens", -1), ("totalSessions", "1"), + ("todayTokensByModel", []), ("todayTokensByModel", {"shared": float("nan")}), + ("recentDays", []), ("recentDays", [{"date": "bad", "messageCount": 20}] * 7), + ("activeDates", [42]), ("activeDays", None), ("modelUsage", {"shared": {}}), + ("modelUsage", {"shared": dict(scanner.empty_bucket(), inputTokens=True)}), + ): + envelope = json.loads(self.db_cache().read_text()) + envelope["stats"] = dict(good, **{field: value}) + invalid.append(envelope) + for value in invalid: + with self.subTest(value=value): + self.db_cache().write_text(json.dumps(value)) + with patch.object(scanner.sqlite3, "connect", wraps=sqlite3.connect) as connect: + self.assertEqual(scanner.scan_opencode_stats(900), good) + connect.assert_called_once() + self.db_cache().write_text('{"stats":') + self.assert_stats(scanner.scan_opencode_stats(900)) + + def test_initial_database_error_does_not_cache_success(self): + self.db.parent.mkdir(parents=True) + self.db.write_bytes(b"not sqlite") + with self.assertRaises(scanner.LocalScanError): + scanner.scan_opencode_stats(900) + self.assertFalse(self.db_cache().exists()) + self.db.write_bytes(b"") + self.database((2,)) + self.assert_stats(scanner.scan_opencode_stats(900)) + + def test_valid_cache_hit_and_force_rescan(self): + self.database() + self.assert_stats(scanner.scan_opencode_stats(0)) + self.database(providers=False) + self.assert_stats(scanner.scan_opencode_stats(900)) + self.assert_stats(scanner.scan_opencode_stats(0), prompts=2, tokens=40) + + def test_future_cache_is_rejected(self): + self.database() + scanner.scan_opencode_stats(0) + cached = json.loads(self.db_cache().read_text()) + cached["stats"]["todayTotalTokens"] = 999 + self.db_cache().write_text(json.dumps(cached)) + future = time.time() + 3600 + os.utime(self.db_cache(), (future, future)) + self.assert_stats(scanner.scan_opencode_stats(900)) + + def test_midnight_and_timezone_invalidate_cache(self): + self.database() + scanner.scan_opencode_stats(0) + tomorrow = dt.date.today() + dt.timedelta(days=1) + days = [(tomorrow - dt.timedelta(days=n)).isoformat() for n in range(6, -1, -1)] + with patch.object(scanner, "local_date_string", return_value=tomorrow.isoformat()), \ + patch.object(scanner, "recent_date_strings", return_value=days): + stats = scanner.scan_opencode_stats(900) + self.assertEqual(stats["todayTotalTokens"], 0) + self.assertEqual(stats["recentDays"][-1]["date"], tomorrow.isoformat()) + scanner.scan_opencode_stats(0) + cached = json.loads(self.db_cache().read_text()) + cached["stats"]["todayTotalTokens"] = 999 + self.db_cache().write_text(json.dumps(cached)) + with patch.dict(os.environ, {"TZ": "Etc/GMT+12"}): + time.tzset() + self.assertNotEqual(scanner.scan_opencode_stats(900)["todayTotalTokens"], 999) + time.tzset() + + def test_cache_unavailable_is_best_effort(self): + self.database() + blocker = self.home / "not-a-directory" + blocker.write_text("block") + with patch.dict(os.environ, {"XDG_CACHE_HOME": str(blocker)}): + self.assert_stats(scanner.scan_opencode_stats(0)) + self.assertIsNone(scanner.scan_pi_usage(0)) + # Permission mocks work even under root and cover both mkdir and writes. + with patch.object(Path, "mkdir", side_effect=PermissionError("read-only cache")): + self.assert_stats(scanner.scan_opencode_stats(0)) + with patch.object(scanner.tempfile, "NamedTemporaryFile", side_effect=PermissionError("unwritable cache")): + self.assert_stats(scanner.scan_opencode_stats(0)) + + def test_atomic_cache_failure_preserves_last_success(self): + self.database() + scanner.scan_opencode_stats(0) + before = self.db_cache().read_bytes() + self.database(providers=False) + with patch.object(scanner.os, "replace", side_effect=PermissionError("read-only destination")): + self.assert_stats(scanner.scan_opencode_stats(0), prompts=2, tokens=40) + self.assertEqual(self.db_cache().read_bytes(), before) + self.assertEqual(list(self.db_cache().parent.glob("*.tmp")), []) + + def session(self, root, name, timestamp, provider="opencode-go", usage=None): + root.mkdir(parents=True, exist_ok=True) + entry = {"type": "message", "id": name, "timestamp": timestamp, "message": { + "role": "assistant", "provider": provider, "model": "shared", + "usage": usage or {"input": 10, "output": 4, "cacheRead": 2, "cacheWrite": 1}, + }} + (root / f"{name}.jsonl").write_text('\n'.join([ + 'not json', '[]', '{"type":"message","message":{"role":"assistant","usage":[]}}', + json.dumps(entry), json.dumps(entry), # Existing same-file ID semantics. + ])) + + def test_unreadable_session_does_not_cache_partial_scan(self): + root = self.home / ".pi/agent/sessions/project" + iso = dt.datetime.now(dt.timezone.utc).isoformat() + self.session(root, "good", iso) + scanner.scan_pi_usage(0) + cache = self.home / "cache/omarchy/agent-usage/opencodego-pi-sessions.json" + before = cache.read_bytes() + self.session(root, "new", iso) + self.session(root, "unreadable", iso) + real_open = Path.open + def open_fixture(path, *args, **kwargs): + if path.name == "unreadable.jsonl": + raise PermissionError("unreadable fixture") + return real_open(path, *args, **kwargs) + with patch.object(Path, "open", open_fixture), self.assertRaises(scanner.LocalScanError): + scanner.scan_pi_usage(0) + self.assertEqual(cache.read_bytes(), before) + self.assertEqual(scanner.scan_pi_usage(0)["totalPrompts"], 3) + + def test_scan_failures_publish_nothing_and_sanitize_stderr(self): + self.database() + iso = dt.datetime.now(dt.timezone.utc).isoformat() + self.session(self.home / ".pi/agent/sessions/project", "good", iso) + secret = "synthetic-secret-path-and-key" + # Cover errors before opening, during discovery, and after partial rows. + real_stat = Path.stat + def stat_fixture(path, *args, **kwargs): + if path == self.db: + raise PermissionError(secret) + return real_stat(path, *args, **kwargs) + class BrokenConnection: + def execute(self, query): + if query == "PRAGMA query_only = ON": + return [] + if "sqlite_master" in query: + return [("message",)] + def rows(): + yield ("session", json.dumps({"role": "assistant", "providerID": "opencode-go", + "tokens": {"input": 10}}), None) + raise sqlite3.OperationalError(secret) + return rows() + def close(self): + pass + self.session(self.home / ".omp/profiles/work/agent/sessions/project", "profile", iso) + real_scandir = scanner.os.scandir + def nested_failure(path): + if str(path).endswith("/project"): + raise PermissionError(secret) + return real_scandir(path) + def session_stat_failure(path, *args, **kwargs): + if path == self.home / ".pi/agent/sessions": + raise PermissionError(secret) + return real_stat(path, *args, **kwargs) + real_open = Path.open + def partial_read(path, *args, **kwargs): + if path.name == "good.jsonl": + def lines(): + yield json.dumps({"type": "message", "message": {"role": "assistant", + "provider": "opencode-go", "usage": {"input": 10}}}) + raise OSError(secret) + return contextlib.nullcontext(lines()) + return real_open(path, *args, **kwargs) + failures = [ + ("database stat", patch.object(Path, "stat", stat_fixture)), + ("partial SQLite cursor", patch.object(scanner.sqlite3, "connect", return_value=BrokenConnection())), + ("session root stat", patch.object(Path, "stat", session_stat_failure)), + ("profile discovery", patch.object(Path, "iterdir", side_effect=PermissionError(secret))), + ("root traversal", patch.object(scanner.os, "scandir", side_effect=PermissionError(secret))), + ("nested traversal", patch.object(scanner.os, "scandir", side_effect=nested_failure)), + ("session open", patch.object(Path, "open", side_effect=PermissionError(secret))), + ("partial session read", patch.object(Path, "open", partial_read)), + ] + for name, failure in failures: + with self.subTest(failure=name), failure: + output, errors = io.StringIO(), io.StringIO() + with patch.object(sys, "argv", ["collector", "--force"]), \ + contextlib.redirect_stdout(output), contextlib.redirect_stderr(errors): + self.assertNotEqual(scanner.main(), 0) + self.assertEqual(output.getvalue(), "") + self.assertEqual(errors.getvalue(), "OpenCode Go local usage scan failed.\n") + self.assert_stats(scanner.scan_opencode_stats(0)) + + def test_custom_pi_roots_replace_default_and_invalidate_only_session_cache(self): + self.database() + scanner.scan_opencode_stats(0) + iso = dt.datetime.now(dt.timezone.utc).isoformat() + self.session(self.home / ".pi/agent/sessions/project", "default", iso) + self.session(self.home / ".omp/agent/sessions/project", "omp", iso) + self.session(self.home / ".omp/profiles/work/agent/sessions/project", "profile", iso) + self.assertEqual(scanner.scan_pi_usage(0)["totalPrompts"], 3) + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("PI_CODING_AGENT_DIR", None) + self.assertEqual(scanner.scan_pi_usage(900)["totalPrompts"], 3) + for name, count in (("custom-a", 1), ("custom-b", 2)): + agent = self.home / name + for n in range(count): + self.session(agent / "sessions/project", str(n), iso) + with patch.dict(os.environ, {"PI_CODING_AGENT_DIR": str(agent)}): + self.assertEqual(scanner.scan_pi_usage(900)["totalPrompts"], count + 2) + with patch.object(scanner.sqlite3, "connect", side_effect=AssertionError("SQLite cache invalidated")): + self.assert_stats(scanner.scan_opencode_stats(900)) + self.assertEqual(scanner.scan_pi_usage(900)["totalPrompts"], 3) + + def test_pi_omp_profiles_and_model_day_merges(self): + self.database() + iso = dt.datetime.now(dt.timezone.utc).isoformat() + yesterday = int((time.time() - 86400) * 1000) + self.session(self.home / ".pi/agent/sessions/project", "pi", iso) + self.session(self.home / ".omp/agent/sessions/project", "omp", self.now) + self.session(self.home / ".omp/profiles/work/agent/sessions/project", "profile", iso, + usage={"totalTokens": 8}) + self.session(self.home / ".pi/agent/sessions/project", "old", yesterday) + self.session(self.home / ".pi/agent/sessions/project", "other", iso, provider="opencode-go-proxy") + pi = scanner.scan_pi_usage(0) + self.assertEqual(pi["totalPrompts"], 4) + self.assertEqual(pi["todayTotalTokens"], 42) + merged = scanner.merge_stats(scanner.scan_opencode_stats(0), pi) + self.assertEqual(merged["totalPrompts"], 5) + self.assertEqual(merged["todayTotalTokens"], 62) + self.assertEqual(merged["todayTokensByModel"], {"shared": 62}) + self.assertEqual(merged["modelUsage"]["shared"]["cacheReadInputTokens"], 8) + self.assertEqual(merged["activeDays"], 2) + self.assertEqual(sum(day["messageCount"] for day in merged["recentDays"]), 79) + self.assertEqual(scanner.scan_pi_usage(900), pi) + cache = self.home / "cache/omarchy/agent-usage/opencodego-pi-sessions.json" + cache.write_text('[]') + self.assertEqual(scanner.scan_pi_usage(900), pi) + + +unittest.main(argv=[sys.argv[0]], verbosity=2) +PY +pass "OpenCode Go scans isolated databases, sessions, and best-effort caches" diff --git a/test/shell.d/agent-usage-update-test.sh b/test/shell.d/agent-usage-update-test.sh index ee3aaf0112e..c36d0755e44 100755 --- a/test/shell.d/agent-usage-update-test.sh +++ b/test/shell.d/agent-usage-update-test.sh @@ -63,3 +63,47 @@ pass "update succeeds when the requested collectors all pass" [[ -e $usage_dir/skipped.json && ! -e $usage_dir/noisy.json ]] || fail "update with agent arguments only runs the named collectors" pass "update with agent arguments only runs the named collectors" + +require_command python3 +PYTHONDONTWRITEBYTECODE=1 python3 - "$ROOT" <<'PY' +import json +import os +import sqlite3 +import subprocess +import sys +import tempfile +from pathlib import Path + +root = Path(sys.argv[1]) +with tempfile.TemporaryDirectory() as scratch: + home = Path(scratch) + fake = home / "omarchy" + (fake / "bin").mkdir(parents=True) + (fake / "bin/omarchy-agent-usage-opencodego").symlink_to(root / "bin/omarchy-agent-usage-opencodego") + env = dict(os.environ, HOME=str(home), OMARCHY_PATH=str(fake), + XDG_DATA_HOME=str(home / "data"), XDG_CACHE_HOME=str(home / "cache"), + XDG_STATE_HOME=str(home / "state"), PI_CODING_AGENT_DIR="", OPENCODE_GO_API_KEY="") + record = home / "state/omarchy/agents/usage/opencodego.json" + record.parent.mkdir(parents=True) + seeded = b'{ "id": "opencodego", "totalPrompts": 91, "limits": [{"percent":0.5}] }\n' + record.write_bytes(seeded) + db = home / "data/opencode/opencode.db" + db.parent.mkdir(parents=True) + db.write_bytes(b"not sqlite") + def update(): + return subprocess.run([str(root / "bin/omarchy-agent-usage-update"), "opencodego", "--force"], + env=env, capture_output=True) + failed = update() + assert failed.returncode != 0, "incomplete collector scan must fail updater" + assert record.read_bytes() == seeded, "failure must preserve complete record byte-for-byte" + db.write_bytes(b"") + with sqlite3.connect(db) as conn: + conn.execute("CREATE TABLE message (session_id text, data text)") + conn.execute("INSERT INTO message VALUES (?, ?)", ("recovered", json.dumps({ + "role": "assistant", "providerID": "opencode-go", "tokens": {"input": 5}}))) + recovered = update() + assert recovered.returncode == 0, recovered.stderr.decode() + assert json.loads(record.read_bytes())["totalPrompts"] == 1, "successful recovery must publish fresh stats" +PY +[[ $? == 0 ]] || fail "update preserves last-good OpenCode Go record and resumes after scan recovery" +pass "update preserves last-good OpenCode Go record and resumes after scan recovery"