commit 4df21f20a0196ff8dd6efa62f2c7ce870b80c7fb Author: Ebeneezer Date: Thu Aug 20 15:31:55 2026 -0700 feat: SuperGrok weekly collector for stock Omarchy agents panel Writes grok.json for omarchy.agents. Yields when official omarchy-agent-usage-grok exists. No bar plugin, no packaged-bin name. diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1af05b3 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.mypy_cache/ +*.json +!tests/**/*.json +auth.json +grok.json +.venv/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..9048094 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 John Morris and Ebeneezer + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..94b9a30 --- /dev/null +++ b/README.md @@ -0,0 +1,86 @@ +# Omarchy Grok usage collector + +Household SuperGrok weekly quota for the stock Omarchy agents panel. Not a +bar plugin. Not a fork of third-party compositor QML. + +The panel already draws whatever JSON lands in: + +```text +~/.local/state/omarchy/agents/usage/.json +``` + +This repo writes `grok.json` there. When Omarchy later ships +`omarchy-agent-usage-grok` under `$OMARCHY_PATH/bin/`, this collector +**stops writing** and leaves that file alone. + +## Requirements + +- Omarchy with the Quattro `omarchy.agents` panel +- Grok Build CLI logged in (`grok login`) so `~/.grok/auth.json` exists +- Python 3.10+ (stdlib only) +- `systemd --user` + +It never stores the Grok token in the usage record, cache, or logs. The +token is read from the CLI auth file and sent only as `Authorization` to +xAI's CLI billing/settings HTTPS endpoints, with redirects refused. + +## Install + +```bash +git clone git@forgejo.jtmorris.net:ebeneezer/omarchy-grok-usage.git +cd omarchy-grok-usage +python3 -m unittest tests.test_collector -v +./install.sh +``` + +`install.sh` copies the collector to `~/.local/bin/omarchy-grok-usage` and +enables a 15-minute user timer. It does not write under `/usr/share/omarchy` +and does not name the binary `omarchy-agent-usage-grok`, so Omarchy's +`omarchy-agent-usage-update` glob will not pick this up. + +Left-click the agents icon on the bar. A Grok tab appears once +`grok.json` is ready (weekly percent is enough; local session charts are +out of scope). + +## Coexistence with official collectors + +`omarchy-agent-usage-update` only runs `$OMARCHY_PATH/bin/omarchy-agent-usage-*`. + +| Situation | What happens | +| --- | --- | +| No official Grok collector | This timer writes `grok.json` | +| Official `omarchy-agent-usage-grok` is executable | `--write` exits 0 and does not touch `grok.json` | +| You want to keep this writer anyway | `OMARCHY_GROK_USAGE_FORCE=1` | + +After an Omarchy update that ships Grok, you can also `./uninstall.sh`. +Uninstall does not delete `grok.json`. + +## Commands + +```bash +omarchy-grok-usage # print the record on stdout +omarchy-grok-usage --write # atomic replace of grok.json, unless yielding +systemctl --user status omarchy-grok-usage.timer +journalctl --user -u omarchy-grok-usage.service -n 50 +``` + +## Security notes + +- Reads `~/.grok/auth.json` (mode 0600, owned by you) +- HTTPS only, no redirect following +- Token never copied into the JSON the panel reads +- No QML in `omarchy-shell` +- No prebuilt binary; you run the Python in this tree + +## Tests + +```bash +python3 -m unittest tests.test_collector -v +``` + +Fixtures use fake tokens. Tests do not touch a live `auth.json` or the +network. + +## License + +MIT. See [LICENSE](LICENSE). diff --git a/bin/omarchy-grok-usage b/bin/omarchy-grok-usage new file mode 100755 index 0000000..95a4436 --- /dev/null +++ b/bin/omarchy-grok-usage @@ -0,0 +1,13 @@ +#!/usr/bin/python3 +"""Repo-checkout launcher. `install.sh` copies src/omarchy_grok_usage.py onto PATH instead.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) +from omarchy_grok_usage import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/DEV.md b/docs/DEV.md new file mode 100644 index 0000000..cdb43a5 --- /dev/null +++ b/docs/DEV.md @@ -0,0 +1,52 @@ +# Agent notes — omarchy-grok-usage + +Household collector. Owner: John Morris. Author on this Forgejo: Ebeneezer. + +## Product lock + +- Fill SuperGrok weekly remaining quota on the stock Omarchy agents panel. +- Do not add a Quickshell / bar plugin. +- Do not name the installed binary `omarchy-agent-usage-grok`. +- Do not patch `/usr/share/omarchy` or `omarchy-agent-usage-update`. +- When official `omarchy-agent-usage-grok` exists, `--write` is a no-op. +- No secrets in git, tests, README examples, or chat. + +## Layout + +| Path | Role | +| --- | --- | +| `src/omarchy_grok_usage.py` | Collector (stdlib Python) | +| `bin/omarchy-grok-usage` | Checkout launcher (`src/` on `sys.path`) | +| `install.sh` / `uninstall.sh` | User systemd timer | +| `systemd/` | `--user` unit + timer | +| `tests/test_collector.py` | unittest, no network | + +Install copies `src/omarchy_grok_usage.py` to `~/.local/bin/omarchy-grok-usage`. + +## Verify + +```bash +python3 -m unittest tests.test_collector -v +``` + +On a machine with `grok login` already done: + +```bash +./install.sh +python3 -c 'import json,pathlib; p=pathlib.Path.home()/".local/state/omarchy/agents/usage/grok.json"; d=json.loads(p.read_text()); assert d["id"]=="grok"; assert "limits" in d; print(d["tierLabel"], d["ready"], [x.get("percent") for x in d["limits"]])' +``` + +Do not print `auth.json`. Do not dump Authorization headers. + +## Git + +- Public repo is fine. Never commit `auth.json`, live `grok.json`, or tokens. +- Forgejo: `ebeneezer/omarchy-grok-usage`. John (`jtmorris`) is an admin collaborator. +- Git author for agent commits: `Ebeneezer `. + +## Out of scope + +- Claude / Codex / Amp / Cursor +- Session token charts from `~/.grok/sessions` +- Chips drawn on the bar itself +- Shipping this as an Omarchy plugin diff --git a/install.sh b/install.sh new file mode 100755 index 0000000..1c2d8a9 --- /dev/null +++ b/install.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Install the household Grok usage collector for the current user. +# Does not touch /usr/share/omarchy and does not install a bar plugin. +set -euo pipefail + +root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +src="$root/src/omarchy_grok_usage.py" +bin_dir="${XDG_BIN_HOME:-$HOME/.local/bin}" +unit_dir="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user" + +if [[ ! -f "$src" ]]; then + echo "missing collector source: $src" >&2 + exit 1 +fi + +mkdir -p "$bin_dir" "$unit_dir" +install -m 0755 "$src" "$bin_dir/omarchy-grok-usage" +install -m 0644 "$root/systemd/omarchy-grok-usage.service" "$unit_dir/omarchy-grok-usage.service" +install -m 0644 "$root/systemd/omarchy-grok-usage.timer" "$unit_dir/omarchy-grok-usage.timer" + +systemctl --user daemon-reload +systemctl --user enable --now omarchy-grok-usage.timer +systemctl --user start omarchy-grok-usage.service + +echo "Installed $bin_dir/omarchy-grok-usage" +echo "Timer: systemctl --user status omarchy-grok-usage.timer" +echo "If an official omarchy-agent-usage-grok later ships, this timer becomes a no-op." diff --git a/src/omarchy_grok_usage.py b/src/omarchy_grok_usage.py new file mode 100755 index 0000000..7e2cf93 --- /dev/null +++ b/src/omarchy_grok_usage.py @@ -0,0 +1,412 @@ +#!/usr/bin/python3 +"""Household SuperGrok weekly-quota collector for Omarchy's agents panel. + +Prints (or writes) one display-ready JSON record with id ``grok``. The stock +``omarchy.agents`` panel already watches +``~/.local/state/omarchy/agents/usage/``. This program does not install a +bar plugin and is not named ``omarchy-agent-usage-grok``, so a future +packaged Omarchy collector wins the glob in ``omarchy-agent-usage-update``. + +When an official ``omarchy-agent-usage-grok`` executable is present, ``--write`` +is a no-op and leaves ``grok.json`` alone. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import json +import os +import re +import sys +import tempfile +import urllib.error +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Mapping + +AGENT_ID = "grok" +AGENT_NAME = "Grok" +AUTH_HELP = "Run `grok login` to restore SuperGrok usage." +BILLING_URL = "https://cli-chat-proxy.grok.com/v1/billing?format=credits" +SETTINGS_URL = "https://cli-chat-proxy.grok.com/v1/settings" +SCHEMA_VERSION = 1 + +HttpGet = Callable[[str, Mapping[str, str]], dict[str, Any]] + + +class HttpError(Exception): + """HTTPS transport, redirect, or protocol failure. Never includes headers.""" + + +@dataclass(frozen=True) +class WriteResult: + action: str + path: Path | None = None + + +def parse_auth(payload: Any) -> tuple[str | None, str | None]: + """Return (token, optional first_name) from Grok CLI ``auth.json``.""" + if not isinstance(payload, dict): + return None, None + for entry in payload.values(): + if not isinstance(entry, dict): + continue + key = entry.get("key") + if isinstance(key, str) and key.strip(): + name = entry.get("first_name") + label = name.strip() if isinstance(name, str) and name.strip() else None + return key, label + return None, None + + +def _numeric(value: Any) -> float | None: + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, dict) and "val" in value: + return _numeric(value.get("val")) + if isinstance(value, str): + try: + return float(value.strip()) + except ValueError: + return None + return None + + +def _iso_reset(value: Any) -> str | None: + if not isinstance(value, str) or not value.strip(): + return None + raw = value.strip().replace("Z", "+00:00") + try: + parsed = dt.datetime.fromisoformat(raw) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=dt.timezone.utc) + return parsed.isoformat() + + +def limits_from_billing(payload: Mapping[str, Any]) -> list[dict[str, Any]]: + config = payload.get("config") if isinstance(payload.get("config"), dict) else payload + if not isinstance(config, dict): + return [] + + percent = _numeric(config.get("creditUsagePercent")) + if percent is None: + used = _numeric(config.get("used")) + limit = _numeric(config.get("monthlyLimit")) + if used is not None and limit and limit > 0: + percent = (used / limit) * 100.0 + if percent is None: + return [] + percent = max(0.0, min(100.0, percent)) + + period = config.get("currentPeriod") + end = None + if isinstance(period, dict): + end = period.get("end") + if end is None: + end = config.get("billingPeriodEnd") + + return [ + { + "label": "Weekly", + "percent": percent, + "resetsAt": _iso_reset(end), + } + ] + + +def _humanize_tier(label: str) -> str: + known = { + "SuperGrokPro": "SuperGrok Pro", + "SuperGrokHeavy": "SuperGrok Heavy", + "SuperGrokPlus": "SuperGrok Plus", + "SuperGrok": "SuperGrok", + } + if label in known: + return known[label] + match = re.match(r"^(.*?)(Pro|Heavy|Plus)$", label) + if match and match.group(1): + return f"{match.group(1)} {match.group(2)}" + return label + + +def tier_from_settings(payload: Mapping[str, Any]) -> str: + display = payload.get("subscription_tier_display") + if isinstance(display, str) and display.strip(): + return display.strip() + tier = payload.get("subscriptionTier") + if isinstance(tier, str) and tier.strip(): + return _humanize_tier(tier.strip()) + return "" + + +def official_grok_collector_present( + env: Mapping[str, str] | None = None, + extra_paths: tuple[str, ...] = (), +) -> bool: + env = os.environ if env is None else env + if str(env.get("OMARCHY_GROK_USAGE_FORCE", "")) == "1": + return False + omarchy = str(env.get("OMARCHY_PATH") or "").rstrip("/") + candidates = [ + f"{omarchy}/bin/omarchy-agent-usage-grok" if omarchy else "", + "/usr/share/omarchy/bin/omarchy-agent-usage-grok", + "/usr/bin/omarchy-agent-usage-grok", + *extra_paths, + ] + seen: set[str] = set() + for raw in candidates: + if not raw or raw in seen: + continue + seen.add(raw) + path = Path(raw) + try: + if path.is_file() and os.access(path, os.X_OK): + return True + except OSError: + continue + return False + + +def empty_stats() -> dict[str, Any]: + return { + "todayPrompts": 0, + "todaySessions": 0, + "todayTotalTokens": 0, + "todayTokensByModel": {}, + "recentDays": [], + "totalPrompts": 0, + "totalSessions": 0, + "activeDays": 0, + "activeDates": [], + "modelUsage": {}, + "hasPromptStats": False, + } + + +def build_record( + *, + limits: list[dict[str, Any]], + tier_label: str, + usage_status_text: str, + auth_help_text: str, + now_iso: str | None = None, +) -> dict[str, Any]: + stamp = now_iso or dt.datetime.now(dt.timezone.utc).isoformat() + record: dict[str, Any] = { + "schemaVersion": SCHEMA_VERSION, + "id": AGENT_ID, + "name": AGENT_NAME, + "updatedAt": stamp, + "ready": len(limits) > 0, + "hasLocalStats": False, + "scope": "account", + "tierLabel": tier_label, + "usageStatusText": usage_status_text, + "authHelpText": auth_help_text, + "limits": limits, + } + record.update(empty_stats()) + return record + + +def maybe_write_record( + *, + record: Mapping[str, Any], + usage_dir: Path, + yield_to_official: bool, +) -> WriteResult: + target = usage_dir / "grok.json" + if yield_to_official: + return WriteResult(action="yielded", path=target) + usage_dir.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp(prefix=".grok.", dir=str(usage_dir), text=True) + tmp_path = Path(tmp_name) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(record, handle, separators=(",", ":"), sort_keys=True) + handle.write("\n") + tmp_path.replace(target) + except Exception: + try: + tmp_path.unlink(missing_ok=True) + except OSError: + pass + raise + return WriteResult(action="wrote", path=target) + + +def _urlopen(request: urllib.request.Request, timeout: int = 10): + opener = urllib.request.build_opener(_NoRedirectHandler) + return opener.open(request, timeout=timeout) + + +class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: ANN001 + raise HttpError(f"redirect refused: {newurl}") + + +def https_get(url: str, headers: Mapping[str, str], timeout: int = 10) -> dict[str, Any]: + if not url.startswith("https://"): + raise HttpError("only https URLs are allowed") + request = urllib.request.Request(url, headers=dict(headers)) + try: + with _urlopen(request, timeout=timeout) as response: + final = response.geturl() + if final.split("#", 1)[0] != url: + raise HttpError(f"redirect refused: {final}") + raw = response.read() + except HttpError: + raise + except urllib.error.HTTPError as error: + raise HttpError(f"http {error.code}") from None + except Exception as error: + raise HttpError("network error") from error + try: + payload = json.loads(raw.decode("utf-8", errors="replace")) + except json.JSONDecodeError as error: + raise HttpError("invalid json") from error + if not isinstance(payload, dict): + raise HttpError("invalid json") + return payload + + +def grok_home(env: Mapping[str, str], home: Path) -> Path: + raw = env.get("GROK_HOME") + if raw: + path = Path(raw).expanduser() + if not path.is_absolute(): + raise ValueError("GROK_HOME must be an absolute path") + return path + return home / ".grok" + + +def load_auth_file(path: Path) -> tuple[str | None, str | None]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None, None + return parse_auth(payload) + + +def _auth_headers(token: str) -> dict[str, str]: + return { + "Authorization": f"Bearer {token}", + "X-XAI-Token-Auth": "xai-grok-cli", + "Accept": "application/json", + "User-Agent": "omarchy-grok-usage/1.0", + } + + +def collect_record( + *, + env: Mapping[str, str] | None = None, + home: Path | None = None, + http_get: HttpGet | None = None, +) -> dict[str, Any]: + env = os.environ if env is None else env + home = Path.home() if home is None else home + getter = https_get if http_get is None else http_get + + try: + auth_path = grok_home(env, home) / "auth.json" + except ValueError: + return build_record( + limits=[], + tier_label="", + usage_status_text="GROK_HOME is invalid.", + auth_help_text=AUTH_HELP, + ) + + token, _label = load_auth_file(auth_path) + if not token: + return build_record( + limits=[], + tier_label="", + usage_status_text="", + auth_help_text=AUTH_HELP, + ) + + headers = _auth_headers(token) + try: + billing = getter(BILLING_URL, headers) + except HttpError as error: + message = str(error) + if "http 401" in message or "http 403" in message: + return build_record( + limits=[], + tier_label="", + usage_status_text="Grok authentication was rejected.", + auth_help_text=AUTH_HELP, + ) + return build_record( + limits=[], + tier_label="", + usage_status_text="Couldn't reach Grok billing. Retrying later.", + auth_help_text=AUTH_HELP, + ) + + limits = limits_from_billing(billing) + tier = "" + try: + settings = getter(SETTINGS_URL, headers) + tier = tier_from_settings(settings) + except HttpError: + tier = "" + if not tier: + tier = "SuperGrok" if limits else "" + + return build_record( + limits=limits, + tier_label=tier, + usage_status_text="" if limits else "Grok billing returned no weekly window.", + auth_help_text=AUTH_HELP, + ) + + +def default_usage_dir(env: Mapping[str, str], home: Path) -> Path: + xdg = env.get("XDG_STATE_HOME") + root = Path(xdg) if xdg else home / ".local" / "state" + return root / "omarchy" / "agents" / "usage" + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Collect SuperGrok weekly usage for the Omarchy agents panel." + ) + parser.add_argument( + "--write", + action="store_true", + help="write grok.json (no-op if an official omarchy-agent-usage-grok exists)", + ) + args = parser.parse_args(argv) + + if official_grok_collector_present() and args.write: + print( + "omarchy-grok-usage: official omarchy-agent-usage-grok present; not writing grok.json", + file=sys.stderr, + ) + return 0 + + record = collect_record() + if args.write: + result = maybe_write_record( + record=record, + usage_dir=default_usage_dir(os.environ, Path.home()), + yield_to_official=False, + ) + if result.path is not None: + print(f"wrote {result.path}", file=sys.stderr) + return 0 + json.dump(record, sys.stdout, separators=(",", ":"), sort_keys=True) + sys.stdout.write("\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/systemd/omarchy-grok-usage.service b/systemd/omarchy-grok-usage.service new file mode 100644 index 0000000..8e49064 --- /dev/null +++ b/systemd/omarchy-grok-usage.service @@ -0,0 +1,12 @@ +[Unit] +Description=Refresh SuperGrok usage for the Omarchy agents panel +After=network-online.target +Wants=network-online.target + +[Service] +Type=oneshot +ExecStart=%h/.local/bin/omarchy-grok-usage --write +Nice=10 + +[Install] +WantedBy=default.target diff --git a/systemd/omarchy-grok-usage.timer b/systemd/omarchy-grok-usage.timer new file mode 100644 index 0000000..5a6655f --- /dev/null +++ b/systemd/omarchy-grok-usage.timer @@ -0,0 +1,12 @@ +[Unit] +Description=Refresh SuperGrok usage every 15 minutes + +[Timer] +OnBootSec=2min +OnUnitActiveSec=15min +AccuracySec=1min +Persistent=true +Unit=omarchy-grok-usage.service + +[Install] +WantedBy=timers.target diff --git a/tests/test_collector.py b/tests/test_collector.py new file mode 100644 index 0000000..74312ab --- /dev/null +++ b/tests/test_collector.py @@ -0,0 +1,275 @@ +#!/usr/bin/env python3 +"""Tests for the household Grok usage collector. + +No live network. No real credentials. Fixtures use obviously fake tokens. +""" + +from __future__ import annotations + +import json +import os +import stat +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + +import omarchy_grok_usage as m # noqa: E402 + +FAKE_TOKEN = "fake-test-token-not-a-real-secret" + + +class TestAuthParse(unittest.TestCase): + def test_extracts_key_from_accounts_xai_entry(self) -> None: + payload = { + "https://accounts.x.ai/sign-in": { + "key": FAKE_TOKEN, + "first_name": "Test", + } + } + token, label = m.parse_auth(payload) + self.assertEqual(token, FAKE_TOKEN) + self.assertEqual(label, "Test") + + def test_extracts_key_from_first_object_with_key(self) -> None: + payload = { + "other": {"nope": 1}, + "grok-com": {"key": FAKE_TOKEN, "first_name": "Ada"}, + } + token, label = m.parse_auth(payload) + self.assertEqual(token, FAKE_TOKEN) + self.assertEqual(label, "Ada") + + def test_missing_key_is_unauthenticated(self) -> None: + token, label = m.parse_auth({"x": {"refresh": "nope"}}) + self.assertIsNone(token) + self.assertIsNone(label) + + def test_empty_key_is_unauthenticated(self) -> None: + token, _label = m.parse_auth({"https://accounts.x.ai/sign-in": {"key": ""}}) + self.assertIsNone(token) + + +class TestBillingMap(unittest.TestCase): + def test_credit_usage_percent_and_period_end(self) -> None: + payload = { + "config": { + "creditUsagePercent": 31.5, + "currentPeriod": {"end": "2026-08-27T00:00:00Z"}, + "monthlyLimit": {"val": 60000}, + "used": {"val": 18900}, + } + } + limits = m.limits_from_billing(payload) + self.assertEqual(len(limits), 1) + self.assertEqual(limits[0]["label"], "Weekly") + self.assertEqual(limits[0]["percent"], 31.5) + self.assertEqual(limits[0]["resetsAt"], "2026-08-27T00:00:00+00:00") + + def test_falls_back_to_used_over_limit(self) -> None: + payload = { + "config": { + "monthlyLimit": {"val": 100}, + "used": {"val": 40}, + "billingPeriodEnd": "2026-09-01T00:00:00+00:00", + } + } + limits = m.limits_from_billing(payload) + self.assertEqual(limits[0]["percent"], 40.0) + self.assertEqual(limits[0]["resetsAt"], "2026-09-01T00:00:00+00:00") + + def test_empty_billing_yields_no_limits(self) -> None: + self.assertEqual(m.limits_from_billing({}), []) + + def test_percent_clamped_to_100(self) -> None: + payload = {"config": {"creditUsagePercent": 140}} + limits = m.limits_from_billing(payload) + self.assertEqual(limits[0]["percent"], 100.0) + + +class TestTierLabel(unittest.TestCase): + def test_prefers_subscription_tier_display(self) -> None: + self.assertEqual( + m.tier_from_settings({"subscription_tier_display": "SuperGrok Heavy"}), + "SuperGrok Heavy", + ) + + def test_splits_camel_subscription_tier(self) -> None: + self.assertEqual( + m.tier_from_settings({"subscriptionTier": "SuperGrokPro"}), + "SuperGrok Pro", + ) + + def test_empty_settings(self) -> None: + self.assertEqual(m.tier_from_settings({}), "") + + +class TestOfficialYield(unittest.TestCase): + def test_yields_when_omarchy_path_collector_exists(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + official = Path(tmp) / "bin" / "omarchy-agent-usage-grok" + official.parent.mkdir(parents=True) + official.write_text("#!/bin/sh\n") + official.chmod(official.stat().st_mode | stat.S_IEXEC) + env = {"OMARCHY_PATH": tmp} + self.assertTrue(m.official_grok_collector_present(env=env)) + + def test_absent_when_no_official_binary(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + env = { + "OMARCHY_PATH": tmp, + "PATH": tmp, + } + self.assertFalse( + m.official_grok_collector_present( + env=env, extra_paths=(str(Path(tmp) / "nope"),) + ) + ) + + def test_force_env_disables_yield(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + official = Path(tmp) / "bin" / "omarchy-agent-usage-grok" + official.parent.mkdir(parents=True) + official.write_text("#!/bin/sh\n") + official.chmod(official.stat().st_mode | stat.S_IEXEC) + env = {"OMARCHY_PATH": tmp, "OMARCHY_GROK_USAGE_FORCE": "1"} + self.assertFalse(m.official_grok_collector_present(env=env)) + + +class TestRecord(unittest.TestCase): + def test_token_never_enters_record(self) -> None: + record = m.build_record( + limits=[{"label": "Weekly", "percent": 10.0, "resetsAt": None}], + tier_label="SuperGrok", + usage_status_text="", + auth_help_text="Run `grok login`.", + now_iso="2026-08-20T00:00:00+00:00", + ) + blob = json.dumps(record) + self.assertNotIn(FAKE_TOKEN, blob) + self.assertNotIn("Authorization", blob) + self.assertEqual(record["id"], "grok") + self.assertEqual(record["schemaVersion"], 1) + self.assertTrue(record["ready"]) + self.assertEqual(record["scope"], "account") + self.assertFalse(record["hasLocalStats"]) + + def test_unauthenticated_record_is_not_ready(self) -> None: + record = m.build_record( + limits=[], + tier_label="", + usage_status_text="", + auth_help_text="Run `grok login` to restore SuperGrok usage.", + now_iso="2026-08-20T00:00:00+00:00", + ) + self.assertFalse(record["ready"]) + self.assertEqual(record["limits"], []) + + +class TestWritePath(unittest.TestCase): + def test_skips_write_when_official_present(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + usage_dir = Path(tmp) / "usage" + usage_dir.mkdir() + existing = usage_dir / "grok.json" + existing.write_text('{"id":"grok","from":"official"}\n') + result = m.maybe_write_record( + record={"id": "grok", "from": "household"}, + usage_dir=usage_dir, + yield_to_official=True, + ) + self.assertEqual(result.action, "yielded") + self.assertEqual(existing.read_text(), '{"id":"grok","from":"official"}\n') + + def test_writes_atomically_when_no_official(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + usage_dir = Path(tmp) / "usage" + result = m.maybe_write_record( + record={"id": "grok", "schemaVersion": 1}, + usage_dir=usage_dir, + yield_to_official=False, + ) + self.assertEqual(result.action, "wrote") + written = json.loads((usage_dir / "grok.json").read_text()) + self.assertEqual(written["id"], "grok") + + def test_https_only_client_rejects_http(self) -> None: + with self.assertRaises(m.HttpError): + m.https_get("http://example.com/billing", headers={}) + + +class TestHttpsClientRedirects(unittest.TestCase): + def test_redirects_are_refused(self) -> None: + class FakeResponse: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self) -> bytes: + return b"{}" + + def geturl(self) -> str: + return "https://evil.example/steal" + + def fake_open(_request, timeout=10): + return FakeResponse() + + with mock.patch.object(m, "_urlopen", fake_open): + with self.assertRaises(m.HttpError) as ctx: + m.https_get( + "https://cli-chat-proxy.grok.com/v1/billing?format=credits", + headers={"Authorization": "Bearer " + FAKE_TOKEN}, + ) + self.assertIn("redirect", str(ctx.exception).lower()) + + +class TestCollect(unittest.TestCase): + def test_collect_record_does_not_embed_token(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + grok = Path(tmp) / ".grok" + grok.mkdir() + (grok / "auth.json").write_text( + json.dumps({"https://accounts.x.ai/sign-in": {"key": FAKE_TOKEN}}), + encoding="utf-8", + ) + + def fake_get(url: str, headers): + auth = headers.get("Authorization", "") + self.assertTrue(auth.startswith("Bearer ")) + if "billing" in url: + return { + "config": { + "creditUsagePercent": 12, + "currentPeriod": {"end": "2026-08-27T00:00:00Z"}, + } + } + if "settings" in url: + return {"subscription_tier_display": "SuperGrok"} + raise AssertionError(url) + + record = m.collect_record( + env={"GROK_HOME": str(grok)}, + home=Path(tmp), + http_get=fake_get, + ) + blob = json.dumps(record) + self.assertNotIn(FAKE_TOKEN, blob) + self.assertEqual(record["limits"][0]["percent"], 12) + self.assertEqual(record["tierLabel"], "SuperGrok") + self.assertTrue(record["ready"]) + + def test_collect_record_without_auth_file(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + record = m.collect_record(env={"GROK_HOME": tmp}, home=Path(tmp)) + self.assertFalse(record["ready"]) + self.assertIn("grok login", record["authHelpText"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/uninstall.sh b/uninstall.sh new file mode 100755 index 0000000..45c5a08 --- /dev/null +++ b/uninstall.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash +# Remove the household Grok usage collector for the current user. +# Leaves grok.json in place so a later official collector can overwrite it. +set -euo pipefail + +bin_dir="${XDG_BIN_HOME:-$HOME/.local/bin}" +unit_dir="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user" + +systemctl --user disable --now omarchy-grok-usage.timer 2>/dev/null || true +systemctl --user disable --now omarchy-grok-usage.service 2>/dev/null || true +rm -f "$unit_dir/omarchy-grok-usage.timer" "$unit_dir/omarchy-grok-usage.service" +rm -f "$bin_dir/omarchy-grok-usage" +systemctl --user daemon-reload 2>/dev/null || true + +echo "Removed omarchy-grok-usage timer and ~/.local/bin/omarchy-grok-usage" +echo "Did not delete ~/.local/state/omarchy/agents/usage/grok.json"