From 24d43bf58a01109365ad8a4a0a234d141b999320 Mon Sep 17 00:00:00 2001 From: Resham Joshi <65915470+iamtoruk@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:25:07 -0700 Subject: [PATCH] fix(codex): read Zstandard-compressed rollouts Codex now moves cold rollouts to rollout-*.jsonl.zst. Discovery, the first-line probe, the parser, fork parent lookup, codex-tps and the context view read them through a zstd decode; a session present in both forms counts once from the plain file. Compressed files are parsed whole and never resumed by byte offset. Without node:zlib zstd (Node < 22.15) they are skipped with one notice. DAILY_CACHE_VERSION 76 so settled days pick up the recovered Codex usage. --- src/codex-fork-replay.ts | 6 +-- src/codex-throughput.ts | 79 +++++++++++++++++++++-------------- src/context-tree-codex.ts | 16 ++++--- src/daily-cache.ts | 6 ++- src/fs-utils.ts | 58 ++++++++++++++++++++++++- src/launcher-homes.ts | 58 +++++++++++++++---------- src/parser.ts | 4 +- src/providers/codex.ts | 28 +++++++------ tests/fs-utils.test.ts | 18 ++++++++ tests/providers/codex.test.ts | 68 ++++++++++++++++++++++++++++++ 10 files changed, 265 insertions(+), 76 deletions(-) diff --git a/src/codex-fork-replay.ts b/src/codex-fork-replay.ts index ac75f93ea..4abb0bf8a 100644 --- a/src/codex-fork-replay.ts +++ b/src/codex-fork-replay.ts @@ -2,7 +2,7 @@ import { stat } from 'fs/promises' import { basename, dirname } from 'path' import { readSessionLines } from './fs-utils.js' -import { walkRolloutFiles } from './launcher-homes.js' +import { rolloutStem, walkRolloutFiles } from './launcher-homes.js' /** Maximum timestamp gap between copied records in Codex's fork replay burst. */ export const CODEX_FORK_REPLAY_MAX_GAP_MS = 1000 @@ -84,7 +84,7 @@ export function codexReplayUsageIdentity(info: { total_token_usage?: CodexUsageC export const codexReplayResponseIdentity = (responseId: string): string => `response:${responseId}` -// Rollout names are `rollout-YYYY-MM-DDTHH-MM-SS-.jsonl`. +// Rollout names are `rollout-YYYY-MM-DDTHH-MM-SS-.jsonl` (or `.jsonl.zst`). const ROLLOUT_ID_OFFSET = 'rollout-YYYY-MM-DDTHH-MM-SS-'.length const rolloutIndexes = new Map>() @@ -92,7 +92,7 @@ function findRollout(home: string, id: string): string | undefined { let index = rolloutIndexes.get(home) if (!index?.has(id)) { index = new Map() - walkRolloutFiles(home, path => index!.set(basename(path).slice(ROLLOUT_ID_OFFSET, -'.jsonl'.length), path)) + walkRolloutFiles(home, path => index!.set(rolloutStem(basename(path)).slice(ROLLOUT_ID_OFFSET), path)) if (!index.has(id)) index.set(id, null) rolloutIndexes.set(home, index) } diff --git a/src/codex-throughput.ts b/src/codex-throughput.ts index d623f700a..8f562dc56 100644 --- a/src/codex-throughput.ts +++ b/src/codex-throughput.ts @@ -1,6 +1,7 @@ import { open, stat } from 'node:fs/promises' import { StringDecoder } from 'node:string_decoder' +import { createTextReadStream } from './fs-utils.js' import { billableOutputTokens } from './models.js' import { isCodexForkReplay, startCodexForkReplay, type CodexForkReplayState } from './codex-fork-replay.js' @@ -448,10 +449,55 @@ export class CodexThroughputReader { if (state.points.length > 10000) state.points.splice(0, state.points.length - 10000) } + private consume(buffer: Buffer): void { + let chunk = this.decoder.write(buffer) + while (chunk.length > 0) { + const newlineIndex = chunk.search(/\r?\n/) + const segment = newlineIndex >= 0 ? chunk.slice(0, newlineIndex) : chunk + this.pending += segment + this.scanDurationSegment(segment) + if (newlineIndex < 0) break + const newlineLength = chunk[newlineIndex] === '\r' ? 2 : 1 + const line = this.pending + const durationOverride = this.pendingDurationMs + this.pending = '' + this.pendingDurationMs = undefined + this.scanDepth = 0 + this.scanPayloadDepth = undefined + this.scanInString = false + this.scanEscape = false + this.scanString = '' + this.scanLastString = '' + this.scanAwaitingColon = false + this.scanCurrentKey = undefined + this.scanCapture = undefined + this.processLine(line, durationOverride) + chunk = chunk.slice(newlineIndex + newlineLength) + } + if (this.pending.length > MAX_PENDING_LINE_CHARS) { + const body = this.pending.startsWith(TRUNCATION_MARKER) + ? this.pending.slice(TRUNCATION_MARKER.length) + : this.pending + this.pending = TRUNCATION_MARKER + body.slice(0, 256 * 1024) + body.slice(-256 * 1024) + } + } + async update(filePath: string, limit = 10, finalize = false): Promise { const info = await stat(filePath) if (info.size < this.offset) this.reset() - const bytesToRead = info.size - this.offset + // A compressed rollout is immutable and its byte offsets are not text + // offsets, so it is decoded whole, once. + const compressed = filePath.endsWith('.zst') + if (compressed && this.offset === 0) { + try { + for await (const buffer of createTextReadStream(filePath)) this.consume(buffer as Buffer) + } catch (err) { + this.reset() + throw err + } + this.offset = info.size + } + const bytesToRead = compressed ? 0 : info.size - this.offset if (bytesToRead > 0) { const file = await open(filePath, 'r') try { @@ -462,36 +508,7 @@ export class CodexThroughputReader { if (bytesRead === 0) break position += bytesRead this.offset = position - let chunk = this.decoder.write(buffer.subarray(0, bytesRead)) - while (chunk.length > 0) { - const newlineIndex = chunk.search(/\r?\n/) - const segment = newlineIndex >= 0 ? chunk.slice(0, newlineIndex) : chunk - this.pending += segment - this.scanDurationSegment(segment) - if (newlineIndex < 0) break - const newlineLength = chunk[newlineIndex] === '\r' ? 2 : 1 - const line = this.pending - const durationOverride = this.pendingDurationMs - this.pending = '' - this.pendingDurationMs = undefined - this.scanDepth = 0 - this.scanPayloadDepth = undefined - this.scanInString = false - this.scanEscape = false - this.scanString = '' - this.scanLastString = '' - this.scanAwaitingColon = false - this.scanCurrentKey = undefined - this.scanCapture = undefined - this.processLine(line, durationOverride) - chunk = chunk.slice(newlineIndex + newlineLength) - } - if (this.pending.length > MAX_PENDING_LINE_CHARS) { - const body = this.pending.startsWith(TRUNCATION_MARKER) - ? this.pending.slice(TRUNCATION_MARKER.length) - : this.pending - this.pending = TRUNCATION_MARKER + body.slice(0, 256 * 1024) + body.slice(-256 * 1024) - } + this.consume(buffer.subarray(0, bytesRead)) } } finally { await file.close() diff --git a/src/context-tree-codex.ts b/src/context-tree-codex.ts index c7d66c06e..1ea627175 100644 --- a/src/context-tree-codex.ts +++ b/src/context-tree-codex.ts @@ -3,7 +3,8 @@ import { existsSync } from 'fs' import { basename, join } from 'path' import { homedir } from 'os' -import { readSessionLines } from './fs-utils.js' +import { readSessionLines, readTextHeadSync } from './fs-utils.js' +import { isRolloutFileName } from './launcher-homes.js' import { estimateTokensFromChars } from './token-estimate.js' import { add, @@ -243,7 +244,7 @@ export async function buildCodexContextTree(session: SessionRef): Promise { return [] } const rollouts: RolloutFile[] = [] + const all = new Set(files) for (const rel of files) { - const match = ROLLOUT_RE.exec(basename(rel)) - if (match) rollouts.push({ filePath: join(root, rel), sessionId: match[1] }) + if (rel.endsWith('.zst') && all.has(rel.slice(0, -4))) continue + const name = basename(rel) + const match = ROLLOUT_RE.exec(name) + if (match && isRolloutFileName(name)) rollouts.push({ filePath: join(root, rel), sessionId: match[1] }) } return rollouts } @@ -299,7 +303,9 @@ export async function findCodexSession(idPrefix: string): Promise { let chunk: string try { - chunk = await readChunk(ref.filePath, 0, 262_144) + chunk = ref.filePath.endsWith('.zst') + ? readTextHeadSync(ref.filePath, 262_144).toString('utf-8') + : await readChunk(ref.filePath, 0, 262_144) } catch { return { project: '', title: '' } } diff --git a/src/daily-cache.ts b/src/daily-cache.ts index 3ded9cc7b..525a29d8f 100644 --- a/src/daily-cache.ts +++ b/src/daily-cache.ts @@ -361,7 +361,11 @@ import type { DateRange, ProjectSummary } from './types.js' // copilotCredits, and a chat session that also ran on the agent host is left // to it. That twin drops calls, so copilot joins // PENDING_REDERIVE_PROVIDER_VERSIONS at 75. -export const DAILY_CACHE_VERSION = 75 +// v76: Codex reads the Zstandard-compressed rollouts (`rollout-*.jsonl.zst`) +// Codex moves cold sessions into, which were skipped, so days settled after +// that move miss their Codex usage. Calls and cost only rise, so no +// PENDING_REDERIVE_PROVIDER_VERSIONS entry is needed. +export const DAILY_CACHE_VERSION = 76 const MIN_SUPPORTED_VERSION = 28 /// Providers whose per-day CALL COUNT means something different at diff --git a/src/fs-utils.ts b/src/fs-utils.ts index 401a54ec4..17030664c 100644 --- a/src/fs-utils.ts +++ b/src/fs-utils.ts @@ -1,5 +1,7 @@ import { readFile, stat } from 'fs/promises' -import { readFileSync, statSync, createReadStream } from 'fs' +import { readFileSync, statSync, createReadStream, openSync, readSync, closeSync } from 'fs' +import { pipeline, type Readable } from 'stream' +import zlib from 'zlib' // Hard cap well below V8's 512 MB string limit. Callers that need line-by-line // processing should use readSessionLines(), which avoids materializing the @@ -23,6 +25,58 @@ function warn(msg: string): void { if (verbose()) process.stderr.write(`codeburn: ${msg}\n`) } +// zstd landed in node:zlib in 22.15 / 23.8; the package floor is 22.13, so the +// exports are optional and callers check zstdSupported before handing over a +// `.zst` path. +const zstd = zlib as unknown as { + createZstdDecompress?: typeof zlib.createZstdDecompress + zstdDecompressSync?: typeof zlib.zstdDecompressSync +} +export const zstdSupported = !!zstd.createZstdDecompress && !!zstd.zstdDecompressSync +const MAX_ZSTD_HEAD_DECODED_BYTES = 64 * 1024 * 1024 + +function requireZstd(filePath: string): void { + if (!zstdSupported) throw new Error(`${filePath} is zstd-compressed; reading it needs Node 22.15+`) +} + +/** Reads a file's text, decompressing a `.zst` file on the fly. */ +export function createTextReadStream(filePath: string, options?: { start?: number; end?: number }): Readable { + if (!filePath.endsWith('.zst')) return createReadStream(filePath, options) + requireZstd(filePath) + const raw = createReadStream(filePath, options) + // pipeline, not pipe: a read error has to fail the decoder too, or its + // consumer waits forever. + return pipeline(raw, zstd.createZstdDecompress!(), () => {}) +} + +/** Up to `maxBytes` of a file's text from its start; for a `.zst` file only the compressed head is read. */ +export function readTextHeadSync(filePath: string, maxBytes: number): Buffer { + const fd = openSync(filePath, 'r') + let head: Buffer + try { + const buf = Buffer.alloc(maxBytes) + head = buf.subarray(0, readSync(fd, buf, 0, maxBytes, 0)) + } finally { + closeSync(fd) + } + if (!filePath.endsWith('.zst') || head.length === 0) return head + requireZstd(filePath) + // ZSTD_e_continue decodes a cut-off frame as far as it goes instead of + // rejecting it as truncated; a head that inflates past the cap is retried + // with less input. + for (;;) { + try { + return zstd.zstdDecompressSync!(head, { + finishFlush: zlib.constants.ZSTD_e_continue, + maxOutputLength: MAX_ZSTD_HEAD_DECODED_BYTES, + }).subarray(0, maxBytes) + } catch (err) { + if ((err as NodeJS.ErrnoException).code !== 'ERR_BUFFER_TOO_LARGE' || head.length < 2) throw err + head = head.subarray(0, head.length >> 1) + } + } +} + // Always surfaced (not verbose-gated): dropping an entire session file silently // understates reported usage with no signal, so oversize skips use this. function notice(msg: string): void { @@ -116,7 +170,7 @@ export async function* readSessionLines( return } - const stream = createReadStream( + const stream = createTextReadStream( filePath, options.startByteOffset !== undefined ? { start: options.startByteOffset } : undefined, ) diff --git a/src/launcher-homes.ts b/src/launcher-homes.ts index dcbb3e571..a807c8db2 100644 --- a/src/launcher-homes.ts +++ b/src/launcher-homes.ts @@ -1,7 +1,9 @@ -import { closeSync, existsSync, openSync, readdirSync, readSync, realpathSync } from 'fs' +import { existsSync, readdirSync, realpathSync } from 'fs' import { homedir } from 'os' import { join, resolve, sep } from 'path' +import { readTextHeadSync, zstdSupported } from './fs-utils.js' + export type LauncherNote = { name: string path: string @@ -49,33 +51,47 @@ export const FIRST_LINE_READ_CAP = 1024 * 1024 * Missing / unreadable files return undefined — callers must not drop those. */ export function rolloutFileSessionId(filePath: string): string | undefined { try { - const fd = openSync(filePath, 'r') - try { - const buf = Buffer.alloc(FIRST_LINE_READ_CAP) - const n = readSync(fd, buf, 0, buf.length, 0) - if (n === 0) return undefined - const slice = buf.subarray(0, n) - const nl = slice.indexOf(0x0a) - // Hit the 1 MiB cap without a newline: incomplete, not a parseable id. - if (nl === -1 && n === FIRST_LINE_READ_CAP) return undefined - const line = (nl === -1 ? slice : slice.subarray(0, nl)).toString('utf8') - if (!line) return undefined - const parsed: unknown = JSON.parse(line) - if (!parsed || typeof parsed !== 'object') return undefined - const payload = (parsed as { payload?: { session_id?: unknown } }).payload - const id = payload?.session_id - return typeof id === 'string' && id.length > 0 ? id : undefined - } finally { - closeSync(fd) - } + const slice = readTextHeadSync(filePath, FIRST_LINE_READ_CAP) + if (slice.length === 0) return undefined + const nl = slice.indexOf(0x0a) + // Hit the 1 MiB cap without a newline: incomplete, not a parseable id. + if (nl === -1 && slice.length === FIRST_LINE_READ_CAP) return undefined + const line = (nl === -1 ? slice : slice.subarray(0, nl)).toString('utf8') + if (!line) return undefined + const parsed: unknown = JSON.parse(line) + if (!parsed || typeof parsed !== 'object') return undefined + const payload = (parsed as { payload?: { session_id?: unknown } }).payload + const id = payload?.session_id + return typeof id === 'string' && id.length > 0 ? id : undefined } catch { return undefined } } +let zstdSkipNoticeSent = false + +/** `rollout-*.jsonl`, or the `.jsonl.zst` form Codex compresses cold rollouts + * into, when this Node can decompress it. */ +export function isRolloutFileName(name: string): boolean { + if (!name.startsWith('rollout-')) return false + if (name.endsWith('.jsonl')) return true + if (!name.endsWith('.jsonl.zst')) return false + if (zstdSupported) return true + if (!zstdSkipNoticeSent) { + zstdSkipNoticeSent = true + process.stderr.write('codeburn: codex: skipping compressed rollouts (.jsonl.zst); counting them needs Node 22.15+ (node:zlib zstd)\n') + } + return false +} + +/** Rollout filename without `.jsonl` / `.jsonl.zst`. */ +export function rolloutStem(name: string): string { + return name.replace(/\.jsonl(\.zst)?$/, '') +} + export function walkRolloutFiles(codexDir: string, visit: (filePath: string) => void): void { const take = (dir: string, file: string) => { - if (file.startsWith('rollout-') && file.endsWith('.jsonl')) visit(join(dir, file)) + if (isRolloutFileName(file)) visit(join(dir, file)) } const sessionsDir = join(codexDir, 'sessions') try { diff --git a/src/parser.ts b/src/parser.ts index 7fc0d5dbd..df7d18c80 100644 --- a/src/parser.ts +++ b/src/parser.ts @@ -3811,7 +3811,9 @@ export async function parseProviderSources( if (await readCachedCodexResults(source.path)) continue workerJobs.push({ kind: 'codex', source }) workerPaths.add(source.path) - pendingBytes += fp.sizeBytes + // The worker memory budget scales with decoded size; real compressed + // rollouts decode to 2.3-4x their size on disk. + pendingBytes += source.path.endsWith('.zst') ? fp.sizeBytes * 4 : fp.sizeBytes } } const decision = workerJobs.length > 0 diff --git a/src/providers/codex.ts b/src/providers/codex.ts index d56130135..8b8049e63 100644 --- a/src/providers/codex.ts +++ b/src/providers/codex.ts @@ -1,10 +1,9 @@ import { readdir, stat } from 'fs/promises' -import { createReadStream } from 'fs' import { createInterface } from 'readline' import { basename, join, resolve } from 'path' import { homedir } from 'os' -import { FS_SCAN_CONCURRENCY, mapWithConcurrency, readSessionLines } from '../fs-utils.js' +import { createTextReadStream, FS_SCAN_CONCURRENCY, mapWithConcurrency, readSessionLines } from '../fs-utils.js' import { billableOutputTokens, calculateCost, getModelCosts, pricingModelAt } from '../models.js' import { readCachedCodexResults, writeCachedCodexResults, getCachedCodexProject, fingerprintFile, type CodexFileFingerprint } from '../codex-cache.js' import { mergeToolIntervals } from '../codex-throughput.js' @@ -14,7 +13,7 @@ import { estimateTokensFromChars } from '../token-estimate.js' import { wslHomes } from '../wsl.js' import type { ToolCall } from '../types.js' import type { Provider, ProbeRoot, SessionSource, SessionParser, ParsedProviderCall } from './types.js' -import { defaultBilledCodexHome, defaultLauncherRoots, FIRST_LINE_READ_CAP, isNestedLauncherCodexHome, listRolloutSessionIds, rolloutFileSessionId, sameCodexHome } from '../launcher-homes.js' +import { defaultBilledCodexHome, defaultLauncherRoots, FIRST_LINE_READ_CAP, isNestedLauncherCodexHome, isRolloutFileName, listRolloutSessionIds, rolloutFileSessionId, rolloutStem, sameCodexHome } from '../launcher-homes.js' const modelDisplayNames: Record = { 'codex-auto-review': 'Codex Auto Review', @@ -203,8 +202,7 @@ async function readFirstLine(filePath: string): Promise { // would miss the trailing newline and reject the session as invalid. // Stream the file via readline so we can read the first line up to // FIRST_LINE_READ_CAP, which keeps memory bounded if the file has no newline. - const stream = createReadStream(filePath, { - encoding: 'utf-8', + const stream = createTextReadStream(filePath, { start: 0, end: FIRST_LINE_READ_CAP - 1, }) @@ -627,7 +625,7 @@ async function discoverSessionsInDir(codexDir: string): Promise const dated = (await mapWithConcurrency(dayDirs, FS_SCAN_CONCURRENCY, async dayDir => (await readdir(dayDir).catch(() => [] as string[])) - .filter(f => f.startsWith('rollout-') && f.endsWith('.jsonl')) + .filter(isRolloutFileName) .map(f => ({ file: f, path: join(dayDir, f) })), )).flat() @@ -637,15 +635,19 @@ async function discoverSessionsInDir(codexDir: string): Promise // archived copies, while basename dedup below prevents double discovery. const archivedDir = join(codexDir, 'archived_sessions') const archived = (await readdir(archivedDir).catch(() => [] as string[])) - .filter(f => f.startsWith('rollout-') && f.endsWith('.jsonl')) + .filter(isRolloutFileName) .map(f => ({ file: f, path: join(archivedDir, f) })) // Dedup before the reads so a basename is opened once, and keep the dated // pass ahead of the archived one — the same precedence the serial walk had. + // While Codex compresses a rollout both forms exist briefly; the plain one + // goes first because it may still be growing. const candidates: Array<{ file: string; path: string }> = [] - for (const entry of [...dated, ...archived]) { - if (seenBasenames.has(entry.file)) continue - seenBasenames.add(entry.file) + const ordered = [...dated, ...archived].sort((a, b) => Number(a.file.endsWith('.zst')) - Number(b.file.endsWith('.zst'))) + for (const entry of ordered) { + const stem = rolloutStem(entry.file) + if (seenBasenames.has(stem)) continue + seenBasenames.add(stem) candidates.push(entry) } const discovered = await mapWithConcurrency(candidates, FS_SCAN_CONCURRENCY, c => discoverSessionFile(c.path)) @@ -1021,7 +1023,7 @@ function createParser(source: SessionSource, seenKeys: Set, capture?: { if (!entry) continue if (entry.type === 'session_meta') { - sessionId = entry.payload?.session_id ?? basename(source.path, '.jsonl') + sessionId = entry.payload?.session_id ?? rolloutStem(basename(source.path)) // Guarded for the same reason as discoverSessionFile: parseCodexLine // hands back an unchecked JSON.parse cast, and a non-string cwd would // ride into projectPath/workingDirectory where the parser's path @@ -1527,7 +1529,9 @@ function createParser(source: SessionSource, seenKeys: Set, capture?: { // Flush the final task, which has no following task_started to trigger it. results.push(...pendingTaskCalls) - const resumeWrite = resumeState ? { offset: resumeOffset, state: resumeState, callCount: resumeCallCount } : undefined + // Offsets into decompressed text cannot pass the cache's raw-byte append + // check, and a compressed rollout is never appended to anyway. + const resumeWrite = resumeState && !source.path.endsWith('.zst') ? { offset: resumeOffset, state: resumeState, callCount: resumeCallCount } : undefined if (capture) { capture.write = { project: source.project, fingerprint: fp, ...(resumeWrite ? { resume: resumeWrite } : {}) } } else { diff --git a/tests/fs-utils.test.ts b/tests/fs-utils.test.ts index 13e746409..4e01809b2 100644 --- a/tests/fs-utils.test.ts +++ b/tests/fs-utils.test.ts @@ -2,12 +2,14 @@ import { describe, it, expect, afterEach, vi } from 'vitest' import { mkdtemp, writeFile, rm } from 'fs/promises' import { tmpdir } from 'os' import { join } from 'path' +import * as zlib from 'zlib' import { MAX_SESSION_FILE_BYTES, MAX_STREAM_SESSION_FILE_BYTES, readSessionFile, readSessionLines, + readTextHeadSync, } from '../src/fs-utils.js' describe('readSessionFile', () => { @@ -193,3 +195,19 @@ describe('readSessionLines', () => { expect(MAX_STREAM_SESSION_FILE_BYTES).toBe(4 * 1024 * 1024 * 1024) }) }) + +describe.skipIf(!('zstdCompressSync' in zlib))('readTextHeadSync', () => { + it('returns the head of a .zst whose compressed head inflates past the decode cap', async () => { + const dir = await mkdtemp(join(tmpdir(), 'codeburn-head-')) + try { + const file = join(dir, 'rollout.jsonl.zst') + const text = Buffer.concat([Buffer.from('{"type":"session_meta"}\n'), Buffer.alloc(80 * 1024 * 1024, 0x61)]) + await writeFile(file, (zlib as unknown as { zstdCompressSync(b: Buffer): Buffer }).zstdCompressSync(text)) + const head = readTextHeadSync(file, 1024 * 1024) + expect(head.length).toBe(1024 * 1024) + expect(head.toString('utf8').split('\n')[0]).toBe('{"type":"session_meta"}') + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) +}) diff --git a/tests/providers/codex.test.ts b/tests/providers/codex.test.ts index 0a98cfb33..3bed041cb 100644 --- a/tests/providers/codex.test.ts +++ b/tests/providers/codex.test.ts @@ -2,9 +2,12 @@ import { describe, it, expect, beforeEach, afterEach } from 'vitest' import { mkdtemp, mkdir, writeFile, rm, stat } from 'fs/promises' import { basename, join } from 'path' import { tmpdir } from 'os' +import zlib from 'zlib' import { createCodexProvider } from '../../src/providers/codex.js' import { clearCodexMemCaches, CODEX_CACHE_VERSION, codexCacheFileName } from '../../src/codex-cache.js' +import { rolloutFileSessionId } from '../../src/launcher-homes.js' +import { readCodexThroughput } from '../../src/codex-throughput.js' import { calculateCost, parseLiteLLMEntry, restorePricingState, snapshotPricingState } from '../../src/models.js' import type { ParsedProviderCall } from '../../src/providers/types.js' @@ -2128,3 +2131,68 @@ describe('codex provider - priority service tier (#1616)', () => { }) }) + +describe.skipIf(typeof (zlib as { zstdCompressSync?: unknown }).zstdCompressSync !== 'function')('codex provider - Zstandard-compressed rollouts', () => { + const lines = [ + // Larger than a stream chunk, so the first line spans several decoded chunks. + JSON.stringify({ + type: 'session_meta', + timestamp: '2026-04-14T10:00:00Z', + payload: { cwd: '/Users/test/zst', originator: 'codex-cli', session_id: 'sess-zst', base_instructions: { text: 'x'.repeat(100_000) } }, + }), + functionCall('exec_command'), + tokenCount({ timestamp: '2026-04-14T10:01:00Z', last: { input: 100, output: 50 }, total: { total: 150 } }), + tokenCount({ timestamp: '2026-04-14T10:02:00Z', last: { input: 300, cached: 100, output: 70 }, total: { total: 620 } }), + ] + + async function writeCompressed(path: string): Promise { + await writeFile(path, zlib.zstdCompressSync(Buffer.from(lines.join('\n') + '\n'))) + return path + } + + async function parseAll(dir: string): Promise { + const provider = createCodexProvider(dir) + const seenKeys = new Set() + const calls: ParsedProviderCall[] = [] + for (const session of await provider.discoverSessions()) { + for await (const call of provider.createSessionParser(session, seenKeys).parse()) calls.push(call) + } + return calls + } + + it('discovers and parses a .jsonl.zst rollout exactly like the plain one', async () => { + const plainHome = join(tmpDir, 'plain') + const zstHome = join(tmpDir, 'zst') + const plainPath = await writeSession(plainHome, '2026-04-14', 'rollout-2026-04-14T10-00-00-zst.jsonl', lines) + const zstPath = await writeCompressed((await writeSession(zstHome, '2026-04-14', 'rollout-2026-04-14T10-00-00-zst.jsonl', lines)) + '.zst') + await rm(zstPath.slice(0, -'.zst'.length)) + + const sessions = await createCodexProvider(zstHome).discoverSessions() + expect(sessions).toEqual([{ path: zstPath, project: 'Users-test-zst', provider: 'codex' }]) + + const plainCalls = await parseAll(plainHome) + expect(plainCalls).toHaveLength(2) + expect(await parseAll(zstHome)).toEqual(plainCalls) + expect(rolloutFileSessionId(zstPath)).toBe('sess-zst') + expect(await readCodexThroughput(zstPath)).toEqual(await readCodexThroughput(plainPath)) + }) + + it('counts a rollout present in both forms once, from the plain file', async () => { + const plainPath = await writeSession(tmpDir, '2026-04-14', 'rollout-2026-04-14T10-00-00-zst.jsonl', lines) + await writeCompressed(plainPath + '.zst') + + const sessions = await createCodexProvider(tmpDir).discoverSessions() + expect(sessions.map(s => s.path)).toEqual([plainPath]) + expect(await parseAll(tmpDir)).toHaveLength(2) + }) + + it('discovers a .jsonl.zst rollout in archived_sessions', async () => { + const archivedDir = join(tmpDir, 'archived_sessions') + await mkdir(archivedDir, { recursive: true }) + const zstPath = await writeCompressed(join(archivedDir, 'rollout-2026-04-14T10-00-00-zst.jsonl.zst')) + + const sessions = await createCodexProvider(tmpDir).discoverSessions() + expect(sessions).toEqual([{ path: zstPath, project: 'Users-test-zst', provider: 'codex' }]) + expect(await parseAll(tmpDir)).toHaveLength(2) + }) +})