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:
@@ -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",
|
||||
},
|
||||
]
|
||||
}
|
||||
],
|
||||
|
||||
Reference in New Issue
Block a user