mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
feat(grok-cli): capture per-session token usage + notional cost
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/<cwd>/<session-id>/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.
This commit is contained in:
@@ -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/<url-encoded-cwd>/
|
||||
<session-id>/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 ``<grok_home>/sessions/<encoded-cwd>/<session-id>/``.
|
||||
"""
|
||||
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())
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user