fix(usage): capture agent token usage from the Claude Code transcript

The token-usage pipeline was fully built — per-session SDK counters,
/usage/status, the orchestrator finalize-fetch that writes token columns and
estimated cost to the spawn-session row, the daily rollup, and the dashboard
— but nothing ever populated the counters. /usage/report had zero callers, so
every session reported zero tokens and the cost dashboard rendered all-zeros.
A redeploy could not fix code that was never written.

Close the loop with the producer that was missing. Claude Code does not pass
token counts to hooks, but it does pass the session transcript path, and each
assistant entry records its API call's usage. Add:

- POST /usage/sync, which parses the transcript and *sets* the cumulative
  totals absolutely (idempotent — re-syncing the same or a grown transcript
  overwrites, never double-counts), with a (size, mtime) short-circuit so an
  unchanged transcript skips the re-parse.
- usage-report-hook.sh, which hands the SDK the transcript path. Registered on
  PostToolUse (keeps mid-run snapshots and reaped-agent sessions accurate) and
  Stop (guarantees a final sync at turn end before finalize reads the totals).

Field mapping verified against a real Claude Code transcript:
message.usage.{input_tokens, output_tokens, cache_read_input_tokens,
cache_creation_input_tokens}. Unit tests cover summation, idempotency, growth,
a missing transcript, and malformed lines.
This commit is contained in:
Renn F
2026-06-11 07:46:19 +02:00
parent e0cd305844
commit cc4ccb7ea3
6 changed files with 324 additions and 2 deletions
+1
View File
@@ -65,6 +65,7 @@ COPY docker/scripts/sdk-startup-hook.sh /app/scripts/sdk-startup-hook.sh
COPY docker/scripts/a2a-check-hook.sh /app/scripts/a2a-check-hook.sh
COPY docker/scripts/bash-guard-hook.sh /app/scripts/bash-guard-hook.sh
COPY docker/scripts/post-tool-budget-hook.sh /app/scripts/post-tool-budget-hook.sh
COPY docker/scripts/usage-report-hook.sh /app/scripts/usage-report-hook.sh
COPY docker/scripts/stop-hook.sh /app/scripts/stop-hook.sh
COPY docker/scripts/user-prompt-hook.sh /app/scripts/user-prompt-hook.sh
COPY docker/scripts/pre-compact-hook.sh /app/scripts/pre-compact-hook.sh
+35
View File
@@ -0,0 +1,35 @@
#!/usr/bin/env bash
# PostToolUse + Stop: sync token usage from the Claude Code transcript.
#
# Claude Code does not pass token counts to hooks, but it does pass the
# path to the session transcript (.jsonl), which records per-message
# `usage`. We hand that path to the SDK server, which parses the transcript
# and SETS the cumulative totals (absolute, idempotent — safe to call after
# every tool and again at Stop). The orchestrator later reads these via
# /usage/status to finalize the spawn-session row and the usage dashboard.
#
# Fire-and-forget — never block Claude on this. Always exit 0.
set -u
SDK_URL="${ROBOCO_SDK_URL:-http://localhost:9000}"
input=$(cat 2>/dev/null || true)
[[ -z "$input" ]] && exit 0
TRANSCRIPT=$(printf '%s' "$input" | python3 - <<'PY'
import json, sys
try:
d = json.loads(sys.stdin.read())
print(d.get("transcript_path", ""))
except Exception:
print("")
PY
)
[[ -z "$TRANSCRIPT" ]] && exit 0
curl -sf -m 3 -X POST "$SDK_URL/usage/sync" \
-H "Content-Type: application/json" \
-d "{\"transcript_path\":\"$TRANSCRIPT\"}" >/dev/null 2>&1 || true
exit 0
+15
View File
@@ -241,3 +241,18 @@ class TokenUsageStatus(BaseModel):
tokens_cache_write: int = Field(
default=0, description="Total cache-write tokens this session"
)
class TranscriptSyncRequest(BaseModel):
"""Payload for POST /usage/sync — hand the SDK the Claude Code transcript.
The hook can't read token counts itself (Claude Code does not pass usage
to hooks), so it passes the ``transcript_path`` instead. The SDK parses
the JSONL transcript, sums the per-message ``usage`` across the session,
and *sets* the cumulative totals absolutely (idempotent — re-syncing the
same or a grown transcript never double-counts).
"""
transcript_path: str = Field(
description="Absolute path to the Claude Code session transcript (.jsonl)"
)
+82 -1
View File
@@ -37,6 +37,7 @@ from roboco.agent_sdk.models import (
TerminalToolRecordRequest,
TokenReportRequest,
TokenUsageStatus,
TranscriptSyncRequest,
VerbAttemptRequest,
VerbCircuitStatus,
)
@@ -422,11 +423,16 @@ class _SessionState:
self.verb_attempts: dict[tuple[str, str | None], deque[float]] = defaultdict(
deque
)
# Cumulative token usage for this session (reported via /usage/report)
# Cumulative token usage for this session. Populated by /usage/sync,
# which parses the Claude Code transcript and *sets* these absolutely
# (the additive /usage/report path remains for explicit deltas).
self.tokens_input: int = 0
self.tokens_output: int = 0
self.tokens_cache_read: int = 0
self.tokens_cache_write: int = 0
# (size, mtime) of the last transcript parsed for /usage/sync, so a
# re-sync of an unchanged transcript skips the re-parse.
self.transcript_fingerprint: tuple[int, float] | None = None
def reset(self) -> None:
self._init_fields()
@@ -735,6 +741,81 @@ def _token_usage_snapshot() -> TokenUsageStatus:
)
def _sum_transcript_usage(path: Path) -> tuple[int, int, int, int]:
"""Sum per-message token usage across a Claude Code JSONL transcript.
Each assistant entry carries a ``message.usage`` block with the token
counts for that API response; summing them yields the session total.
Returns ``(input, output, cache_read, cache_write)``. Malformed lines are
skipped a single bad line must never lose the whole count.
"""
tin = tout = tcr = tcw = 0
with path.open("r", encoding="utf-8", errors="ignore") as fh:
for raw in fh:
stripped = raw.strip()
if not stripped:
continue
try:
entry = json.loads(stripped)
except (ValueError, TypeError):
continue
message = entry.get("message")
if not isinstance(message, dict):
continue
usage = message.get("usage")
if not isinstance(usage, dict):
continue
tin += int(usage.get("input_tokens", 0) or 0)
tout += int(usage.get("output_tokens", 0) or 0)
tcr += int(usage.get("cache_read_input_tokens", 0) or 0)
tcw += int(usage.get("cache_creation_input_tokens", 0) or 0)
return tin, tout, tcr, tcw
@app.post("/usage/sync", response_model=TokenUsageStatus)
async def usage_sync(req: TranscriptSyncRequest) -> TokenUsageStatus:
"""Parse the Claude Code transcript and *set* cumulative token totals.
Claude Code does not pass token usage to hooks, so the usage-report hook
hands us the transcript path and we derive the totals ourselves. The set
is absolute and idempotent: re-syncing the same or a grown transcript
overwrites rather than accumulates, so it is safe to call after every
tool use and again at Stop. An unchanged transcript (same size + mtime)
short-circuits without re-parsing.
"""
path = Path(req.transcript_path)
try:
stat = path.stat()
except OSError:
# Transcript not written yet (very first turn) — nothing to sync.
return _token_usage_snapshot()
fingerprint = (stat.st_size, stat.st_mtime)
if fingerprint == _state.transcript_fingerprint:
return _token_usage_snapshot()
try:
tin, tout, tcr, tcw = _sum_transcript_usage(path)
except OSError as exc:
logger.warning("Transcript usage sync failed", error=str(exc))
return _token_usage_snapshot()
_state.tokens_input = tin
_state.tokens_output = tout
_state.tokens_cache_read = tcr
_state.tokens_cache_write = tcw
_state.transcript_fingerprint = fingerprint
logger.debug(
"Token usage synced from transcript",
tokens_input=tin,
tokens_output=tout,
tokens_cache_read=tcr,
tokens_cache_write=tcw,
)
return _token_usage_snapshot()
@app.post("/journal/post_mortem")
async def journal_post_mortem(req: PostMortemRequest) -> dict[str, str]:
"""SessionEnd hook submits a post-mortem; we log it and flush to the main API."""
+21 -1
View File
@@ -969,6 +969,19 @@ class AgentOrchestrator:
}
],
},
# Sync token usage from the transcript so /usage/status
# (and the cost dashboard) reflect real spend. Idempotent
# absolute set — running it per tool keeps mid-run
# snapshots and reaped-agent sessions accurate.
{
"matcher": "*",
"hooks": [
{
"type": "command",
"command": "/app/scripts/usage-report-hook.sh",
}
],
},
],
# Stop guard: refuse silent exits unless a terminal tool was
# just called (idle/substitute/escalate/pause/...). Second
@@ -979,7 +992,14 @@ class AgentOrchestrator:
{
"type": "command",
"command": "/app/scripts/stop-hook.sh",
}
},
# Final token-usage sync at turn end — guarantees
# the session total is captured before the agent
# idles and the orchestrator finalizes the row.
{
"type": "command",
"command": "/app/scripts/usage-report-hook.sh",
},
]
}
],
+170
View File
@@ -0,0 +1,170 @@
"""Token-usage capture — /usage/sync parses the transcript and sets totals.
The agent SDK exposes /usage/report (additive) and /usage/status (read),
but nothing ever fed token counts in, so every session reported zero and the
cost dashboard rendered all-zeros. The fix: the usage-report hook hands the
SDK the Claude Code transcript path; /usage/sync parses the per-message
``usage`` blocks and *sets* the cumulative totals absolutely. These tests pin
that contract correct summation, idempotency (no double-count on re-sync),
graceful handling of a missing/partial transcript, and growth on re-sync.
Expected totals are derived from the input rows (no magic literals), so the
assertions track whatever the fixtures declare.
"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING
import pytest
import roboco.agent_sdk.server as srv
from fastapi.testclient import TestClient
if TYPE_CHECKING:
from collections.abc import Iterator, Sequence
from pathlib import Path
_OK = 200
# Each row is (input, output, cache_read, cache_write).
_UsageRow = tuple[int, int, int, int]
@pytest.fixture(autouse=True)
def _reset_state() -> Iterator[None]:
srv._state.reset()
yield
srv._state.reset()
@pytest.fixture
def client() -> TestClient:
return TestClient(srv.app)
def _assistant_line(row: _UsageRow) -> str:
inp, out, cread, cwrite = row
return json.dumps(
{
"type": "assistant",
"message": {
"role": "assistant",
"usage": {
"input_tokens": inp,
"output_tokens": out,
"cache_read_input_tokens": cread,
"cache_creation_input_tokens": cwrite,
},
},
}
)
def _write(path: Path, *lines: str) -> None:
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
def _expected(rows: Sequence[_UsageRow]) -> dict[str, int]:
return {
"tokens_input": sum(r[0] for r in rows),
"tokens_output": sum(r[1] for r in rows),
"tokens_cache_read": sum(r[2] for r in rows),
"tokens_cache_write": sum(r[3] for r in rows),
}
def test_sums_usage_across_assistant_messages(
client: TestClient, tmp_path: Path
) -> None:
rows: list[_UsageRow] = [(100, 20, 5, 3), (50, 10, 2, 1)]
transcript = tmp_path / "session.jsonl"
_write(transcript, *(_assistant_line(r) for r in rows))
resp = client.post("/usage/sync", json={"transcript_path": str(transcript)})
assert resp.status_code == _OK
body = resp.json()
for key, value in _expected(rows).items():
assert body[key] == value
def test_status_reflects_synced_totals(client: TestClient, tmp_path: Path) -> None:
rows: list[_UsageRow] = [(200, 40, 0, 0)]
transcript = tmp_path / "session.jsonl"
_write(transcript, *(_assistant_line(r) for r in rows))
client.post("/usage/sync", json={"transcript_path": str(transcript)})
status = client.get("/usage/status").json()
for key, value in _expected(rows).items():
assert status[key] == value
def test_resync_is_idempotent_not_additive(client: TestClient, tmp_path: Path) -> None:
"""The set is absolute — syncing the same transcript twice must not double."""
rows: list[_UsageRow] = [(100, 20, 0, 0)]
transcript = tmp_path / "session.jsonl"
_write(transcript, *(_assistant_line(r) for r in rows))
client.post("/usage/sync", json={"transcript_path": str(transcript)})
client.post("/usage/sync", json={"transcript_path": str(transcript)})
status = client.get("/usage/status").json()
for key, value in _expected(rows).items():
assert status[key] == value
def test_resync_after_growth_overwrites_with_new_total(
client: TestClient, tmp_path: Path
) -> None:
first: list[_UsageRow] = [(100, 20, 0, 0)]
grown: list[_UsageRow] = [(100, 20, 0, 0), (80, 15, 0, 0)]
transcript = tmp_path / "session.jsonl"
_write(transcript, *(_assistant_line(r) for r in first))
client.post("/usage/sync", json={"transcript_path": str(transcript)})
# The transcript grows as the turn continues.
_write(transcript, *(_assistant_line(r) for r in grown))
client.post("/usage/sync", json={"transcript_path": str(transcript)})
status = client.get("/usage/status").json()
for key, value in _expected(grown).items():
assert status[key] == value
def test_missing_transcript_returns_zero_without_error(
client: TestClient, tmp_path: Path
) -> None:
resp = client.post(
"/usage/sync", json={"transcript_path": str(tmp_path / "nope.jsonl")}
)
assert resp.status_code == _OK
assert resp.json() == _expected([])
def test_malformed_lines_are_skipped(client: TestClient, tmp_path: Path) -> None:
rows: list[_UsageRow] = [(100, 20, 0, 0), (50, 10, 0, 0)]
transcript = tmp_path / "session.jsonl"
_write(
transcript,
"not json at all",
_assistant_line(rows[0]),
json.dumps({"type": "user", "message": {"role": "user"}}), # no usage
"{ broken",
_assistant_line(rows[1]),
)
body = client.post("/usage/sync", json={"transcript_path": str(transcript)}).json()
exp = _expected(rows)
assert body["tokens_input"] == exp["tokens_input"]
assert body["tokens_output"] == exp["tokens_output"]
def test_parser_handles_entries_without_message(tmp_path: Path) -> None:
rows: list[_UsageRow] = [(10, 5, 0, 0)]
transcript = tmp_path / "session.jsonl"
_write(
transcript,
json.dumps({"type": "system", "subtype": "init"}),
_assistant_line(rows[0]),
)
tin, tout, cread, cwrite = srv._sum_transcript_usage(transcript)
exp = _expected(rows)
assert (tin, tout, cread, cwrite) == (
exp["tokens_input"],
exp["tokens_output"],
exp["tokens_cache_read"],
exp["tokens_cache_write"],
)