diff --git a/apps/framework/scripts/upload-braintrust.test.ts b/apps/framework/scripts/upload-braintrust.test.ts index a33a154c..ef6cb4ec 100644 --- a/apps/framework/scripts/upload-braintrust.test.ts +++ b/apps/framework/scripts/upload-braintrust.test.ts @@ -1,7 +1,10 @@ import { describe, expect, it } from 'vitest'; import { experimentName, + logTranscript, runViewUrl, + unwrapShell, + type SpanSink, tokenMetrics, utcStamp, } from './upload-braintrust.js'; @@ -89,3 +92,157 @@ describe('runViewUrl', () => { ); }); }); + +interface RecordedSpan { + name?: string; + start?: number; + end?: number; + metrics?: unknown; + input?: unknown; + children: RecordedSpan[]; +} + +function recorder(node: RecordedSpan): SpanSink { + return { + startSpan(args) { + const child: RecordedSpan = { + name: args?.name, + start: args?.startTime, + children: [], + }; + node.children.push(child); + return recorder(child); + }, + log(event) { + if (event.metrics) { + node.metrics = event.metrics; + } + if (event.input !== undefined && node.name === 'm') { + node.input = event.input; + } + }, + end(args) { + node.end = args?.endTime; + }, + }; +} + +describe('logTranscript', () => { + it('logs back-to-back LLM spans per request, with tool calls as siblings under task', () => { + const root: RecordedSpan = { children: [] }; + const ms = (s: number) => (100 + s) * 1000; + logTranscript(recorder(root), { + prompt: 'go', + agentReport: 'done', + checks: [], + passed: true, + modelId: 'm', + startTime: 100, + endTime: 170, + toolLabels: [], + transcript: [ + { + type: 'message', + role: 'assistant', + content: 'a', + ts: ms(4), + requestId: 'r1', + usage: { + inputTokens: 10, + cacheReadInputTokens: 0, + cacheWriteInputTokens: 0, + outputTokens: 2, + }, + }, + { + type: 'tool_call', + name: 'Skill', + input: {}, + output: 'loaded', + id: 't1', + ts: ms(4), + resultTs: ms(5), + requestId: 'r1', + }, + { + type: 'tool_call', + name: 'Bash', + input: {}, + output: 'ok', + id: 't2', + ts: ms(8), + resultTs: ms(31), + requestId: 'r2', + }, + { + type: 'message', + role: 'assistant', + content: 'b', + ts: ms(38), + requestId: 'r3', + }, + ], + }); + const span = ( + name: string, + start: number, + end: number, + ...children: RecordedSpan[] + ) => ({ + name, + start: 100 + start, + end: 100 + end, + children, + }); + const call = (id: string, name: string) => ({ + id, + type: 'function', + function: { name, arguments: '{}' }, + }); + const prompt = { role: 'user', content: 'go' }; + const first = { + role: 'assistant', + content: 'a', + tool_calls: [call('t1', 'Skill')], + }; + const skillResult = { role: 'tool', tool_call_id: 't1', content: 'loaded' }; + const second = { + role: 'assistant', + content: '', + tool_calls: [call('t2', 'Bash')], + }; + const bashResult = { role: 'tool', tool_call_id: 't2', content: 'ok' }; + expect(root.children).toEqual([ + span( + 'task', + 0, + 38, + { + ...span('m', 0, 4), + input: [prompt], + metrics: { prompt_tokens: 10, completion_tokens: 2, tokens: 12 }, + }, + span('Skill', 4, 5), + { ...span('m', 5, 8), input: [prompt, first, skillResult] }, + span('Bash', 8, 31), + { + ...span('m', 31, 38), + input: [prompt, first, skillResult, second, bashResult], + } + ), + span('passed', 70, 70), + ]); + }); +}); + +describe('unwrapShell', () => { + it("drops Codex's bash -lc wrapper and leaves other commands alone", () => { + expect(unwrapShell(`/bin/bash -lc "ls -la && cat 'a b'"`)).toBe( + "ls -la && cat 'a b'" + ); + expect(unwrapShell(`/bin/bash -lc "apply_patch <<'PATCH' x PATCH'`)).toBe( + "apply_patch <<'PATCH' x PATCH" + ); + expect(unwrapShell('supabase status')).toBe('supabase status'); + }); +}); diff --git a/apps/framework/scripts/upload-braintrust.ts b/apps/framework/scripts/upload-braintrust.ts index a88ec419..b0e99576 100644 --- a/apps/framework/scripts/upload-braintrust.ts +++ b/apps/framework/scripts/upload-braintrust.ts @@ -10,7 +10,10 @@ import { appendFile, stat } from 'node:fs/promises'; import { basename, dirname, join, resolve } from 'node:path'; import { fileURLToPath } from 'node:url'; import { z } from 'zod'; -import type { AgentUsage } from '@supabase-evals/core/eval-metadata'; +import { + modelUsageSchema, + type AgentUsage, +} from '@supabase-evals/core/eval-metadata'; import { normalizeExperimentName, readFlag, @@ -37,11 +40,15 @@ const DRY = rawArgs.includes('--dry'); // `eval:upload` passes this cutoff to select files written by its run. const SINCE = Number(readFlag(rawArgs, 'since') ?? 0); +const requestUsageSchema = modelUsageSchema.omit({ model: true }); const transcriptPartSchema = z.discriminatedUnion('type', [ z.object({ type: z.literal('message'), role: z.enum(['system', 'user', 'assistant']), content: z.string(), + ts: z.number().optional(), + requestId: z.string().optional(), + usage: requestUsageSchema.optional(), }), z.object({ type: z.literal('tool_call'), @@ -49,6 +56,11 @@ const transcriptPartSchema = z.discriminatedUnion('type', [ input: z.record(z.string(), z.unknown()).default({}), output: z.unknown().optional(), error: z.string().optional(), + ts: z.number().optional(), + resultTs: z.number().optional(), + id: z.string().optional(), + requestId: z.string().optional(), + usage: requestUsageSchema.optional(), }), ]); const transcriptSchema = z.array(transcriptPartSchema).catch([]); @@ -61,6 +73,7 @@ interface PendingRow { passed: boolean; checks: unknown; modelId?: string; + modelProvider?: string; transcript: TranscriptPart[]; toolLabels: (string | undefined)[]; metadata: Record; @@ -198,10 +211,18 @@ function toolLabels(toolCalls: unknown): (string | undefined)[] { if (path) { return basename(path); } - return command ? summarize(command) : url; + return command ? summarize(unwrapShell(command)) : url; }); } +// Codex runs every command as `/bin/bash -lc ""`, sometimes joining +// mixed quote segments (`"apply_patch <<'PATCH' … PATCH'`), so only the outer +// quote on each side is dropped. +export function unwrapShell(command: string): string { + const match = command.match(/^\/bin\/(?:ba|z)?sh -lc ([\s\S]*)$/); + return match ? match[1].replace(/^["']|["']$/g, '') : command; +} + function summarize(command: string): string { const line = command.trim().replace(/\s+/g, ' '); return line.length > 60 ? `${line.slice(0, 59)}…` : line; @@ -252,6 +273,7 @@ async function collectRows( continue; } const endTime = durationMs ? mtimeMs / 1000 : undefined; + const transcript = transcriptSchema.parse(result.transcript); const sessionArchivePath = join( dirname(absolutePath), 'session-archive.tar.gz' @@ -265,7 +287,8 @@ async function collectRows( passed: result.passed === true, checks: result.checks, modelId: display?.modelId, - transcript: transcriptSchema.parse(result.transcript), + modelProvider: display?.modelProvider, + transcript, toolLabels: toolLabels(result.toolCalls), sessionArchivePath: existsSync(sessionArchivePath) ? sessionArchivePath @@ -300,9 +323,12 @@ async function collectRows( ? { tool_call_count: result.toolCallCount } : {}), }, - ...(endTime && durationMs - ? { startTime: endTime - durationMs / 1000, endTime } - : {}), + ...runWindow( + transcript, + endTime && durationMs + ? { startTime: endTime - durationMs / 1000, endTime } + : {} + ), }; const rows = byExperiment.get(result.experiment) ?? []; @@ -312,67 +338,231 @@ async function collectRows( return byExperiment; } -function logTranscript(parent: Span, row: PendingRow): void { - // Parsers lack per-entry timestamps, so keep child spans at the run start. - const at = row.startTime ? { startTime: row.startTime } : {}; - const ended = row.startTime ? { endTime: row.startTime } : undefined; +/** + * Widens the mtime-derived window to cover transcript timestamps, since mtime + * lands after scoring and the agent clock can drift from the host's. + */ +function runWindow( + transcript: TranscriptPart[], + derived: { startTime?: number; endTime?: number } +): { startTime?: number; endTime?: number } { + const times = transcript + .flatMap((part) => [ + part.ts, + part.type === 'tool_call' ? part.resultTs : undefined, + ]) + .filter((ms) => ms !== undefined) + .map((ms) => ms / 1000); + if (!times.length) { + return derived; + } + return { + startTime: Math.min(...times, derived.startTime ?? Infinity), + endTime: Math.max(...times, derived.endTime ?? -Infinity), + }; +} + +type StartSpanArgs = Parameters[0]; +type LogEvent = Parameters[0]; +/** The slice of `Span` we use, so tests can record spans. */ +export interface SpanSink { + startSpan(args: StartSpanArgs): SpanSink; + log(event: LogEvent): void; + end(args?: { endTime?: number }): unknown; +} + +/** + * Mirrors Braintrust's Claude Code tracing under the Eval() `task` span: one + * LLM span per model request, with its tool calls as siblings, and the score + * span as a sibling of `task`. Each LLM span's input is the conversation so + * far, which is what the Thread view renders. + * https://github.com/braintrustdata/braintrust-coding-agent-plugins/blob/c0346dcdb16ae9f136b3abf83efb6458191f9300/bt-daemon/src/translate/claude.rs#L1173-L1300 + * https://github.com/braintrustdata/braintrust-spec/blob/b068e39112e081e45b6070e035877f1e2e83f9b7/skills/instrumentation-spec/references/features/eval-spans.md#trace-shape + * + * An LLM span starts where the previous span ended, since the agent only + * writes a request's timestamps once its blocks finish generating. Braintrust's + * Claude Code translator starts it at the first block instead, which leaves the + * generation time as a gap. Their Codex translator does what we do. + * https://github.com/braintrustdata/braintrust-coding-agent-plugins/blob/c0346dcdb16ae9f136b3abf83efb6458191f9300/bt-daemon/src/translate/codex.rs#L929-L932 + * + * task 0s → 38s + * ├─ llm (text + Skill) 0s → 4s + * ├─ Skill 4s → 5s + * ├─ llm (Bash only) 5s → 8s + * ├─ Bash 8s → 31s + * └─ llm (text) 31s → 38s + * + * Parts without a `requestId` each get their own LLM span. Transcripts without + * timestamps collapse every span to the run start. + */ +export function logTranscript( + parent: SpanSink, + row: Pick< + PendingRow, + | 'prompt' + | 'agentReport' + | 'checks' + | 'passed' + | 'modelId' + | 'modelProvider' + | 'transcript' + | 'toolLabels' + | 'startTime' + | 'endTime' + > +): void { + const sec = (ms: number | undefined) => (ms ? ms / 1000 : undefined); + const task = parent.startSpan({ + name: 'task', + type: 'task', + startTime: row.startTime, + }); + task.log({ input: row.prompt, output: row.agentReport }); + let latestTime = row.startTime; let toolIndex = 0; - // Seed the next assistant input with the prompt omitted from the transcript. - let pendingInput = row.prompt ? [{ role: 'user', content: row.prompt }] : []; + // Seeded with the prompt, which the transcript omits. + const history: Record[] = row.prompt + ? [{ role: 'user', content: row.prompt }] + : []; + let trailingUser: Record[] = []; + interface Request { + id?: string; + span: SpanSink; + endTime?: number; + inputLength: number; + message: { role: 'assistant'; content: string; tool_calls?: unknown[] }; + usage?: z.infer; + } + const requests: Request[] = []; + + const requestFor = (part: TranscriptPart, at: number | undefined) => { + const id = part.requestId; + const open = id ? requests.find((r) => r.id === id) : undefined; + if (open) { + open.endTime = latest(open.endTime, at); + open.usage ??= part.usage; + return open; + } + const startTime = + latestTime === undefined || at === undefined + ? (latestTime ?? at) + : Math.min(latestTime, at); + const request: Request = { + id, + span: task.startSpan({ + name: row.modelId ?? 'assistant', + type: 'llm', + startTime, + }), + endTime: at, + inputLength: history.length, + message: { role: 'assistant', content: '' }, + usage: part.usage, + }; + history.push(request.message); + trailingUser = []; + requests.push(request); + return request; + }; for (const part of row.transcript) { + const at = sec(part.ts) ?? latestTime; if (part.type === 'message') { if (part.role === 'user') { - pendingInput.push({ role: 'user', content: part.content }); - continue; - } - if (part.role === 'system') { - const span = parent.startSpan({ name: 'system', type: 'task', ...at }); + const message = { role: 'user', content: part.content }; + history.push(message); + trailingUser.push(message); + } else if (part.role === 'system') { + const span = task.startSpan({ + name: 'system', + type: 'task', + startTime: at, + }); span.log({ output: part.content }); - span.end(ended); - continue; + span.end({ endTime: at }); + } else { + const { message } = requestFor(part, at); + message.content = message.content + ? `${message.content}\n${part.content}` + : part.content; } - const span = parent.startSpan({ - name: row.modelId ?? 'assistant', - type: 'llm', - ...at, - }); - span.log({ - ...(pendingInput.length ? { input: pendingInput } : {}), - output: [{ role: 'assistant', content: part.content }], - ...(row.modelId ? { metadata: { model: row.modelId } } : {}), - }); - span.end(ended); - pendingInput = []; + latestTime = latest(latestTime, at); continue; } + if (part.requestId) { + const { message } = requestFor(part, at); + message.tool_calls = [ + ...(message.tool_calls ?? []), + { + id: part.id, + type: 'function', + function: { name: part.name, arguments: JSON.stringify(part.input) }, + }, + ]; + history.push({ + role: 'tool', + tool_call_id: part.id, + content: part.error ?? part.output ?? null, + }); + } const label = row.toolLabels[toolIndex++]; - const span = parent.startSpan({ + const endTime = sec(part.resultTs) ?? at; + latestTime = latest(latestTime, endTime); + const span = task.startSpan({ name: label ? `${part.name}: ${label}` : part.name, type: 'tool', - ...at, + startTime: at, }); span.log({ input: part.input, + // The span name adds a label, so keep the raw name filterable. + // https://github.com/braintrustdata/braintrust-spec/blob/b068e39112e081e45b6070e035877f1e2e83f9b7/skills/instrumentation-spec/references/features/skill-load-metadata.md + metadata: { tool_name: part.name }, ...(part.output !== undefined ? { output: part.output } : {}), ...(part.error ? { error: part.error } : {}), }); - span.end(ended); + span.end({ endTime }); + } + + for (const { span, endTime, inputLength, message, usage } of requests) { + span.log({ + input: history.slice(0, inputLength), + output: [message], + ...(usage + ? { metrics: tokenMetrics([{ model: row.modelId ?? '', ...usage }]) } + : {}), + metadata: { model: row.modelId, provider: row.modelProvider }, + }); + span.end({ endTime }); } // Preserve a trailing user message with no assistant response. - if (pendingInput.length) { - const span = parent.startSpan({ name: 'user', type: 'task', ...at }); - span.log({ output: pendingInput }); - span.end(ended); + if (trailingUser.length) { + const span = task.startSpan({ + name: 'user', + type: 'task', + startTime: latestTime, + }); + span.log({ output: trailingUser }); + span.end({ endTime: latestTime }); } + task.end({ endTime: latestTime }); - const scorer = parent.startSpan({ name: 'passed', type: 'score', ...at }); + const scoredAt = row.endTime ?? latestTime; + const scorer = parent.startSpan({ + name: 'passed', + type: 'score', + startTime: scoredAt, + }); scorer.log({ output: row.checks, scores: { passed: row.passed ? 1 : 0 }, }); - scorer.end(ended); + scorer.end({ endTime: scoredAt }); +} + +function latest(a: number | undefined, b: number | undefined) { + return a === undefined || b === undefined ? (a ?? b) : Math.max(a, b); } async function main() { diff --git a/packages/core/src/agents/claude-code/parser.test.ts b/packages/core/src/agents/claude-code/parser.test.ts index 0cec8505..7dd0e95d 100644 --- a/packages/core/src/agents/claude-code/parser.test.ts +++ b/packages/core/src/agents/claude-code/parser.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from 'vitest'; -import { claudeCodeParser } from './parser.js'; +import { claudeCodeParser, sessionRequestUsage } from './parser.js'; import { adaptTranscript } from '../../parsers/adapt.js'; /** A representative Claude Code `--print` JSONL session. */ @@ -11,6 +11,7 @@ const SESSION = [ type: 'assistant', timestamp: '2026-06-18T10:00:00.000Z', message: { + id: 'msg_1', role: 'assistant', content: [ { type: 'text', text: 'Let me list the files.' }, @@ -211,13 +212,23 @@ describe('adaptTranscript', () => { it('renders a scorer-facing transcript (messages + tool calls, raw args preserved)', () => { expect(adapted.transcript).toEqual([ - { type: 'message', role: 'assistant', content: 'Let me list the files.' }, + { + type: 'message', + role: 'assistant', + content: 'Let me list the files.', + ts: Date.parse('2026-06-18T10:00:00.000Z'), + requestId: 'msg_1', + }, { type: 'tool_call', name: 'Bash', input: { command: 'ls -la' }, output: 'file1\nfile2', error: undefined, + ts: Date.parse('2026-06-18T10:00:00.000Z'), + resultTs: Date.parse('2026-06-18T10:00:01.000Z'), + id: 'toolu_1', + requestId: 'msg_1', }, { type: 'tool_call', @@ -225,6 +236,7 @@ describe('adaptTranscript', () => { input: { query: 'rls' }, output: undefined, error: 'boom', + id: 'toolu_2', }, { type: 'message', @@ -234,3 +246,34 @@ describe('adaptTranscript', () => { ]); }); }); + +describe('sessionRequestUsage', () => { + it('keys final usage by message id and folds cache buckets into input', () => { + const line = (outputTokens: number) => + JSON.stringify({ + type: 'assistant', + message: { + id: 'msg_1', + usage: { + input_tokens: 2, + cache_read_input_tokens: 5, + cache_creation_input_tokens: 3, + output_tokens: outputTokens, + }, + }, + }); + expect(sessionRequestUsage(`${line(259)}\n${line(259)}`)).toEqual( + new Map([ + [ + 'msg_1', + { + inputTokens: 10, + cacheReadInputTokens: 5, + cacheWriteInputTokens: 3, + outputTokens: 259, + }, + ], + ]) + ); + }); +}); diff --git a/packages/core/src/agents/claude-code/parser.ts b/packages/core/src/agents/claude-code/parser.ts index dcc9249c..cac6a784 100644 --- a/packages/core/src/agents/claude-code/parser.ts +++ b/packages/core/src/agents/claude-code/parser.ts @@ -13,6 +13,7 @@ import type { ParsedTranscript, + RequestUsage, ToolCall, TranscriptEvent, } from '../../transcript/types.js'; @@ -192,6 +193,38 @@ function loadedSkillsFromClaudeCodeCall( return []; } +// Claude Code writes one line per content block, all sharing `message.id`. +function assistantRequestId(data: Record): string | undefined { + const id = isRecord(data.message) ? data.message.id : undefined; + if (typeof id === 'string') return id; + return typeof data.requestId === 'string' ? data.requestId : undefined; +} + +/** + * Final usage per request from Claude Code's session JSONL. Stream-json stdout + * repeats each request's `message_start` usage, so its output count stays at + * the first few tokens. Anthropic's input count excludes both cache buckets, + * so they're added back to match `ModelUsage`. + */ +export function sessionRequestUsage(jsonl: string): Map { + const byRequest = new Map(); + for (const record of parseJsonlRecords(jsonl).records) { + if (record.type !== 'assistant') continue; + const id = assistantRequestId(record); + const usage = isRecord(record.message) ? record.message.usage : undefined; + if (!id || !isRecord(usage)) continue; + const cacheRead = Number(usage.cache_read_input_tokens) || 0; + const cacheWrite = Number(usage.cache_creation_input_tokens) || 0; + byRequest.set(id, { + inputTokens: (Number(usage.input_tokens) || 0) + cacheRead + cacheWrite, + cacheReadInputTokens: cacheRead, + cacheWriteInputTokens: cacheWrite, + outputTokens: Number(usage.output_tokens) || 0, + }); + } + return byRequest; +} + function recordToEvents(data: Record): TranscriptEvent[] { const events: TranscriptEvent[] = []; const timestamp = @@ -232,10 +265,12 @@ function recordToEvents(data: Record): TranscriptEvent[] { } } } else if (type === 'assistant' || data.role === 'assistant') { + const requestId = assistantRequestId(data); const content = extractText(data); if (content) { events.push({ timestamp, + requestId, type: 'message', role: 'assistant', content, @@ -255,6 +290,7 @@ function recordToEvents(data: Record): TranscriptEvent[] { events.push( enrich({ timestamp, + requestId, type: 'tool_call', tool: { name: normalizeToolName(use.name, CLAUDE_CODE_TOOLS), diff --git a/packages/core/src/agents/claude-code/runner.ts b/packages/core/src/agents/claude-code/runner.ts index ee43603e..9246aed0 100644 --- a/packages/core/src/agents/claude-code/runner.ts +++ b/packages/core/src/agents/claude-code/runner.ts @@ -8,6 +8,7 @@ import type { Model as AnthropicModel } from '@anthropic-ai/sdk/resources/messag import type { McpServerConfig } from '../../index.js'; import { isRecord, parseJsonlRecords } from '../../json.js'; import type { AgentRunner } from '../types.js'; +import { sessionRequestUsage } from './parser.js'; import { npmGlobalBin, npmInstallGlobal, @@ -29,6 +30,16 @@ export const claudeCodeRunner: AgentRunner = { defaultModel: 'claude-sonnet-4-6', sessionDir: '"$HOME/.claude/projects"', + async enrichEvents(sandbox, events) { + const session = await sandbox.exec(`cat ${this.sessionDir}/*/*.jsonl`); + if (!session.ok) return; + const usage = sessionRequestUsage(session.stdout); + for (const event of events) { + const request = event.requestId && usage.get(event.requestId); + if (request) event.usage = request; + } + }, + async install(sandbox, version) { await npmInstallGlobal( sandbox, diff --git a/packages/core/src/agents/codex/parser.test.ts b/packages/core/src/agents/codex/parser.test.ts index 532e5f4e..33bec7ea 100644 --- a/packages/core/src/agents/codex/parser.test.ts +++ b/packages/core/src/agents/codex/parser.test.ts @@ -1,7 +1,8 @@ import { describe, expect, it } from 'vitest'; -import { codexParser } from './parser.js'; +import { codexParser, enrichFromRollout } from './parser.js'; import { codexRunner } from './runner.js'; import { adaptTranscript } from '../../parsers/adapt.js'; +import type { TranscriptEvent } from '../../transcript/types.js'; /** A representative `codex exec --json` stream (shapes captured from CLI 0.138). */ const SESSION = [ @@ -282,3 +283,278 @@ describe('codexRunner.deriveStopReason', () => { expect(codexRunner.deriveStopReason!('', timedOut)).toBe('timeout'); }); }); + +describe('enrichFromRollout', () => { + it('pairs stdout events with rollout records by position', () => { + const line = (timestamp: string, type: string, payload: object) => + JSON.stringify({ timestamp, type, payload }); + const rollout = [ + line('t1', 'response_item', { type: 'message', role: 'assistant' }), + line('t2', 'response_item', { type: 'function_call', call_id: 'c1' }), + line('t3', 'token_usage_record', { + response_id: 'r1', + usage: { input_tokens: 10, cached_input_tokens: 4, output_tokens: 2 }, + }), + line('t4', 'event_msg', { + type: 'item_completed', + item: { + type: 'CommandExecution', + id: 'c1', + command: ['/bin/bash', '-lc', 'ls'], + }, + }), + line('t5', 'response_item', { + type: 'function_call_output', + call_id: 'c1', + }), + ].join('\n'); + const events: TranscriptEvent[] = [ + { type: 'message', role: 'assistant', content: 'hi' }, + { + type: 'tool_call', + tool: { + name: 'shell', + originalName: 'x', + id: 'item_1', + command: '/bin/bash -lc ls', + }, + }, + { + type: 'tool_result', + tool: { name: 'shell', originalName: 'x', id: 'item_1' }, + }, + ]; + enrichFromRollout(events, rollout); + const usage = { + inputTokens: 10, + cacheReadInputTokens: 4, + cacheWriteInputTokens: 0, + outputTokens: 2, + }; + expect(events.map((e) => [e.timestamp, e.requestId, e.usage])).toEqual([ + ['t1', 'r1', usage], + ['t2', 'r1', usage], + ['t5', undefined, undefined], + ]); + }); + + it('adds an empty message for a request with no events', () => { + const line = (timestamp: string, type: string, payload: object) => + JSON.stringify({ timestamp, type, payload }); + const usageRecord = (at: string, id: string) => + line(at, 'token_usage_record', { response_id: id, usage: {} }); + const rollout = [ + line('t1', 'response_item', { type: 'message', role: 'assistant' }), + usageRecord('t2', 'r1'), + // A compaction call, which leaves no response item. + usageRecord('t3', 'r2'), + line('t4', 'response_item', { type: 'message', role: 'assistant' }), + usageRecord('t5', 'r3'), + ].join('\n'); + const events: TranscriptEvent[] = [ + { type: 'message', role: 'assistant', content: 'a' }, + { type: 'message', role: 'assistant', content: 'b' }, + ]; + enrichFromRollout(events, rollout); + expect(events.map((e) => [e.content, e.timestamp, e.requestId])).toEqual([ + ['a', 't1', 'r1'], + ['', 't3', 'r2'], + ['b', 't4', 'r3'], + ]); + }); + + it('adds only itemless empty messages when messages do not pair', () => { + const line = (timestamp: string, type: string, payload: object) => + JSON.stringify({ timestamp, type, payload }); + const usageRecord = (at: string, id: string) => + line(at, 'token_usage_record', { response_id: id, usage: {} }); + const rollout = [ + line('t1', 'response_item', { type: 'message', role: 'assistant' }), + usageRecord('t2', 'r1'), + usageRecord('t3', 'r2'), + line('t4', 'response_item', { type: 'message', role: 'assistant' }), + usageRecord('t5', 'r3'), + ].join('\n'); + const events: TranscriptEvent[] = [ + { type: 'message', role: 'assistant', content: 'a' }, + ]; + enrichFromRollout(events, rollout); + expect(events.map((e) => [e.content, e.timestamp, e.requestId])).toEqual([ + ['a', undefined, undefined], + ['', 't3', 'r2'], + ]); + }); +}); + +describe('enrichFromRollout tool pairing', () => { + const line = (timestamp: string, type: string, payload: object) => + JSON.stringify({ timestamp, type, payload }); + const issued = (id: string, at: string) => + line(at, 'response_item', { type: 'function_call', call_id: id }); + const completed = (id: string, script: string) => + line('', 'event_msg', { + type: 'item_completed', + item: { + type: 'CommandExecution', + id, + command: ['/bin/bash', '-lc', script], + }, + }); + const finished = (id: string, at: string) => + line(at, 'response_item', { type: 'function_call_output', call_id: id }); + const call = (id: string, command: string): TranscriptEvent[] => [ + { + type: 'tool_call', + tool: { name: 'shell', originalName: 'x', id, command }, + }, + { type: 'tool_result', tool: { name: 'shell', originalName: 'x', id } }, + ]; + const times = (events: TranscriptEvent[], rollout: string[]) => { + enrichFromRollout(events, rollout.join('\n')); + return events.map((e) => e.timestamp); + }; + + it('matches shell-quoted commands to the argv', () => { + expect( + times(call('item_1', `/bin/bash -lc 'cat "a b.txt"'`), [ + issued('c1', 't1'), + completed('c1', 'cat "a b.txt"'), + finished('c1', 't2'), + ]) + ).toEqual(['t1', 't2']); + }); + + it('skips all pairing when one command has extra arguments', () => { + expect( + times( + [ + ...call('item_1', '/bin/bash -lc pwd'), + ...call('item_2', '/bin/bash -lc ls -la'), + ], + [ + issued('c1', 't1'), + completed('c1', 'pwd'), + finished('c1', 't2'), + issued('c2', 't3'), + completed('c2', 'ls'), + finished('c2', 't4'), + ] + ) + ).toEqual([undefined, undefined, undefined, undefined]); + }); + + it('matches a redacted secret', () => { + expect( + times( + call('item_1', `/bin/bash -lc 'login --password [REDACTED_SECRET]'`), + [ + issued('c1', 't1'), + completed('c1', 'login --password hunter2'), + finished('c1', 't2'), + ] + ) + ).toEqual(['t1', 't2']); + }); + + it('rejects different text around a redacted secret', () => { + expect( + times( + call('item_1', `/bin/bash -lc 'logout --password [REDACTED_SECRET]'`), + [ + issued('c1', 't1'), + completed('c1', 'login --password hunter2'), + finished('c1', 't2'), + ] + ) + ).toEqual([undefined, undefined]); + }); + + it('skips all pairing when the rollout is missing a command', () => { + expect( + times( + [ + ...call('item_1', '/bin/bash -lc pwd'), + ...call('item_2', '/bin/bash -lc ls'), + ], + [issued('c1', 't1'), completed('c1', 'pwd'), finished('c1', 't2')] + ) + ).toEqual([undefined, undefined, undefined, undefined]); + }); + + it('adds only itemless empty messages when tools do not pair', () => { + const events = [ + ...call('item_1', '/bin/bash -lc pwd'), + ...call('item_2', '/bin/bash -lc ls'), + ]; + const usageRecord = (at: string, id: string) => + line(at, 'token_usage_record', { response_id: id, usage: {} }); + enrichFromRollout( + events, + [ + issued('c1', 't1'), + completed('c1', 'pwd'), + usageRecord('t2', 'r1'), + finished('c1', 't3'), + usageRecord('t4', 'r2'), + ].join('\n') + ); + expect(events.map((e) => [e.type, e.requestId])).toEqual([ + ['tool_call', undefined], + ['tool_result', undefined], + ['tool_call', undefined], + ['tool_result', undefined], + ['message', 'r2'], + ]); + }); + + it('ignores non-tool rollout items', () => { + expect( + times(call('item_1', '/bin/bash -lc pwd'), [ + issued('c1', 't1'), + line('', 'event_msg', { + type: 'item_completed', + item: { type: 'ContextCompaction', id: 'cc_1' }, + }), + completed('c1', 'pwd'), + finished('c1', 't2'), + ]) + ).toEqual(['t1', 't2']); + }); + + it('pairs repeated commands in order', () => { + const command = `/bin/bash -lc 'supabase status'`; + expect( + times( + [...call('item_1', command), ...call('item_2', command)], + [ + issued('c1', 't1'), + completed('c1', 'supabase status'), + finished('c1', 't2'), + issued('c2', 't3'), + completed('c2', 'supabase status'), + finished('c2', 't4'), + ] + ) + ).toEqual(['t1', 't2', 't3', 't4']); + }); + + it('pairs parallel calls in completion order', () => { + // c1 starts first, but both streams list c2 first because it finished first. + expect( + times( + [ + ...call('item_2', `/bin/bash -lc 'echo hi'`), + ...call('item_1', `/bin/bash -lc 'sleep 2'`), + ], + [ + issued('c1', 't1'), + issued('c2', 't2'), + completed('c2', 'echo hi'), + finished('c2', 't3'), + completed('c1', 'sleep 2'), + finished('c1', 't4'), + ] + ) + ).toEqual(['t2', 't3', 't1', 't4']); + }); +}); diff --git a/packages/core/src/agents/codex/parser.ts b/packages/core/src/agents/codex/parser.ts index f02afe10..46eb274d 100644 --- a/packages/core/src/agents/codex/parser.ts +++ b/packages/core/src/agents/codex/parser.ts @@ -21,6 +21,7 @@ import { isRecord, parseJsonlRecords } from '../../json.js'; import type { ParsedTranscript, + RequestUsage, ToolCall, TranscriptEvent, } from '../../transcript/types.js'; @@ -238,3 +239,227 @@ export const codexParser: AgentTranscriptParser = { return { events, errors }; }, }; + +/** + * Rollout names for the stdout items that `itemToEvents` makes tool calls from. + * Counting any other rollout item would shift the pairing. + * https://github.com/openai/codex/blob/rust-v0.154.0/codex-rs/protocol/src/items.rs#L45-L77 + */ +const ROLLOUT_TOOL_ITEMS = new Set([ + 'CommandExecution', + 'FileChange', + 'McpToolCall', + 'WebSearch', +]); + +/** + * Fills in event times, model requests, and per-request usage from the session + * rollout, which the `--json` stream lacks. The two streams list tool items and + * assistant messages in the same order but under different ids, so they're + * paired by position. Tool pairs must also agree on the command or MCP tool, + * and any count or content mismatch leaves those messages or tool calls + * untouched. A request no event came from, like a compaction call, gets an + * empty message. If the messages or tool calls didn't pair, only a request with + * no `response_item` records gets one, since the others may just be missing + * their requestId. + * + * rollout: reasoning → message → function_call(c1) → token_usage_record(r1) + * → function_call_output(c1) + * events: message → tool_call → tool_result, each tagged requestId r1 + */ +export function enrichFromRollout( + events: TranscriptEvent[], + rollout: string +): void { + interface Request { + id?: string; + usage?: RequestUsage; + /** No response item before its usage record, like a compaction call. */ + itemless?: boolean; + } + const callStarts = new Map(); + const callEnds = new Map(); + const messages: { at?: string; request: Request }[] = []; + const toolItems: Record[] = []; + const requests: { at?: string; request: Request }[] = []; + let open: Request | undefined; + + for (const record of parseJsonlRecords(rollout).records) { + const at = str(record.timestamp); + const payload = isRecord(record.payload) ? record.payload : {}; + if (record.type === 'token_usage_record') { + const request = open ?? { itemless: true }; + request.id = str(payload.response_id); + request.usage = rolloutUsage(payload.usage); + requests.push({ at, request }); + open = undefined; + } else if (record.type === 'response_item') { + const callId = str(payload.call_id); + if ( + payload.type === 'function_call_output' || + payload.type === 'custom_tool_call_output' + ) { + if (callId) callEnds.set(callId, at); + continue; + } + if (payload.type === 'message' && payload.role !== 'assistant') continue; + open ??= {}; + if (payload.type === 'message') messages.push({ at, request: open }); + if (callId) callStarts.set(callId, { at, request: open }); + } else if ( + record.type === 'event_msg' && + payload.type === 'item_completed' && + isRecord(payload.item) + ) { + const { type, id } = payload.item; + if (typeof id === 'string' && ROLLOUT_TOOL_ITEMS.has(String(type))) { + toolItems.push(payload.item); + } + } + } + + const tag = (event: TranscriptEvent, request: Request) => { + if (request.id) event.requestId = request.id; + if (request.usage) event.usage = request.usage; + }; + const assistant = events.filter( + (e) => e.type === 'message' && e.role === 'assistant' + ); + const messagesPaired = assistant.length === messages.length; + if (messagesPaired) { + assistant.forEach((event, i) => { + event.timestamp = messages[i].at ?? event.timestamp; + tag(event, messages[i].request); + }); + } + const calls = events.filter((e) => e.type === 'tool_call'); + const toolsPaired = + calls.length === toolItems.length && + calls.every((event, i) => sameCall(event, toolItems[i])); + if (toolsPaired) { + const callIdByItem = new Map(); + calls.forEach((event, i) => { + const itemId = String(toolItems[i].id); + const start = callStarts.get(itemId); + if (event.tool?.id) callIdByItem.set(event.tool.id, itemId); + if (!start) return; + event.timestamp = start.at ?? event.timestamp; + tag(event, start.request); + }); + for (const event of events) { + if (event.type !== 'tool_result' || !event.tool?.id) continue; + const callId = callIdByItem.get(event.tool.id); + const end = callId ? callEnds.get(callId) : undefined; + if (end) event.timestamp = end; + } + } + const tagged = new Set(events.map((e) => e.requestId)); + for (const { at, request } of requests) { + if (!request.id || tagged.has(request.id)) continue; + // Unpaired events have no requestId, so their requests only look silent. + if (!(messagesPaired && toolsPaired) && !request.itemless) continue; + const silent: TranscriptEvent = { + type: 'message', + role: 'assistant', + content: '', + timestamp: at, + }; + tag(silent, request); + // Rollout timestamps share one ISO format, so they sort as strings. + const next = events.findIndex((e) => at && e.timestamp && e.timestamp > at); + events.splice(next === -1 ? events.length : next, 0, silent); + } +} + +/** + * A command must split into exactly the rollout's argv. Stdout redacts secrets + * the rollout keeps, so `[REDACTED_SECRET]` matches any text in its word. + * https://github.com/openai/codex/blob/5fa5aaf0fffd6593ca60ed974af3e87e29f3be28/codex-rs/secrets/src/sanitizer.rs#L15-L22 + * + * stdout: [bash, -lc, "login --password [REDACTED_SECRET]"] + * rollout: [bash, -lc, "login --password hunter2"] → same call + */ +function sameCall(event: TranscriptEvent, item: Record) { + if (item.type === 'CommandExecution' && Array.isArray(item.command)) { + const words = event.tool?.command + ? splitShellWords(event.tool.command) + : undefined; + const argv = item.command.map(String); + return ( + !!words && + words.length === argv.length && + words.every((word, i) => sameWord(word, argv[i])) + ); + } + if (item.type === 'McpToolCall') { + return event.tool?.call?.toolName === item.tool; + } + return true; +} + +const REDACTED = '[REDACTED_SECRET]'; + +function sameWord(word: string, raw: string): boolean { + if (!word.includes(REDACTED)) return word === raw; + const escape = (text: string) => text.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); + const pattern = word.split(REDACTED).map(escape).join('[\\s\\S]*'); + return new RegExp(`^${pattern}$`).test(raw); +} + +/** + * POSIX shell word splitting (quotes and backslashes, no expansion), enough to + * recover the argv Codex shell-joined into a `command_execution` command. + * Undefined on an unterminated quote or trailing backslash. + */ +function splitShellWords(command: string): string[] | undefined { + const words: string[] = []; + let word: string | undefined; + for (let i = 0; i < command.length; i += 1) { + const ch = command[i]; + if (ch === "'") { + const end = command.indexOf("'", i + 1); + if (end === -1) return undefined; + word = (word ?? '') + command.slice(i + 1, end); + i = end; + } else if (ch === '"') { + let quoted = ''; + i += 1; + for (; i < command.length && command[i] !== '"'; i += 1) { + const next = command[i + 1]; + if ( + command[i] === '\\' && + next !== undefined && + '$`"\\\n'.includes(next) + ) { + i += 1; + if (next === '\n') continue; + } + quoted += command[i]; + } + if (i >= command.length) return undefined; + word = (word ?? '') + quoted; + } else if (ch === '\\') { + i += 1; + if (i >= command.length) return undefined; + if (command[i] !== '\n') word = (word ?? '') + command[i]; + } else if (/\s/.test(ch)) { + if (word !== undefined) words.push(word); + word = undefined; + } else { + word = (word ?? '') + ch; + } + } + if (word !== undefined) words.push(word); + return words; +} + +/** Codex's `input_tokens` already includes both cache buckets. */ +function rolloutUsage(usage: unknown): RequestUsage | undefined { + if (!isRecord(usage)) return undefined; + return { + inputTokens: Number(usage.input_tokens) || 0, + cacheReadInputTokens: Number(usage.cached_input_tokens) || 0, + cacheWriteInputTokens: Number(usage.cache_write_input_tokens) || 0, + outputTokens: Number(usage.output_tokens) || 0, + }; +} diff --git a/packages/core/src/agents/codex/runner.ts b/packages/core/src/agents/codex/runner.ts index c3c17e99..7fabe2bd 100644 --- a/packages/core/src/agents/codex/runner.ts +++ b/packages/core/src/agents/codex/runner.ts @@ -11,6 +11,7 @@ import type { ChatModel } from 'openai/resources/shared'; import type { McpServerConfig } from '../../index.js'; import { isRecord, parseJsonlRecords } from '../../json.js'; import type { AgentRunner } from '../types.js'; +import { enrichFromRollout } from './parser.js'; import { npmGlobalBin, npmInstallGlobal, @@ -99,16 +100,18 @@ export const codexRunner: AgentRunner = { `cat ${userPromptPath} | ${codex} ${flags}`, { timeoutMs: timeoutSec * 1000, env: { OPENAI_API_KEY: apiKey } } ); - // The --json stream has no per-response boundary, but the session rollout - // Codex writes to disk logs one token_count event per model response. + return { command, raw: command.stdout }; + }, + + async enrichEvents(sandbox, events) { const rollout = await sandbox.exec( `cat "$(ls -t "$HOME"/.codex/sessions/*/*/*/rollout-*.jsonl 2>/dev/null | head -1)"` ); - return { - command, - raw: command.stdout, - stepCount: countModelResponses(rollout.stdout), - }; + if (!rollout.ok) return; + enrichFromRollout(events, rollout.stdout); + // The --json stream has no per-response boundary, but the rollout logs one + // token_count event per model response. + return { stepCount: countModelResponses(rollout.stdout) }; }, deriveStopReason(raw, command) { diff --git a/packages/core/src/agents/engine.ts b/packages/core/src/agents/engine.ts index 757e8401..742570c6 100644 --- a/packages/core/src/agents/engine.ts +++ b/packages/core/src/agents/engine.ts @@ -110,7 +110,7 @@ export function createCliAgent( await writeSandboxFile(sandbox, USER_PROMPT_PATH, args.userPrompt); const start = Date.now(); - const { command, raw, stepCount } = await runner.exec({ + const { command, raw } = await runner.exec({ sandbox, model: options.model, apiKey, @@ -128,22 +128,28 @@ export function createCliAgent( mcpServerNames: Object.keys(args.mcpServers ?? {}), }) : { events: [] }; - const adapted = adaptTranscript(events); - - if (runner.sessionDir && args.sessionArchivePath) { + const archiveSession = async () => { + const archivePath = args.sessionArchivePath; + if (!runner.sessionDir || !archivePath) return; // Keeps a failed archive from leaving a --force rerun's old one behind. - rmSync(args.sessionArchivePath, { force: true }); - const staged = `${SESSION_ARCHIVE_STAGING_DIR}/${basename(args.sessionArchivePath)}`; + rmSync(archivePath, { force: true }); + const staged = `${SESSION_ARCHIVE_STAGING_DIR}/${basename(archivePath)}`; const tar = await sandbox.exec( `mkdir -p ${SESSION_ARCHIVE_STAGING_DIR} && tar -czf ${staged} --exclude=auth.json -C ${runner.sessionDir} .` ); if (tar.ok) { await sandbox.copyToHost( SESSION_ARCHIVE_STAGING_DIR, - dirname(args.sessionArchivePath) + dirname(archivePath) ); } - } + }; + // Both only read the session files, so neither waits on the other. + const [enriched] = await Promise.all([ + runner.enrichEvents?.(sandbox, events), + archiveSession(), + ]); + const adapted = adaptTranscript(events); // Surface run failures that would otherwise be invisible in results // (visible under --debug): the CLI's own error events, or a run that @@ -167,7 +173,9 @@ export function createCliAgent( stoppedReason: runner.deriveStopReason?.(raw, command) ?? processStopReason(command), usage: runner.extractUsage?.(raw, options.model), - stepCount: stepCount ?? runner.extractStepCount?.(raw), + stepCount: + (enriched ? enriched.stepCount : undefined) ?? + runner.extractStepCount?.(raw), durationMs, }; }, diff --git a/packages/core/src/agents/grok/parser.test.ts b/packages/core/src/agents/grok/parser.test.ts index 198cf38c..725e44b7 100644 --- a/packages/core/src/agents/grok/parser.test.ts +++ b/packages/core/src/agents/grok/parser.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from 'vitest'; -import { grokParser } from './parser.js'; +import { enrichFromUpdates, grokParser } from './parser.js'; import { buildGrokConfig, grokRunner } from './runner.js'; import type { AgentSandbox } from '../types.js'; @@ -406,3 +406,61 @@ describe('grok runner: GROK_HOME', () => { expect(warmup.env).not.toHaveProperty('GROK_HOME'); }); }); + +describe('grok request timing and usage', () => { + it('numbers steps by usage events and times events from updates.jsonl', () => { + const { events } = grokParser.parseTranscript( + [ + // Tool calls follow `usage`, as in Grok's real stdout. + JSON.stringify({ type: 'text', data: 'Looking.' }), + JSON.stringify({ + type: 'usage', + usage: { + input_tokens: 5, + cache_read_input_tokens: 3, + output_tokens: 1, + }, + }), + JSON.stringify({ + type: 'tool_call', + toolCallId: 'c1', + toolName: 'read_file', + rawInput: {}, + }), + JSON.stringify({ + type: 'tool_call_update', + toolCallId: 'c1', + status: 'completed', + rawOutput: {}, + }), + ].join('\n') + ); + const update = (ms: number, update: object) => + JSON.stringify({ params: { _meta: { agentTimestampMs: ms }, update } }); + enrichFromUpdates( + events, + [ + update(1000, { sessionUpdate: 'agent_message_chunk' }), + update(1500, { sessionUpdate: 'agent_message_chunk' }), + update(2000, { sessionUpdate: 'tool_call', toolCallId: 'c1' }), + update(3000, { + sessionUpdate: 'tool_call_update', + toolCallId: 'c1', + status: 'completed', + }), + ].join('\n') + ); + const iso = (ms: number) => new Date(ms).toISOString(); + expect(events.map((e) => [e.type, e.timestamp, e.requestId])).toEqual([ + ['message', iso(1500), 'step-0'], + ['tool_call', iso(2000), 'step-0'], + ['tool_result', iso(3000), undefined], + ]); + expect(events[0].usage).toEqual({ + inputTokens: 8, + cacheReadInputTokens: 3, + cacheWriteInputTokens: 0, + outputTokens: 1, + }); + }); +}); diff --git a/packages/core/src/agents/grok/parser.ts b/packages/core/src/agents/grok/parser.ts index e81afe58..b7b223a8 100644 --- a/packages/core/src/agents/grok/parser.ts +++ b/packages/core/src/agents/grok/parser.ts @@ -36,6 +36,7 @@ import { isRecord, parseJsonlRecords } from '../../json.js'; import type { ParsedTranscript, + RequestUsage, ToolCall, TranscriptEvent, } from '../../transcript/types.js'; @@ -227,10 +228,83 @@ function loadedSkills( return inline ? extractLoadedSkillsFromText(inline) : []; } +/** Grok counts cache reads out of `input_tokens`, so they're added back. */ +function stepUsage(usage: unknown): RequestUsage | undefined { + if (!isRecord(usage)) return undefined; + const cacheRead = Number(usage.cache_read_input_tokens) || 0; + const cacheWrite = Number(usage.cache_creation_input_tokens) || 0; + return { + inputTokens: (Number(usage.input_tokens) || 0) + cacheRead + cacheWrite, + cacheReadInputTokens: cacheRead, + cacheWriteInputTokens: cacheWrite, + outputTokens: Number(usage.output_tokens) || 0, + }; +} + +/** + * Fills in event times from the session's `updates.jsonl`, which the stdout + * stream lacks. Tool calls match by `toolCallId`. Assistant messages pair by + * position with runs of `agent_message_chunk`, and take the run's last time. + */ +export function enrichFromUpdates( + events: TranscriptEvent[], + updates: string +): void { + const callStarts = new Map(); + const callEnds = new Map(); + const messageEnds: number[] = []; + let inMessage = false; + for (const record of parseJsonlRecords(updates).records) { + const params = isRecord(record.params) ? record.params : {}; + const meta = isRecord(params._meta) ? params._meta : {}; + const update = isRecord(params.update) ? params.update : {}; + const at = Number(meta.agentTimestampMs) || undefined; + const kind = update.sessionUpdate; + const id = str(update.toolCallId); + if (kind === 'agent_message_chunk') { + if (!inMessage) messageEnds.push(0); + if (at) messageEnds[messageEnds.length - 1] = at; + inMessage = true; + continue; + } + inMessage = false; + if (!id || !at) continue; + if (kind === 'tool_call') callStarts.set(id, at); + if ( + kind === 'tool_call_update' && + statusSuccess(update.status) !== undefined + ) { + callEnds.set(id, at); + } + } + const iso = (ms: number | undefined) => + ms ? new Date(ms).toISOString() : undefined; + for (const event of events) { + const id = event.tool?.id; + if (!id) continue; + const at = + event.type === 'tool_call' ? callStarts.get(id) : callEnds.get(id); + event.timestamp = iso(at) ?? event.timestamp; + } + const messages = events.filter( + (e) => e.type === 'message' && e.role === 'assistant' + ); + if (messages.length !== messageEnds.length) return; + messages.forEach((event, i) => { + event.timestamp = iso(messageEnds[i]) ?? event.timestamp; + }); +} + export const grokParser: AgentTranscriptParser = { parseTranscript(raw: string): ParsedTranscript { const { records, errors } = parseJsonlRecords(raw); const events: TranscriptEvent[] = []; + // Each model call ends with a `usage` event, so steps number the requests. + let step = 0; + const stepUsages = new Map(); + const requestId = () => `step-${step}`; + // Grok emits a request's tool calls after its `usage` event. + const toolRequestId = () => `step-${Math.max(step - 1, 0)}`; // Join the adjacent events of one type into one event. let textBuffer = ''; @@ -239,7 +313,12 @@ export const grokParser: AgentTranscriptParser = { const content = textBuffer.trim(); textBuffer = ''; if (content) { - events.push({ type: 'message', role: 'assistant', content }); + events.push({ + type: 'message', + role: 'assistant', + content, + requestId: requestId(), + }); } }; const flushThought = () => { @@ -316,7 +395,11 @@ export const grokParser: AgentTranscriptParser = { const skills = loadedSkills(args, normalized); if (skills.length > 0) tool.loadedSkills = skills; - events.push({ type: 'tool_call', tool }); + events.push({ + type: 'tool_call', + tool, + requestId: toolRequestId(), + }); pending.set(id, { id, originalName, call, args }); break; } @@ -347,9 +430,14 @@ export const grokParser: AgentTranscriptParser = { flushAll(); break; } - // The `available_commands` and `usage` events make no transcript - // event. The `extractUsage` function in the runner reads the token - // data directly from the stream. + case 'usage': { + flushAll(); + const usage = stepUsage(record.usage); + if (usage) stepUsages.set(requestId(), usage); + step += 1; + break; + } + // The `available_commands` event makes no transcript event. default: break; } @@ -362,6 +450,10 @@ export const grokParser: AgentTranscriptParser = { // A run that stops early can leave a call with no last update. Send these // calls with no status, thus the record contains them. for (const call of pending.values()) emitResult(call); + for (const event of events) { + const usage = event.requestId && stepUsages.get(event.requestId); + if (usage) event.usage = usage; + } return { events, errors }; }, diff --git a/packages/core/src/agents/grok/runner.ts b/packages/core/src/agents/grok/runner.ts index 105d291b..38e5275a 100644 --- a/packages/core/src/agents/grok/runner.ts +++ b/packages/core/src/agents/grok/runner.ts @@ -16,6 +16,7 @@ import type { McpServerConfig } from '../../index.js'; import { isRecord, parseJsonlRecords } from '../../json.js'; import type { AgentRunner } from '../types.js'; +import { enrichFromUpdates } from './parser.js'; import { SCRATCH, npmGlobalBin, @@ -55,6 +56,13 @@ export const grokRunner: AgentRunner = { defaultModel: DEFAULT_GROK_MODEL, sessionDir: `${GROK_HOME}/sessions`, + async enrichEvents(sandbox, events) { + const updates = await sandbox.exec( + `cat ${this.sessionDir}/*/*/updates.jsonl` + ); + if (updates.ok) enrichFromUpdates(events, updates.stdout); + }, + async install(sandbox, version, apiKey) { await npmInstallGlobal( sandbox, diff --git a/packages/core/src/agents/opencode/parser.test.ts b/packages/core/src/agents/opencode/parser.test.ts index cb7faa64..218c97a9 100644 --- a/packages/core/src/agents/opencode/parser.test.ts +++ b/packages/core/src/agents/opencode/parser.test.ts @@ -245,3 +245,67 @@ describe('opencodeParser', () => { expect(events[0].content).toBe('MessageOutputLengthError'); }); }); + +describe('opencode request timing and usage', () => { + it('times tools from state.time and tags events with their step usage', () => { + const stream = [ + JSON.stringify({ + type: 'tool_use', + timestamp: 1000, + part: { + type: 'tool', + tool: 'bash', + callID: 't1', + messageID: 'm1', + state: { + status: 'completed', + input: { command: 'ls' }, + output: '', + time: { start: 400, end: 900 }, + }, + }, + }), + JSON.stringify({ + type: 'step_finish', + part: { + type: 'step-finish', + messageID: 'm1', + tokens: { + input: 5, + output: 2, + reasoning: 1, + cache: { read: 10, write: 0 }, + }, + }, + }), + ].join('\n'); + const [call, result] = opencodeParser.parseTranscript(stream).events; + expect([call.timestamp, result.timestamp]).toEqual([ + new Date(400).toISOString(), + new Date(900).toISOString(), + ]); + expect(call.usage).toEqual({ + inputTokens: 15, + cacheReadInputTokens: 10, + cacheWriteInputTokens: 0, + outputTokens: 3, + }); + }); + + it('emits an empty message for a step with no parts', () => { + const stream = JSON.stringify({ + type: 'step_finish', + timestamp: 2000, + part: { type: 'step-finish', messageID: 'm2', tokens: { output: 4 } }, + }); + const [event] = opencodeParser.parseTranscript(stream).events; + expect(event).toMatchObject({ + type: 'message', + role: 'assistant', + content: '', + requestId: 'm2', + timestamp: new Date(2000).toISOString(), + usage: { outputTokens: 4 }, + }); + }); +}); diff --git a/packages/core/src/agents/opencode/parser.ts b/packages/core/src/agents/opencode/parser.ts index 0a2e2909..4fd3e809 100644 --- a/packages/core/src/agents/opencode/parser.ts +++ b/packages/core/src/agents/opencode/parser.ts @@ -13,9 +13,9 @@ * {"type":"step_finish","part":{"type":"step-finish","reason":"stop","tokens":{…}}} * * A `tool_use` record is self-contained (input + output + status), so it yields - * a paired tool_call + tool_result correlated by `part.callID`. Step records - * carry token/finish info and produce no transcript event (the runner reads the - * terminal `step_finish` reason for the stop reason). + * a paired tool_call + tool_result correlated by `part.callID`. A `step_finish` + * gives its step's events their usage, or an empty message when the step had + * no other parts. The runner reads the last one's reason for the stop reason. * * Adapted from `@supabase/agent-evals` (packages/agent-eval/src/parsers). */ @@ -23,6 +23,7 @@ import { isRecord, parseJsonlRecords } from '../../json.js'; import type { ParsedTranscript, + RequestUsage, ToolCall, TranscriptEvent, } from '../../transcript/types.js'; @@ -113,6 +114,29 @@ const OPENCODE_ARG_FIELDS: ArgFieldMap = { url: ['url'], }; +/** + * A `step_finish` record's tokens as `ModelUsage` fields. opencode keeps cache + * out of `input` and reasoning out of `output`. + * https://github.com/anomalyco/opencode/blob/dev/packages/opencode/src/acp/usage.ts + */ +export function stepUsage( + record: Record +): RequestUsage | undefined { + if (record.type !== 'step_finish' || !isRecord(record.part)) return undefined; + const tokens = record.part.tokens; + if (!isRecord(tokens)) return undefined; + const cache = isRecord(tokens.cache) ? tokens.cache : undefined; + const cacheRead = Number(cache?.read) || 0; + const cacheWrite = Number(cache?.write) || 0; + return { + inputTokens: (Number(tokens.input) || 0) + cacheRead + cacheWrite, + cacheReadInputTokens: cacheRead, + cacheWriteInputTokens: cacheWrite, + outputTokens: + (Number(tokens.output) || 0) + (Number(tokens.reasoning) || 0), + }; +} + /** Epoch-ms (or pass-through ISO) → ISO string. */ function toISO(value: unknown): string | undefined { if (typeof value === 'number') return new Date(value).toISOString(); @@ -153,6 +177,7 @@ function partToEvents( ? [ { timestamp, + requestId: str(part.messageID), type: 'message', role: 'assistant', content: text, @@ -172,6 +197,8 @@ function partToEvents( const args = isRecord(state.input) ? state.input : {}; const status = str(state.status); const metadata = isRecord(state.metadata) ? state.metadata : undefined; + // The record lands when the call settles, so its own time is the end. + const time = isRecord(state.time) ? state.time : undefined; // The builtin tool set is fully enumerated in OPENCODE_TOOLS, so any // unmapped name is an MCP/custom tool (`_`). Builtins are // `other`; unmapped names are attributed to a configured MCP server by @@ -198,12 +225,18 @@ function partToEvents( if (loadedSkills.length > 0) tool.loadedSkills = loadedSkills; const events: TranscriptEvent[] = [ - { timestamp, type: 'tool_call', tool, raw }, + { + timestamp: toISO(time?.start) ?? timestamp, + requestId: str(part.messageID), + type: 'tool_call', + tool, + raw, + }, ]; // The result is in the same record; emit it only once the call completed. if (status && status !== 'running' && status !== 'pending') { events.push({ - timestamp, + timestamp: toISO(time?.end) ?? timestamp, type: 'tool_result', tool: { name, @@ -275,13 +308,34 @@ export const opencodeParser: AgentTranscriptParser = { const { records, errors } = parseJsonlRecords(raw); const mcpServerNames = ctx?.mcpServerNames ?? []; const events: TranscriptEvent[] = []; + const usageByMessage = new Map(); for (const record of records) { + const part = isRecord(record.part) ? record.part : undefined; + const messageId = str(part?.messageID); + const usage = stepUsage(record); + if (messageId && usage) { + usageByMessage.set(messageId, usage); + // A step that emitted no text or tool part still gets an LLM span. + if (!events.some((e) => e.requestId === messageId)) { + events.push({ + timestamp: toISO(record.timestamp), + requestId: messageId, + type: 'message', + role: 'assistant', + content: '', + }); + } + } try { events.push(...recordToEvents(record, mcpServerNames)); } catch (e) { errors.push(e instanceof Error ? e.message : String(e)); } } + for (const event of events) { + const usage = event.requestId && usageByMessage.get(event.requestId); + if (usage) event.usage = usage; + } return { events, errors }; }, }; diff --git a/packages/core/src/agents/opencode/runner.ts b/packages/core/src/agents/opencode/runner.ts index 92695129..e1b68369 100644 --- a/packages/core/src/agents/opencode/runner.ts +++ b/packages/core/src/agents/opencode/runner.ts @@ -20,15 +20,16 @@ // opencode's own config schema (type-only; pinned to the installed CLI version // via the catalog). The transcript stream is deliberately NOT typed from this -// SDK: `run --format json` emits a reduced, differently-shaped record than the -// SDK's server-API `Part`/`Event` entities (no id/sessionID/messageID; different -// discriminants), so the parser stays schema-defensive — see ./parser.ts. +// SDK: `run --format json` emits differently shaped records than the SDK's +// server-API `Part`/`Event` entities, so the parser stays schema-defensive. +// See ./parser.ts. import type { Config, McpLocalConfig } from '@opencode-ai/sdk'; import type { McpServerConfig } from '../../index.js'; import type { ModelProvider } from '../../eval-metadata.js'; import { modelProviderSchema } from '../../eval-metadata.js'; import { isRecord, parseJsonlRecords } from '../../json.js'; import type { AgentRunner } from '../types.js'; +import { stepUsage } from './parser.js'; import { SCRATCH, npmGlobalBin, @@ -193,10 +194,7 @@ export function createOpencodeRunner( }, extractUsage(raw, model) { - // opencode keeps cache out of `input` and reasoning out of `output`. - // https://github.com/anomalyco/opencode/blob/dev/packages/opencode/src/acp/usage.ts if (!raw) return undefined; - const { records } = parseJsonlRecords(raw); let sawUsage = false; const usage = { model, @@ -205,20 +203,14 @@ export function createOpencodeRunner( cacheWriteInputTokens: 0, outputTokens: 0, }; - for (const record of records) { - if (record.type !== 'step_finish' || !isRecord(record.part)) continue; - const tokens = record.part.tokens; - if (!isRecord(tokens)) continue; + for (const record of parseJsonlRecords(raw).records) { + const step = stepUsage(record); + if (!step) continue; sawUsage = true; - const cache = isRecord(tokens.cache) ? tokens.cache : undefined; - const cacheRead = Number(cache?.read) || 0; - const cacheWrite = Number(cache?.write) || 0; - usage.inputTokens += - (Number(tokens.input) || 0) + cacheRead + cacheWrite; - usage.cacheReadInputTokens += cacheRead; - usage.cacheWriteInputTokens += cacheWrite; - usage.outputTokens += - (Number(tokens.output) || 0) + (Number(tokens.reasoning) || 0); + usage.inputTokens += step.inputTokens; + usage.cacheReadInputTokens += step.cacheReadInputTokens; + usage.cacheWriteInputTokens += step.cacheWriteInputTokens; + usage.outputTokens += step.outputTokens; } return sawUsage ? [usage] : undefined; }, diff --git a/packages/core/src/agents/types.ts b/packages/core/src/agents/types.ts index 1db6ae62..de42c50f 100644 --- a/packages/core/src/agents/types.ts +++ b/packages/core/src/agents/types.ts @@ -19,6 +19,7 @@ import type { ReasoningEffortLevel, } from '../eval-metadata.js'; import type { AgentTranscriptParser } from '../parsers/types.js'; +import type { TranscriptEvent } from '../transcript/types.js'; export type AgentMetadata = { agent: AgentHarnessId; @@ -75,12 +76,6 @@ export interface RunnerExecResult { command: CommandResult; /** Raw transcript (JSONL) — stdout for streaming CLIs, or read from disk. */ raw?: string; - /** - * Model responses in the run, for runners that read them from somewhere - * other than `raw` (Codex uses its session rollout). See `stepCount` in - * eval-metadata.ts for what counts as one. - */ - stepCount?: number; } /** A CLI coding agent's execution strategy. `M` is its SDK model-id type. */ @@ -124,6 +119,14 @@ export interface AgentRunner { extractStepCount?(raw: string | undefined): number | undefined; /** Sandbox dir (shell expression) where the CLI writes its own session files. */ sessionDir?: string; + /** + * Fills in event `timestamp`, `requestId`, and `usage` from the CLI's session + * files, since stdout lacks them. Returns `stepCount` if the files have it. + */ + enrichEvents?( + sandbox: AgentSandbox, + events: TranscriptEvent[] + ): Promise<{ stepCount?: number } | void>; } /** diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 76966770..d11edfbf 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -14,7 +14,7 @@ import { import { basename, dirname, join, resolve } from 'node:path'; import { tmpdir } from 'node:os'; import { fileURLToPath } from 'node:url'; -import type { ToolCall, ToolName } from './transcript/types.js'; +import type { RequestUsage, ToolCall, ToolName } from './transcript/types.js'; import { createClient, type SupabaseClient } from '@supabase/supabase-js'; import { createMCPClient } from '@ai-sdk/mcp'; import { Experimental_StdioMCPTransport as StdioMCPTransport } from '@ai-sdk/mcp/mcp-stdio'; @@ -159,11 +159,19 @@ export interface ScoreResult { checks?: CheckResult[]; } +/** + * `ts` and `resultTs` are epoch ms, set when the agent records event times. + * `requestId` groups the parts one model request emitted, and `usage` is that + * request's tokens (see `TranscriptEvent`). + */ export type TranscriptPart = | { type: 'message'; role: 'system' | 'user' | 'assistant'; content: string; + ts?: number; + requestId?: string; + usage?: RequestUsage; } | { type: 'tool_call'; @@ -171,6 +179,12 @@ export type TranscriptPart = input: Record; output?: unknown; error?: string; + ts?: number; + resultTs?: number; + /** Correlates the call with its result, as `TranscriptEvent.tool.id`. */ + id?: string; + requestId?: string; + usage?: RequestUsage; }; export type TranscriptSerializationOptions = { diff --git a/packages/core/src/parsers/adapt.ts b/packages/core/src/parsers/adapt.ts index 58e916ca..525963b3 100644 --- a/packages/core/src/parsers/adapt.ts +++ b/packages/core/src/parsers/adapt.ts @@ -23,6 +23,7 @@ export interface AdaptedTranscript { interface ResolvedResult { result?: unknown; error?: string; + ts?: number; } export function adaptTranscript(events: TranscriptEvent[]): AdaptedTranscript { @@ -31,10 +32,10 @@ export function adaptTranscript(events: TranscriptEvent[]): AdaptedTranscript { const resultsById = new Map(); for (const event of events) { if (event.type !== 'tool_result' || !event.tool?.id) continue; - resultsById.set( - event.tool.id, - toResolved(event.tool.result, event.tool.success) - ); + resultsById.set(event.tool.id, { + ...toResolved(event.tool.result, event.tool.success), + ...timed(event.timestamp), + }); } const transcript: TranscriptPart[] = []; @@ -43,11 +44,19 @@ export function adaptTranscript(events: TranscriptEvent[]): AdaptedTranscript { let steps = 0; for (const event of events) { - if (event.type === 'message' && event.role && event.content) { - const content = event.content.trim(); - if (!content) continue; - transcript.push({ type: 'message', role: event.role, content }); - if (event.role === 'assistant') { + if (event.type === 'message' && event.role) { + const content = event.content?.trim() ?? ''; + // An empty message still marks a model request that emitted nothing. + if (!content && !event.requestId) continue; + transcript.push({ + type: 'message', + role: event.role, + content, + ...timed(event.timestamp), + ...(event.requestId ? { requestId: event.requestId } : {}), + ...(event.usage ? { usage: event.usage } : {}), + }); + if (event.role === 'assistant' && content) { agentReport = content; steps += 1; } @@ -67,6 +76,11 @@ export function adaptTranscript(events: TranscriptEvent[]): AdaptedTranscript { input: body, output: resolved?.error === undefined ? resolved?.result : undefined, error: resolved?.error, + ...timed(event.timestamp), + ...(resolved?.ts ? { resultTs: resolved.ts } : {}), + ...(event.tool.id ? { id: event.tool.id } : {}), + ...(event.requestId ? { requestId: event.requestId } : {}), + ...(event.usage ? { usage: event.usage } : {}), }); toolCalls.push({ tool: call, @@ -99,6 +113,11 @@ function toResolved( return { result }; } +function timed(timestamp: string | undefined): { ts?: number } { + const ts = parseTs(timestamp); + return ts ? { ts } : {}; +} + function parseTs(timestamp: string | undefined): number { if (!timestamp) return 0; const ms = Date.parse(timestamp); diff --git a/packages/core/src/transcript/types.ts b/packages/core/src/transcript/types.ts index abca11e9..c0582674 100644 --- a/packages/core/src/transcript/types.ts +++ b/packages/core/src/transcript/types.ts @@ -43,10 +43,21 @@ export type ToolCall = | { kind: 'mcp'; server: string; toolName: string } | { kind: 'other'; toolName: string }; +import type { ModelUsage } from '../eval-metadata.js'; + +/** Token usage of the model request that emitted an event. */ +export type RequestUsage = Omit; + /** A single normalized event in an agent transcript. */ export interface TranscriptEvent { /** ISO timestamp of the event, when the agent records one. */ timestamp?: string; + /** + * Model request that emitted this event (Claude Code `message.id`, OpenCode + * `part.messageID`). One request can emit text plus several tool calls. + */ + requestId?: string; + usage?: RequestUsage; /** Event kind. */ type: 'message' | 'tool_call' | 'tool_result' | 'thinking' | 'error'; /** For `message` events: the speaker. */