feat(usage): track tokens, cost, and turns from Claude Code session JSONL#97
Closed
manzil-infinity180 wants to merge 11 commits into
Closed
feat(usage): track tokens, cost, and turns from Claude Code session JSONL#97manzil-infinity180 wants to merge 11 commits into
manzil-infinity180 wants to merge 11 commits into
Conversation
There was a problem hiding this comment.
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_pathin session state so usage can be re-read later (e.g., after interrupted sessions). - Add
internal/usageJSONL 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 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 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 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 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 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 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 |
Contributor
Author
Contributor
Author
|
We are getting completely different in/out token + Cost :( |
Contributor
Author
Contributor
Author
|
Check #111 |
…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>
b927096 to
84887a2
Compare
Docs Preview
|
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>
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.




Closes #96. Rebased onto
dev.aflock's MCP transport carries no token usage or turn-boundary signal, so policies with
maxSpendUSD/maxTokens*/maxTurnswere silently bypassed. Claude Code's session JSONL has the data andinternal/usagenow parses it incrementally with byte-offset memoization — called per tool call on the MCP path, atSessionEndto settle the trailing turn, and at verify time when metrics look stale.Dedup follows CodexBar (
steipete/CodexBar—CostUsageScanner+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.