mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
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:
@@ -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
|
||||
|
||||
Executable
+35
@@ -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
|
||||
@@ -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)"
|
||||
)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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",
|
||||
},
|
||||
]
|
||||
}
|
||||
],
|
||||
|
||||
@@ -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"],
|
||||
)
|
||||
Reference in New Issue
Block a user