From e6e27c354ca908d364e696edefe9bc3cd35ae3f3 Mon Sep 17 00:00:00 2001 From: Edward Abrams <102749+zeroaltitude@users.noreply.github.com> Date: Wed, 6 May 2026 15:34:46 -0700 Subject: [PATCH] perf(usage-cost-cache): throttle full-cache rewrites during refresh MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit refreshCostUsageCache rewrote the entire .usage-cost-cache.json after every single scanned session. With ~8190 stale session files and a 108MB cache, that was O(N * cacheSize) work — sustained CPU burn from repeated 100MB+ JSON.stringify + atomic-replace cycles every refresh. Checkpoint policy now batches durable writes: - At most one rewrite per 256 scanned files - Or one rewrite per 5s of wall time - Final write only when something actually changed (no-op refresh on a fully-fresh cache no longer rewrites the file) Crash safety is preserved: an interrupted refresh still has a recent checkpoint on disk, and the next run rescans only the unfinished tail (file size + mtime + pricingFingerprint match). Validation: - pnpm vitest run src/infra/session-cost-usage.test.ts (39/39 pass) - New test 'throttles cache writes during a large stale refresh' confirms cache renames stay below sessionCount/4 (was ~sessionCount+1) and that a no-op refresh issues zero cache writes. - pnpm check:changed (clean) Beads: openclaw-0zr --- scripts/proof-usage-cost-cache-refresh.ts | 114 ++++++++++++++++++++++ src/infra/session-cost-usage.test.ts | 72 ++++++++++++++ src/infra/session-cost-usage.ts | 33 ++++++- 3 files changed, 216 insertions(+), 3 deletions(-) create mode 100644 scripts/proof-usage-cost-cache-refresh.ts diff --git a/scripts/proof-usage-cost-cache-refresh.ts b/scripts/proof-usage-cost-cache-refresh.ts new file mode 100644 index 000000000000..d2add3edf9dd --- /dev/null +++ b/scripts/proof-usage-cost-cache-refresh.ts @@ -0,0 +1,114 @@ +/** + * Real-runtime proof for usage-cost cache refresh batching. + * + * Drives the production `refreshCostUsageCache` and `loadCostUsageSummaryFromCache` + * code paths against an on-disk OPENCLAW_STATE_DIR. It creates a synthetic session + * corpus, performs a cold refresh, appends to every transcript so the cache entries + * are stale, and refreshes again. The assertions pin that the aggregate cache remains + * fresh and correct after many stale files are processed in one refresh. + * + * Run with: pnpm tsx scripts/proof-usage-cost-cache-refresh.ts + */ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { performance } from "node:perf_hooks"; +import { + loadCostUsageSummaryFromCache, + refreshCostUsageCache, +} from "../src/infra/session-cost-usage.js"; + +const sessionCount = Number.parseInt(process.env.OPENCLAW_USAGE_COST_PROOF_SESSIONS ?? "400", 10); +const root = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-usage-cost-proof-")); +const previousStateDir = process.env.OPENCLAW_STATE_DIR; +process.env.OPENCLAW_STATE_DIR = root; + +try { + const sessionsDir = path.join(root, "agents", "main", "sessions"); + await fs.mkdir(sessionsDir, { recursive: true }); + + const firstTimestamp = "2026-02-05T12:00:00.000Z"; + const secondTimestamp = "2026-02-05T12:01:00.000Z"; + const makeEntry = (sessionId: string, timestamp: string, totalTokens: number) => + JSON.stringify({ + type: "message", + timestamp, + sessionId, + message: { + role: "assistant", + provider: "openai", + model: "gpt-5.5", + usage: { + input: totalTokens, + output: 0, + totalTokens, + cost: { total: totalTokens / 1000 }, + }, + }, + }); + + for (let index = 0; index < sessionCount; index += 1) { + const sessionId = `usage-cost-proof-${index}`; + await fs.writeFile( + path.join(sessionsDir, `${sessionId}.jsonl`), + `${JSON.stringify({ type: "session", version: 1, id: sessionId })}\n${makeEntry(sessionId, firstTimestamp, 1)}\n`, + "utf-8", + ); + } + + const coldStart = performance.now(); + await refreshCostUsageCache(); + const coldMs = performance.now() - coldStart; + + for (let index = 0; index < sessionCount; index += 1) { + const sessionId = `usage-cost-proof-${index}`; + await fs.appendFile( + path.join(sessionsDir, `${sessionId}.jsonl`), + `${makeEntry(sessionId, secondTimestamp, 2)}\n`, + "utf-8", + ); + } + + const refreshStart = performance.now(); + await refreshCostUsageCache(); + const staleRefreshMs = performance.now() - refreshStart; + + const summary = await loadCostUsageSummaryFromCache({ + startMs: Date.UTC(2026, 1, 5), + endMs: Date.UTC(2026, 1, 5) + 24 * 60 * 60 * 1000 - 1, + requestRefresh: false, + }); + + const expectedTokens = sessionCount * 3; + if (summary.totals.totalTokens !== expectedTokens) { + throw new Error(`expected ${expectedTokens} tokens, got ${summary.totals.totalTokens}`); + } + if (summary.cacheStatus?.status !== "fresh") { + throw new Error(`expected fresh cache, got ${summary.cacheStatus?.status ?? "missing"}`); + } + + const cachePath = path.join(sessionsDir, ".usage-cost-cache.json"); + const cacheStats = await fs.stat(cachePath); + console.log( + JSON.stringify( + { + sessionCount, + coldMs: Math.round(coldMs), + staleRefreshMs: Math.round(staleRefreshMs), + cacheBytes: cacheStats.size, + totalTokens: summary.totals.totalTokens, + cacheStatus: summary.cacheStatus?.status, + }, + null, + 2, + ), + ); + console.log("All runtime assertions passed."); +} finally { + if (previousStateDir === undefined) { + delete process.env.OPENCLAW_STATE_DIR; + } else { + process.env.OPENCLAW_STATE_DIR = previousStateDir; + } + await fs.rm(root, { recursive: true, force: true }); +} diff --git a/src/infra/session-cost-usage.test.ts b/src/infra/session-cost-usage.test.ts index b6f9817b2b79..22bce31c7611 100644 --- a/src/infra/session-cost-usage.test.ts +++ b/src/infra/session-cost-usage.test.ts @@ -1410,6 +1410,78 @@ describe("session cost usage", () => { }); }); + it("throttles cache writes during a large stale refresh and skips writes when nothing changed", async () => { + const root = await makeSessionCostRoot("cost-cache-throttle"); + const sessionsDir = path.join(root, "agents", "main", "sessions"); + await fs.mkdir(sessionsDir, { recursive: true }); + const cachePath = path.join(sessionsDir, ".usage-cost-cache.json"); + + const sessionCount = 300; // > USAGE_COST_CACHE_CHECKPOINT_FILES (256) + const baseTimestamp = "2026-02-05T12:00:00.000Z"; + const makeEntry = (totalTokens: number, timestamp: string) => + JSON.stringify({ + type: "message", + timestamp, + message: { + role: "assistant", + provider: "openai", + model: "gpt-5.4", + usage: { + input: totalTokens, + output: 0, + totalTokens, + cost: { total: totalTokens / 1000 }, + }, + }, + }); + + for (let index = 0; index < sessionCount; index += 1) { + const sessionId = `sess-throttle-${index}`; + await fs.writeFile( + path.join(sessionsDir, `${sessionId}.jsonl`), + `${JSON.stringify({ type: "session", version: 1, id: sessionId })}\n${makeEntry(1, baseTimestamp)}\n`, + "utf-8", + ); + } + + await withStateDir(root, async () => { + const renameSpy = vi.spyOn(fs, "rename"); + const cacheRenamesBefore = renameSpy.mock.calls.length; + try { + await refreshCostUsageCache(); + const cacheRenamesAfterCold = renameSpy.mock.calls.filter( + ([, dest]) => dest === cachePath, + ).length; + + // Without throttling this would be ~sessionCount + 1 writes (601). + // With checkpointing it must be far fewer than the file count. + expect(cacheRenamesAfterCold).toBeGreaterThan(0); + expect(cacheRenamesAfterCold).toBeLessThan(sessionCount / 4); + + const summary = await loadCostUsageSummaryFromCache({ + startMs: Date.UTC(2026, 1, 5), + endMs: Date.UTC(2026, 1, 5) + 24 * 60 * 60 * 1000 - 1, + requestRefresh: false, + }); + expect(summary.totals.totalTokens).toBe(sessionCount); + expect(summary.cacheStatus?.status).toBe("fresh"); + + // No-op refresh: nothing stale, nothing deleted -> no cache rewrite. + const renamesBeforeNoOp = renameSpy.mock.calls.filter( + ([, dest]) => dest === cachePath, + ).length; + await refreshCostUsageCache(); + const renamesAfterNoOp = renameSpy.mock.calls.filter( + ([, dest]) => dest === cachePath, + ).length; + expect(renamesAfterNoOp).toBe(renamesBeforeNoOp); + } finally { + renameSpy.mockRestore(); + void cacheRenamesBefore; + } + }); + }); + it("batches stale session summary refreshes for the same agent", async () => { const root = await makeSessionCostRoot("cost-cache-session-batch"); const sessionsDir = path.join(root, "agents", "main", "sessions"); diff --git a/src/infra/session-cost-usage.ts b/src/infra/session-cost-usage.ts index f86d8ba0240c..0293d15ffe9c 100644 --- a/src/infra/session-cost-usage.ts +++ b/src/infra/session-cost-usage.ts @@ -91,6 +91,11 @@ const USAGE_COST_CACHE_VERSION = 4; const USAGE_COST_CACHE_FILE = ".usage-cost-cache.json"; const USAGE_COST_CACHE_LOCK_WRITE_GRACE_MS = 10_000; const USAGE_COST_TRANSCRIPT_STAT_CONCURRENCY = 32; +// Checkpoint policy for refreshCostUsageCache: bound the cost of full cache +// serialization when scanning thousands of session files. Smaller of the two +// limits triggers the next durable write. +const USAGE_COST_CACHE_CHECKPOINT_FILES = 256; +const USAGE_COST_CACHE_CHECKPOINT_INTERVAL_MS = 5_000; const logger = createSubsystemLogger("usage-cost-cache"); type UsageCostRefreshState = { @@ -1527,9 +1532,11 @@ async function refreshCostUsageCacheForPath(params?: { ? files : files.filter((file) => file.mtimeMs >= refreshStartMs); const livePaths = new Set(files.map((file) => file.filePath)); + let cacheMutated = false; for (const filePath of Object.keys(cache.files)) { if (!livePaths.has(filePath)) { delete cache.files[filePath]; + cacheMutated = true; } } @@ -1551,6 +1558,14 @@ async function refreshCostUsageCacheForPath(params?: { .slice(0, maxFiles); const resolveCost = createUsageCostResolver(params?.config); + // Throttle full cache rewrites: writing a 100MB+ JSON cache after every + // single scanned session balloons CPU/IO into O(N * cacheSize). Instead, + // checkpoint at most once every USAGE_COST_CACHE_CHECKPOINT_INTERVAL_MS + // (or every USAGE_COST_CACHE_CHECKPOINT_FILES files) so an interrupted + // refresh still makes durable forward progress while a normal refresh of + // thousands of files only pays the serialization cost a handful of times. + let dirtyCount = 0; + let lastCheckpointMs = Date.now(); for (const file of staleFiles) { cache.files[file.filePath] = await scanUsageFileForCache({ file, @@ -1559,12 +1574,24 @@ async function refreshCostUsageCacheForPath(params?: { previous: cache.files[file.filePath], includeSessionSummary: sessionSummaryFiles.has(file.filePath), }); + dirtyCount += 1; + cacheMutated = true; + const now = Date.now(); + if ( + dirtyCount >= USAGE_COST_CACHE_CHECKPOINT_FILES || + now - lastCheckpointMs >= USAGE_COST_CACHE_CHECKPOINT_INTERVAL_MS + ) { + cache.updatedAt = now; + await writeUsageCostCache(cachePath, cache); + dirtyCount = 0; + lastCheckpointMs = Date.now(); + } + } + + if (cacheMutated || dirtyCount > 0) { cache.updatedAt = Date.now(); await writeUsageCostCache(cachePath, cache); } - - cache.updatedAt = Date.now(); - await writeUsageCostCache(cachePath, cache); return "refreshed"; } finally { await lock.release();