diff --git a/bin/omarchy-agent-usage-codex b/bin/omarchy-agent-usage-codex index 972c164dee9..22b8956b050 100755 --- a/bin/omarchy-agent-usage-codex +++ b/bin/omarchy-agent-usage-codex @@ -29,6 +29,11 @@ AGENT_ID = "codex" AGENT_NAME = "Codex" AUTH_HELP = "Run `codex login` to authenticate." +# session_meta.model_provider for the built-in OpenAI backend, which is the +# only one this subscription pays for. Any other id is a custom provider in +# config.toml pointed at some other endpoint. +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 @@ -307,6 +312,19 @@ def scan_native_codex_sessions(): entry = json.loads(raw) except Exception: continue + if entry.get("type") == "session_meta": + # 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. The field is written once, on the first record. + 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) diff --git a/bin/omarchy-agent-usage-local b/bin/omarchy-agent-usage-local new file mode 100755 index 00000000000..44a96685626 --- /dev/null +++ b/bin/omarchy-agent-usage-local @@ -0,0 +1,533 @@ +#!/usr/bin/python3 +# omarchy:summary=Print the self-hosted model usage record as JSON +# omarchy:args=[--force] [--limits-only] +# omarchy:hidden=true +"""Collect self-hosted model usage into one display-ready JSON record. + +Every other collector answers "what did this subscription cost me". This one +answers the question none of them can: how much work ran on hardware you own. +It counts the turns the subscription collectors deliberately skip — Codex CLI +pointed at Ollama, opencode on a local gateway, pi on llama.cpp — whether the +endpoint is this machine or a box down the hall. + +An endpoint is self-hosted when its URL resolves to a loopback, private, or +link-local address, or to a bare/LAN hostname (`server.home`, `gpu-box`). +Provider ids with no URL to resolve fall back to a list of well-known local +runtimes. Neither is exhaustive, so `~/.config/omarchy/agents/local.json` can +name more: + + {"providers": ["homelab", "my-gateway"]} + +There are no limits and no balance to report: the meter on self-hosted work is +the electricity bill. +""" + +import argparse +import fcntl +import hashlib +import ipaddress +import json +import os +import sqlite3 +import sys +import tempfile +import time +import tomllib +from datetime import datetime, timedelta +from pathlib import Path +from urllib.parse import urlparse + +AGENT_ID = "local" +AGENT_NAME = "Local" +TIER_LABEL = "Self-hosted" + +# Runtimes that serve models from your own hardware. Used when a provider id +# carries no URL to classify — opencode and pi record the id but not the +# endpoint, and Codex's built-in `--oss` provider has no config.toml entry. +KNOWN_LOCAL_PROVIDERS = { + "gpt4all", "jan", "koboldcpp", "llama-cpp", "llamacpp", "llamafile", + "lm-studio", "lmstudio", "local", "localai", "mlx", "ollama", + "openai-local", "oss", "sglang", "tabby", "text-generation-webui", "vllm", +} + +# Session files are scanned back this far, matching the Codex collector's +# window. Pi transcripts carry no rollover, so they are read in full. +NATIVE_WINDOW_DAYS = 30 + +# Same reuse windows as the subscription collectors: a no-flag run reuses a +# scan only long enough to dedup the concurrent collectors the update command +# backgrounds, while --limits-only promises nothing about stats freshness and +# may reuse one for far longer. This collector has no limits to fetch, so +# --limits-only is purely a cheap way to redraw. +SCAN_REUSE_SECONDS = 20 +LIMITS_ONLY_REUSE_SECONDS = 900 + +today = datetime.now().strftime("%Y-%m-%d") + + +def local_day(value): + if value is None: + return datetime.now().strftime("%Y-%m-%d") + if isinstance(value, (int, float)): + # pi message timestamps are milliseconds; Codex timestamps are seconds. + if value > 10_000_000_000: + value = value / 1000 + return datetime.fromtimestamp(value).strftime("%Y-%m-%d") + text = str(value) + try: + if text.endswith("Z"): + dt = datetime.fromisoformat(text[:-1] + "+00:00") + else: + dt = datetime.fromisoformat(text) + if dt.tzinfo is not None: + dt = dt.astimezone() + return dt.strftime("%Y-%m-%d") + except Exception: + return datetime.now().strftime("%Y-%m-%d") + + +def number(value): + try: + return max(0, int(value or 0)) + except Exception: + return 0 + + +def model_name(value): + text = str(value or "").strip() + return text or "local" + + +def is_self_hosted(url): + """True when a base URL points at hardware the user controls.""" + try: + host = urlparse(str(url or "")).hostname + except Exception: + return False + if not host: + return False + host = host.lower().rstrip(".") + if host == "localhost": + return True + try: + ip = ipaddress.ip_address(host) + return ip.is_loopback or ip.is_private or ip.is_link_local + except ValueError: + pass + # A bare hostname is a LAN name by construction; the rest are the suffixes + # home routers and mDNS hand out. + if "." not in host: + return True + return host.endswith((".home", ".lan", ".local", ".internal", ".localdomain", ".home.arpa")) + + +def configured_providers(): + path = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "omarchy" / "agents" / "local.json" + try: + data = json.loads(path.read_text()) + except Exception: + return set() + names = data.get("providers") if isinstance(data, dict) else None + if not isinstance(names, list): + return set() + return {str(name).strip().lower() for name in names if str(name).strip()} + + +def codex_endpoints(): + """provider id -> base_url, from the Codex CLI's own config.""" + codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + try: + with (codex_home / "config.toml").open("rb") as handle: + config = tomllib.load(handle) + except Exception: + return {} + providers = config.get("model_providers") + if not isinstance(providers, dict): + return {} + endpoints = {} + for name, entry in providers.items(): + if isinstance(entry, dict): + endpoints[str(name).lower()] = str(entry.get("base_url") or "") + return endpoints + + +class Scan: + def __init__(self): + self.extra_providers = configured_providers() + self.endpoints = codex_endpoints() + self.today = today + self.recent_days = [ + (datetime.now() - timedelta(days=offset)).strftime("%Y-%m-%d") for offset in range(6, -1, -1) + ] + self.recent = {day: 0 for day in self.recent_days} + self.today_prompts = 0 + self.today_sessions = set() + self.today_tokens = 0 + self.today_by_model = {} + self.total_prompts = 0 + self.total_sessions = set() + self.active_days = set() + self.model_usage = {} + self.providers_seen = set() + self.endpoints_seen = set() + self.seen_pi_messages = set() + self.complete = True + + def provider_is_local(self, provider): + pid = str(provider or "").strip().lower() + if not pid: + return False + if pid in self.extra_providers or pid in KNOWN_LOCAL_PROVIDERS: + return True + url = self.endpoints.get(pid) + return bool(url) and is_self_hosted(url) + + def local_provider_ids(self): + """Every id this scan would accept, for pushing filters down into SQL.""" + ids = set(KNOWN_LOCAL_PROVIDERS) | set(self.extra_providers) + for pid, url in self.endpoints.items(): + if url and is_self_hosted(url): + ids.add(pid) + return ids + + def note_provider(self, provider): + pid = str(provider or "").strip().lower() + self.providers_seen.add(pid) + url = self.endpoints.get(pid) + if url: + host = urlparse(url).hostname + if host: + self.endpoints_seen.add(host.lower()) + + def add(self, day, session_key, model, input_tokens, output_tokens, cache_read, cache_write): + total = input_tokens + output_tokens + cache_read + cache_write + if total <= 0: + return + self.total_prompts += 1 + self.total_sessions.add(session_key) + self.active_days.add(day) + if day in self.recent: + self.recent[day] += total + if day == self.today: + self.today_prompts += 1 + self.today_sessions.add(session_key) + self.today_tokens += total + self.today_by_model[model] = self.today_by_model.get(model, 0) + total + bucket = self.model_usage.setdefault( + model, + {"inputTokens": 0, "outputTokens": 0, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}, + ) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + + # ------------------------------------------------------------ native codex + def scan_codex(self): + codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + cutoff = time.time() - NATIVE_WINDOW_DAYS * 24 * 60 * 60 + for root in (codex_home / "sessions", codex_home / "archived_sessions"): + if not root.exists(): + continue + for path in root.rglob("*.jsonl"): + try: + if path.stat().st_mtime < cutoff: + continue + except OSError: + continue + self._scan_codex_file(path) + + def _scan_codex_file(self, path): + current_model = "local" + try: + with path.open(errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except Exception: + continue + if entry.get("type") == "session_meta": + provider = (entry.get("payload") or {}).get("model_provider") + # The subscription collector counts these; this one must not. + if not self.provider_is_local(provider): + return + self.note_provider(provider) + continue + if entry.get("type") == "turn_context": + payload = entry.get("payload") or {} + current_model = model_name(payload.get("model") or payload.get("model_slug") or current_model) + continue + payload = entry.get("payload") or entry + if entry.get("type") == "response_item" and isinstance(payload, dict): + payload = payload.get("payload") or payload + if not isinstance(payload, dict) or payload.get("type") != "token_count": + continue + usage = (payload.get("info") or {}).get("last_token_usage") or {} + cache_read = number(usage.get("cached_input_tokens")) + cache_write = number(usage.get("cache_write_input_tokens")) + input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write) + output_tokens = number(usage.get("output_tokens")) + day = local_day(entry.get("timestamp") or path.stat().st_mtime) + self.add(day, str(path), current_model, input_tokens, output_tokens, cache_read, cache_write) + except Exception: + return + + # ----------------------------------------------------------------- pi / omp + def scan_pi(self): + for root in (Path.home() / ".pi" / "agent" / "sessions", Path.home() / ".omp" / "agent" / "sessions"): + if not root.exists(): + continue + for path in root.rglob("*.jsonl"): + try: + with path.open(errors="replace") as handle: + for raw in handle: + try: + entry = json.loads(raw) + except Exception: + continue + message = entry.get("message") or {} + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + provider = message.get("provider") or message.get("api") + if not self.provider_is_local(provider): + continue + key = str(entry.get("id") or entry.get("uuid") or f"{path}:{raw[:64]}") + if key in self.seen_pi_messages: + continue + self.seen_pi_messages.add(key) + self.note_provider(provider) + usage = message.get("usage") or {} + day = local_day(entry.get("timestamp") or message.get("timestamp")) + self.add( + day, str(path), model_name(message.get("model")), + number(usage.get("input")), number(usage.get("output")), + number(usage.get("cacheRead")), number(usage.get("cacheWrite")), + ) + except Exception: + continue + + # ----------------------------------------------------------------- opencode + def scan_opencode(self): + db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + if not db.exists(): + return + try: + conn = sqlite3.connect(f"file:{db}?mode=ro&immutable=1", uri=True) + except Exception: + return + # OpenCode DBs hold every historical message with its JSON, so the + # provider filter belongs in SQL rather than in a Python loop over every + # row. The candidate ids are known before the query runs: the well-known + # runtimes, whatever the user configured, and any Codex provider whose + # endpoint resolves to hardware they host. json_valid() guards the + # extract so one malformed row cannot abort the scan. + candidates = sorted(self.local_provider_ids()) + if not candidates: + return + placeholders = ",".join("?" * len(candidates)) + try: + for session_id, raw in conn.execute( + "SELECT session_id, data FROM message" + " WHERE data LIKE '%\"role\"%:%\"assistant\"%'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.role') END = 'assistant'" + " AND CASE WHEN json_valid(data) THEN json_extract(data, '$.providerID') END" + f" IN ({placeholders})", + candidates, + ): + try: + entry = json.loads(raw) + except Exception: + continue + # Exact match again in Python: the SQL narrowed the rows, it did not + # vet them. + if not isinstance(entry, dict) or not self.provider_is_local(entry.get("providerID")): + continue + self.note_provider(entry.get("providerID")) + tokens = entry.get("tokens") or {} + cache = tokens.get("cache") or {} + created = (entry.get("time") or {}).get("created") + # opencode keeps thinking tokens out of output; both are generated. + self.add( + local_day(created), "opencode:" + str(session_id), model_name(entry.get("modelID")), + number(tokens.get("input")), number(tokens.get("output")) + number(tokens.get("reasoning")), + number(cache.get("read")), number(cache.get("write")), + ) + except Exception: + # A scan cut short by a database error must not be cached as the whole + # story, or the missing usage stays suppressed until the cache expires. + self.complete = False + return + finally: + try: + conn.close() + except Exception: + pass + + def stats(self): + """Snapshot the aggregated local usage into the record's stats dict.""" + return { + "todayPrompts": self.today_prompts, + "todaySessions": len(self.today_sessions), + "todayTotalTokens": self.today_tokens, + "todayTokensByModel": self.today_by_model, + "recentDays": [{"date": day, "messageCount": self.recent[day]} for day in self.recent_days], + "totalPrompts": self.total_prompts, + "totalSessions": len(self.total_sessions), + "activeDays": len(self.active_days), + "activeDates": sorted(self.active_days), + "modelUsage": self.model_usage, + } + + +def cache_root(): + root = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "omarchy" / "agent-usage" + root.mkdir(parents=True, exist_ok=True) + return root + + +def scan_cache_paths(): + codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex")) + db = Path(os.environ.get("XDG_DATA_HOME") or (Path.home() / ".local" / "share")) / "opencode" / "opencode.db" + config = Path(os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")) / "omarchy" / "agents" / "local.json" + # The digest covers every path the scan reads: the codex session roots and + # its config.toml, the opencode DB, the provider config, and (via + # Path.home()) the pi/omp session roots. + digest = hashlib.sha1( + "\n".join(str(part) for part in (Path.home(), codex_home, db, config)).encode("utf-8") + ).hexdigest()[:16] + root = cache_root() + return root / f"local-scan-{digest}.json", root / f"local-scan-{digest}.lock" + + +def read_fresh_json(path, max_age_seconds): + if max_age_seconds <= 0 or not path.exists(): + return None + try: + # A negative age means the mtime is in the future: the clock moved + # backwards since the write, so the cache's freshness cannot be trusted. + age = time.time() - path.stat().st_mtime + if 0 <= age <= max_age_seconds: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + return None + + +def write_json(path, payload): + # A temp name unique to this writer, not derived from the target: several + # collectors can run at once, and a shared temp path means the second + # replace finds the first one's file already moved away. + handle_fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp") + tmp = Path(tmp_name) + try: + with os.fdopen(handle_fd, "w", encoding="utf-8") as handle: + handle.write(json.dumps(payload, separators=(",", ":")) + "\n") + # mkstemp opens at 0600; nothing in the cache is sensitive, so open it + # up to the usual 0644. + tmp.chmod(0o644) + tmp.replace(path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +# The cache payload is a versioned envelope around the stats dict, so a +# corrupted or foreign-shaped file is a cache miss (rescan + rewrite) instead +# of a crash or a garbage record. +def read_cached_stats(cache_file, max_age_seconds): + cached = read_fresh_json(cache_file, max_age_seconds) + if not isinstance(cached, dict) or cached.get("schemaVersion") != 1: + return None + # today* fields only mean "today" on the day they were scanned. A cache + # from another local date (midnight passed, or the clock moved) is a miss, + # not merely old, whatever its mtime says. + if cached.get("scanDate") != today: + return None + stats = cached.get("stats") + if not isinstance(stats, dict): + return None + if not all(key in stats for key in ("todayPrompts", "todayTotalTokens", "recentDays", "activeDates", "modelUsage")): + return None + return stats + + +def write_cached_stats(cache_file, stats): + try: + write_json(cache_file, {"schemaVersion": 1, "scanDate": today, "stats": stats}) + except Exception as exc: + print(f"omarchy-agent-usage-local: could not write usage cache ({exc})", file=sys.stderr) + + +def run_local_scan(): + scan = Scan() + scan.scan_codex() + scan.scan_pi() + scan.scan_opencode() + return scan.stats(), scan.complete + + +def cached_stats(max_age): + """Stats, with the cache as a pure optimization. + + The cache must never take the collector down: any cache-layer failure + (unwritable cache root, lock errors, disk full) degrades to a direct scan + and a warning on stderr. The JSON record is the contract; the cache is not. + """ + try: + return _cached_stats(max_age) + except Exception as exc: + print(f"omarchy-agent-usage-local: cache unavailable ({exc}); scanning directly", file=sys.stderr) + stats, _ = run_local_scan() + return stats + + +def _cached_stats(max_age): + cache_file, lock_file = scan_cache_paths() + + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + + with lock_file.open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + cached = read_cached_stats(cache_file, max_age) + if cached is not None: + return cached + stats, complete = run_local_scan() + # An interrupted scan still serves this run, but caching it would + # suppress the missing usage for every reader until the cache expires. + if complete: + write_cached_stats(cache_file, stats) + return stats + + +def main(): + parser = argparse.ArgumentParser() + # Accepted for parity with the other collectors: the update command passes + # them to everything it runs. There are no limits to fetch here, so + # --limits-only just means "a redraw may reuse an older scan". + parser.add_argument("--force", action="store_true") + parser.add_argument("--limits-only", action="store_true") + args = parser.parse_args() + + max_age = 0 if args.force else (LIMITS_ONLY_REUSE_SECONDS if args.limits_only else SCAN_REUSE_SECONDS) + + record = { + "schemaVersion": 1, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": datetime.now().astimezone().isoformat(), + "ready": True, + "hasLocalStats": True, + "tierLabel": TIER_LABEL, + "usageStatusText": "", + "authHelpText": "", + # Nothing meters self-hosted work, so the panel draws no limit rows. + "limits": [], + } + record.update(cached_stats(max_age)) + print(json.dumps(record, separators=(",", ":"))) + + +if __name__ == "__main__": + main() diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 51ad4e43c2d..d9c5315edcd 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -53,8 +53,17 @@ light surfaces — and 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 | native Codex CLI session files (plus pi and opencode sessions) | +| `codex` | The Codex app-server RPC | native Codex CLI session files (plus pi and opencode sessions), limited to turns the subscription actually served | | `fireworks` | Estimated prepaid balance: configured funding minus rated account costs | Fireworks billing API, grouped by day and model for the last 30 days | +| `local` | None — self-hosted work is not metered | The turns the subscription collectors skip: Codex, pi, and opencode sessions served by a self-hosted endpoint | + +Every collector counts only the turns its own subscription paid for. Codex +CLI in particular 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 real ones. +They are skipped by `session_meta.model_provider`, so a local model never +shows up as subscription usage. 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 @@ -64,6 +73,33 @@ falls back to local stats only. A non-default Claude directory is honored via opencode stores in `~/.local/share/opencode/auth.json` when Fireworks is signed in there. +### Local + +The `local` collector is the counterpart to the subscription ones. They answer +"what did this cost me"; it answers "how much ran on hardware I own" — Codex +CLI pointed at Ollama, opencode on a local gateway, pi on llama.cpp. Every turn +the other collectors skip as foreign lands here instead, so nothing a coding +agent does goes uncounted. + +An endpoint counts as self-hosted when its URL resolves to a loopback, +private, or link-local address, or to a bare or LAN hostname — `localhost`, +`10.0.0.50`, `gpu-box`, `server.home` all qualify, so a model served by +another machine on your network counts the same as one on this one. Provider +ids with no URL to resolve (opencode and pi record the id but not the +endpoint) fall back to a list of well-known runtimes: Ollama, LM Studio, +llama.cpp, vLLM, LocalAI and friends. Neither list is exhaustive, so a custom +gateway can be named in `~/.config/omarchy/agents/local.json`: + +```json +{"providers": ["homelab", "my-gateway"]} +``` + +A third-party API that is neither this subscription nor your own hardware — +OpenRouter, say — is counted by no collector on purpose. + +There are no limits and no balance: the only meter on self-hosted work is the +electricity bill, so the panel draws the token history and nothing else. + ### Fireworks balance The collector first asks the account's `:getBalance` endpoint for the real diff --git a/shell/plugins/agents/manifest.json b/shell/plugins/agents/manifest.json index 1743a798214..0e7a3957c7a 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, and Fireworks usage, limits, and pace in a native Omarchy bar panel.", + "description": "Claude Code, Codex, Fireworks, and self-hosted model usage, limits, and pace in a native Omarchy bar panel.", "kinds": ["bar-widget"], "activation": "on-demand", "entryPoints": { @@ -21,7 +21,8 @@ "providers": { "claude": { "enabled": true }, "codex": { "enabled": true }, - "fireworks": { "enabled": true } + "fireworks": { "enabled": true }, + "local": { "enabled": true } }, "refreshIntervalSec": 900, "syncMode": "Off", diff --git a/test/shell.d/agent-usage-codex-scanner-test.sh b/test/shell.d/agent-usage-codex-scanner-test.sh index 3ce4883cdbf..4e4def4bb95 100644 --- a/test/shell.d/agent-usage-codex-scanner-test.sh +++ b/test/shell.d/agent-usage-codex-scanner-test.sh @@ -603,3 +603,39 @@ result=$(HOME="$INTERRUPTED_HOME" CODEX_HOME="$INTERRUPTED_HOME/.codex" XDG_CACH [[ $(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" + +# 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) +trap 'rm -rf "$TEST_HOME" "$PI_HOME" "$OPENCODE_HOME" "$CACHE_HOME" "$FRESH_HOME" "$MALFORMED_HOME" "$UNWRITABLE_HOME" "$INTERRUPTED_HOME" "$PROVIDER_HOME"' EXIT +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" <"$LOCAL_HOME/.codex/config.toml" <<'EOF' +[model_providers.homelab] +base_url = "http://server.home:11434/v1" + +[model_providers.openrouter] +base_url = "https://openrouter.ai/api/v1" +EOF + +timestamp="$(date +%Y-%m-%d)T12:00:00Z" +rollout() { + local name=$1 provider=$2 model=$3 + cat >"$LOCAL_HOME/.codex/sessions/$(date +%Y/%m/%d)/rollout-$name.jsonl" <"$LOCAL_HOME/.config/omarchy/agents/local.json" +rollout mystery mystery qwen-custom +# --force: the config changed inside the scan-reuse window, so a no-flag run +# would legitimately serve the previous scan. +local_result=$(env "${env_args[@]}" "$ROOT/bin/omarchy-agent-usage-local" --force) +[[ $(jq -r '.modelUsage["qwen-custom"].inputTokens' <<<"$local_result") == "1000" ]] || + fail "Local collector honors configured provider ids" "$local_result" +pass "Local collector honors configured provider ids" + +# The local and Codex collectors partition the same sessions directory: every +# rollout is counted by exactly one of them, so the fix cannot lose a turn and +# the new tab cannot double-count one. +cat >"$LOCAL_HOME/codex" <<'EOF' +#!/bin/bash +while read -r request; do + id=$(jq -r '.id // empty' <<<"$request") + case "$(jq -r '.method // empty' <<<"$request")" in + initialize) jq -cn --argjson id "$id" '{id: $id, result: {}}' ;; + account/read) jq -cn --argjson id "$id" '{id: $id, result: {account: {}}}' ;; + account/rateLimits/read) jq -cn --argjson id "$id" '{id: $id, result: {rateLimits: {}}}' ;; + esac +done +EOF +chmod +x "$LOCAL_HOME/codex" + +codex_result=$(env "${env_args[@]}" PATH="$LOCAL_HOME:$PATH" "$ROOT/bin/omarchy-agent-usage-codex" --force) + +[[ $(jq -r '.modelUsage | keys | join(",")' <<<"$codex_result") == "gpt-sub" ]] || + fail "Codex collector keeps only the subscription's own turns" "$codex_result" + +overlap=$(jq -n --argjson a "$(jq -c '.modelUsage | keys' <<<"$local_result")" \ + --argjson b "$(jq -c '.modelUsage | keys' <<<"$codex_result")" \ + '$a - ($a - $b) | length') +[[ $overlap == "0" ]] || + fail "Local and Codex collectors must not count the same model twice" "$local_result$codex_result" +pass "Local and Codex collectors partition sessions without gaps or overlap" + +# opencode records the provider id but not the endpoint, so the well-known +# runtimes and configured ids carry the filter - pushed into SQL, since these +# databases keep every historical message with its JSON. +OC_HOME=$(mktemp -d) +trap 'rm -rf "$LOCAL_HOME" "$OC_HOME"' EXIT +mkdir -p "$OC_HOME/.config/omarchy/agents" +echo '{"providers":["homelab"]}' >"$OC_HOME/.config/omarchy/agents/local.json" + +python3 - "$OC_HOME/.local/share/opencode/opencode.db" <<'PY' +import json, sqlite3, sys, time +from pathlib import Path + +db = Path(sys.argv[1]) +db.parent.mkdir(parents=True, exist_ok=True) +conn = sqlite3.connect(db) +conn.execute("CREATE TABLE message (id text PRIMARY KEY, session_id text NOT NULL, time_created integer NOT NULL, time_updated integer NOT NULL, data text NOT NULL)") +now_ms = int(time.time() * 1000) + +def message(mid, provider, model, role="assistant", inp=0, out=0, reasoning=0): + return (mid, "ses_1", now_ms, now_ms, json.dumps({ + "role": role, "providerID": provider, "modelID": model, + "tokens": {"input": inp, "output": out, "reasoning": reasoning, "cache": {"read": 0, "write": 0}}, + "time": {"created": now_ms}, + })) + +conn.executemany("INSERT INTO message VALUES (?, ?, ?, ?, ?)", [ + message("m1", "ollama", "qwen-oc", inp=100, out=10, reasoning=5), + message("m2", "homelab", "qwen-cfg", inp=50, out=5), + message("m3", "anthropic", "claude-paid", inp=999, out=999), + message("m4", "openai", "gpt-sub", inp=999, out=999), + message("m5", "ollama", "qwen-oc", role="user", inp=999), +]) +conn.execute("INSERT INTO message VALUES ('m6', 'ses_1', ?, ?, 'not json at all')", (now_ms, now_ms)) +conn.commit(); conn.close() +PY + +oc_env=(HOME="$OC_HOME" CODEX_HOME="$OC_HOME/.codex" XDG_CACHE_HOME="$OC_HOME/.cache" + XDG_DATA_HOME="$OC_HOME/.local/share" XDG_CONFIG_HOME="$OC_HOME/.config") +oc=$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local") + +# 115 = 100 + 10 + 5 reasoning; 55 = 50 + 5. Paid providers and the user +# message contribute nothing, and the malformed row does not abort the scan. +[[ $(jq -r '.todayTotalTokens' <<<"$oc") == "170" ]] || + fail "Local collector counts opencode turns on self-hosted providers" "$oc" +[[ $(jq -r '.modelUsage | keys | sort | join(",")' <<<"$oc") == "qwen-cfg,qwen-oc" ]] || + fail "Local collector filters opencode rows to self-hosted providers" "$oc" +[[ $(jq -r '.modelUsage["qwen-oc"].outputTokens' <<<"$oc") == "15" ]] || + fail "Local collector counts opencode reasoning tokens as generated" "$oc" +pass "Local collector counts self-hosted opencode turns past paid and malformed rows" + +# The scan is cached like every other collector's: one envelope per data-path +# set, readable, versioned, and bypassed by --force. +cache_file=$(ls "$OC_HOME/.cache/omarchy/agent-usage/"local-scan-*.json 2>/dev/null | head -n 1) +[[ -n $cache_file && $(stat -c %a "$cache_file") == "644" ]] || + fail "Local collector leaves a readable cache file behind" "$oc" +[[ $(jq -r '.schemaVersion' "$cache_file") == "1" && $(jq -r '.stats.todayTotalTokens' "$cache_file") == "170" ]] || + fail "Local collector writes a versioned cache envelope" "$oc" +pass "Local collector caches its scan" + +python3 - "$OC_HOME/.local/share/opencode/opencode.db" <<'PY' +import json, sqlite3, sys, time +conn = sqlite3.connect(sys.argv[1]); now_ms = int(time.time() * 1000) +conn.execute("INSERT INTO message VALUES (?, ?, ?, ?, ?)", ("m7", "ses_1", now_ms, now_ms, json.dumps({ + "role": "assistant", "providerID": "ollama", "modelID": "qwen-oc", + "tokens": {"input": 30, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + "time": {"created": now_ms}}))) +conn.commit(); conn.close() +PY + +[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local")") == "170" ]] || + fail "Local collector reuses a seconds-old scan" "$oc" +[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --force)") == "200" ]] || + fail "Local collector --force rescans past the cache" "$oc" +pass "Local collector --force rescans past the cache" + +# A cache stamped with another local date holds another day's today* stats. +jq -c '.scanDate = "1999-01-01"' "$cache_file" >"$cache_file.tmp" && mv "$cache_file.tmp" "$cache_file" +[[ $(jq -r '.scanDate' <<<"$(cat "$cache_file")") == "1999-01-01" ]] && + [[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --limits-only)") == "200" ]] || + fail "Local collector treats a cache from another day as a miss" "$oc" +pass "Local collector treats a cache from another day as a miss" + +# A corrupt-but-parseable cache is a miss, not a garbage record. +printf '[]' >"$cache_file" +[[ $(jq -r '.todayTotalTokens' <<<"$(env "${oc_env[@]}" "$ROOT/bin/omarchy-agent-usage-local" --limits-only)") == "200" ]] || + fail "Local collector recovers from a corrupt cache file" "$oc" +pass "Local collector recovers from a corrupt cache file" + +# An unwritable cache must not kill the collector: the record is the contract. +touch "$OC_HOME/not-a-dir" +broken=$(env "${oc_env[@]}" XDG_CACHE_HOME="$OC_HOME/not-a-dir" "$ROOT/bin/omarchy-agent-usage-local" 2>/dev/null) +[[ $(jq -r '.todayTotalTokens' <<<"$broken") == "200" ]] || + fail "Local collector still prints a record when the cache is unwritable" "$broken" +pass "Local collector survives an unwritable cache" + +# A scan cut short by a database error must not be cached as the whole story. +BROKEN_HOME=$(mktemp -d) +trap 'rm -rf "$LOCAL_HOME" "$OC_HOME" "$BROKEN_HOME"' EXIT +python3 - "$BROKEN_HOME/.local/share/opencode/opencode.db" <<'PY' +import sqlite3, sys +from pathlib import Path +db = Path(sys.argv[1]); db.parent.mkdir(parents=True, exist_ok=True) +conn = sqlite3.connect(db); conn.execute("CREATE TABLE unrelated (id text PRIMARY KEY)") +conn.commit(); conn.close() +PY +env HOME="$BROKEN_HOME" CODEX_HOME="$BROKEN_HOME/.codex" XDG_CACHE_HOME="$BROKEN_HOME/.cache" \ + XDG_DATA_HOME="$BROKEN_HOME/.local/share" XDG_CONFIG_HOME="$BROKEN_HOME/.config" \ + "$ROOT/bin/omarchy-agent-usage-local" >/dev/null +[[ -z $(ls "$BROKEN_HOME/.cache/omarchy/agent-usage/"local-scan-*.json 2>/dev/null) ]] || + fail "Local collector must not cache an interrupted scan" "broken-db" +pass "Local collector does not cache an interrupted opencode scan"