diff --git a/apps/server/src/usage/usageTranscriptReader.ts b/apps/server/src/usage/usageTranscriptReader.ts index 5b92f48356aa4..ddfe10901ad19 100644 --- a/apps/server/src/usage/usageTranscriptReader.ts +++ b/apps/server/src/usage/usageTranscriptReader.ts @@ -394,32 +394,69 @@ export async function readTranscriptRecords( let scanOffset = start; let pendingChunks: Buffer[] = []; let pendingBytes = 0; + let oversized = false; let streaming: ReturnType | undefined; let decoder: NodeStringDecoder.StringDecoder | undefined; + let carriesUsage = false; + let markerTail = Buffer.alloc(0); + const markerOverlap = + Math.max(...format.usageMarkers.map((marker) => Buffer.byteLength(marker))) - 1; const selectPath = selectUsageFields(format.selectFields); + const checkMarkers = (segment: Buffer) => { + if (carriesUsage) return; + const boundary = Buffer.concat([markerTail, segment.subarray(0, markerOverlap)]); + carriesUsage = mightCarryUsage(segment) || mightCarryUsage(boundary); + markerTail = Buffer.from( + segment.length >= markerOverlap + ? segment.subarray(segment.length - markerOverlap) + : Buffer.concat([markerTail, segment]).subarray(-markerOverlap), + ); + }; const append = (segment: Buffer) => { - if (!streaming && pendingBytes + segment.length <= streamingThresholdBytes) { + if (!oversized && pendingBytes + segment.length <= streamingThresholdBytes) { // Copy: the next read overwrites the chunk this segment points into. if (segment.length > 0) pendingChunks.push(Buffer.from(segment)); pendingBytes += segment.length; return; } - if (!streaming) { - // Usage has no import-history budget; retain selected fields regardless of surrounding tool payload size. - streaming = createTranscriptJsonReader(() => {}, selectPath, { maxDepth: Infinity }); - decoder = new NodeStringDecoder.StringDecoder("utf8"); - for (const pending of pendingChunks) streaming.write(decoder.write(pending)); + if (!oversized) { + oversized = true; + if (compressed) { + // Usage has no import-history budget; retain selected fields regardless of surrounding tool payload size. + streaming = createTranscriptJsonReader(() => {}, selectPath, { maxDepth: Infinity }); + decoder = new NodeStringDecoder.StringDecoder("utf8"); + for (const pending of pendingChunks) streaming.write(decoder.write(pending)); + } else { + for (const pending of pendingChunks) checkMarkers(pending); + } pendingChunks = []; pendingBytes = 0; } - streaming.write(decoder!.write(segment)); + if (streaming) streaming.write(decoder!.write(segment)); + else checkMarkers(segment); }; - const finish = (lineState: State, out: UsageRecord[]) => { + const finish = async (lineState: State, out: UsageRecord[], end: number) => { if (streaming) { streaming.write(decoder!.end()); const projected = streaming.finish(); for (const record of format.parseProjected(projected, lineState)) out.push(record); + } else if (oversized) { + if (carriesUsage) { + // Replay only usage-bearing oversized lines; tool output can otherwise cost minutes of JSON tokenization. + const replay = createTranscriptJsonReader(() => {}, selectPath, { maxDepth: Infinity }); + const replayDecoder = new NodeStringDecoder.StringDecoder("utf8"); + const line = NodeFS.createReadStream(filePath, { + fd: handle.fd, + start: resumeOffset, + end: end - 1, + autoClose: false, + highWaterMark: 1024 * 1024, + }) as AsyncIterable; + for await (const chunk of line) replay.write(replayDecoder.write(chunk)); + replay.write(replayDecoder.end()); + for (const record of format.parseProjected(replay.finish(), lineState)) out.push(record); + } } else if (pendingBytes > 0) { const line = pendingChunks.length === 1 @@ -429,8 +466,11 @@ export async function readTranscriptRecords( } pendingChunks = []; pendingBytes = 0; + oversized = false; streaming = undefined; decoder = undefined; + carriesUsage = false; + markerTail = Buffer.alloc(0); }; for await (const chunk of readChunks(handle, start, compressed)) { let lineStart = 0; @@ -440,11 +480,15 @@ export async function readTranscriptRecords( append(chunk.subarray(lineStart)); break; } - if (!streaming && pendingBytes === 0) { + if ( + !oversized && + pendingBytes === 0 && + newlineIndex - lineStart <= streamingThresholdBytes + ) { parseLine(chunk.subarray(lineStart, newlineIndex), state, records); } else { append(chunk.subarray(lineStart, newlineIndex)); - finish(state, records); + await finish(state, records, scanOffset + newlineIndex); } lineStart = newlineIndex + 1; resumeOffset = scanOffset + lineStart; @@ -457,7 +501,7 @@ export async function readTranscriptRecords( const tailRecords: UsageRecord[] = []; const tailState = decodeState === undefined ? state : Option.getOrThrow(decodeState(encodedState)); - finish(tailState, tailRecords); + await finish(tailState, tailRecords, scanOffset); const guardLength = compressed ? 0 : Math.min(GUARD_LENGTH, resumeOffset); let guardHash = 0; diff --git a/apps/server/src/usage/usageTranscriptStreaming.test.ts b/apps/server/src/usage/usageTranscriptStreaming.test.ts index 635b5ae4e6bec..1416ea3df3b1d 100644 --- a/apps/server/src/usage/usageTranscriptStreaming.test.ts +++ b/apps/server/src/usage/usageTranscriptStreaming.test.ts @@ -2,8 +2,10 @@ import * as NodeFSP from "node:fs/promises"; import * as NodeOS from "node:os"; import * as NodePath from "node:path"; +import * as NodeZlib from "node:zlib"; import { afterEach, beforeEach, describe, expect, it } from "@effect/vitest"; +import { vi } from "vite-plus/test"; import { TEST_FORMATS } from "./usageTestFormats.ts"; import { @@ -108,6 +110,80 @@ async function scan( } describe("large usage records", () => { + it("skips JSON projection for oversized lines without usage markers", async () => { + const path = NodePath.join(dir, "unrelated.jsonl"); + await NodeFSP.writeFile( + path, + JSON.stringify({ type: "response_item", payload: { content } }) + + "\n" + + codex.map((line) => JSON.stringify(line)).join("\n") + + "\n", + ); + const parseProjected = vi.fn(TEST_FORMATS.codex.parseProjected); + const result = await readWithDefaultThreshold( + path, + { ...TEST_FORMATS.codex, parseProjected }, + undefined, + { streamingThresholdBytes: 256 * 1024 }, + ); + expect(result?.records.map((record) => record.totals)).toEqual([ + { + uncachedInputTokens: 75, + cachedInputTokens: 20, + cacheCreationTokens: 5, + outputTokens: 99, + reasoningTokens: 10, + }, + ]); + expect(parseProjected).not.toHaveBeenCalled(); + }); + + it("replays an oversized usage line whose marker spans a read boundary", async () => { + const path = NodePath.join(dir, "split-marker.jsonl"); + const head = codex + .slice(0, 2) + .map((line) => JSON.stringify(line) + "\n") + .join(""); + const usageLine = (pad: number) => + JSON.stringify({ pad: "x".repeat(pad), ...codex[2], tail: "y".repeat(16 * 1024) }) + "\n"; + const markerOffset = Buffer.from(usageLine(0)).indexOf('"token_count"'); + const readBoundary = 256 * 1024; + const pad = readBoundary - 6 - Buffer.byteLength(head) - markerOffset; + await NodeFSP.writeFile(path, head + usageLine(pad)); + const parseProjected = vi.fn(TEST_FORMATS.codex.parseProjected); + const result = await readWithDefaultThreshold( + path, + { ...TEST_FORMATS.codex, parseProjected }, + undefined, + { streamingThresholdBytes: readBoundary }, + ); + expect(result?.records.map((record) => record.totals)).toEqual([ + { + uncachedInputTokens: 75, + cachedInputTokens: 20, + cacheCreationTokens: 5, + outputTokens: 99, + reasoningTokens: 10, + }, + ]); + expect(parseProjected).toHaveBeenCalledTimes(1); + }); + + it("projects oversized usage lines in compressed transcripts", async () => { + const path = NodePath.join(dir, "archived.jsonl.zst"); + await NodeFSP.writeFile(path, NodeZlib.zstdCompressSync(JSON.stringify(claude()) + "\n")); + const result = await readTranscriptRecords(path, "claude"); + expect(result?.records.map((record) => record.totals)).toEqual([ + { + uncachedInputTokens: 100, + cachedInputTokens: 20, + cacheCreationTokens: 5, + outputTokens: 99, + reasoningTokens: 0, + }, + ]); + }); + it("streams Muse completion envelopes and resumes a partial appended event", async () => { const entry = { schema_version: 1,