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();