mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
Renn F
579dfb997b
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.
99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
"""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
|