diff --git a/.claude/settings.json b/.claude/settings.json index ddbdce134..a5c951163 100644 --- a/.claude/settings.json +++ b/.claude/settings.json @@ -13,7 +13,7 @@ { "hooks": [ { - "command": "chainloop trace hook claude post-tool-use", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook claude post-tool-use'", "type": "command" } ], @@ -24,7 +24,7 @@ { "hooks": [ { - "command": "chainloop trace hook claude post-tool-use", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook claude post-tool-use'", "type": "command" } ], @@ -35,7 +35,7 @@ { "hooks": [ { - "command": "chainloop trace hook claude pre-tool-use", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook claude pre-tool-use'", "type": "command" } ], @@ -46,7 +46,7 @@ { "hooks": [ { - "command": "chainloop trace hook claude session-end", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook claude session-end'", "type": "command" } ] @@ -56,7 +56,7 @@ { "hooks": [ { - "command": "chainloop trace hook claude session-start", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook claude session-start'", "type": "command" } ] @@ -66,7 +66,7 @@ { "hooks": [ { - "command": "chainloop trace hook claude user-prompt-submit", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook claude user-prompt-submit'", "type": "command" } ] diff --git a/.cursor/hooks.json b/.cursor/hooks.json index 7fd0d796a..defd6e51c 100644 --- a/.cursor/hooks.json +++ b/.cursor/hooks.json @@ -2,33 +2,33 @@ "hooks": { "afterFileEdit": [ { - "command": "chainloop trace hook cursor after-file-edit", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook cursor after-file-edit'", "timeout": 30 } ], "postToolUse": [ { - "command": "chainloop trace hook cursor post-tool-use", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook cursor post-tool-use'", "matcher": "Shell", "timeout": 30 } ], "preToolUse": [ { - "command": "chainloop trace hook cursor pre-tool-use", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then echo \"{\\\"permission\\\":\\\"allow\\\"}\"; exit 0; fi; exec chainloop trace hook cursor pre-tool-use'", "matcher": "Shell", "timeout": 30 } ], "sessionEnd": [ { - "command": "chainloop trace hook cursor session-end", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook cursor session-end'", "timeout": 30 } ], "sessionStart": [ { - "command": "chainloop trace hook cursor session-start", + "command": "sh -c 'if ! command -v chainloop \u003e/dev/null 2\u003e\u00261; then exit 0; fi; exec chainloop trace hook cursor session-start'", "timeout": 30 } ] diff --git a/.opencode/plugins/chainloop-trace.ts b/.opencode/plugins/chainloop-trace.ts index 71f2a458e..2471f07fb 100644 --- a/.opencode/plugins/chainloop-trace.ts +++ b/.opencode/plugins/chainloop-trace.ts @@ -5,20 +5,29 @@ const commandTools = ["bash","shell"] const skillTool = "skill" // skillDirFrom returns the folder of the skill that a skill tool call loaded: -// the dir field of the result metadata, or else the base directory line of -// the result text. Each OpenCode major gives the result in its own shape, so -// each candidate is tried in turn. +// the dir (1.x) or directory (2) field of the result metadata, or else the +// base directory line of the result text. Each OpenCode major gives the +// result in its own shape, so each candidate is tried in turn. function skillDirFrom(...results: any[]): string { for (const r of results) { - const dir = r?.metadata?.dir - if (typeof dir === "string" && dir) return dir - const text = typeof r === "string" ? r : typeof r?.output === "string" ? r.output : "" - const m = /^Base directory for this skill:\s*(.+)$/m.exec(text) + for (const dir of [r?.metadata?.dir, r?.metadata?.directory]) { + if (typeof dir === "string" && dir) return dir + } + const m = /^Base directory for this skill:\s*(.+)$/m.exec(resultText(r)) if (m) return m[1].trim() } return "" } +// resultText returns the text of a tool result: the result itself, its +// output (1.x), or the text items of its content (2). +function resultText(r: any): string { + if (typeof r === "string") return r + if (typeof r?.output === "string") return r.output + if (Array.isArray(r?.content)) return r.content.map((c: any) => (typeof c?.text === "string" ? c.text : "")).join("\n") + return "" +} + function filePathsFromArgs(args: any): string[] { if (args?.filePath) return [args.filePath] if (args?.path) return [args.path] diff --git a/.pi/extensions/chainloop-trace.ts b/.pi/extensions/chainloop-trace.ts new file mode 100644 index 000000000..d68c59b29 --- /dev/null +++ b/.pi/extensions/chainloop-trace.ts @@ -0,0 +1,246 @@ +// chainloop-trace: managed Pi extension +import { spawn } from "node:child_process" +import { homedir } from "node:os" +import { resolve } from "node:path" +import { fileURLToPath } from "node:url" + +const maxOutputBytes = 64 * 1024 +const hookTimeoutMs = 60 * 1000 +const killGraceMs = 5 * 1000 +const instructionMarker = "chainloop-trace-instruction:" +const reminderMarker = "chainloop-trace-spec-reminder" +const treeMarker = "chainloop-trace-tree-marker" +const fileWritingTools = ["edit","write"] +const commandTools = ["bash"] +const sessionEndReasons = new Set(["quit", "new", "resume", "fork"]) +const unicodeSpaces = /[\u00A0\u2000-\u200A\u202F\u205F\u3000]/g + +function resolveToolPath(input, cwd) { + let normalized = input.replace(unicodeSpaces, " ") + if (normalized.startsWith("@")) normalized = normalized.slice(1) + if (normalized === "~") return homedir() + if (normalized.startsWith("~/") || (process.platform === "win32" && normalized.startsWith("~\\"))) { + normalized = resolve(homedir(), normalized.slice(2)) + } + if (/^file:\/\//.test(normalized)) normalized = fileURLToPath(normalized) + return resolve(cwd, normalized) +} + +function sessionID(ctx) { + return ctx.sessionManager.getHeader()?.id ?? ctx.sessionManager.getSessionId() +} + +function hookPayload(ctx, hookEventName) { + const transcriptPath = ctx.sessionManager.getSessionFile() + return { + session_id: sessionID(ctx), + hook_event_name: hookEventName, + cwd: ctx.cwd, + ...(transcriptPath ? { transcript_path: transcriptPath } : {}), + } +} + +function runHook(ctx, event, payload) { + return new Promise((resolveHook) => { + let settled = false + let failed = false + let stdoutBytes = 0 + let stderrBytes = 0 + const stdout = [] + const stderr = [] + let killTimer + let timeoutTimer + let exited = false + let child + + const finish = (response) => { + if (settled) return + settled = true + if (timeoutTimer) clearTimeout(timeoutTimer) + if (killTimer) clearTimeout(killTimer) + resolveHook(response) + } + + const closeStreams = () => { + try { child?.stdin.destroy() } catch {} + try { child?.stdout.destroy() } catch {} + try { child?.stderr.destroy() } catch {} + } + + try { + child = spawn("chainloop", ["trace", "hook", "pi", event], { + cwd: ctx.cwd, + stdio: ["pipe", "pipe", "pipe"], + }) + + const stop = () => { + if (failed) return + failed = true + if (!exited) { + try { child.kill("SIGTERM") } catch {} + } + if (!killTimer) { + killTimer = setTimeout(() => { + if (!exited) { + try { child.kill("SIGKILL") } catch {} + } + closeStreams() + finish() + }, killGraceMs) + } + } + + const capture = (chunks, chunk, current) => { + const remaining = maxOutputBytes - current + if (remaining > 0) chunks.push(chunk.subarray(0, remaining)) + if (chunk.length > remaining) stop() + return Math.min(maxOutputBytes, current + chunk.length) + } + + child.stdout.on("data", (chunk) => { stdoutBytes = capture(stdout, chunk, stdoutBytes) }) + child.stderr.on("data", (chunk) => { stderrBytes = capture(stderr, chunk, stderrBytes) }) + child.on("error", () => { + exited = true + finish() + }) + child.on("exit", () => { exited = true }) + child.on("close", (code) => { + exited = true + if (failed || code !== 0) return finish() + const output = Buffer.concat(stdout).toString("utf8").trim() + if (!output) return finish() + try { + finish(JSON.parse(output)) + } catch { + finish() + } + }) + child.stdin.on("error", () => {}) + child.stdin.end(JSON.stringify(payload)) + timeoutTimer = setTimeout(stop, hookTimeoutMs) + } catch { + finish() + } + }) +} + +function showBanner(ctx, response) { + try { + if (ctx.hasUI && response?.banner) ctx.ui.notify(response.banner, "info") + } catch {} +} + +function showMessage(ctx, response) { + if (!response?.message) return + if (ctx.hasUI) ctx.ui.notify(response.message, "info") + else process.stderr.write(response.message + "\n") +} + +export default function (pi) { + const starts = new Map() + const instructed = new Set() + + const start = async (ctx, hookEventName) => { + const id = sessionID(ctx) + const response = await runHook(ctx, "session-start", hookPayload(ctx, hookEventName)) + if (response?.instruction) starts.set(id, response) + showBanner(ctx, response) + return response + } + + pi.on("session_start", async (event, ctx) => { + try { + await start(ctx, event.type + ":" + event.reason) + } catch {} + }) + + pi.on("session_shutdown", async (event, ctx) => { + if (!sessionEndReasons.has(event.reason)) return + const id = sessionID(ctx) + try { + await runHook(ctx, "session-end", hookPayload(ctx, event.type + ":" + event.reason)) + } catch { + } finally { + starts.delete(id) + instructed.delete(id) + } + }) + + pi.on("session_tree", (event, _ctx) => { + try { + pi.appendEntry(treeMarker, { selected: event.newLeafId }) + } catch {} + }) + + pi.on("before_agent_start", async (event, ctx) => { + try { + const id = sessionID(ctx) + const marker = instructionMarker + id + const persisted = ctx.sessionManager.getEntries().some((entry) => + entry.type === "custom_message" && entry.customType === marker, + ) + const alreadyInstructed = instructed.has(id) || persisted + if (persisted) instructed.add(id) + + const response = starts.get(id) ?? await start(ctx, event.type) + const reminder = await runHook(ctx, "user-prompt-submit", hookPayload(ctx, event.type)) + const content = [] + let includesInstruction = false + + if (!alreadyInstructed && response?.instruction) { + content.push(response.instruction) + includesInstruction = true + } + if (reminder?.instruction) content.push(reminder.instruction) + if (content.length === 0) return + + if (includesInstruction) instructed.add(id) + return { + message: { + customType: includesInstruction ? marker : reminderMarker, + content: content.join("\n\n"), + display: false, + details: { sessionId: id }, + }, + } + } catch { + return + } + }) + + const toolPayload = (ctx, event) => { + if (!fileWritingTools.includes(event.toolName) && !commandTools.includes(event.toolName)) return + const payload = { + ...hookPayload(ctx, event.type), + tool_name: event.toolName, + tool_use_id: event.toolCallId, + ...(event.type === "tool_result" && event.isError ? { tool_failed: true } : {}), + } + if (fileWritingTools.includes(event.toolName)) { + const path = event.input?.path + if (typeof path !== "string" || path.length === 0) return + payload.file_path = resolveToolPath(path, ctx.cwd) + } + if (commandTools.includes(event.toolName) && typeof event.input?.command === "string") { + payload.command = event.input.command + } + return payload + } + + pi.on("tool_call", async (event, ctx) => { + try { + const payload = toolPayload(ctx, event) + if (payload) await runHook(ctx, "pre-tool-use", payload) + } catch {} + return undefined + }) + + pi.on("tool_result", async (event, ctx) => { + try { + const payload = toolPayload(ctx, event) + if (!payload) return + const response = await runHook(ctx, "post-tool-use", payload) + showMessage(ctx, response) + } catch {} + }) +} diff --git a/app/cli/internal/trace/README.md b/app/cli/internal/trace/README.md index 1936c7d0d..f1f648639 100644 --- a/app/cli/internal/trace/README.md +++ b/app/cli/internal/trace/README.md @@ -12,7 +12,9 @@ Pi uses the project-local `.pi/extensions/chainloop-trace.ts` extension. Run `ch Pi supports the JSONL v3 format used by Pi 0.80.10. The compatibility suite pins Pi 1.3.0 with `openrouter/anthropic/claude-haiku-5.5`. Pi cells begin at **Unknown** and move only from end-to-end results or an explicit manual check. -Pi 1.3.0 passed 25 feature checks in the local end-to-end suite. Pi 0.80.10 also completed targeted TUI and JSON-headless smokes that exercised session start, `write`, `edit`, `bash`, commit attribution and session shutdown; the TUI showed the local-mode welcome message. Subagents, skills tracking, spec source pointers and pasted image pointers are not implemented. The headless suite cannot verify the UI-only welcome message, cache accounting without cache activity, the control-plane credentials warning, the session link after a control-plane push or control-plane attestation. The multi-agent feature has no registered check. +Pi evidence preserves native `input`, `output`, `cacheRead`, and `cacheWrite` as separate categories. Evidence `total_tokens` is recomputed as base input + output, excluding both cache categories, regardless of supplied `totalTokens`. The 100/20/10/5 fixture therefore records total 120, not 135. Supplied `cost.total` remains authoritative, without repricing. Invalid or overflowing counters are ignored with deterministic warnings; if input + output overflows, the evidence total is omitted with a warning while individual counters remain. Missing usage contributes zero. Only the selected branch is counted, including its history after resume, not abandoned branches or complete session spend. The Chainloop evidence field `schema_version` remains `v1`; it is not a Pi release or JSONL format version. Earlier Chainloop parsers used supplied totals, which may include cache, under the same evidence version; normalize historical evidence from raw records before combining totals. + +Pi 1.3.0 passed 26 feature checks in the local end-to-end suite, including a cache-active accounting run. Pi 0.80.10 also completed targeted TUI and JSON-headless smokes that exercised session start, `write`, `edit`, `bash`, commit attribution and session shutdown; the TUI showed the local-mode welcome message. Subagents, skills tracking, spec source pointers and pasted image pointers are not implemented. The headless suite cannot verify the UI-only welcome message, the control-plane credentials warning, the session link after a control-plane push or control-plane attestation. The multi-agent feature has no registered check. Codex uses native hooks in `.codex/hooks.json` (Codex CLI 0.163.0). Run `chainloop trace init --codex` or `chainloop trace run --local --codex -- codex exec …`. Review and trust the installed hooks in Codex before starting a session. See the [official hooks documentation](https://learn.chatgpt.com/docs/hooks). Installed hooks have a `Chainloop Trace` status label in Codex. Re-running init adds labels to existing Chainloop hooks. Changing the hooks requires a new trust review. @@ -49,7 +51,7 @@ What the session evidence holds about the conversation. |---------|---------------|-------------|--------|--------------|------------|---|---| | Transcript | The evidence holds the conversation (`raw_session`): the messages and the tool calls of the agent. | Yes | Yes | Partial | Partial | Yes | Yes | | Token usage and cost | The evidence holds the tokens that the session used and their cost in USD. | Yes | No | Yes | Yes | Yes | Yes | -| Cache token accounting | Cache reads and writes stay in separate evidence fields. Base input and total exclude cache. The suite compares stored token counts with the raw usage records. | Unknown | No | Unknown | Unknown | Unknown | Unknown | +| Cache token accounting | Cache reads and writes stay in separate evidence fields. Base input and total exclude cache. The suite compares stored token counts with the raw usage records. | Unknown | No | Unknown | Unknown | Yes | Unknown | | Model detection | The evidence names the primary model and the other models that the session used. | Yes | Partial | Yes | Yes | Yes | Yes | | Subagents | The evidence holds the conversations of the subagents, and their token usage, as part of the session. | Yes | No | No | No | No | No | | Missed session start | A session whose session-start hook did not fire is still recorded, from its first tool hook. | Yes | Unknown | Unknown | Yes | Yes | Unknown | @@ -69,7 +71,7 @@ Why not Yes: - **Subagents, OpenCode:** a subagent runs in a child session with its own export. It is not part of the evidence of its parent session. - **Missed session start, Cursor and OpenCode 1.x:** not verified yet. The end-to-end suite finds the values (spec issue-3635). -The cache accounting check reuses the edit-and-commit session, so it adds no agent runs. It compares the stored evidence with source usage, including unique Codex responses and Claude subagents, and fails if input or total includes cache or if a cache category is lost. The Codex check also rejects inflated usage by comparing with deduplicated responses or accepted cumulative snapshots. The Claude check sums all assistant records, so matching duplicated records in both source and evidence will pass; Claude deduplication remains a gap. A session with no cache activity is reported as not covered. Regression fixtures cover the compaction overcount, duplicate response IDs, resumed sessions, cumulative fallbacks, reasoning, and cache pricing. The live agent matrix must be rerun to record the new check's result. +The cache accounting check reuses the edit-and-commit session, so it adds no agent runs. It compares the stored evidence with source usage, including unique Codex responses and Claude subagents, and fails if input or total includes cache or if a cache category is lost. The Codex check also rejects inflated usage by comparing with deduplicated responses or accepted cumulative snapshots. The Claude check sums all assistant records, so matching duplicated records in both source and evidence will pass; Claude deduplication remains a gap. A session with no cache activity is reported as not covered. Regression fixtures cover the compaction overcount, duplicate response IDs, resumed sessions, cumulative fallbacks, reasoning, and cache pricing. A live cache-active Pi 1.3.0 run passed with `openrouter/anthropic/claude-haiku-5.5`. ## Code attribution diff --git a/app/cli/internal/trace/codex/TOKEN_ACCOUNTING.md b/app/cli/internal/trace/codex/TOKEN_ACCOUNTING.md index 11a60e147..79d96d9a0 100644 --- a/app/cli/internal/trace/codex/TOKEN_ACCOUNTING.md +++ b/app/cli/internal/trace/codex/TOKEN_ACCOUNTING.md @@ -1,6 +1,7 @@ # Codex token accounting -Codex and Claude Code evidence use disjoint token categories: +The current Codex, Pi, and Claude Code parsers use disjoint token categories +in Chainloop evidence: - `input_tokens`: base input, excluding cache reads and cache writes. - `cache_read_input_tokens`: input served from cache. @@ -78,32 +79,34 @@ total tokens, $900.8926 estimated cost, and 147 false reset warnings. ## Other provider parsers -This table describes the current code. This fix changes Codex only. +This table describes the current parsers, not agent release versions. | Provider | Input/cache fields | Total and cost | Duplicate and reasoning handling | | --- | --- | --- | --- | | Claude Code | Base input, cache reads, and cache writes are separate. | Total is base input + output. Cost includes cache reads and both 5-minute and 1-hour writes. | Sums each assistant usage record, including subagents, without message-ID deduplication. No separate reasoning count is added. Duplicate usage records can inflate counts and cost. | | Cursor | The collected transcript supplies no usage. | Usage and cost are unavailable, with a warning. | No token accounting to deduplicate or normalize. | -| Pi parser | Copies input, output, cache reads, and cache writes separately. | Uses the supplied `totalTokens` and `cost.total` without reconciling categories. Its fixture total includes cache. | Rejects duplicate entry IDs and reads only the selected branch. No separate reasoning count is added. | +| Pi parser | Copies native input, output, cache reads, and cache writes separately. | Recomputes total as base input + output, excluding cache; supplied `cost.total` remains authoritative. Earlier Chainloop parsers used supplied totals, which may include cache. | Rejects duplicate entry IDs and reads only the selected branch, including resumed history. No separate reasoning count is added. Independent E2E source comparison has regression coverage and passed a live cache-active Pi 1.3.0 run. | | OpenCode | Copies input, output, and cache read/write fields separately. | Total is input + output. Prefers session counters/cost over message sums. | Message sums have no ID deduplication. A separate source `reasoning` counter is not represented in the parser's token types. | -Pi's selected-branch scope and source totals differ from a complete session +Pi's selected-branch scope differs from a complete session spend report. OpenCode's separate reasoning counter needs a follow-up decision and normalization if it is to be included in output. Those behaviors and Claude's missing deduplication remain outside this Codex fix. ## Existing evidence and display -Codex evidence now sets `data.schema_version` to `v2`. The JSON structure and -schema URL stay the same; the data version identifies the accounting semantics. -Consumers must check both `data.agent.name` and `data.schema_version`: +`data.schema_version` is a Chainloop evidence field, not a Codex or Pi release +version or the Pi JSONL session format version. Codex already emits `v2`; +Pi continues to emit `v1`. This Pi accounting fix changes neither the schema +nor its version. The JSON structure and schema URL stay the same. | Agent and data version | Meaning of input and total | Inclusive processed-token total | | --- | --- | --- | | Codex `v1` | Input includes cache; total includes input and output. | Use `total_tokens` without adding the cache fields. Counts may still contain the old overcount. | | Codex `v2` | Input and total exclude cache, matching the existing Claude convention. | Add both cache fields to `total_tokens` exactly once. | +| Pi `v1` | Input is separate; the fixed parser recomputes total as input + output. Earlier parsers used supplied totals, which may include cache. | Recompute from native input, output, and cache fields when handling historical evidence; the evidence version alone does not distinguish the two behaviors. | -Other providers retain their existing versions and semantics. An unknown +Providers other than Codex and Pi retain their existing versions and semantics. An unknown version requires an explicit compatibility decision; do not infer its meaning from whether cache fields are present. No new evidence fields are needed. Consumers must not treat inclusive input as base input for pricing, or aggregate diff --git a/app/cli/internal/trace/e2e/checks_tokens.go b/app/cli/internal/trace/e2e/checks_tokens.go index 3f0e88387..bb73c4690 100644 --- a/app/cli/internal/trace/e2e/checks_tokens.go +++ b/app/cli/internal/trace/e2e/checks_tokens.go @@ -18,9 +18,11 @@ import ( "encoding/json" "errors" "fmt" + "math/big" "github.com/chainloop-dev/chainloop/app/cli/internal/trace/claude" "github.com/chainloop-dev/chainloop/app/cli/internal/trace/codex" + "github.com/chainloop-dev/chainloop/app/cli/internal/trace/pi" "github.com/chainloop-dev/chainloop/pkg/attestation/crafter/materials/aicodingsession" ) @@ -40,6 +42,40 @@ func checkCacheTokenAccounting(s *Session) error { if err != nil { return err } + case pi.Name: + // Pi raw_session contains only the selected branch. Read native fields + // independently of the provider parser; source totalTokens may include cache. + for _, raw := range d.RawSession["main"] { + var r struct { + Type string `json:"type"` + Message struct { + Role string `json:"role"` + Usage map[string]json.RawMessage `json:"usage"` + } `json:"message"` + } + if json.Unmarshal(raw, &r) != nil || r.Type != "message" || r.Message.Role != "assistant" { + continue + } + for _, field := range []struct { + name string + total *int + }{ + {"input", &want.InputTokens}, {"output", &want.OutputTokens}, + {"cacheRead", &want.CacheReadInputTokens}, {"cacheWrite", &want.CacheCreationInputTokens}, + } { + value, ok := new(big.Rat).SetString(string(r.Message.Usage[field.name])) + if !ok || value.Sign() < 0 || !value.IsInt() || !value.Num().IsInt64() { + continue + } + count := value.Num().Int64() + if count <= int64(int(^uint(0)>>1)-*field.total) { + *field.total += int(count) + } + } + } + if want.OutputTokens > int(^uint(0)>>1)-want.InputTokens { + return errors.New("Pi source input + output token total overflows") + } case claude.Name: for _, records := range d.RawSession { for _, raw := range records { diff --git a/app/cli/internal/trace/e2e/checks_tokens_test.go b/app/cli/internal/trace/e2e/checks_tokens_test.go index 26632297c..1fcc893c9 100644 --- a/app/cli/internal/trace/e2e/checks_tokens_test.go +++ b/app/cli/internal/trace/e2e/checks_tokens_test.go @@ -23,6 +23,7 @@ import ( "github.com/chainloop-dev/chainloop/app/cli/internal/trace/claude" "github.com/chainloop-dev/chainloop/app/cli/internal/trace/codex" + "github.com/chainloop-dev/chainloop/app/cli/internal/trace/pi" "github.com/chainloop-dev/chainloop/pkg/attestation/crafter/materials/aicodingsession" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -110,12 +111,78 @@ func TestClaudeCacheComparisonDoesNotDeduplicate(t *testing.T) { func TestCacheAccountingNotCovered(t *testing.T) { for _, s := range []*Session{ cacheSession("cursor"), + cacheSession(pi.Name, `{"type":"message","message":{"role":"assistant","usage":{"input":100,"output":50}}}`), cacheSession(codex.Name, `{"type":"token_usage_record","payload":{"response_id":"a","usage":{"input_tokens":1000,"output_tokens":50}}}`), } { require.ErrorIs(t, checkCacheTokenAccounting(s), errNotCovered) } } +func TestPiCacheAccountingSourceComparison(t *testing.T) { + for _, tc := range []struct { + name string + change func(*aicodingsession.Data) + }{ + {name: "native disjoint counters"}, + {name: "inclusive total", change: func(d *aicodingsession.Data) { d.Usage.TotalTokens = 1050 }}, + {name: "cache subtracted from base input", change: func(d *aicodingsession.Data) { + d.Usage.InputTokens = 0 + d.Usage.TotalTokens = 50 + }}, + {name: "missing cache writes", change: func(d *aicodingsession.Data) { d.Usage.CacheCreationInputTokens = 0 }}, + {name: "inflated consistent counts", change: func(d *aicodingsession.Data) { + d.Usage.InputTokens = 200 + d.Usage.TotalTokens = 250 + }}, + } { + t.Run(tc.name, func(t *testing.T) { + // Given: selected raw records include a continuation and non-assistant usage. + raw := `{"type":"message","message":{"role":"assistant","usage":{"input":50,"output":25,"cacheRead":400,"cacheWrite":50,"totalTokens":525}}}` + s := cacheSession(pi.Name, raw, raw, + `{"type":"message","message":{"role":"user","usage":{"input":999}}}`, + `{"type":"compaction","message":{"role":"assistant","usage":{"input":999}}}`) + if tc.change != nil { + tc.change(&s.Stored.Evidence.Data) + } + // When + err := checkCacheTokenAccounting(s) + // Then + if tc.change == nil { + require.NoError(t, err) + } else { + require.Error(t, err) + require.NotErrorIs(t, err, errNotCovered) + } + }) + } +} + +func TestPiCacheAccountingIgnoresEvidenceVersion(t *testing.T) { + // Given: Pi accounting comes from native usage, not the evidence version. + s := cacheSession(pi.Name, + `{"type":"message","message":{"role":"assistant","usage":{"input":100,"output":50,"cacheRead":800,"cacheWrite":100,"totalTokens":1050}}}`) + evidence(s).SchemaVersion = "v2" + + // When + err := checkCacheTokenAccounting(s) + + // Then + require.NoError(t, err) +} + +func TestPiCacheAccountingSelectedBranchFixture(t *testing.T) { + // Given: stored evidence retains selected raw records with invalid usage, + // compaction, and both cache categories from the parser fixture. + raw, err := os.ReadFile("../pi/testdata/session-v3.golden.json") + require.NoError(t, err) + s := &Session{Stored: &StoredSession{}} + require.NoError(t, json.Unmarshal(raw, &s.Stored.Evidence)) + // When + err = checkCacheTokenAccounting(s) + // Then + require.NoError(t, err) +} + func TestCacheAccountingCompactedSession(t *testing.T) { raw, err := os.ReadFile("../codex/testdata/compacted-session.jsonl") require.NoError(t, err) diff --git a/app/cli/internal/trace/pi/parse.go b/app/cli/internal/trace/pi/parse.go index b188b0b56..4954227a9 100644 --- a/app/cli/internal/trace/pi/parse.go +++ b/app/cli/internal/trace/pi/parse.go @@ -262,7 +262,6 @@ type evidenceAccumulator struct { outputTokens int cacheRead int cacheWrite int - totalTokens int totalCost float64 models map[string]struct{} primaryModel string @@ -307,6 +306,12 @@ func buildEvidence(parsed *parsedSession, agentVersion string) *aicodingsession. if parsed.TruncatedLine > 0 { acc.warnings = append(acc.warnings, fmt.Sprintf("ignored truncated final Pi session line %d", parsed.TruncatedLine)) } + totalTokens := 0 + if acc.outputTokens > int(^uint(0)>>1)-acc.inputTokens { + acc.warnings = append(acc.warnings, "Pi input + output token total overflows; total_tokens omitted") + } else { + totalTokens = acc.inputTokens + acc.outputTokens + } endedAt := parsed.StartedAt if len(parsed.Selected) > 0 { @@ -336,7 +341,7 @@ func buildEvidence(parsed *parsedSession, agentVersion string) *aicodingsession. Usage: &aicodingsession.Usage{ InputTokens: acc.inputTokens, OutputTokens: acc.outputTokens, - TotalTokens: acc.totalTokens, + TotalTokens: totalTokens, CacheReadInputTokens: acc.cacheRead, CacheCreationInputTokens: acc.cacheWrite, EstimatedCostUSD: roundedCost(acc.totalCost), @@ -449,7 +454,10 @@ func (a *evidenceAccumulator) addUsage(line int, message messageView) { a.addInt(line, "output", usage["output"], &a.outputTokens) a.addInt(line, "cacheRead", usage["cacheRead"], &a.cacheRead) a.addInt(line, "cacheWrite", usage["cacheWrite"], &a.cacheWrite) - a.addInt(line, "totalTokens", usage["totalTokens"], &a.totalTokens) + // Retain diagnostics for invalid source totals, but never use them for + // evidence: Pi totals can include both disjoint cache categories. + var suppliedTotal int + a.addInt(line, "totalTokens", usage["totalTokens"], &suppliedTotal) costRaw, exists := usage["cost"] if !exists || isNull(costRaw) { diff --git a/app/cli/internal/trace/pi/parse_test.go b/app/cli/internal/trace/pi/parse_test.go index 27d777f31..f77c918a4 100644 --- a/app/cli/internal/trace/pi/parse_test.go +++ b/app/cli/internal/trace/pi/parse_test.go @@ -49,6 +49,7 @@ func TestParseSessionEvidence(t *testing.T) { require.NoError(t, err) assert.Equal(t, "pi", result.Data.Agent.Name) + assert.Equal(t, "v1", result.Data.SchemaVersion) assert.Equal(t, "0.80.10", result.Data.Agent.Version) assert.Equal(t, testSessionID, result.Data.Session.ID) assert.Equal(t, "2026-01-02T03:04:00Z", result.Data.Session.StartedAt) @@ -63,7 +64,7 @@ func TestParseSessionEvidence(t *testing.T) { require.NotNil(t, result.Data.Usage) assert.Equal(t, 300, result.Data.Usage.InputTokens) assert.Equal(t, 60, result.Data.Usage.OutputTokens) - assert.Equal(t, 395, result.Data.Usage.TotalTokens) + assert.Equal(t, 360, result.Data.Usage.TotalTokens) assert.Equal(t, 30, result.Data.Usage.CacheReadInputTokens) assert.Equal(t, 7, result.Data.Usage.CacheCreationInputTokens) assert.InDelta(t, 0.3, result.Data.Usage.EstimatedCostUSD, 0.0000001) @@ -209,13 +210,53 @@ func TestParseSessionUsageOverflow(t *testing.T) { require.NoError(t, err) assert.Equal(t, maxInt, result.Data.Usage.InputTokens) assert.Equal(t, 1000, result.Data.Usage.OutputTokens) + assert.Zero(t, result.Data.Usage.TotalTokens) assert.InDelta(t, 1.7e308, result.Data.Usage.EstimatedCostUSD, 1e292) assert.Equal(t, []string{ "Pi entry on line 3 has invalid usage.input; ignored", "Pi entry on line 3 has invalid usage.cost.total; ignored", + "Pi input + output token total overflows; total_tokens omitted", }, result.Data.Warnings) } +func TestParseSessionCacheExclusiveTotals(t *testing.T) { + for _, total := range []string{`135`, `0`, `null`, `"bad"`, `1e400`} { + t.Run(total, func(t *testing.T) { + // Given: native counters are disjoint, regardless of the supplied total. + content := fmt.Sprintf(`{"type":"session","version":3,"id":"pi-session-1","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"} +{"type":"message","id":"first001","parentId":null,"timestamp":"2026-01-02T03:04:01Z","message":{"role":"assistant","usage":{"input":100,"output":20,"cacheRead":10,"cacheWrite":5,"totalTokens":%s,"cost":{"total":0.1234}}}} +`, total) + // When + result, err := parseBytes(t, []byte(content)) + // Then + require.NoError(t, err) + assert.Equal(t, "v1", result.Data.SchemaVersion) + assert.Equal(t, &aicodingsession.Usage{ + InputTokens: 100, OutputTokens: 20, TotalTokens: 120, + CacheReadInputTokens: 10, CacheCreationInputTokens: 5, EstimatedCostUSD: 0.1234, + }, result.Data.Usage) + assert.NoError(t, validateEvidenceData(result.Data)) + }) + } +} + +func TestParseSessionResumedUsage(t *testing.T) { + // Given: a resumed transcript extends the selected branch. + content := `{"type":"session","version":3,"id":"pi-session-1","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"} +{"type":"message","id":"first001","parentId":null,"timestamp":"2026-01-02T03:04:01Z","message":{"role":"assistant","usage":{"input":100,"output":20,"cacheRead":10,"cacheWrite":5,"cost":{"total":0.1}}}} +{"type":"message","id":"resume01","parentId":"first001","timestamp":"2026-01-03T03:04:01Z","message":{"role":"assistant","usage":{"input":30,"output":4,"cacheRead":50,"cacheWrite":2,"totalTokens":86,"cost":{"total":0.2}}}} +` + // When + result, err := parseBytes(t, []byte(content)) + // Then + require.NoError(t, err) + assert.Equal(t, &aicodingsession.Usage{ + InputTokens: 130, OutputTokens: 24, TotalTokens: 154, + CacheReadInputTokens: 60, CacheCreationInputTokens: 7, EstimatedCostUSD: 0.3, + }, result.Data.Usage) + assert.Empty(t, result.Data.Warnings) +} + func parseBytes(t *testing.T, content []byte) (*aicodingsession.Evidence, error) { t.Helper() rawDir := t.TempDir() diff --git a/app/cli/internal/trace/pi/testdata/session-v3.golden.json b/app/cli/internal/trace/pi/testdata/session-v3.golden.json index eff98a6e2..586c7db39 100644 --- a/app/cli/internal/trace/pi/testdata/session-v3.golden.json +++ b/app/cli/internal/trace/pi/testdata/session-v3.golden.json @@ -25,7 +25,7 @@ "usage": { "input_tokens": 300, "output_tokens": 60, - "total_tokens": 395, + "total_tokens": 360, "cache_read_input_tokens": 30, "cache_creation_input_tokens": 7, "estimated_cost_usd": 0.3