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.
This commit is contained in:
Ebeneezer 2026-08-20 15:31:55 -07:00
commit 4df21f20a0
11 changed files with 935 additions and 0 deletions

9
.gitignore vendored Normal file
View file

@ -0,0 +1,9 @@
__pycache__/
*.py[cod]
.pytest_cache/
.mypy_cache/
*.json
!tests/**/*.json
auth.json
grok.json
.venv/

21
LICENSE Normal file
View file

@ -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.

86
README.md Normal file
View file

@ -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/<id>.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).

13
bin/omarchy-grok-usage Executable file
View file

@ -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())

52
docs/DEV.md Normal file
View file

@ -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 <ebeneezer@jtmorris.net>`.
## 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

27
install.sh Executable file
View file

@ -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."

412
src/omarchy_grok_usage.py Executable file
View file

@ -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())

View file

@ -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

View file

@ -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

275
tests/test_collector.py Normal file
View file

@ -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()

16
uninstall.sh Executable file
View file

@ -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"