From 52b618339f94aaea6701a38deb4e18d3a277313d Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Thu, 12 Nov 2026 02:55:43 -0800 Subject: [PATCH] Stream protected usage caches without whole-file buffers (#41478) * Stream usage source caches without whole-document strings * Keep the validated source-array type through migration * Narrow migrated source records without an assertion * Assert cache preservation without runtime-specific error copy * Stream protected usage cache persistence and verification * Assert bounded usage cache input before gzip compression * Retain compressed usage cache publication bounds * Verify persisted output in usage cache streaming contract --- .../claude-usage/persisted-source-cache.ts | 52 +++--- .../persisted-source-migration.test.ts | 57 ++++++ .../claude-usage/persisted-usage-report.ts | 13 ++ .../source-cache-compression.test.ts | 172 +++++++++++++++--- .../claude-usage/source-cache-compression.ts | 145 ++++++++++++--- src/main/claude-usage/store.ts | 7 +- src/main/usage/cache/usage-cache-json.ts | 16 +- .../usage/cache/usage-cache-protected-json.ts | 85 +++++++++ .../usage-source-cache-streaming.test.ts | 58 ++++++ .../usage-provider-store-contract.ts | 1 + .../usage-provider-store-lifecycle.ts | 24 +-- 11 files changed, 524 insertions(+), 106 deletions(-) create mode 100644 src/main/claude-usage/persisted-source-migration.test.ts create mode 100644 src/main/usage/cache/usage-cache-protected-json.ts create mode 100644 src/main/usage/cache/usage-source-cache-streaming.test.ts diff --git a/src/main/claude-usage/persisted-source-cache.ts b/src/main/claude-usage/persisted-source-cache.ts index 1dbe3f126ef..12fb712fc83 100644 --- a/src/main/claude-usage/persisted-source-cache.ts +++ b/src/main/claude-usage/persisted-source-cache.ts @@ -1,20 +1,18 @@ -import { readFile } from 'node:fs/promises' import type { ClaudeUsagePersistedFile, ClaudeUsagePersistedState } from './types' import { encodeClaudeUsagePersistedFiles } from './persisted-token-columns' -import { - compressClaudeUsageSourceText, - decodeClaudeUsageSourceText -} from './source-cache-compression' +import { writeClaudeUsageJson, readClaudeUsageJson } from './source-cache-compression' import { hasClaudeUsageProtectedCheckpoint, normalizeClaudeUsageSourceFiles, type ClaudeUsageVerifiedSources } from './persisted-projection-validation' -import { parseClaudeUsageReport, serializeClaudeUsageReport } from './persisted-usage-report' -import { sealUsageCacheJson, verifyUsageCacheJson } from '../usage/cache/usage-cache-json-integrity' +import { + parseClaudeUsageReport, + readClaudeUsageReportFile, + serializeClaudeUsageReport +} from './persisted-usage-report' import { usageSourceCachePath, - writeUsageSourceCacheData, type UsageCacheSplitRequest, type UsageCacheSplitResult, type UsageSourceCacheRef @@ -39,9 +37,12 @@ export async function readClaudeUsageSourceCache( return { sources: [] } } try { - const text = await decodeClaudeUsageSourceText(await readFile(cacheRef.path)) - const parsed = JSON.parse(text) - const verified = verifyUsageCacheJson(text, parsed?.usageIntegrity, SOURCE_DOMAIN) + const { parsed, verified } = await readClaudeUsageJson<{ + usageIntegrity?: string + schemaVersion: number + worktreeFingerprint: string | null + sources: ClaudeUsagePersistedFile[] + }>(cacheRef.path, SOURCE_DOMAIN) if ( (parsed?.schemaVersion !== cacheRef.schemaVersion && !(cacheRef.schemaVersion === 7 && parsed?.schemaVersion === 6)) || @@ -100,41 +101,40 @@ export async function writeClaudeUsageSourceCache( sources: readonly ClaudeUsagePersistedFile[] ): Promise { const cacheRef = { ...ref } - const material = JSON.stringify({ - schemaVersion: cacheRef.schemaVersion, - worktreeFingerprint: cacheRef.worktreeFingerprint, - sources: encodeClaudeUsagePersistedFiles(sources) - }) - await writeUsageSourceCacheData( + await writeClaudeUsageJson( cacheRef, - await compressClaudeUsageSourceText( - cacheRef.schemaVersion === 6 ? material : sealUsageCacheJson(material, SOURCE_DOMAIN) - ) + { + schemaVersion: cacheRef.schemaVersion, + worktreeFingerprint: cacheRef.worktreeFingerprint, + sources: encodeClaudeUsagePersistedFiles(sources) + }, + SOURCE_DOMAIN, + cacheRef.schemaVersion !== 6 ) } export async function splitClaudeUsageCacheFile( request: UsageCacheSplitRequest ): Promise { - let text: string + let loaded: Awaited> try { - text = await readFile(request.cacheFile, 'utf8') + loaded = await readClaudeUsageReportFile(request.cacheFile) } catch (error) { if (error instanceof Error && 'code' in error && error.code === 'ENOENT') { return { reportText: null, migrated: false } } throw error } - const parsed: ClaudeUsagePersistedState = JSON.parse(text) + const parsed: ClaudeUsagePersistedState = loaded.state if (!Array.isArray(parsed.processedFiles) || parsed.processedFiles.length === 0) { - const validated = await parseClaudeUsageReport(text, parsed) + const validated = await parseClaudeUsageReport('', parsed, loaded.integrityVerified) return { - reportText: validated === parsed ? text : serializeClaudeUsageReport(validated), + reportText: serializeClaudeUsageReport(validated), migrated: false, reportIntegrityVerified: validated.schemaVersion === 7 } } - const validated = await parseClaudeUsageReport(text, parsed) + const validated = await parseClaudeUsageReport('', parsed, loaded.integrityVerified) const protectedInline = parsed.processedFiles.some(hasClaudeUsageProtectedCheckpoint) const state = protectedInline ? { ...validated, schemaVersion: 7 } : validated const ref = { diff --git a/src/main/claude-usage/persisted-source-migration.test.ts b/src/main/claude-usage/persisted-source-migration.test.ts new file mode 100644 index 00000000000..2fc2840cb1f --- /dev/null +++ b/src/main/claude-usage/persisted-source-migration.test.ts @@ -0,0 +1,57 @@ +import { sealUsageCacheJson } from '../usage/cache/usage-cache-json-integrity' +import { mkdir, readFile, rm, writeFile } from 'node:fs/promises' +import { afterEach, beforeEach, expect, it, vi } from 'vitest' +import { + createPersistedUsageTestFixture, + type PersistedUsageTestFixture +} from './persisted-usage-test-fixture' + +vi.mock('./transcript-file-discovery', () => ({ listClaudeTranscriptFiles: vi.fn() })) +vi.mock('../usage/scan-worker/usage-scan-worker-spawn', () => ({ + scanClaudeUsageFilesViaWorker: vi.fn(), + splitUsageCacheFileViaWorker: vi.fn() +})) +let fixture: PersistedUsageTestFixture +beforeEach(async () => { + fixture = await createPersistedUsageTestFixture() +}) +afterEach(async () => { + await fixture.cleanup() +}) + +it('preserves protected inline ownership through failed migration, toggle and retry', async () => { + const source = await fixture.writeResumable() + await writeFile(fixture.reportPath, JSON.stringify(fixture.state([source]))) + await mkdir(fixture.sourceRef.path) + vi.spyOn(console, 'warn').mockImplementation(() => {}) + const first = fixture.createStore() + await first.whenLoaded() + await first.setEnabled(false) + await first.flush() + const retained = JSON.parse(await readFile(fixture.reportPath, 'utf8')) + expect(retained.processedFiles).toEqual([source]) + expect(retained.scanState.enabled).toBe(false) + await rm(fixture.sourceRef.path, { recursive: true }) + const retried = fixture.createStore() + await retried.whenLoaded() + await retried.flush() + expect(retried.getScanState()).toMatchObject({ enabled: false, hasAnyClaudeData: true }) + expect(await fixture.decodedSources()).toEqual([source]) + expect(JSON.parse(await readFile(fixture.reportPath, 'utf8')).processedFiles).toBeUndefined() +}) + +it('does not trust a canonical reconstruction when original report verification failed', async () => { + const source = await fixture.writeResumable() + const { processedFiles: _files, ...report } = fixture.state([source], 7) + const canonical = sealUsageCacheJson( + JSON.stringify({ ...report, padding: 'a'.repeat(8 * 1024 * 1024) }), + 'claude-usage-report-v1' + ) + await writeFile(fixture.reportPath, canonical.replace('{"usageIntegrity"', '{ "usageIntegrity"')) + fixture.splitWorker.mockRejectedValueOnce(new Error('worker unavailable')) + vi.spyOn(console, 'warn').mockImplementation(() => {}) + const store = fixture.createStore() + await store.whenLoaded() + expect(store.getScanState()).toMatchObject({ hasAnyClaudeData: false, lastScanCompletedAt: null }) + await store.flush() +}) diff --git a/src/main/claude-usage/persisted-usage-report.ts b/src/main/claude-usage/persisted-usage-report.ts index 29e661494bd..955b6789cc6 100644 --- a/src/main/claude-usage/persisted-usage-report.ts +++ b/src/main/claude-usage/persisted-usage-report.ts @@ -1,3 +1,4 @@ +import { readClaudeUsageJson } from './source-cache-compression' import type { ClaudeUsagePersistedState } from './types' import { hasClaudeUsageProtectedCheckpoint, @@ -54,3 +55,15 @@ export function parseClaudeUsageReport( return invalidateReport(state) } } + +export async function readClaudeUsageReportFile(path: string): Promise<{ + state: ClaudeUsagePersistedState + integrityVerified: boolean +}> { + const { parsed, verified } = await readClaudeUsageJson( + path, + REPORT_DOMAIN, + false + ) + return { state: parsed, integrityVerified: verified } +} diff --git a/src/main/claude-usage/source-cache-compression.test.ts b/src/main/claude-usage/source-cache-compression.test.ts index 29e9f12c0d3..7fbfb9a2c91 100644 --- a/src/main/claude-usage/source-cache-compression.test.ts +++ b/src/main/claude-usage/source-cache-compression.test.ts @@ -1,36 +1,150 @@ -import { gzipSync } from 'node:zlib' -import { expect, it } from 'vitest' -import { - compressClaudeUsageSourceText, - decodeClaudeUsageSourceText -} from './source-cache-compression' +import { mkdtemp, readFile, readdir, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { Readable } from 'node:stream' +import { gzipSync, gunzipSync } from 'node:zlib' +import { afterEach, beforeEach, expect, it, vi } from 'vitest' +import * as durable from '../filesystem/durable-file-write' +import { sealUsageCacheJson, verifyUsageCacheJson } from '../usage/cache/usage-cache-json-integrity' +import { readClaudeUsageJson, writeClaudeUsageJson } from './source-cache-compression' -const source = JSON.stringify({ sources: ['Unicode: 日本語 🌊'.repeat(10_000)] }) - -it('keeps small and existing plain caches readable without changing them', async () => { - const small = '{"schemaVersion":7,"sources":[]}' - expect(await compressClaudeUsageSourceText(small)).toBe(small) - expect(await decodeClaudeUsageSourceText(Buffer.from(small))).toBe(small) - expect(await decodeClaudeUsageSourceText(source)).toBe(source) +let directory: string +const domain = 'claude-usage-sources-v1' +const source = { + schemaVersion: 7, + sources: [{ path: `${'a'.repeat(4095)}😀é\ud800`, ownedDedupeKeys: ['日本語 🌊'.repeat(10_000)] }] +} +const ref = () => ({ + path: join(directory, 'sources.json'), + schemaVersion: 7, + worktreeFingerprint: null, + reuse: true +}) +beforeEach(async () => { + directory = await mkdtemp(join(tmpdir(), 'orca-claude-stream-')) +}) +afterEach(async () => { + await rm(directory, { recursive: true, force: true }) }) -it.each([0, 70_001, 129 * 1024 * 1024])( - 'rejects inconsistent or excessive decoded length %i', - async (declared) => { - const saved = gzipSync('x'.repeat(70_000)) - saved.writeUInt32LE(declared, saved.length - 4) - await expect(decodeClaudeUsageSourceText(saved)).rejects.toThrow() +it('cross-reads existing protected plain and gzip text and emits the same verified material', async () => { + const original = sealUsageCacheJson(JSON.stringify(source), domain) + for (const saved of [Buffer.from(original), gzipSync(original)]) { + await writeFile(ref().path, saved) + const loaded = await readClaudeUsageJson(ref().path, domain) + expect(loaded.verified).toBe(true) + expect(loaded.parsed).toMatchObject(source) } -) - -it('bounds decompression to a forged smaller length', async () => { - const saved = gzipSync('x'.repeat(500_000)) - saved.writeUInt32LE(70_000, saved.length - 4) - await expect(decodeClaudeUsageSourceText(saved)).rejects.toThrow() + await writeClaudeUsageJson(ref(), source, domain, true) + const saved = await readFile(ref().path) + const text = saved[0] === 0x1f ? gunzipSync(saved).toString('utf8') : saved.toString('utf8') + const parsed = JSON.parse(text) + expect(verifyUsageCacheJson(text, parsed.usageIntegrity, domain)).toBe(true) + expect(parsed.sources).toEqual(source.sources) + expect((await readClaudeUsageJson(ref().path, domain)).parsed).toEqual(parsed) }) -it('rejects concatenated streams rather than accepting a partial generation', async () => { - const first = gzipSync('x'.repeat(70_000)) - const second = gzipSync('y'.repeat(70_000)) - await expect(decodeClaudeUsageSourceText(Buffer.concat([first, second]))).rejects.toThrow() +it('keeps small unsigned schema6 and sealed schema7 material plain and readable', async () => { + for (const schemaVersion of [6, 7]) { + await writeClaudeUsageJson(ref(), { schemaVersion, sources: [] }, domain, schemaVersion === 7) + const saved = await readFile(ref().path, 'utf8') + expect(saved.startsWith('{')).toBe(true) + expect(await readClaudeUsageJson(ref().path, domain)).toEqual({ + parsed: JSON.parse(saved), + verified: schemaVersion === 7 + }) + } +}) + +it('rejects corruption and partial generations before returning sources', async () => { + const original = sealUsageCacheJson(JSON.stringify(source), domain) + const gzip = gzipSync(original) + const forged = Buffer.from(gzip) + forged.writeUInt32LE(70_000, forged.length - 4) + for (const saved of [ + Buffer.from(original.replace('日本語', '日木語')), + Buffer.from(`${original} null`), + gzip.subarray(0, -7), + forged, + Buffer.concat([gzip, gzip]) + ]) { + await writeFile(ref().path, saved) + await expect(readClaudeUsageJson(ref().path, domain)).rejects.toThrow() + } +}) + +it('removes only its staging file when compressed publication fails', async () => { + const { mkdir } = await import('node:fs/promises') + await mkdir(ref().path) + await expect(writeClaudeUsageJson(ref(), source, domain, true)).rejects.toThrow() + expect(await readdir(directory)).toEqual(['sources.json']) +}) + +it.each([0, 129 * 1024 * 1024])('rejects out-of-range gzip footer %i', async (declared) => { + const saved = gzipSync(JSON.stringify(source)) + saved.writeUInt32LE(declared, saved.length - 4) + await writeFile(ref().path, saved) + await expect(readClaudeUsageJson(ref().path, domain)).rejects.toThrow() +}) + +it('streams bounded gzip inputs and publication at N and 10N source histories', async () => { + type Receipt = { iterable: boolean; maximum: number; bytes: number; chunks: number } + const inputs: Receipt[] = [] + const outputs: Receipt[] = [] + async function* measure(payload: durable.DurableFilePayload, receipts: Receipt[]) { + const receipt = { + iterable: typeof payload !== 'string' && !(payload instanceof Uint8Array), + maximum: 0, + bytes: 0, + chunks: 0 + } + receipts.push(receipt) + const iterable = + typeof payload === 'string' || payload instanceof Uint8Array ? [payload] : payload + for await (const chunk of iterable) { + const bytes = typeof chunk === 'string' ? Buffer.byteLength(chunk) : chunk.byteLength + receipt.maximum = Math.max(receipt.maximum, bytes) + receipt.bytes += bytes + receipt.chunks++ + yield chunk + } + } + const originalFrom = Readable.from + const inputSpy = vi + .spyOn(Readable, 'from') + .mockImplementation((input, options) => + originalFrom.call(Readable, measure(input, inputs), options) + ) + const originalWrite = durable.writeTempFileDurable + const outputSpy = vi + .spyOn(durable, 'writeTempFileDurable') + .mockImplementation((path, payload, mode) => + originalWrite(path, measure(payload, outputs), mode) + ) + try { + for (const events of [10_000, 100_000]) { + const value = { + schemaVersion: 7, + sources: [ + { ownedDedupeKeys: Array.from({ length: events }, (_, i) => `request-${i}|message-${i}`) } + ] + } + await writeClaudeUsageJson(ref(), value, domain, true) + const loaded = await readClaudeUsageJson(ref().path, domain) + expect(loaded.verified).toBe(true) + expect(loaded.parsed).toMatchObject(value) + } + for (const receipt of [...inputs, ...outputs]) { + expect(receipt.iterable).toBe(true) + expect(receipt.maximum).toBeLessThanOrEqual(256 * 1024) + } + for (const receipt of inputs) { + expect(receipt.chunks).toBeLessThanOrEqual(Math.ceil(receipt.bytes / (32 * 1024)) + 1) + } + expect(inputs).toHaveLength(2) + expect(outputs).toHaveLength(2) + } finally { + inputSpy.mockRestore() + outputSpy.mockRestore() + } }) diff --git a/src/main/claude-usage/source-cache-compression.ts b/src/main/claude-usage/source-cache-compression.ts index 71b667dd3df..d34566f4e0f 100644 --- a/src/main/claude-usage/source-cache-compression.ts +++ b/src/main/claude-usage/source-cache-compression.ts @@ -1,41 +1,126 @@ -import { promisify } from 'node:util' -import { gunzip, gzip } from 'node:zlib' +import { open, mkdir, rm, stat } from 'node:fs/promises' +import { threadId } from 'node:worker_threads' +import { dirname } from 'node:path' +import { Readable } from 'node:stream' +import { StringDecoder } from 'node:string_decoder' +import { createGzip, createGunzip } from 'node:zlib' +import { + durableWriteTempPath, + renameDurable, + removeStaleDurableWriteTempFiles, + writeTempFileDurable +} from '../filesystem/durable-file-write' +import { + protectedUsageCacheJson, + readProtectedUsageCacheJson +} from '../usage/cache/usage-cache-protected-json' +import { + writeUsageSourceCacheData, + type UsageSourceCacheRef +} from '../usage/cache/usage-source-cache-file' -const gzipAsync = promisify(gzip) -const gunzipAsync = promisify(gunzip) const MIN_COMPRESSED_BYTES = 64 * 1024 const MAX_COMPRESSED_BYTES = 128 * 1024 * 1024 -const ZLIB_CHUNK_BYTES = 256 * 1024 +const CHUNK_BYTES = 256 * 1024 -export async function compressClaudeUsageSourceText(text: string): Promise { - const bytes = Buffer.byteLength(text) - if (bytes < MIN_COMPRESSED_BYTES || bytes > MAX_COMPRESSED_BYTES) { - return text +/** Keep magic, declared length and body on the same opened generation. */ +async function* decodedSourceText(path: string): AsyncGenerator { + const handle = await open(path, 'r') + try { + const magic = Buffer.alloc(2) + await handle.read(magic, 0, 2, 0) + const gzip = magic[0] === 0x1f && magic[1] === 0x8b + let expected = 0 + if (gzip) { + const { size } = await handle.stat() + if (size < 18) { + throw new Error('Saved Claude usage gzip stream is truncated.') + } + const footer = Buffer.alloc(4) + await handle.read(footer, 0, 4, size - 4) + expected = footer.readUInt32LE(0) + if (expected < MIN_COMPRESSED_BYTES || expected > MAX_COMPRESSED_BYTES) { + throw new Error('Saved Claude usage gzip length is invalid.') + } + } + const input = handle.createReadStream({ + start: 0, + autoClose: false, + highWaterMark: CHUNK_BYTES + }) + const expanded = gzip ? input.pipe(createGunzip({ chunkSize: CHUNK_BYTES })) : input + input.on('error', (error) => expanded.destroy(error)) + const decoder = new StringDecoder('utf8') + let bytes = 0 + try { + for await (const chunk of expanded) { + bytes += chunk.byteLength + if (gzip && (bytes > expected || bytes > MAX_COMPRESSED_BYTES)) { + throw new Error('Saved Claude usage compression length is invalid.') + } + yield decoder.write(chunk) + } + if (gzip && bytes !== expected) { + throw new Error('Saved Claude usage compression length is invalid.') + } + const remaining = decoder.end() + if (remaining) { + yield remaining + } + } finally { + input.destroy() + expanded.destroy() + } + } finally { + await handle.close() } - const compressed = await gzipAsync(text, { level: 1, chunkSize: ZLIB_CHUNK_BYTES }) - return compressed.byteLength < bytes ? compressed : text } -export async function decodeClaudeUsageSourceText(saved: string | Buffer): Promise { - if (typeof saved === 'string') { - return saved +export async function readClaudeUsageJson( + path: string, + domain: string, + strict = true +): Promise<{ parsed: T; verified: boolean }> { + return readProtectedUsageCacheJson(decodedSourceText(path), domain, strict) +} + +export async function writeClaudeUsageJson( + ref: UsageSourceCacheRef, + value: unknown, + domain: string, + sealed: boolean +): Promise { + const material = protectedUsageCacheJson(value, domain, sealed) + if (material.bytes < MIN_COMPRESSED_BYTES || material.bytes > MAX_COMPRESSED_BYTES) { + await writeUsageSourceCacheData(ref, material.encode()) + return } - if (saved[0] !== 0x1f || saved[1] !== 0x8b) { - return saved.toString('utf8') + await mkdir(dirname(ref.path), { recursive: true }) + await removeStaleDurableWriteTempFiles(ref.path) + const candidate = durableWriteTempPath(ref.path, threadId > 0 ? String(threadId) : undefined) + async function* gzipChunks(): AsyncGenerator { + const input = Readable.from(material.encode()) + const compressed = input.pipe(createGzip({ level: 1, chunkSize: CHUNK_BYTES })) + input.on('error', (error) => compressed.destroy(error)) + try { + yield* compressed + } finally { + input.destroy() + compressed.destroy() + } } - if (saved.length < 18) { - throw new Error('Saved Claude usage gzip stream is truncated.') + let published = false + try { + await writeTempFileDurable(candidate, gzipChunks(), 0o600) + if ((await stat(candidate)).size < material.bytes) { + await renameDurable(candidate, ref.path) + published = true + } else { + await writeUsageSourceCacheData(ref, material.encode()) + } + } finally { + if (!published) { + await rm(candidate, { force: true }).catch(() => {}) + } } - const bytes = saved.readUInt32LE(saved.length - 4) - if (bytes < MIN_COMPRESSED_BYTES || bytes > MAX_COMPRESSED_BYTES) { - throw new Error('Saved Claude usage gzip length is invalid.') - } - const decoded = await gunzipAsync(saved, { - maxOutputLength: bytes, - chunkSize: ZLIB_CHUNK_BYTES - }) - if (decoded.byteLength !== bytes) { - throw new Error('Saved Claude usage compression length is invalid.') - } - return decoded.toString('utf8') } diff --git a/src/main/claude-usage/store.ts b/src/main/claude-usage/store.ts index a6cda7d260b..e634443983c 100644 --- a/src/main/claude-usage/store.ts +++ b/src/main/claude-usage/store.ts @@ -22,7 +22,11 @@ import { buildBreakdown, buildDaily, buildSummary } from './claude-usage-report- import { buildRecentSessions } from './claude-usage-session-rows' import type { AutomationUsageLookupInput } from './claude-usage-automation-attribution' import { resolveAutomationRunUsage } from './claude-usage-automation-attribution' -import { parseClaudeUsageReport, serializeClaudeUsageReport } from './persisted-usage-report' +import { + parseClaudeUsageReport, + readClaudeUsageReportFile, + serializeClaudeUsageReport +} from './persisted-usage-report' // Pre-v6 totals lack stable fork ownership or the 1-hour cache-write split and must be rescanned. // v7 stores protected column checkpoints on the worker and a separately protected report. @@ -75,6 +79,7 @@ export class ClaudeUsageStore extends UsageProviderStoreLifecycle< createDefaultState: getDefaultState, normalizeState: normalizePersistedState, parseReport: parseClaudeUsageReport, + readFallbackCache: readClaudeUsageReportFile, serializeReport: serializeClaudeUsageReport, providerId: 'claude', sourceKey: 'processedFiles', diff --git a/src/main/usage/cache/usage-cache-json.ts b/src/main/usage/cache/usage-cache-json.ts index 862aceff858..520cc71ea27 100644 --- a/src/main/usage/cache/usage-cache-json.ts +++ b/src/main/usage/cache/usage-cache-json.ts @@ -134,15 +134,21 @@ export function* encodeUsageCacheJson(value: unknown): Generator { } /** EOF validation precedes exposing any of the privately assembled records. */ -export async function readUsageCacheJson(path: string): Promise { +export async function parseUsageCacheJson(chunks: AsyncIterable): Promise { const assembler = new Assembler() const reader = createJsonTokenReader((token) => assembler.consume(token)) - for await (const text of createReadStream(path, { - encoding: 'utf8', - highWaterMark: TEXT_BATCH_CHARS - })) { + for await (const text of chunks) { reader.write(text) } reader.end() return assembler.current } + +export function readUsageCacheJson(path: string): Promise { + return parseUsageCacheJson( + createReadStream(path, { + encoding: 'utf8', + highWaterMark: TEXT_BATCH_CHARS + }) + ) +} diff --git a/src/main/usage/cache/usage-cache-protected-json.ts b/src/main/usage/cache/usage-cache-protected-json.ts new file mode 100644 index 00000000000..1e91c55469b --- /dev/null +++ b/src/main/usage/cache/usage-cache-protected-json.ts @@ -0,0 +1,85 @@ +import { createHash } from 'node:crypto' +import { encodeUsageCacheJson, parseUsageCacheJson } from './usage-cache-json' + +const INTEGRITY_HEADER = /^\{"usageIntegrity":"([a-f0-9]{64})",/ +const HEADER_CHARS = '{"usageIntegrity":"'.length + 64 + '",'.length + +export function protectedUsageCacheJson( + value: unknown, + domain: string, + sealed = true +): { + bytes: number + encode: () => Generator +} { + const hash = createHash('sha256').update(domain).update('\n') + let bytes = 0 + for (const text of encodeUsageCacheJson(value)) { + hash.update(text) + bytes += Buffer.byteLength(text) + } + const digest = hash.digest('hex') + const prefix = `{"usageIntegrity":"${digest}",` + return { + bytes: bytes + (sealed ? Buffer.byteLength(prefix) - 1 : 0), + encode: function* () { + let first = true + for (const text of encodeUsageCacheJson(value)) { + if (first && sealed) { + if (!text.startsWith('{')) { + throw new Error('Usage cache integrity requires a JSON object.') + } + yield prefix + yield text.slice(1) + } else { + yield text + } + first = false + } + } + } +} + +/** Verify the original decoded text, never an equivalent re-encoding. */ +export async function readProtectedUsageCacheJson( + chunks: AsyncIterable, + domain: string, + strict = true +): Promise<{ parsed: T; verified: boolean }> { + const hash = createHash('sha256').update(domain).update('\n{') + let leading = '' + let header: RegExpExecArray | null | undefined + async function* inspect(): AsyncGenerator { + for await (const text of chunks) { + if (header === undefined) { + const taken = Math.min(HEADER_CHARS - leading.length, text.length) + leading += text.slice(0, taken) + if (leading.length === HEADER_CHARS) { + header = INTEGRITY_HEADER.exec(leading) + if (header) { + hash.update(text.slice(taken)) + } + } + } else if (header) { + hash.update(text) + } + yield text + } + } + const parsed = await parseUsageCacheJson(inspect()) + if (!parsed || typeof parsed !== 'object') { + throw new Error('Saved usage cache must be a JSON object.') + } + header ??= INTEGRITY_HEADER.exec(leading) + if (!header) { + if (strict && parsed.usageIntegrity !== undefined) { + throw new Error('Saved usage cache integrity framing is invalid.') + } + return { parsed, verified: false } + } + const verified = parsed.usageIntegrity === header[1] && hash.digest('hex') === header[1] + if (strict && !verified) { + throw new Error('Saved usage cache integrity does not match its contents.') + } + return { parsed, verified } +} diff --git a/src/main/usage/cache/usage-source-cache-streaming.test.ts b/src/main/usage/cache/usage-source-cache-streaming.test.ts new file mode 100644 index 00000000000..68e9cc88ca0 --- /dev/null +++ b/src/main/usage/cache/usage-source-cache-streaming.test.ts @@ -0,0 +1,58 @@ +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { expect, it, vi } from 'vitest' +import * as durable from '../../filesystem/durable-file-write' +import { readUsageSourceCache, writeUsageSourceCache } from './usage-source-cache-file' + +it('bounds source-cache write batches as one source grows from N to 10N events', async () => { + const directory = mkdtempSync(join(tmpdir(), 'orca-usage-stream-contract-')) + const receipts: { bytes: number; maximum: number; chunks: number }[] = [] + const originalWrite = durable.writeFileDurable + const write = vi + .spyOn(durable, 'writeFileDurable') + .mockImplementation(async (tmp, final, data, mode) => { + const receipt = { bytes: 0, maximum: 0, chunks: 0 } + async function* chunks() { + const iterable = typeof data === 'string' || data instanceof Uint8Array ? [data] : data + for await (const chunk of iterable) { + const bytes = typeof chunk === 'string' ? Buffer.byteLength(chunk) : chunk.byteLength + receipt.bytes += bytes + receipt.maximum = Math.max(receipt.maximum, bytes) + receipt.chunks++ + yield chunk + } + } + await originalWrite(tmp, final, chunks(), mode) + receipts.push(receipt) + }) + try { + for (const events of [10_000, 100_000]) { + const ref = { + path: join(directory, 'sources.json'), + schemaVersion: 7, + worktreeFingerprint: null, + reuse: true + } + const sources = [ + { + ownedEventKeys: Array.from( + { length: events }, + (_, index) => `turn-${index}|100000,50000,1000,0,101000|1000,500,10,0,1010` + ) + } + ] + await writeUsageSourceCache(ref, sources) + expect(await readUsageSourceCache(ref)).toEqual(sources) + } + for (const receipt of receipts) { + // Output-sized work, with bounded batches and no awaited write per event. + expect(receipt.maximum).toBeLessThanOrEqual(256 * 1024) + expect(receipt.chunks).toBeLessThanOrEqual(Math.ceil(receipt.bytes / (32 * 1024)) + 1) + } + expect(receipts).toHaveLength(2) + } finally { + write.mockRestore() + rmSync(directory, { recursive: true, force: true }) + } +}) diff --git a/src/main/usage/provider-lifecycle/usage-provider-store-contract.ts b/src/main/usage/provider-lifecycle/usage-provider-store-contract.ts index 586f64a70e8..184e3ef0d54 100644 --- a/src/main/usage/provider-lifecycle/usage-provider-store-contract.ts +++ b/src/main/usage/provider-lifecycle/usage-provider-store-contract.ts @@ -37,6 +37,7 @@ export type UsageProviderStoreLifecycleConfig< integrityVerified?: boolean ) => State | Promise serializeReport?: (state: State) => string + readFallbackCache?: (path: string) => Promise<{ state: State; integrityVerified: boolean }> providerId?: 'claude' sourceKey: SourceKey dataPresenceKey: DataPresenceKey diff --git a/src/main/usage/provider-lifecycle/usage-provider-store-lifecycle.ts b/src/main/usage/provider-lifecycle/usage-provider-store-lifecycle.ts index e90a1aa0f3c..a76137d9dd1 100644 --- a/src/main/usage/provider-lifecycle/usage-provider-store-lifecycle.ts +++ b/src/main/usage/provider-lifecycle/usage-provider-store-lifecycle.ts @@ -1,5 +1,4 @@ import { readFileSync, statSync } from 'node:fs' -import { readFile } from 'node:fs/promises' import { AgentTokenUsageReporter } from '../reporting/token-usage/agent-token-usage-reporter' import { join, parse } from 'node:path' import { @@ -194,22 +193,17 @@ export abstract class UsageProviderStoreLifecycle< // Preserve usage history when the worker cannot start. console.warn(`${this.config.logTag} Reading the usage cache on the main thread:`, error) fallback = true - if (!this.config.providerId) { - fallbackState = await readUsageCacheJson(cacheFile) - if (!fallbackState) { - throw new Error('Saved usage cache is invalid') - } - const { [this.config.sourceKey]: _sources, ...report } = fallbackState - return { - reportText: JSON.stringify(report), - migrated: false, - reportIntegrityVerified: false - } + const loaded = this.config.readFallbackCache + ? await this.config.readFallbackCache(cacheFile) + : { state: await readUsageCacheJson(cacheFile), integrityVerified: false } + fallbackState = loaded.state + if (!fallbackState) { + throw new Error('Saved usage cache is invalid') } return { - reportText: await readFile(cacheFile, 'utf-8'), + reportText: '', migrated: false, - reportIntegrityVerified: false + reportIntegrityVerified: loaded.integrityVerified } }) if (reportText === null) { @@ -221,7 +215,7 @@ export abstract class UsageProviderStoreLifecycle< !('usageIntegrity' in parsed) && Array.isArray(parsed[this.config.sourceKey]) && parsed[this.config.sourceKey].length > 0 - await this.applyReport(reportText, parsed, !fallback && reportIntegrityVerified === true) + await this.applyReport(reportText, parsed, reportIntegrityVerified === true) if (migrated) { // Future launches can read the compact report synchronously. await this.writeToDisk().catch(() => {})