From 579dfb997bf986d7bdc0069225e614f54a631843 Mon Sep 17 00:00:00 2001 From: Renn F Date: Fri, 19 Jun 2026 03:28:39 +0200 Subject: [PATCH] feat(grok-cli): capture per-session token usage + notional cost MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Grok runs on the SuperGrok subscription, but — exactly like Claude on Max — we still record per-agent tokens and a notional cost for the dashboard. The grok CLI writes a cumulative totalTokens per turn into ~/.grok/sessions///updates.jsonl (the grok analogue of the Claude transcript / old opencode.db); the max is the session total. This reader locates that file (url-encoded cwd), extracts the total, and prices it at the output rate (no input/output split from the CLI; conservative + matches the reasoning-at-output convention). Validated against a real grok-build session (18253 tokens -> $0.0365). Entrypoint + finalize wiring follows. --- roboco/llm/providers/grok_cli_usage.py | 110 ++++++++++++++++++ .../unit/llm/providers/test_grok_cli_usage.py | 98 ++++++++++++++++ 2 files changed, 208 insertions(+) create mode 100644 roboco/llm/providers/grok_cli_usage.py create mode 100644 tests/unit/llm/providers/test_grok_cli_usage.py diff --git a/roboco/llm/providers/grok_cli_usage.py b/roboco/llm/providers/grok_cli_usage.py new file mode 100644 index 00000000..d138b4b0 --- /dev/null +++ b/roboco/llm/providers/grok_cli_usage.py @@ -0,0 +1,110 @@ +"""Capture token usage from a Grok CLI session for the usage / cost dashboard. + +The grok CLI persists each session under ``~/.grok/sessions// +/updates.jsonl``; every update carries a cumulative +``params._meta.totalTokens``, so the maximum across the file is the session's +total token count. This is the grok analogue of the Claude Code transcript and +the old opencode.db — Grok runs on the SuperGrok subscription, but (exactly like +Claude on Max) we still record per-agent tokens and a notional cost. + +The grok CLI reports a single ``totalTokens`` with no input/output split, so the +notional cost prices the whole total at the output rate (the higher rate — +conservative, and consistent with reasoning tokens billing at the output rate). + +The agent entrypoint runs ``python -m roboco.llm.providers.grok_cli_usage`` after +the run to write a small ``usage.json`` (``{model, total_tokens, cost_usd}``) +into a per-agent dir the orchestrator reads back at finalize. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any +from urllib.parse import quote + +from roboco.billing.pricing import calculate_cost + +# Where the entrypoint writes the captured usage for the orchestrator to read. +USAGE_OUT_PATH = Path( + os.environ.get("ROBOCO_GROK_USAGE_FILE", "/tmp/roboco-grok-usage.json") +) + + +def total_tokens_from_updates(updates_path: Path) -> int: + """Return the max cumulative ``totalTokens`` in a grok ``updates.jsonl``. + + Each line is a ``session/update`` JSON-RPC event whose ``params._meta`` (or a + top-level field on older formats) carries a cumulative ``totalTokens``. The + maximum is the session total. Returns 0 for a missing / empty / unparseable + file (best-effort — usage capture never fails a run). + """ + best = 0 + try: + with updates_path.open(encoding="utf-8") as fh: + for raw in fh: + text = raw.strip() + if not text: + continue + try: + event = json.loads(text) + except json.JSONDecodeError: + continue + best = max(best, _extract_total_tokens(event)) + except OSError: + return 0 + return best + + +def _extract_total_tokens(event: dict[str, Any]) -> int: + """Pull ``totalTokens`` from an update event (nested ``_meta`` or top-level).""" + meta = (event.get("params") or {}).get("_meta") or {} + value = meta.get("totalTokens", event.get("totalTokens", 0)) + return int(value) if isinstance(value, (int, float)) else 0 + + +def find_updates_path(grok_home: Path, cwd: str, session_id: str) -> Path | None: + """Locate ``updates.jsonl`` for a session, or ``None`` if not present. + + grok keys the session dir by the url-encoded working directory (``/`` → + ``%2F``) under ``/sessions///``. + """ + if not (cwd and session_id): + return None + encoded = quote(cwd, safe="") + candidate = grok_home / "sessions" / encoded / session_id / "updates.jsonl" + return candidate if candidate.is_file() else None + + +def usage_and_cost(model: str, total_tokens: int) -> tuple[int, float]: + """Return ``(total_tokens, notional_cost_usd)`` for a grok session. + + No input/output split is available, so the whole total is priced at the + output rate (folded into ``tokens_output``). + """ + return total_tokens, calculate_cost( + model, tokens_input=0, tokens_output=total_tokens + ) + + +def main() -> int: + """Entrypoint: write ``usage.json`` (model, total_tokens, cost) for the run.""" + grok_home = Path(os.environ.get("GROK_HOME", str(Path.home() / ".grok"))) + cwd = os.environ.get("ROBOCO_GROK_RUN_CWD", str(Path.cwd())) + session_id = os.environ.get("ROBOCO_AGENT_SESSION_ID", "") + model = os.environ.get("ROBOCO_AGENT_MODEL", "grok-build") + + updates = find_updates_path(grok_home, cwd, session_id) + total = total_tokens_from_updates(updates) if updates else 0 + tokens, cost = usage_and_cost(model, total) + + USAGE_OUT_PATH.write_text( + json.dumps({"model": model, "total_tokens": tokens, "cost_usd": cost}), + encoding="utf-8", + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/llm/providers/test_grok_cli_usage.py b/tests/unit/llm/providers/test_grok_cli_usage.py new file mode 100644 index 00000000..51dc8107 --- /dev/null +++ b/tests/unit/llm/providers/test_grok_cli_usage.py @@ -0,0 +1,98 @@ +"""grok_cli_usage — capture token usage from a Grok CLI session store.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +from roboco.llm.providers import grok_cli_usage as gu + +if TYPE_CHECKING: + from pathlib import Path + + +def _write_updates(path: Path, totals: list[int]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + lines = [] + for i, t in enumerate(totals): + lines.append( + json.dumps( + { + "method": "session/update", + "params": { + "sessionId": "s1", + "update": {"sessionUpdate": "agent_message_chunk"}, + "_meta": {"totalTokens": t, "chunkId": i}, + }, + } + ) + ) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def test_total_tokens_is_the_running_max(tmp_path: Path) -> None: + upd = tmp_path / "updates.jsonl" + _write_updates(upd, [3863, 18133, 18220, 18253, 18220]) + assert gu.total_tokens_from_updates(upd) == 18253 # noqa: PLR2004 + + +def test_total_tokens_zero_for_missing_or_empty(tmp_path: Path) -> None: + assert gu.total_tokens_from_updates(tmp_path / "nope.jsonl") == 0 + empty = tmp_path / "empty.jsonl" + empty.write_text("\n \n", encoding="utf-8") + assert gu.total_tokens_from_updates(empty) == 0 + + +def test_total_tokens_tolerates_bad_lines(tmp_path: Path) -> None: + upd = tmp_path / "updates.jsonl" + upd.write_text( + 'not json\n{"params":{"_meta":{"totalTokens":42}}}\n{"x":1}\n', + encoding="utf-8", + ) + assert gu.total_tokens_from_updates(upd) == 42 # noqa: PLR2004 + + +def test_find_updates_path_encodes_cwd(tmp_path: Path) -> None: + home = tmp_path / ".grok" + cwd = "/data/workspaces/roboco/backend/be-dev-1" + sid = "019edd59-bc7b-7920" + target = ( + home / "sessions" / "%2Fdata%2Fworkspaces%2Froboco%2Fbackend%2Fbe-dev-1" / sid + ) + _write_updates(target / "updates.jsonl", [10]) + found = gu.find_updates_path(home, cwd, sid) + assert found is not None + assert found == target / "updates.jsonl" + + +def test_find_updates_path_none_when_absent(tmp_path: Path) -> None: + home = tmp_path / ".grok" + assert gu.find_updates_path(home, "/x", "sid") is None + assert gu.find_updates_path(home, "", "sid") is None # missing cwd + assert gu.find_updates_path(home, "/x", "") is None # missing session id + + +def test_usage_and_cost_prices_total_at_output_rate() -> None: + # grok-build output rate is $2.00/1M → 1M tokens = $2.00. + tokens, cost = gu.usage_and_cost("grok-build", 1_000_000) + assert tokens == 1_000_000 # noqa: PLR2004 + assert abs(cost - 2.00) < 1e-6 # noqa: PLR2004 + + +def test_main_writes_usage_file(tmp_path: Path, monkeypatch) -> None: # type: ignore[no-untyped-def] + home = tmp_path / ".grok" + cwd = "/ws/be-dev-1" + sid = "sid-1" + target = home / "sessions" / "%2Fws%2Fbe-dev-1" / sid + _write_updates(target / "updates.jsonl", [1234]) + out = tmp_path / "usage.json" + monkeypatch.setattr(gu, "USAGE_OUT_PATH", out) + monkeypatch.setenv("GROK_HOME", str(home)) + monkeypatch.setenv("ROBOCO_GROK_RUN_CWD", cwd) + monkeypatch.setenv("ROBOCO_AGENT_SESSION_ID", sid) + monkeypatch.setenv("ROBOCO_AGENT_MODEL", "grok-build") + assert gu.main() == 0 + data = json.loads(out.read_text()) + assert data["total_tokens"] == 1234 # noqa: PLR2004 + assert data["model"] == "grok-build" + assert data["cost_usd"] > 0.0