diff --git a/scripts/trace-model-usage-e2e.mjs b/scripts/trace-model-usage-e2e.mjs index b3f1389d..7496bbfc 100644 --- a/scripts/trace-model-usage-e2e.mjs +++ b/scripts/trace-model-usage-e2e.mjs @@ -38,7 +38,7 @@ const TABLE = process.env.HIVEMIND_SESSIONS_TABLE ?? "sessions_modeltest"; const RUN_TAG = "modeltest-" + Date.now().toString(36); const MAX_FILES_PER_AGENT = Number(process.env.MAX_FILES ?? 300); -const { parseClaudeTurnMeta, parseCodexTurnMeta, sdkTurnMeta } = await import( +const { parseClaudeTurnMeta, parseCodexTurnMeta, scanClaudeTurnForCapture, sdkTurnMeta } = await import( join(ROOT, "dist/src/notifications/model-usage.js") ); const { loadConfig } = await import(join(ROOT, "dist/src/config.js")); @@ -143,6 +143,50 @@ function runHook(entry, payload) { return r.status === 0; } +/** + * The transcript's newest non-sidechain assistant text — what a live Stop + * hook would receive as `last_assistant_message` for that transcript. + */ +function lastClaudeAssistantText(file) { + let raw; + try { + raw = readFileSync(file, "utf-8"); + } catch { + return null; + } + const lines = raw.split("\n"); + for (let i = lines.length - 1; i >= 0; i--) { + const t = lines[i].trim(); + if (!t) continue; + let e; + try { e = JSON.parse(t); } catch { continue; } + // Mirror the live scanner: the newest text-bearing, non-sidechain + // assistant record — with NO usage requirement, because that is the + // record the hook will correlate against. + if (e?.type !== "assistant" || e.isSidechain === true) continue; + const c = e.message?.content; + const text = typeof c === "string" + ? c + : Array.isArray(c) ? c.map((b) => (b?.type === "text" && typeof b.text === "string" ? b.text : "")).join("") : ""; + if (text.trim()) return text; + } + return null; +} + +/** + * Live-faithful Claude parse for transcript selection: only pick transcripts + * whose FINAL turn would actually enrich under the hook's correlation + * semantics (newest text-bearing record matches + carries usage). Replaying a + * transcript whose final turn has no usage would correctly produce an + * unenriched row — useless for this harness's per-model assertions. + */ +function parseClaudeLive(file) { + const text = lastClaudeAssistantText(file); + if (!text) return null; + const scan = scanClaudeTurnForCapture(file, text); + return scan.status === "match" ? scan.meta : null; +} + async function main() { const config = loadConfig(); if (!config) { @@ -153,7 +197,7 @@ async function main() { const claudeModels = pickOnePerModel( listTranscripts(join(homedir(), ".claude", "projects"), MAX_FILES_PER_AGENT), - parseClaudeTurnMeta, + parseClaudeLive, ); const codexModels = pickOnePerModel( listTranscripts(join(homedir(), ".codex", "sessions"), MAX_FILES_PER_AGENT), @@ -176,10 +220,15 @@ async function main() { let n = 0; for (const [model, file] of claudeModels) { const sid = `${RUN_TAG}-cc-${n++}`; + // The capture hook correlates last_assistant_message with the transcript's + // newest assistant record (off-by-one guard), so a synthetic message would + // make every enrichment come back null — replay the transcript's REAL + // final assistant text. + const lastText = lastClaudeAssistantText(file); jobs.push({ agent: "claude_code", model, run: () => runHook("dist/src/hooks/capture.js", { session_id: sid, transcript_path: file, cwd: ROOT, - hook_event_name: "Stop", last_assistant_message: `e2e trace for ${model}`, + hook_event_name: "Stop", last_assistant_message: lastText ?? `e2e trace for ${model}`, }) }); } for (const [model, file] of codexModels) { diff --git a/src/hooks/capture.ts b/src/hooks/capture.ts index 04575fdb..c0406e6e 100644 --- a/src/hooks/capture.ts +++ b/src/hooks/capture.ts @@ -15,7 +15,7 @@ import { DeeplakeApi } from "../deeplake-api.js"; import { projectNameFromCwd } from "../utils/project-name.js"; import { log as _log } from "../utils/debug.js"; import { buildSessionPath } from "../utils/session-path.js"; -import { parseClaudeTurnMeta } from "../notifications/model-usage.js"; +import { parseClaudeTurnMetaLive } from "../notifications/model-usage.js"; import { bumpTotalCount, loadTriggerConfig, @@ -139,7 +139,23 @@ async function main(): Promise { // last assistant turn (best-effort; null on any read/parse failure). On // SubagentStop, last_assistant_message belongs to the subagent transcript; // transcript_path points at the parent session, so prefer the agent one. - const modelMeta = parseClaudeTurnMeta(input.agent_transcript_path ?? input.transcript_path); + // The Stop event races the transcript writer (the turn's final assistant + // record may not be flushed yet), so this re-reads across a short backoff + // and correlates the record with last_assistant_message — never attributing + // a PREVIOUS turn's tokens to this one. A sidechain record is a subagent's + // on a main-session Stop, but IS the turn on SubagentStop. + // An empty last_assistant_message can't be correlated with a transcript + // record — skipping the correlation would silently accept the PREVIOUS + // turn's record, so skip the enrichment instead. + const expectText = typeof input.last_assistant_message === "string" ? input.last_assistant_message.trim() : ""; + if (!expectText) log("empty last_assistant_message — skipping turn-meta enrichment"); + const modelMeta = expectText + ? await parseClaudeTurnMetaLive( + input.agent_transcript_path ?? input.transcript_path, + expectText, + Boolean(input.agent_transcript_path), + ) + : null; entry = { id: crypto.randomUUID(), ...meta, diff --git a/src/notifications/model-usage.ts b/src/notifications/model-usage.ts index 0399513a..b06c8ff6 100644 --- a/src/notifications/model-usage.ts +++ b/src/notifications/model-usage.ts @@ -272,7 +272,8 @@ interface ClaudeUsage { interface ClaudeLine { type?: string; cwd?: unknown; - message?: { model?: unknown; usage?: ClaudeUsage; stop_reason?: unknown }; + isSidechain?: unknown; + message?: { model?: unknown; usage?: ClaudeUsage; stop_reason?: unknown; content?: unknown }; } /** @@ -399,6 +400,165 @@ export function parseClaudeTurnMeta(transcriptPath?: string): TraceModelMeta | n } } +/** Flatten a Claude message `content` (string or block array) to its text. */ +function flattenClaudeText(content: unknown): string { + if (typeof content === "string") return content; + if (!Array.isArray(content)) return ""; + return content + .map((b) => (b && typeof b === "object" && (b as { type?: unknown }).type === "text" && typeof (b as { text?: unknown }).text === "string" ? (b as { text: string }).text : "")) + .join(""); +} + +/** + * Whether the transcript's assistant text plausibly IS the hook's + * `last_assistant_message`. Exact match after trimming, or an ANCHORED + * prefix with enough shared length that a coincidental hit is implausible — + * an unanchored substring test would let a short reply ("Done") match the + * PREVIOUS turn ("Done — details…") and resurrect the off-by-one this + * correlation exists to prevent. + */ +const MIN_TRUNCATION_MATCH_LEN = 32; + +function textMatches(entryText: string, expect: string): boolean { + const a = entryText.trim(); + const b = expect.trim(); + if (!a || !b) return false; + if (a === b) return true; + const shorter = a.length <= b.length ? a : b; + const longer = a.length <= b.length ? b : a; + return shorter.length >= MIN_TRUNCATION_MATCH_LEN && longer.startsWith(shorter); +} + +/** + * One point-in-time scan of a Claude transcript for the CURRENT turn's + * assistant metadata. + * + * `"match"` — the newest relevant assistant record is complete (has usage) and, + * when `expectText` is given, its text corresponds to the captured + * `last_assistant_message`. `meta` is set. + * + * `"retry"` — the transcript hasn't caught up with the turn being captured: + * the file is missing/unreadable, ends in a partial JSONL record, has no + * relevant assistant record yet, or its newest relevant assistant record lacks + * usage or doesn't correspond to `expectText`. The caller should re-read after + * a short backoff. Crucially this NEVER falls back to an older usage-bearing + * record — that would silently attribute the previous turn's tokens to this + * one (the off-by-one this scan exists to prevent). + * + * `includeSidechain` — on SubagentStop the agent transcript IS the sidechain, + * so its records must not be filtered; on a main-session Stop, sidechain + * records belong to subagents and must be skipped. + */ +export interface ClaudeTurnScan { + status: "match" | "retry"; + meta?: TraceModelMeta; +} + +export function scanClaudeTurnForCapture( + transcriptPath: string | undefined, + expectText?: string, + includeSidechain = false, +): ClaudeTurnScan { + if (!transcriptPath) return { status: "retry" }; + const r = openTranscript(transcriptPath); + if (!r) return { status: "retry" }; + try { + // The tail window can slice the newest records away only when a single + // record exceeds 1 MiB — scan the whole snapshot when the tail finds + // nothing relevant. + return scanClaudeTurn(r.lines, expectText, includeSidechain) ?? + scanClaudeTurn(r.readWhole(), expectText, includeSidechain) ?? + { status: "retry" }; + } finally { + r.close(); + } +} + +/** Reverse-scan for the newest relevant assistant record. Null = nothing relevant in these lines. */ +function scanClaudeTurn( + lines: string[] | null, + expectText: string | undefined, + includeSidechain: boolean, +): ClaudeTurnScan | null { + if (!lines) return null; + let sawNonBlank = false; + for (let i = lines.length - 1; i >= 0; i--) { + const trimmed = lines[i].trim(); + if (!trimmed) continue; + let entry: ClaudeLine; + try { + entry = JSON.parse(trimmed) as ClaudeLine; + } catch { + // A malformed FINAL record is an in-progress append — wait for it to + // complete. Malformed records further up are just skipped. + if (!sawNonBlank) { + log("trailing partial record — retry"); + return { status: "retry" }; + } + sawNonBlank = true; + continue; + } + sawNonBlank = true; + if (!entry || typeof entry !== "object") continue; + const msg = entry.message; + if (entry.type !== "assistant" || !msg) continue; + if (!includeSidechain && entry.isSidechain === true) continue; + if (expectText !== undefined) { + // Only a TEXT-bearing record can be the one behind + // last_assistant_message — a turn's trailing tool_use/thinking records + // have no text and must not decide the correlation. + const text = flattenClaudeText(msg.content); + if (!text.trim()) continue; + if (!textMatches(text, expectText)) { + log("newest assistant record does not match last_assistant_message — retry"); + return { status: "retry" }; + } + } + // Newest relevant assistant record found — it alone decides the outcome. + if (!msg.usage) { + log("newest assistant record has no usage — retry"); + return { status: "retry" }; + } + const meta: TraceModelMeta = { + model: typeof msg.model === "string" ? msg.model : undefined, + reasoning_effort: readClaudeEffortLevel(typeof entry.cwd === "string" ? entry.cwd : undefined) ?? null, + token_usage: normalizeClaudeUsage(msg.usage), + }; + if (typeof msg.stop_reason === "string") meta.stop_reason = msg.stop_reason; + const extra = claudeUsageExtra(msg.usage); + if (extra) meta.usage_extra = extra; + return { status: "match", meta }; + } + return null; +} + +/** + * Live-capture variant of {@link parseClaudeTurnMeta}: the Stop hook often + * fires before Claude has flushed the turn's final assistant record to the + * transcript, so a single read races the writer (observed in real sessions: + * turn 1 captured with null model/usage while the record landed ~ms later). + * Re-opens and re-scans the file across a short bounded backoff and returns + * null — capture proceeds unenriched — if the record never materializes. + */ +export async function parseClaudeTurnMetaLive( + transcriptPath: string | undefined, + expectText?: string, + includeSidechain = false, + backoffMs: readonly number[] = [25, 50, 100], +): Promise { + let scan = scanClaudeTurnForCapture(transcriptPath, expectText, includeSidechain); + for (const delay of backoffMs) { + if (scan.status === "match") break; + await new Promise((resolve) => setTimeout(resolve, delay)); + scan = scanClaudeTurnForCapture(transcriptPath, expectText, includeSidechain); + } + if (scan.status !== "match") { + log(`turn meta unavailable after ${backoffMs.length + 1} reads — writing unenriched`); + return null; + } + return scan.meta ?? null; +} + // --------------------------------------------------------------------------- // Codex // --------------------------------------------------------------------------- diff --git a/tests/claude-code/capture-hook.test.ts b/tests/claude-code/capture-hook.test.ts index 79eb62f4..3b699511 100644 --- a/tests/claude-code/capture-hook.test.ts +++ b/tests/claude-code/capture-hook.test.ts @@ -262,6 +262,9 @@ describe("capture hook — event-type branches", () => { last_assistant_message: "reply text", }); await runHook(); + // The Stop path re-reads the (here: absent) transcript across a bounded + // backoff before inserting unenriched — wait for the INSERT, not a tick. + await vi.waitFor(() => expect(queryMock).toHaveBeenCalled(), { timeout: 2000 }); const sql = queryMock.mock.calls[0][0] as string; expect(sql).toContain('"type":"assistant_message"'); expect(sql).toContain('"content":"reply text"'); @@ -277,6 +280,7 @@ describe("capture hook — event-type branches", () => { agent_transcript_path: "/tmp/agent.jsonl", }); await runHook(); + await vi.waitFor(() => expect(queryMock).toHaveBeenCalled(), { timeout: 2000 }); const sql = queryMock.mock.calls[0][0] as string; expect(sql).toContain('"agent_transcript_path":"/tmp/agent.jsonl"'); }); diff --git a/tests/shared/notifications-model-usage.test.ts b/tests/shared/notifications-model-usage.test.ts index 28b83901..82320b83 100644 --- a/tests/shared/notifications-model-usage.test.ts +++ b/tests/shared/notifications-model-usage.test.ts @@ -6,8 +6,10 @@ import { join } from "node:path"; import { mkdirSync } from "node:fs"; import { parseClaudeTurnMeta, + parseClaudeTurnMetaLive, parseCodexTurnMeta, normalizeSdkUsage, + scanClaudeTurnForCapture, sdkTurnMeta, readClaudeEffortLevel, readTailLines, @@ -541,3 +543,133 @@ describe("parseCodexTurnMeta — usage_extra quota", () => { expect(meta).not.toHaveProperty("usage_extra"); }); }); + +// --------------------------------------------------------------------------- +// Stop-capture race: scanClaudeTurnForCapture / parseClaudeTurnMetaLive +// --------------------------------------------------------------------------- + +const USAGE = { input_tokens: 10, output_tokens: 42 }; + +function claudeAssistantText(model: string, usage: Record | undefined, text: string, extra: object = {}) { + return { + type: "assistant", + message: { model, role: "assistant", content: [{ type: "text", text }], ...(usage ? { usage } : {}) }, + ...extra, + }; +} + +describe("scanClaudeTurnForCapture", () => { + it("matches the newest assistant record when it carries usage and the expected text", () => { + const file = writeTranscript([ + { type: "user" }, + claudeAssistantText("claude-fable-5", USAGE, "final answer"), + ]); + const scan = scanClaudeTurnForCapture(file, "final answer"); + expect(scan.status).toBe("match"); + expect(scan.meta?.model).toBe("claude-fable-5"); + expect(scan.meta?.token_usage).toEqual({ input_tokens: 10, output_tokens: 42 }); + }); + + it("retries instead of falling back to the PREVIOUS turn when the current record is absent", () => { + // Only turn 1 is on disk; the Stop being captured is turn 2. + const file = writeTranscript([ + claudeAssistantText("claude-fable-5", { input_tokens: 1, output_tokens: 1 }, "turn one answer"), + ]); + const scan = scanClaudeTurnForCapture(file, "turn two answer"); + expect(scan.status).toBe("retry"); + expect(scan.meta).toBeUndefined(); + }); + + it("retries when the newest assistant record has no usage yet", () => { + const file = writeTranscript([ + claudeAssistantText("claude-fable-5", undefined, "final answer"), + ]); + expect(scanClaudeTurnForCapture(file, "final answer").status).toBe("retry"); + }); + + it("retries on a trailing partial JSONL record", () => { + const file = writeTranscript([claudeAssistantText("claude-fable-5", USAGE, "final answer")]); + writeFileSync(file, '{"type":"assistant","message":{"mod', { flag: "a" }); + expect(scanClaudeTurnForCapture(file, "final answer").status).toBe("retry"); + }); + + it("retries on a missing file and on an empty transcript", () => { + expect(scanClaudeTurnForCapture(join(TEMP_DIR, "absent.jsonl")).status).toBe("retry"); + const empty = join(TEMP_DIR, "empty.jsonl"); + writeFileSync(empty, ""); + expect(scanClaudeTurnForCapture(empty).status).toBe("retry"); + }); + + it("skips sidechain records on a main-session Stop but not on SubagentStop", () => { + const file = writeTranscript([ + claudeAssistantText("claude-fable-5", USAGE, "main answer"), + claudeAssistantText("claude-haiku-4-5", { input_tokens: 2, output_tokens: 3 }, "subagent answer", { isSidechain: true }), + ]); + const main = scanClaudeTurnForCapture(file, "main answer", false); + expect(main.status).toBe("match"); + expect(main.meta?.model).toBe("claude-fable-5"); + const sub = scanClaudeTurnForCapture(file, "subagent answer", true); + expect(sub.status).toBe("match"); + expect(sub.meta?.model).toBe("claude-haiku-4-5"); + }); + + it("matches without text correlation when no expected text is given", () => { + const file = writeTranscript([claudeAssistantText("claude-fable-5", USAGE, "whatever")]); + expect(scanClaudeTurnForCapture(file).status).toBe("match"); + }); + + it("tolerates one-sided truncation of the expected text only with an anchored, long-enough prefix", () => { + const long = "a sufficiently long final answer that exceeds the anchored-prefix minimum length"; + const file = writeTranscript([claudeAssistantText("claude-fable-5", USAGE, long + " with extra tail detail")]); + expect(scanClaudeTurnForCapture(file, long).status).toBe("match"); + // A SHORT shared prefix is not evidence of identity — "Done" must not + // match the previous turn's "Done — details…" (the off-by-one guard). + const file2 = writeTranscript([claudeAssistantText("claude-fable-5", USAGE, "Done — full details follow here")]); + expect(scanClaudeTurnForCapture(file2, "Done").status).toBe("retry"); + // Unanchored containment must not match either. + const file3 = writeTranscript([claudeAssistantText("claude-fable-5", USAGE, "prefix junk then a sufficiently long final answer that exceeds the anchored-prefix minimum length")]); + expect(scanClaudeTurnForCapture(file3, long).status).toBe("retry"); + }); + + it("skips textless trailing records (tool_use) when correlating, without falling back across turns", () => { + // The turn's final flush order can put a textless record last; the + // text-bearing record of the SAME turn must still match… + const file = writeTranscript([ + claudeAssistantText("claude-fable-5", USAGE, "the actual final answer text here padded long"), + { type: "assistant", message: { model: "claude-fable-5", role: "assistant", content: [{ type: "tool_use", id: "t1" }], usage: USAGE } }, + ]); + expect(scanClaudeTurnForCapture(file, "the actual final answer text here padded long").status).toBe("match"); + // …but a PREVIOUS turn's text must not satisfy the current turn. + expect(scanClaudeTurnForCapture(file, "a different current-turn reply").status).toBe("retry"); + }); +}); + +describe("parseClaudeTurnMetaLive", () => { + it("recovers the turn's meta when the record lands between reads (the real flush race)", async () => { + const file = join(TEMP_DIR, "grow.jsonl"); + writeFileSync(file, JSON.stringify({ type: "user" }) + "\n"); + const pending = parseClaudeTurnMetaLive(file, "late answer", false, [10, 10, 200]); + setTimeout(() => { + writeFileSync(file, JSON.stringify(claudeAssistantText("claude-fable-5", USAGE, "late answer")) + "\n", { flag: "a" }); + }, 15); + const meta = await pending; + expect(meta?.model).toBe("claude-fable-5"); + expect(meta?.token_usage).toEqual({ input_tokens: 10, output_tokens: 42 }); + }); + + it("returns null (unenriched) instead of the previous turn's usage when the record never lands", async () => { + const file = writeTranscript([ + claudeAssistantText("claude-fable-5", { input_tokens: 999, output_tokens: 999 }, "turn one answer"), + ]); + const meta = await parseClaudeTurnMetaLive(file, "turn two answer", false, [1, 1, 1]); + expect(meta).toBeNull(); + }); + + it("resolves immediately on first read when the record is already complete", async () => { + const file = writeTranscript([claudeAssistantText("claude-fable-5", USAGE, "done")]); + const started = Date.now(); + const meta = await parseClaudeTurnMetaLive(file, "done", false, [500, 500, 500]); + expect(meta?.model).toBe("claude-fable-5"); + expect(Date.now() - started).toBeLessThan(400); + }); +});