diff --git a/src/codex-cache.ts b/src/codex-cache.ts index a7f2f7c1..ca3f306a 100644 --- a/src/codex-cache.ts +++ b/src/codex-cache.ts @@ -22,7 +22,8 @@ import type { ParsedProviderCall } from './providers/types.js' const CODEX_CACHE_VERSION = 8 const CACHE_FILE = 'codex-results.json' -type FileFingerprint = { dev: number; ino: number; mtimeMs: number; sizeBytes: number } +export type CodexFileFingerprint = { dev: number; ino: number; mtimeMs: number; sizeBytes: number } +type FileFingerprint = CodexFileFingerprint type FileEntry = { // Absent on entries written before the resume support landed. diff --git a/src/parse-worker.ts b/src/parse-worker.ts index da0d8a99..3329ebd7 100644 --- a/src/parse-worker.ts +++ b/src/parse-worker.ts @@ -1,6 +1,8 @@ import { parentPort, workerData } from 'worker_threads' import { restorePricingState, type PricingSnapshot } from './models.js' +import type { ParseJob } from './parse-workers.js' import { parseClaudeFileFull } from './parser.js' +import { parseCodexFileFull } from './providers/codex.js' const port = parentPort if (!port) throw new Error('parse-worker must be started as a worker thread') @@ -10,14 +12,21 @@ restorePricingState((workerData as { pricing: PricingSnapshot }).pricing) // The parsed turns go back as a JSON string rather than as a live object graph: // structured-cloning a whole corpus of turns costs more than the parallel parse // saves, while a string is a single copy the parent re-parses at memcpy speed. -// `msgIds` is every streaming message id this file claimed; the parent uses it to -// prove no earlier file already owned one before installing the result. -port.on('message', (msg: { filePath: string }) => { +// `msgIds` / `keys` is every dedup key this file claimed; the parent uses it to +// prove no earlier file already owned one before installing the result. A Codex +// job also carries back the cache entry it would have written, because the cache +// module's per-directory state belongs to the parent, not to a thread. +port.on('message', (msg: ParseJob) => { void (async () => { try { - const seenMsgIds = new Set() - const parsed = await parseClaudeFileFull(msg.filePath, seenMsgIds) - port.postMessage({ json: parsed === null ? null : JSON.stringify({ ...parsed, msgIds: [...seenMsgIds] }) }) + const seen = new Set() + if (msg.kind === 'codex') { + const parsed = await parseCodexFileFull(msg.source, seen) + port.postMessage({ json: JSON.stringify({ ...parsed, keys: [...seen] }) }) + return + } + const parsed = await parseClaudeFileFull(msg.filePath, seen) + port.postMessage({ json: parsed === null ? null : JSON.stringify({ ...parsed, msgIds: [...seen] }) }) } catch (err) { port.postMessage({ error: err instanceof Error ? err.message : String(err) }) } diff --git a/src/parse-workers.ts b/src/parse-workers.ts index 873c4f6a..842592ae 100644 --- a/src/parse-workers.ts +++ b/src/parse-workers.ts @@ -2,6 +2,7 @@ import { availableParallelism, totalmem } from 'os' import { Worker } from 'worker_threads' import { snapshotPricingState } from './models.js' import type { ClaudeFileParse } from './parser.js' +import type { SessionSource } from './providers/types.js' // Each worker holds one file's entry list plus its serialized result. 256 MB is // the observed high-water mark for the largest real session files; going over the @@ -10,8 +11,12 @@ const PER_WORKER_RSS_BYTES = 256 * 1024 * 1024 const MEMORY_BUDGET_CAP_BYTES = 2 * 1024 * 1024 * 1024 const MIN_AVAILABLE_BYTES = 4 * 1024 * 1024 * 1024 const MIN_FILES_PER_WORKER = 50 -// Below these, a parse is warm/incremental and the thread startup + result -// transfer costs more than the parallelism buys. +const MIN_BYTES_PER_WORKER = 200 * 1024 * 1024 +// Below BOTH of these, a parse is warm/incremental and the thread startup + +// result transfer costs more than the parallelism buys. Either one on its own +// qualifies: a Codex corpus is a few hundred rollouts of which a handful carry +// most of the gigabytes, so a file-count-only gate leaves the biggest workload +// there is (150 huge rollouts) parsing serially. const MIN_PENDING_FILES = 200 const MIN_PENDING_BYTES = 200 * 1024 * 1024 @@ -53,16 +58,23 @@ export function decideParseWorkers( // Workload gates first, so a warm run's log line says "warm", not whatever the // machine happened to look like at that moment. - if (pending.files < MIN_PENDING_FILES) return { workers: 0, reason: `below ${MIN_PENDING_FILES} pending files; ${inputs}` } - if (pending.bytes < MIN_PENDING_BYTES) return { workers: 0, reason: `below ${Math.round(MIN_PENDING_BYTES / 1e6)} MB pending; ${inputs}` } + if (pending.files < MIN_PENDING_FILES && pending.bytes < MIN_PENDING_BYTES) { + return { workers: 0, reason: `below ${MIN_PENDING_FILES} pending files and ${Math.round(MIN_PENDING_BYTES / 1e6)} MB pending; ${inputs}` } + } if (sys.cores <= 2) return { workers: 0, reason: `too few cores; ${inputs}` } if (sys.availableBytes < MIN_AVAILABLE_BYTES) return { workers: 0, reason: `below ${Math.round(MIN_AVAILABLE_BYTES / 1e9)} GB available memory; ${inputs}` } const memoryBudget = Math.min(0.25 * sys.availableBytes, MEMORY_BUDGET_CAP_BYTES) + // Files and bytes each earn threads on their own: a few hundred huge rollouts + // are as parallelisable as a few thousand small transcripts, and gating the + // count on files alone would hand a 6 GB / 60-file workload a single thread. const workers = Math.min( sys.cores - 1, Math.floor(memoryBudget / PER_WORKER_RSS_BYTES), - Math.floor(pending.files / MIN_FILES_PER_WORKER), + Math.max( + Math.floor(pending.files / MIN_FILES_PER_WORKER), + Math.floor(pending.bytes / MIN_BYTES_PER_WORKER), + ), ) return { workers, reason: inputs } } @@ -90,11 +102,19 @@ function workerEntryUrl(): string { return new URL(`./parse-worker${ext}`, import.meta.url).href } -export type ClaudeWorkerResult = - | { ok: true; parsed: (ClaudeFileParse & { msgIds: string[] }) | null } +/// One whole-file parse for a worker to run. Both kinds carry exactly what the +/// serial per-file parse takes, so the worker can run that same function. +export type ParseJob = + | { kind: 'claude'; filePath: string } + | { kind: 'codex'; source: SessionSource } + +export type ParseWorkerResult = + | { ok: true; parsed: T | null } | { ok: false; error: string } -type Task = { filePath: string; resolve: (r: ClaudeWorkerResult) => void } +export type ClaudeWorkerParse = ClaudeFileParse & { msgIds: string[] } + +type Task = { job: ParseJob; resolve: (r: ParseWorkerResult) => void } type WorkerMessage = { json?: string | null; error?: string } @@ -132,13 +152,13 @@ export class ParseWorkerPool { /// Parse one file off-thread. Never rejects: a worker-side failure (or a dead /// pool) comes back as `ok: false` so the caller can fall back to an in-process /// parse and never lose a file to a crashed thread. - submit(filePath: string): Promise { - return new Promise((resolve) => { + submit(job: ParseJob): Promise> { + return new Promise>((resolve) => { if (this.closed || this.workers.length === 0) { resolve({ ok: false, error: 'parse worker pool unavailable' }) return } - this.queue.push({ filePath, resolve }) + this.queue.push({ job, resolve: resolve as (r: ParseWorkerResult) => void }) this.pump() }) } @@ -159,7 +179,7 @@ export class ParseWorkerPool { const worker = this.idle.pop()! const task = this.queue.shift()! this.inflight.set(worker, task) - worker.postMessage({ filePath: task.filePath }) + worker.postMessage(task.job) } } @@ -197,22 +217,22 @@ export class ParseWorkerPool { } } -/// Yield results for `filePaths` in the SAME order they were given, no matter -/// which worker finishes first. Keeps exactly `pool.size` files in flight, so at -/// most that many parsed results are buffered while the caller installs one. -export async function* parseFilesInOrder( +/// Yield results for `jobs` in the SAME order they were given, no matter which +/// worker finishes first. Keeps exactly `pool.size` files in flight, so at most +/// that many parsed results are buffered while the caller installs one. +export async function* parseFilesInOrder( pool: ParseWorkerPool, - filePaths: readonly string[], -): AsyncGenerator { - const inflight: Array> = [] + jobs: readonly ParseJob[], +): AsyncGenerator, void, void> { + const inflight: Array>> = [] let next = 0 const fill = (): void => { - while (inflight.length < Math.max(1, pool.size) && next < filePaths.length) { - inflight.push(pool.submit(filePaths[next++]!)) + while (inflight.length < Math.max(1, pool.size) && next < jobs.length) { + inflight.push(pool.submit(jobs[next++]!)) } } fill() - for (let i = 0; i < filePaths.length; i++) { + for (let i = 0; i < jobs.length; i++) { const result = await inflight.shift()! fill() yield result diff --git a/src/parser.ts b/src/parser.ts index 8cbd579b..f7bf7167 100644 --- a/src/parser.ts +++ b/src/parser.ts @@ -6,7 +6,7 @@ import { calculateCost, calculateLocalModelSavings, getShortModelName, isProxied import { resolveSubagentAttribution, sessionIdentity } from './sessions-report.js' import { normalizeContentBlocks, flatSlice, flatString } from './content-utils.js' import { discoverAllSessions, getProvider } from './providers/index.js' -import { flushCodexCache, withCodexCacheDirectory } from './codex-cache.js' +import { flushCodexCache, readCachedCodexResults, withCodexCacheDirectory, writeCachedCodexResults } from './codex-cache.js' import { antigravityCascadeIdFromPath, flushAntigravityCache, shouldReparseAntigravitySource } from './providers/antigravity.js' import { getClaudeConfigDirs, getDesktopSessionsDirs } from './providers/claude.js' import { isSqliteBusyError } from './sqlite.js' @@ -31,7 +31,8 @@ import { saveCache, } from './session-cache.js' import { acquireCacheRefreshLock, type RefreshLockHandle } from './cache-refresh-lock.js' -import { decideParseWorkers, parseFilesInOrder, ParseWorkerPool } from './parse-workers.js' +import { decideParseWorkers, parseFilesInOrder, ParseWorkerPool, type ClaudeWorkerParse, type ParseJob } from './parse-workers.js' +import type { CodexFullParse } from './providers/codex.js' import { dateKey } from './day-aggregator.js' import type { ParsedProviderCall, SessionSource } from './providers/types.js' import type { @@ -2032,7 +2033,9 @@ async function scanProjectDirs( process.stderr.write(`codeburn: parse workers unavailable, parsing serially (${err instanceof Error ? err.message : String(err)})\n`) } } - const offThread = pool ? parseFilesInOrder(pool, fullReparsePaths) : null + const offThread = pool + ? parseFilesInOrder(pool, fullReparsePaths.map(filePath => ({ kind: 'claude', filePath }))) + : null // Files whose worker result had to be thrown away because an earlier file had // already claimed one of its message ids. Expected to stay near zero; a large // count means the corpus is full of resumed sessions and the pool is doing @@ -3078,6 +3081,45 @@ async function parseProviderSources( } } + // Codex rollouts are the bulk of a cold parse (multi-GB against Claude's + // hundreds of MB), so whole-file decodes go to worker threads. A file the + // codex cache can serve exactly, or resume into from a byte offset, stays + // in-process: it reads a few KB, and it is the codex cache's own per-directory + // state that a thread must never own. The eligible list is built with the same + // filters (and in the same order) the parse loop applies, so the Nth result + // parseFilesInOrder yields is the Nth file that reaches the worker branch. + // The decision is per provider rather than pooled across Claude+Codex because + // the two scans run one after the other — at most one pool is ever alive — and + // a per-provider count is what the verbose line can honestly report. + const workerJobs: ParseJob[] = [] + const workerPaths = new Set() + let workerDiscards = 0 + let pendingBytes = 0 + if (providerName === 'codex' && !readOnly) { + for (const { source, fp } of changedSources) { + if (dateRange && fp.mtimeMs < dateRange.start.getTime()) continue + if (await readCachedCodexResults(source.path)) continue + workerJobs.push({ kind: 'codex', source }) + workerPaths.add(source.path) + pendingBytes += fp.sizeBytes + } + } + const decision = workerJobs.length > 0 + ? decideParseWorkers({ files: workerJobs.length, bytes: pendingBytes }) + : { workers: 0, reason: 'no full parses pending' } + if (providerName === 'codex' && !readOnly && process.env['CODEBURN_VERBOSE'] === '1') { + process.stderr.write(`codeburn: codex parse workers=${decision.workers} (${decision.reason})\n`) + } + let pool: ParseWorkerPool | null = null + if (decision.workers > 0) { + try { + pool = new ParseWorkerPool(decision.workers) + } catch (err) { + process.stderr.write(`codeburn: parse workers unavailable, parsing serially (${err instanceof Error ? err.message : String(err)})\n`) + } + } + const offThread = pool ? parseFilesInOrder(pool, workerJobs) : null + // Parse changed files, update cache let didParse = false // Track which paths have already been cleared this pass so that subsequent @@ -3101,12 +3143,36 @@ async function parseProviderSources( clearedPaths.add(source.path) } - const parser = provider.createSessionParser(source, parserDedup, dateRange) - try { - const providerCalls: ParsedProviderCall[] = [] - for await (const call of parser.parse()) { - providerCalls.push(call) + // Off-thread results arrive in this order, so the Nth eligible file here + // is the Nth yielded result. A worker decodes against an EMPTY dedup set, + // so an EMPTY key intersection is the proof that a serial parse would + // have dropped nothing either — that, and only that, makes the result + // installable. On any overlap (a forked rollout replaying its parent's + // token_count history is exactly this) the WHOLE file is discarded and + // re-parsed in-process against the real dedup set. + let providerCalls: ParsedProviderCall[] | undefined + if (offThread && workerPaths.has(source.path)) { + const result = (await offThread.next()).value + if (result?.ok && result.parsed) { + if (result.parsed.keys.some(k => parserDedup.has(k))) { + workerDiscards++ + } else { + for (const k of result.parsed.keys) parserDedup.add(k) + providerCalls = result.parsed.calls + // The worker never touches the codex cache; publish its entry here, + // in install order, so flushCodexCache writes what serial would. + const write = result.parsed.write + if (write) await writeCachedCodexResults(source.path, write.project, providerCalls, write.fingerprint, write.resume) + } + } + } + if (!providerCalls) { + const parser = provider.createSessionParser(source, parserDedup, dateRange) + providerCalls = [] + for await (const call of parser.parse()) { + providerCalls.push(call) + } } const canonicalCalls = await Promise.all(providerCalls.map(canonicalizeProviderCallProject)) const turns = providerCallsToCachedTurns(canonicalCalls) @@ -3162,6 +3228,10 @@ async function parseProviderSources( } } } finally { + await pool?.close() + if (pool && process.env['CODEBURN_VERBOSE'] === '1') { + process.stderr.write(`codeburn: codex parse workers done, ${workerDiscards}/${workerJobs.length} results re-parsed in-process on id overlap\n`) + } if (didParse && providerName === 'codex') await flushCodexCache() if (didParse && providerName === 'antigravity') { const liveIds = new Set(sources.map(s => antigravityCascadeIdFromPath(s.path))) diff --git a/src/providers/codex.ts b/src/providers/codex.ts index 41d2a4df..3328a3ef 100644 --- a/src/providers/codex.ts +++ b/src/providers/codex.ts @@ -6,7 +6,7 @@ import { homedir } from 'os' import { readSessionLines } from '../fs-utils.js' import { calculateCost } from '../models.js' -import { readCachedCodexResults, writeCachedCodexResults, getCachedCodexProject, fingerprintFile } from '../codex-cache.js' +import { readCachedCodexResults, writeCachedCodexResults, getCachedCodexProject, fingerprintFile, type CodexFileFingerprint } from '../codex-cache.js' import { normalizeContentBlocks } from '../content-utils.js' import { estimateTokensFromChars } from '../token-estimate.js' import type { ToolCall } from '../types.js' @@ -624,10 +624,22 @@ function isResumeState(value: unknown): value is CodexResumeState { && typeof v['currentTurnId'] === 'string' } -function createParser(source: SessionSource, seenKeys: Set): SessionParser { +/** What the serial path would have written to the codex cache for one file. */ +export type CodexCacheWrite = { + project: string + fingerprint: CodexFileFingerprint + resume?: { offset: number; state: unknown; callCount: number } +} + +// When `capture` is passed the parse is a whole-file decode that never touches +// the codex cache: no hit lookup (so no resume), and the entry it would have +// written comes back through `capture` for the caller to install. That is what +// lets a worker thread run this exact decode without owning the cache module's +// per-directory state. +function createParser(source: SessionSource, seenKeys: Set, capture?: { write?: CodexCacheWrite }): SessionParser { return { async *parse(): AsyncGenerator { - const hit = await readCachedCodexResults(source.path) + const hit = capture ? null : await readCachedCodexResults(source.path) if (hit?.kind === 'exact') { for (const call of hit.calls) { if (seenKeys.has(call.deduplicationKey)) continue @@ -1112,13 +1124,12 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars // Flush the final task, which has no following task_started to trigger it. results.push(...pendingTaskCalls) - await writeCachedCodexResults( - source.path, - source.project, - results, - fp, - resumeState ? { offset: resumeOffset, state: resumeState, callCount: resumeCallCount } : undefined, - ) + const resumeWrite = resumeState ? { offset: resumeOffset, state: resumeState, callCount: resumeCallCount } : undefined + if (capture) { + capture.write = { project: source.project, fingerprint: fp, ...(resumeWrite ? { resume: resumeWrite } : {}) } + } else { + await writeCachedCodexResults(source.path, source.project, results, fp, resumeWrite) + } for (const call of results) { yield call @@ -1127,6 +1138,20 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars } } +export type CodexFullParse = { calls: ParsedProviderCall[]; write?: CodexCacheWrite } + +/// Decode one rollout end to end, exactly as the serial path does for a file +/// with no cache entry, without reading or writing the codex cache. `seenKeys` +/// is the dedup set the decode runs against — pass an empty one off-thread and +/// let the caller prove no earlier file claimed any of the keys before +/// installing the result. +export async function parseCodexFileFull(source: SessionSource, seenKeys: Set): Promise { + const capture: { write?: CodexCacheWrite } = {} + const calls: ParsedProviderCall[] = [] + for await (const call of createParser(source, seenKeys, capture).parse()) calls.push(call) + return { calls, ...(capture.write ? { write: capture.write } : {}) } +} + export function createCodexProvider(codexDir?: string): Provider { const dir = getCodexDir(codexDir)