From c139e2d0171d6ee9e4d5782b1223b6c6add06a93 Mon Sep 17 00:00:00 2001
From: Renn F <rennf93@users.noreply.github.com>
Date: Fri, 19 Jun 2026 03:32:32 +0200
Subject: [PATCH] feat(grok-cli): wire usage capture into the run (session id +
 post-run extract)

The provider pins a fixed session id (ROBOCO_AGENT_SESSION_ID, reused from the
agent session id as on the Claude path); the entrypoint passes it to
'grok -p -s <id>' so the run's session store is locatable, then runs the usage
reader post-run (best-effort) to write the captured tokens + cost. The
orchestrator-side finalize that reads that file follows.
---
 docker/scripts/grok-cli-agent-entrypoint.sh | 12 ++++++++++++
 roboco/llm/providers/grok.py                |  5 +++++
 tests/unit/llm/test_providers.py            |  2 ++
 3 files changed, 19 insertions(+)

diff --git a/docker/scripts/grok-cli-agent-entrypoint.sh b/docker/scripts/grok-cli-agent-entrypoint.sh
index e1b6b361..36c1ea6e 100755
--- a/docker/scripts/grok-cli-agent-entrypoint.sh
+++ b/docker/scripts/grok-cli-agent-entrypoint.sh
@@ -36,11 +36,16 @@ fi
 RUN_LOG="/tmp/grok-run.json"
 ERR_LOG="/tmp/grok-run.err"
 WORKSPACE="${ROBOCO_WORKSPACE:-$PWD}"
+# A fixed session id (set by the provider) makes the run's session store
+# locatable for usage capture below; absent, grok generates its own.
+SESSION_ARGS=()
+[ -n "${ROBOCO_AGENT_SESSION_ID:-}" ] && SESSION_ARGS=(-s "${ROBOCO_AGENT_SESSION_ID}")
 set +e
 grok -p "${ROBOCO_INITIAL_PROMPT:-}" \
   -m "${ROBOCO_AGENT_MODEL:-grok-build}" \
   --cwd "$WORKSPACE" \
   --output-format json \
+  "${SESSION_ARGS[@]}" \
   "${GROK_ARGS[@]}" \
   < /dev/null > "$RUN_LOG" 2> "$ERR_LOG"
 run_rc=$?
@@ -49,6 +54,13 @@ set -e
 cat "$RUN_LOG"
 [ -s "$ERR_LOG" ] && cat "$ERR_LOG" >&2
 
+# Capture token usage from the grok session store (~/.grok/sessions). The
+# orchestrator reads the written usage file back at finalize — the grok analogue
+# of the Claude transcript. Best-effort; never fails the run. Run from /app for
+# the same module-resolution reason as the render above.
+( cd /app && ROBOCO_GROK_RUN_CWD="$WORKSPACE" \
+    python -m roboco.llm.providers.grok_cli_usage ) || true
+
 # Rate-limit detection (parity with the opencode B4 path): an xAI 429 / quota
 # error ends the run without a terminal verb. Detect it from the run output and
 # exit 75 (EX_TEMPFAIL) so the orchestrator PARKS the grok provider instead of
diff --git a/roboco/llm/providers/grok.py b/roboco/llm/providers/grok.py
index d52504ff..50943a7e 100644
--- a/roboco/llm/providers/grok.py
+++ b/roboco/llm/providers/grok.py
@@ -175,6 +175,11 @@ class GrokCliProvider(AgentProvider):
                 f"ROBOCO_INITIAL_PROMPT={initial_prompt or ''}",
             ]
         )
+        if config.claude_session_id:
+            # A fixed session id (reused as the generic agent session id, as on
+            # the Claude path) lets the entrypoint pin `grok -p -s <id>` so the
+            # run's session store is locatable for token-usage capture.
+            cmd.extend(["-e", f"ROBOCO_AGENT_SESSION_ID={config.claude_session_id}"])
 
     async def stop(self, instance_id: str, graceful: bool = True) -> None:
         await stop_container(instance_id, graceful)
diff --git a/tests/unit/llm/test_providers.py b/tests/unit/llm/test_providers.py
index 2c3149c6..acf7d959 100644
--- a/tests/unit/llm/test_providers.py
+++ b/tests/unit/llm/test_providers.py
@@ -213,6 +213,8 @@ async def test_grok_spawn_wires_gateway_env_and_image_last() -> None:
     assert "ROBOCO_MCP_CONFIG=/app/mcp-config.json" in cmd
     assert "ROBOCO_AGENT_ID=be-dev-1" in cmd  # renderer computes per-role flags
     assert "ROBOCO_AGENT_MODEL=grok-build" in cmd
+    # Fixed session id so usage capture can locate the run's session store.
+    assert "ROBOCO_AGENT_SESSION_ID=sess-1" in cmd
     # Identity wiring from the shared host helpers is present.
     assert "ROBOCO_AGENT_TOKEN=hmac-be-dev-1" in cmd
     # The image is the final docker-run argument.