From 7713daf57ff53aa6e5ddf231cfe273cfb79285ca Mon Sep 17 00:00:00 2001 From: Renn F Date: Fri, 19 Jun 2026 07:54:59 +0200 Subject: [PATCH] fix(grok): stream one-shot output live + capture real token usage Two gaps the buffered run hid, both verified in the real image with mounted SuperGrok auth: - Observability: the entrypoint buffered grok's output to a temp file and only cat it after the run, so `docker logs` was blank while the agent worked. Switch the one-shot to --output-format streaming-json piped through tee: grok flushes each thought/text event incrementally (confirmed token-by-token live in-container), so the agent's reasoning shows in docker logs in real time, parity with the Claude stream-json path. Read the session id back from the NDJSON run log (the terminal `end` event) since -s does not pin it. - Usage: total_tokens read 0 for every grok run. grok nests the cumulative totalTokens on params.update._meta, but the reader looked at params._meta (which only holds event ids); the unit fixture had the same wrong shape, so the tests masked it. Read the real path (with params._meta / top-level fallbacks) and fix the fixture to the real grok shape. Verified live: usage.json now reports total_tokens=3262, cost_usd=0.006524 (was 0). --- docker/scripts/grok-cli-agent-entrypoint.sh | 18 +++-- roboco/llm/providers/grok_cli_usage.py | 70 ++++++++++++++----- .../unit/llm/providers/test_grok_cli_usage.py | 21 +++++- 3 files changed, 83 insertions(+), 26 deletions(-) diff --git a/docker/scripts/grok-cli-agent-entrypoint.sh b/docker/scripts/grok-cli-agent-entrypoint.sh index a9e7fe59..5ee9ca3d 100755 --- a/docker/scripts/grok-cli-agent-entrypoint.sh +++ b/docker/scripts/grok-cli-agent-entrypoint.sh @@ -42,18 +42,24 @@ WORKSPACE="${ROBOCO_WORKSPACE:-$PWD}" # render step above (grok_cli_config) wrote it from the mounted system prompt. # NOTE: grok generates its own session id and ignores a requested one (`-s` does # not pin it), so we do NOT pass a session id in; usage capture below reads the -# real id back out of the JSON run log instead. +# real id back out of the run log instead. +# +# `--output-format streaming-json` + `tee` streams the run to the container's +# stdout LIVE (so `docker logs` shows the agent reasoning/answering in real time, +# parity with the Claude path's stream-json) while ALSO capturing it to RUN_LOG +# for the session-id / usage read below. Without this the run is invisible until +# it ends (the buffered-to-a-file black box). stderr (grok's tool calls / +# diagnostics) goes to ERR_LOG and is surfaced after the run. set +e grok -p "${ROBOCO_INITIAL_PROMPT:-}" \ -m "${ROBOCO_AGENT_MODEL:-grok-build}" \ --cwd "$WORKSPACE" \ - --output-format json \ + --output-format streaming-json \ "${GROK_ARGS[@]}" \ - < /dev/null > "$RUN_LOG" 2> "$ERR_LOG" -run_rc=$? + < /dev/null 2> "$ERR_LOG" | tee "$RUN_LOG" +run_rc=${PIPESTATUS[0]} set -e -# Surface the run output + any stderr into the agent log. -cat "$RUN_LOG" +# stdout already streamed live via tee; surface stderr (tool calls / errors) too. [ -s "$ERR_LOG" ] && cat "$ERR_LOG" >&2 # Capture token usage from the grok session store (~/.grok/sessions). The reader diff --git a/roboco/llm/providers/grok_cli_usage.py b/roboco/llm/providers/grok_cli_usage.py index 1edf0d9f..f74d62cb 100644 --- a/roboco/llm/providers/grok_cli_usage.py +++ b/roboco/llm/providers/grok_cli_usage.py @@ -25,6 +25,7 @@ falling back to ``ROBOCO_AGENT_SESSION_ID`` only when no log is given. from __future__ import annotations +import contextlib import json import os from pathlib import Path @@ -42,10 +43,11 @@ USAGE_OUT_PATH = Path( 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). + Each line is a ``session/update`` JSON-RPC event. grok carries the cumulative + ``totalTokens`` on the inner ``params.update._meta`` (the per-chunk metadata); + older / alternate shapes put it on ``params._meta`` or the top level. The + maximum across the file is the session total. Returns 0 for a missing / empty + / unparseable file (best-effort — usage capture never fails a run). """ best = 0 try: @@ -65,9 +67,21 @@ def total_tokens_from_updates(updates_path: Path) -> int: 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)) + """Pull the cumulative ``totalTokens`` from a grok ``session/update`` event. + + grok nests the counter on ``params.update._meta`` (the per-chunk metadata); + ``params._meta`` there only carries event/timing ids. Fall back to + ``params._meta`` and a top-level field for older / alternate shapes. Returns 0 + when no recognised field is present. + """ + params = event.get("params") or {} + update_meta = (params.get("update") or {}).get("_meta") or {} + params_meta = params.get("_meta") or {} + for meta in (update_meta, params_meta): + if isinstance(meta, dict) and "totalTokens" in meta: + value = meta["totalTokens"] + return int(value) if isinstance(value, (int, float)) else 0 + value = event.get("totalTokens", 0) return int(value) if isinstance(value, (int, float)) else 0 @@ -125,21 +139,41 @@ def capture_session_usage( return 0 -def session_id_from_run_log(run_log: Path) -> str | None: - """Read the ``sessionId`` grok generated from its ``--output-format json`` log. +def _sid_from_obj(obj: object) -> str | None: + """A non-empty ``sessionId`` / ``session_id`` from a parsed event, else None.""" + if not isinstance(obj, dict): + return None + sid = obj.get("sessionId") or obj.get("session_id") + return sid if isinstance(sid, str) and sid else None - ``grok -p`` does not honour a requested session id, so the entrypoint hands - us the run's JSON output and we read the real id back. Returns ``None`` for a - missing / non-JSON / id-less log. + +def session_id_from_run_log(run_log: Path) -> str | None: + """Read the ``sessionId`` grok generated from the run's output log. + + ``grok -p`` does not honour a requested session id, so the entrypoint hands us + the run's output and we read the real id back. Handles BOTH the single-object + ``--output-format json`` log and the NDJSON ``--output-format streaming-json`` + log (the id rides on the terminal ``end`` event). Returns ``None`` for a + missing / id-less log. """ try: - payload = json.loads(run_log.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): + text = run_log.read_text(encoding="utf-8") + except OSError: return None - if not isinstance(payload, dict): - return None - sid = payload.get("sessionId") or payload.get("session_id") - return sid if isinstance(sid, str) and sid else None + # A single (possibly pretty-printed multi-line) JSON object first. + with contextlib.suppress(json.JSONDecodeError): + sid = _sid_from_obj(json.loads(text)) + if sid: + return sid + # Else scan NDJSON lines (streaming-json); the last sessionId wins. + found: str | None = None + for raw in text.splitlines(): + stripped = raw.strip() + if not stripped: + continue + with contextlib.suppress(json.JSONDecodeError): + found = _sid_from_obj(json.loads(stripped)) or found + return found def main() -> int: diff --git a/tests/unit/llm/providers/test_grok_cli_usage.py b/tests/unit/llm/providers/test_grok_cli_usage.py index b8e4e63b..95b62f38 100644 --- a/tests/unit/llm/providers/test_grok_cli_usage.py +++ b/tests/unit/llm/providers/test_grok_cli_usage.py @@ -17,14 +17,19 @@ def _write_updates(path: Path, totals: list[int]) -> None: path.parent.mkdir(parents=True, exist_ok=True) lines = [] for i, t in enumerate(totals): + # Real grok shape: the cumulative totalTokens rides on the INNER + # params.update._meta, not params._meta (which only holds event ids). lines.append( json.dumps( { "method": "session/update", "params": { "sessionId": "s1", - "update": {"sessionUpdate": "agent_message_chunk"}, - "_meta": {"totalTokens": t, "chunkId": i}, + "update": { + "sessionUpdate": "agent_message_chunk", + "_meta": {"totalTokens": t, "chunkId": i}, + }, + "_meta": {"eventId": f"s1-{i}"}, }, } ) @@ -142,6 +147,18 @@ def test_session_id_from_run_log_reads_the_real_id(tmp_path: Path) -> None: assert gu.session_id_from_run_log(log) == "019edd9d-real" +def test_session_id_from_run_log_reads_streaming_ndjson(tmp_path: Path) -> None: + # streaming-json: the id rides on the terminal `end` event line. + log = tmp_path / "run.ndjson" + log.write_text( + '{"type":"thought","data":"hm"}\n' + '{"type":"text","data":"ok"}\n' + '{"type":"end","stopReason":"EndTurn","sessionId":"019stream-real"}\n', + encoding="utf-8", + ) + assert gu.session_id_from_run_log(log) == "019stream-real" + + def test_session_id_from_run_log_none_for_bad_log(tmp_path: Path) -> None: assert gu.session_id_from_run_log(tmp_path / "absent.json") is None bad = tmp_path / "bad.json"