Skip to content

feat(usage): track tokens, cost, and turns from Claude Code session JSONL#97

Closed
manzil-infinity180 wants to merge 11 commits into
devfrom
fix/jsonl-token-tracking-96
Closed

feat(usage): track tokens, cost, and turns from Claude Code session JSONL#97
manzil-infinity180 wants to merge 11 commits into
devfrom
fix/jsonl-token-tracking-96

Conversation

@manzil-infinity180

@manzil-infinity180 manzil-infinity180 commented May 6, 2026

Copy link
Copy Markdown
Contributor

Closes #96. Rebased onto dev.

aflock's MCP transport carries no token usage or turn-boundary signal, so policies with maxSpendUSD / maxTokens* / maxTurns were silently bypassed. Claude Code's session JSONL has the data and internal/usage now parses it incrementally with byte-offset memoization — called per tool call on the MCP path, at SessionEnd to settle the trailing turn, and at verify time when metrics look stale.

Dedup follows CodexBar (steipete/CodexBarCostUsageScanner+Claude.swift): key is (message.id, requestId) with overwrite-last semantics so cumulative streaming chunks converge on the final snapshot instead of the first chunk's smaller usage. Rows missing either ID stay as separate entries.

Pricing: hardcoded for opus/sonnet/haiku with AFLOCK_PRICE_* overrides; cache-tier weights applied; unknown model warns and zeroes.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds token, cost (USD), and assistant-turn tracking for MCP sessions by incrementally parsing Claude Code’s session JSONL transcript, enabling enforcement and post-hoc verification of policy limits like maxSpendUSD, maxTokensIn/Out, and maxTurns (Closes #96).

Changes:

  • Persist Claude Code transcript_path in session state so usage can be re-read later (e.g., after interrupted sessions).
  • Add internal/usage JSONL tracker + pricing logic (with env overrides) to compute tokens/turns/cost from transcript deltas.
  • Integrate usage delta application into the MCP server action-record path and add best-effort “final settle” at SessionEnd, plus a verify-time stale-metrics repair path.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 8 comments.

Show a summary per file
File Description
pkg/aflock/types.go Persists TranscriptPath in session state for later usage repair/verification.
internal/mcp/server.go Initializes and applies JSONL usage deltas into session metrics during MCP tool execution flow.
internal/hooks/handler.go Settles trailing JSONL usage at SessionEnd before post-hoc limit evaluation.
internal/verify/verifier.go Attempts to repair stale metrics by re-reading transcript JSONL during aflock verify.
internal/usage/jsonl.go New incremental JSONL parser/tracker with byte-offset memoization.
internal/usage/pricing.go New pricing table + env overrides and cost computation helpers.
internal/usage/jsonl_test.go Tests for delta reading, offset persistence, truncation handling, corrupt lines, etc.
internal/usage/pricing_test.go Tests for pricing math and env override behavior.
internal/usage/testdata/sample.jsonl Sample Claude Code JSONL fixture used by tracker tests.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread internal/usage/jsonl.go
Comment on lines +198 to +215
entry := jsonlEntry{}
if err := json.Unmarshal(line, &entry); err != nil {
continue
}
if entry.Type != "assistant" {
continue
}
if entry.Message == nil || entry.Message.Usage == nil {
continue
}
c.InputTokens += entry.Message.Usage.InputTokens
c.OutputTokens += entry.Message.Usage.OutputTokens
c.CacheReadInputTokens += entry.Message.Usage.CacheReadInputTokens
c.CacheCreationInputTokens += entry.Message.Usage.CacheCreationInputTokens
c.AssistantTurns++
if entry.Message.Model != "" {
c.Model = entry.Message.Model
}
Comment thread internal/usage/jsonl.go
Comment on lines +189 to +215
scanner := bufio.NewScanner(bytes.NewReader(b))
// Claude messages can be long (full prompts + responses) — give the
// scanner enough buffer to handle them.
scanner.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
for scanner.Scan() {
line := scanner.Bytes()
if len(bytes.TrimSpace(line)) == 0 {
continue
}
entry := jsonlEntry{}
if err := json.Unmarshal(line, &entry); err != nil {
continue
}
if entry.Type != "assistant" {
continue
}
if entry.Message == nil || entry.Message.Usage == nil {
continue
}
c.InputTokens += entry.Message.Usage.InputTokens
c.OutputTokens += entry.Message.Usage.OutputTokens
c.CacheReadInputTokens += entry.Message.Usage.CacheReadInputTokens
c.CacheCreationInputTokens += entry.Message.Usage.CacheCreationInputTokens
c.AssistantTurns++
if entry.Message.Model != "" {
c.Model = entry.Message.Model
}
Comment thread internal/usage/jsonl.go
Comment on lines +99 to +106
func (t *Tracker) saveOffset() {
if t.offsetPath == "" {
return
}
// Write best-effort. Failure is non-fatal — at worst we re-process
// already-consumed lines on restart, which is correct (just wasteful).
_ = os.WriteFile(t.offsetPath, []byte(strconv.FormatInt(t.offset, 10)), 0o600)
}
Comment on lines +401 to +417
tracker := usage.NewTracker(sessionState.TranscriptPath, v.stateManager.SessionDir(sessionState.SessionID))
delta, err := tracker.ReadDelta()
if err != nil {
return
}
if delta.AssistantTurns == 0 && delta.InputTokens == 0 && delta.OutputTokens == 0 {
return
}
cost := usage.ComputeCostUSD(delta)
v.stateManager.UpdateMetrics(sessionState, delta.InputTokens, delta.OutputTokens, cost)
for i := 0; i < delta.AssistantTurns; i++ {
v.stateManager.IncrementTurns(sessionState)
}
// Persist the repair so future verify runs (or status commands) see the
// fixed metrics. Best-effort.
_ = v.stateManager.Save(sessionState)
}
Comment thread internal/hooks/handler.go
Comment on lines +872 to +876
if input == nil || input.TranscriptPath == "" || sessionState == nil {
return
}
sessionDir := h.stateManager.SessionDir(input.SessionID)
tracker := usage.NewTracker(input.TranscriptPath, sessionDir)
Comment thread internal/mcp/server.go
Comment on lines +1248 to +1261
// applyUsageDelta polls the JSONL tracker for new token/turn data since the
// last call and folds it into sessionState.Metrics. No-op when no tracker is
// active (non-Claude MCP client, or pre-discovery). Errors are logged once
// at debug-ish level and never block tool dispatch — usage tracking is
// best-effort instrumentation, not a security gate.
func (s *Server) applyUsageDelta(sessionState *aflock.SessionState) {
if s.usageTracker == nil || sessionState == nil {
return
}
delta, err := s.usageTracker.ReadDelta()
if err != nil {
fmt.Fprintf(os.Stderr, "[aflock] Warning: usage tracker read failed: %v\n", err)
return
}
Comment on lines +1 to +3
{"parentUuid":null,"isSidechain":false,"promptId":"42a2aba2-15e9-4d02-8317-4e9be46b8e68","type":"user","message":{"role":"user","content":"Use the aflock-uds tool get_token to mint a JWT for this session and show me the token"},"uuid":"97f71efb-900a-4b0b-b9ba-d6ad1ecd53a8","timestamp":"2026-05-05T14:30:23.727Z","permissionMode":"default","userType":"external","entrypoint":"cli","cwd":"/Users/rahulxf/work-dir/aflock-example/peer-cred-pr88/workspace/run-mcp-spire","sessionId":"3512a4a5-e0ec-4cb1-b0c5-e3855c059b48","version":"2.1.128","gitBranch":"main"}
{"parentUuid":"afc8931f-cf1d-416d-aed3-95b6d1511c09","isSidechain":false,"message":{"model":"claude-opus-4-7","id":"msg_01Q1FKGMvvi7gQNfrGFdjXQA","type":"message","role":"assistant","content":[{"type":"thinking","thinking":"","signature":"ErcCClkIDRgCKkBxW39xIsS3s+vAKPPHg+jLxMIXOrl0TyP2XMYkOsxSoVc6v9anGbnY7nA5TFP8KjiWWJV8xTn03hdxc6xjZz5JMg9jbGF1ZGUtb3B1cy00LTc4ABIMQjIjt9ooBZ7uPSkNGgwTrWZsy4+aC4y/Qp4iMEpdvYiaSKxWijWky/eN0j7HpDh+aiAIq3tNS/RSix7wO6FR+0wZojUAillCf8KU2SqLAb1YOGNhWC7UuZotDJVIXB1r4fbjtz61zuiFMi9hIk9Xj3VVwRyz5ry6JMDn7kmzbub+rc0XklcAMKYcKv/h0GALfEkz9WDLhaArCfxGHOjA9dAvCQsWtfLKRZX6PvwVOLKdKj6oyHxHnBIp9q1uP3unnj6/enhlSooaT75hkoaBW1qgSUN77zJHjCwYAQ=="}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":8791,"cache_read_input_tokens":11595,"output_tokens":171,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":8791,"ephemeral_5m_input_tokens":0},"inference_geo":"","iterations":[{"input_tokens":6,"output_tokens":171,"cache_read_input_tokens":11595,"cache_creation_input_tokens":8791,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":8791},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011CajajMPHiCdZqrhDdQnHD","type":"assistant","uuid":"3f990e43-98c9-43d4-9348-6c62ac629500","timestamp":"2026-05-05T14:30:26.729Z","userType":"external","entrypoint":"cli","cwd":"/Users/rahulxf/work-dir/aflock-example/peer-cred-pr88/workspace/run-mcp-spire","sessionId":"3512a4a5-e0ec-4cb1-b0c5-e3855c059b48","version":"2.1.128","gitBranch":"main"}
{"parentUuid":"3f990e43-98c9-43d4-9348-6c62ac629500","isSidechain":false,"message":{"model":"claude-opus-4-7","id":"msg_01Q1FKGMvvi7gQNfrGFdjXQA","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01UdgrhkjgtFfZCLGGPaysvP","name":"ToolSearch","input":{"query":"select:mcp__aflock-uds__get_token","max_results":1},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":8791,"cache_read_input_tokens":11595,"output_tokens":171,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":8791,"ephemeral_5m_input_tokens":0},"inference_geo":"","iterations":[{"input_tokens":6,"output_tokens":171,"cache_read_input_tokens":11595,"cache_creation_input_tokens":8791,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":8791},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011CajajMPHiCdZqrhDdQnHD","type":"assistant","uuid":"c7a4d0fe-6f6a-4a72-8128-0ad5d4bfd675","timestamp":"2026-05-05T14:30:27.015Z","userType":"external","entrypoint":"cli","cwd":"/Users/rahulxf/work-dir/aflock-example/peer-cred-pr88/workspace/run-mcp-spire","sessionId":"3512a4a5-e0ec-4cb1-b0c5-e3855c059b48","version":"2.1.128","gitBranch":"main"}
Comment thread internal/usage/pricing.go
Comment on lines +92 to +117
// ComputeCostUSD returns the USD cost of c using the model's pricing.
// If c.Model is unknown (no defaults, no env overrides), logs a one-time
// WARNING and returns 0 — we deliberately do not guess.
func ComputeCostUSD(c Cumulative) float64 {
if c.Model == "" {
return 0
}
p, known := ResolvePricing(c.Model)
if !known {
unknownLogMu.Lock()
if !loggedUnknown[c.Model] {
loggedUnknown[c.Model] = true
fmt.Fprintf(os.Stderr,
"[aflock] WARNING: unknown model %q in pricing table — costUSD will be 0. "+
"Set AFLOCK_PRICE_INPUT_%s / AFLOCK_PRICE_OUTPUT_%s to override.\n",
c.Model, modelSlug(c.Model), modelSlug(c.Model))
}
unknownLogMu.Unlock()
return 0
}
cost := 0.0
cost += float64(c.InputTokens) * p.InputPerMTok / 1_000_000
cost += float64(c.OutputTokens) * p.OutputPerMTok / 1_000_000
cost += float64(c.CacheCreationInputTokens) * p.InputPerMTok * p.CacheWriteMult / 1_000_000
cost += float64(c.CacheReadInputTokens) * p.InputPerMTok * p.CacheReadMult / 1_000_000
return cost
@manzil-infinity180

manzil-infinity180 commented May 6, 2026

Copy link
Copy Markdown
Contributor Author

Aflock TUI

Claude /usage


You could see some big difference here in Cost. Need to do some investigation for it
Actually i am reading from .jsonl even after their is getting some difference

@manzil-infinity180

manzil-infinity180 commented May 6, 2026

Copy link
Copy Markdown
Contributor Author

We are getting completely different in/out token + Cost :(

@manzil-infinity180

Copy link
Copy Markdown
Contributor Author

Aflock TUI

## Claude /usage You could see some big difference here in Cost. Need to do some investigation for it Actually i am reading from `.jsonl ` even after their is getting some difference

Tested live. Token counts now match the JSONL exactly (input 33 in both aflock and /usage). Dedup by message.id, 5m/1h cache tier split, and mtime-filtered JSONL discovery all working.

Cost still diverges ($0.94 vs /usage $0.24) — claude-code makes internal LLM calls (haiku title gen, compaction) that aren't in the JSONL, and Max subscription seems to use rates we don't have.

Holding off on merging. Fixes are real wins but want to think more before shipping with the cost-divergence as a known limitation.

@manzil-infinity180

Copy link
Copy Markdown
Contributor Author

Check #111

manzil-infinity180 and others added 7 commits May 17, 2026 11:00
…SONL

aflock's MCP transport doesn't carry the agent's token usage or any
turn-boundary signal, so policies with maxSpendUSD/maxTokens*/maxTurns
were silently bypassed: counters stayed at zero, limits never fired.

Claude Code writes the same data to its session JSONL (the file
aflock already locates at startup), with message.usage.{input_tokens,
output_tokens, cache_*} per assistant turn. New internal/usage parses
that JSONL incrementally with offset memoization (sidecar persistence
for crash-restart correctness), called from three places: per tool
call on the MCP path, on SessionEnd to settle the trailing assistant
turn, and at verify time when state metrics look stale (e.g. Ctrl-C
bypassed SessionEnd).

Pricing is hardcoded for opus/sonnet/haiku with AFLOCK_PRICE_*_<MODEL>
env overrides, applies cache-tier weights (writes 1.25x, reads 0.10x),
and warns + zeroes when the model is unknown rather than guessing.

Signed-off-by: Rahul Vishwakarma <rahulvs2809@gmail.com>
Resolve struct field alignment formatting that golangci-lint flagged on
the PR — gofmt prefers tighter padding when the surrounding lines
become uniform.

Signed-off-by: Rahul Vishwakarma <rahulvs2809@gmail.com>
Claude Code emits the same assistant message as multiple JSONL lines
(different parentUuid, identical message.id and identical usage block)
when the message participates in a tool-use chain. Counting each line
double-counted tokens by ~40%.

Tracker now keeps a per-instance set of seen message.ids and skips
duplicates. Parity against jq-deduped ground truth is exact on a real
session: 10 turns / 35 input / 1659 output (was 14 / 54 / 2586).

Signed-off-by: Rahul Vishwakarma <rahulvs2809@gmail.com>
Anthropic prices the two prompt-cache write tiers differently:
5-minute ephemeral at 1.25x input, 1-hour ephemeral at 2.00x input.
We were applying a single 1.25x multiplier to both, undercharging
1-hour cache creations by 60%.

Pricing struct now carries CacheWrite5mMult and CacheWrite1hMult.
ComputeCostUSD reads the per-tier breakdown when available, falling
back to the scalar (credited to 5m, Anthropic's pre-1h-tier default)
on older JSONLs. Env overrides exist for both new tiers; the legacy
AFLOCK_PRICE_CACHE_WRITE_MULT_<MODEL> still works and sets both.

Signed-off-by: Rahul Vishwakarma <rahulvs2809@gmail.com>
Eager discovery in initAgentIdentity often runs before claude has
written its first JSONL line — claude spawns MCP servers eagerly but
writes the JSONL only after the first user prompt. Result: tracker
stayed nil for the whole session and metrics never populated.

applyUsageDelta now retries discovery on every tool call until a
JSONL appears, then propagates the path to sessionState.TranscriptPath
for verify-side fallback.

repairStaleMetricsFromJSONL no longer gates on tokensIn==0 — that
heuristic missed sessions where mid-session metrics had populated
but the trailing assistant turn arrived after the last tool call.
ReadDelta is idempotent via byte-offset memoization, so re-running
on already-consumed JSONL is a no-op anyway.

Signed-off-by: Rahul Vishwakarma <rahulvs2809@gmail.com>
Init-time discovery was caching a stale JSONL path from a previous
claude session. Filter to files modified at or after this session's
StartedAt so we lock onto the current session.

Signed-off-by: Rahul Vishwakarma <rahulvs2809@gmail.com>
PR #97's original parser keyed dedup on message.id alone with skip-first
semantics, which under-counted when claude-code emits the SAME assistant
turn as multiple streaming chunks where each chunk's usage block is
cumulative. The first row was kept and all growth dropped.

Now follows CodexBar's approach (steipete/CodexBar —
CostUsageScanner+Claude.swift): dedup key is (message.id, requestId), and
the map is overwritten on every occurrence so the final cumulative
snapshot wins. Rows missing either ID stay as separate entries to avoid
silently dropping usage.

ReadDelta diffs the post-merge total against a baseline so callers fold
only the growth into session metrics when a streaming row is updated
across multiple reads — no double-counting.

New tests cover: streaming chunks where output grows across reads, same
message.id with different requestId (must both count), and a streaming
row updated across two ReadDelta calls.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown

Docs Preview

Status Deployed
Branch fix/jsonl-token-tracking-96
Preview URL https://fix-jsonl-token-tracking-96.aflock-d0m.pages.dev

@manzil-infinity180

Copy link
Copy Markdown
Contributor Author

manzil-infinity180 and others added 4 commits May 17, 2026 11:11
Validating against four real ~/.claude/projects/ JSONLs showed every
token aggregate matched jq's independent group-by exactly (313 deduped
assistant turns across the four files), but Cumulative.Model surfaced
"<synthetic>" on a file that contained a single claude-code-injected
compaction line alongside 332 real claude-opus-4-7 turns.

Map iteration order in computeTotal is non-deterministic, so the
synthetic line could win the last-write. Cumulative.Model feeds pricing
lookup, so propagating the sentinel would either fail the lookup or
yield a nonsensical cost. parseRows now normalizes the sentinel to ""
so only real inference models reach the cumulative.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
handleSessionEnd called settleFinalUsage to fold the trailing assistant
turn into sessionState.Metrics in memory, but Save fired only when a
post-hoc limit was exceeded — so sessions that stayed within limits
dropped their JSONL-derived token / cost / turn counts on the floor.
A live test on /tmp/aflock-usage-live-test/ exposed it: usage.offset
advanced to EOF (the Tracker persists its own offset inside ReadDelta)
but state.json kept the pre-SessionEnd zero snapshot.

Now: always Save after settleFinalUsage and the post-hoc check. The
limit-exceeded branch no longer needs its own Save.

Also persist input.TranscriptPath into sessionState.TranscriptPath in
hooks mode. MCP mode already sets it in server.go; hooks mode left it
empty, so verify-time repairStaleMetricsFromJSONL could not act on a
Ctrl-C'd session.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
In MCP-only deployments (no aflock hooks configured), applyUsageDelta
fires inside recordAction once per tool call. After the LAST tool call,
claude's wrap-up assistant turns land in the JSONL but no tool call
follows to trigger another read — so state.json silently undercounts.

A live MCP+SPIRE test exposed this: state.json reported 12 turns / 37
input / 2517 output while jq over the same JSONL showed 14 / 44 / 2730.
The missing rows were post-last-tool-call wrap-up messages.

Adds settleFinalUsageOnExit and defers it from Serve, ServeUnix, and
ServeHTTP so transport close picks up the trailing rows. Idempotent
with handler.settleFinalUsage at SessionEnd via the Tracker's byte-
offset memoization — running both in hooks+MCP deployments is safe.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Live MCP+SPIRE test exposed a 1-turn undercount: claude-code wrote a
"No response requested." row with model="<synthetic>" and zero usage,
34 seconds AFTER the last real assistant turn — landing in the JSONL
after the user's /exit and after aflock's transport-close defer had
already fired.

Synthetic rows are claude-code internal placeholders, not real LLM
calls. They carry zero tokens and zero cost, so counting them as turns
inflates maxTurns enforcement on noise without contributing to billing.
Skip them in parseRows so:

  - turns reflects real API calls (correct semantics for maxTurns)
  - the post-/exit race becomes moot — the synthetic row wouldn't count
    even if we DID catch it
  - normalizeModel can be removed: synthetic never reaches row creation

Updates the synthetic-row test to assert zero contribution. Existing
real-JSONL fixture (1.7MB, 333 assistant rows, 1 synthetic) still
matches jq when both filter synthetic.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: implement token/cost/turn tracking via Claude Code session JSONL

2 participants