mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
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).
This commit is contained in:
@@ -42,18 +42,24 @@ WORKSPACE="${ROBOCO_WORKSPACE:-$PWD}"
|
|||||||
# render step above (grok_cli_config) wrote it from the mounted system prompt.
|
# 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
|
# 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
|
# 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
|
set +e
|
||||||
grok -p "${ROBOCO_INITIAL_PROMPT:-}" \
|
grok -p "${ROBOCO_INITIAL_PROMPT:-}" \
|
||||||
-m "${ROBOCO_AGENT_MODEL:-grok-build}" \
|
-m "${ROBOCO_AGENT_MODEL:-grok-build}" \
|
||||||
--cwd "$WORKSPACE" \
|
--cwd "$WORKSPACE" \
|
||||||
--output-format json \
|
--output-format streaming-json \
|
||||||
"${GROK_ARGS[@]}" \
|
"${GROK_ARGS[@]}" \
|
||||||
< /dev/null > "$RUN_LOG" 2> "$ERR_LOG"
|
< /dev/null 2> "$ERR_LOG" | tee "$RUN_LOG"
|
||||||
run_rc=$?
|
run_rc=${PIPESTATUS[0]}
|
||||||
set -e
|
set -e
|
||||||
# Surface the run output + any stderr into the agent log.
|
# stdout already streamed live via tee; surface stderr (tool calls / errors) too.
|
||||||
cat "$RUN_LOG"
|
|
||||||
[ -s "$ERR_LOG" ] && cat "$ERR_LOG" >&2
|
[ -s "$ERR_LOG" ] && cat "$ERR_LOG" >&2
|
||||||
|
|
||||||
# Capture token usage from the grok session store (~/.grok/sessions). The reader
|
# Capture token usage from the grok session store (~/.grok/sessions). The reader
|
||||||
|
|||||||
@@ -25,6 +25,7 @@ falling back to ``ROBOCO_AGENT_SESSION_ID`` only when no log is given.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import contextlib
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -42,10 +43,11 @@ USAGE_OUT_PATH = Path(
|
|||||||
def total_tokens_from_updates(updates_path: Path) -> int:
|
def total_tokens_from_updates(updates_path: Path) -> int:
|
||||||
"""Return the max cumulative ``totalTokens`` in a grok ``updates.jsonl``.
|
"""Return the max cumulative ``totalTokens`` in a grok ``updates.jsonl``.
|
||||||
|
|
||||||
Each line is a ``session/update`` JSON-RPC event whose ``params._meta`` (or a
|
Each line is a ``session/update`` JSON-RPC event. grok carries the cumulative
|
||||||
top-level field on older formats) carries a cumulative ``totalTokens``. The
|
``totalTokens`` on the inner ``params.update._meta`` (the per-chunk metadata);
|
||||||
maximum is the session total. Returns 0 for a missing / empty / unparseable
|
older / alternate shapes put it on ``params._meta`` or the top level. The
|
||||||
file (best-effort — usage capture never fails a run).
|
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
|
best = 0
|
||||||
try:
|
try:
|
||||||
@@ -65,9 +67,21 @@ def total_tokens_from_updates(updates_path: Path) -> int:
|
|||||||
|
|
||||||
|
|
||||||
def _extract_total_tokens(event: dict[str, Any]) -> int:
|
def _extract_total_tokens(event: dict[str, Any]) -> int:
|
||||||
"""Pull ``totalTokens`` from an update event (nested ``_meta`` or top-level)."""
|
"""Pull the cumulative ``totalTokens`` from a grok ``session/update`` event.
|
||||||
meta = (event.get("params") or {}).get("_meta") or {}
|
|
||||||
value = meta.get("totalTokens", event.get("totalTokens", 0))
|
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
|
return int(value) if isinstance(value, (int, float)) else 0
|
||||||
|
|
||||||
|
|
||||||
@@ -125,21 +139,41 @@ def capture_session_usage(
|
|||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|
||||||
def session_id_from_run_log(run_log: Path) -> str | None:
|
def _sid_from_obj(obj: object) -> str | None:
|
||||||
"""Read the ``sessionId`` grok generated from its ``--output-format json`` log.
|
"""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
|
def session_id_from_run_log(run_log: Path) -> str | None:
|
||||||
missing / non-JSON / id-less log.
|
"""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:
|
try:
|
||||||
payload = json.loads(run_log.read_text(encoding="utf-8"))
|
text = run_log.read_text(encoding="utf-8")
|
||||||
except (OSError, json.JSONDecodeError):
|
except OSError:
|
||||||
return None
|
return None
|
||||||
if not isinstance(payload, dict):
|
# A single (possibly pretty-printed multi-line) JSON object first.
|
||||||
return None
|
with contextlib.suppress(json.JSONDecodeError):
|
||||||
sid = payload.get("sessionId") or payload.get("session_id")
|
sid = _sid_from_obj(json.loads(text))
|
||||||
return sid if isinstance(sid, str) and sid else None
|
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:
|
def main() -> int:
|
||||||
|
|||||||
@@ -17,14 +17,19 @@ def _write_updates(path: Path, totals: list[int]) -> None:
|
|||||||
path.parent.mkdir(parents=True, exist_ok=True)
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
lines = []
|
lines = []
|
||||||
for i, t in enumerate(totals):
|
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(
|
lines.append(
|
||||||
json.dumps(
|
json.dumps(
|
||||||
{
|
{
|
||||||
"method": "session/update",
|
"method": "session/update",
|
||||||
"params": {
|
"params": {
|
||||||
"sessionId": "s1",
|
"sessionId": "s1",
|
||||||
"update": {"sessionUpdate": "agent_message_chunk"},
|
"update": {
|
||||||
"_meta": {"totalTokens": t, "chunkId": i},
|
"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"
|
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:
|
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
|
assert gu.session_id_from_run_log(tmp_path / "absent.json") is None
|
||||||
bad = tmp_path / "bad.json"
|
bad = tmp_path / "bad.json"
|
||||||
|
|||||||
Reference in New Issue
Block a user