From 2e6ba44706cdb0eaedd228f201c1c271fa151bf5 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Tue, 26 May 2026 03:27:54 +0200 Subject: [PATCH] fix(perf): bound session transcript stat fanout --- src/infra/session-cost-usage.test.ts | 62 +++++++++++++++++++ src/infra/session-cost-usage.ts | 90 ++++++++++++---------------- 2 files changed, 101 insertions(+), 51 deletions(-) diff --git a/src/infra/session-cost-usage.test.ts b/src/infra/session-cost-usage.test.ts index 944e8594466e..23f6ab283372 100644 --- a/src/infra/session-cost-usage.test.ts +++ b/src/infra/session-cost-usage.test.ts @@ -561,6 +561,68 @@ describe("session cost usage", () => { }); }); + it("limits transcript stat fanout when listing durable cost inputs", async () => { + const root = await makeSessionCostRoot("cost-cache-stat-fanout"); + const sessionsDir = path.join(root, "agents", "main", "sessions"); + await fs.mkdir(sessionsDir, { recursive: true }); + await Promise.all( + Array.from({ length: 48 }, async (_, index) => { + const sessionId = `sess-stat-fanout-${index}`; + await fs.writeFile( + path.join(sessionsDir, `${sessionId}.jsonl`), + transcriptText(sessionId, { + type: "message", + timestamp: "2026-02-05T12:00:00.000Z", + message: { + role: "assistant", + usage: { + input: 1, + output: 1, + totalTokens: 2, + cost: { total: 0.002 }, + }, + }, + }), + "utf-8", + ); + }), + ); + + const originalStat = nodeFs.promises.stat.bind(nodeFs.promises); + let activeStats = 0; + let maxActiveStats = 0; + const statSpy = vi.spyOn(nodeFs.promises, "stat").mockImplementation(async (target) => { + const targetPath = String(target); + if (targetPath.startsWith(sessionsDir) && targetPath.endsWith(".jsonl")) { + activeStats += 1; + maxActiveStats = Math.max(maxActiveStats, activeStats); + try { + await new Promise((resolve) => setTimeout(resolve, 2)); + return await originalStat(target); + } finally { + activeStats -= 1; + } + } + return await originalStat(target); + }); + + await withStateDir(root, async () => { + try { + 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.cacheStatus?.status).toBe("stale"); + } finally { + statSpy.mockRestore(); + } + }); + + expect(maxActiveStats).toBeGreaterThan(1); + expect(maxActiveStats).toBeLessThanOrEqual(32); + }); + it("invalidates durable aggregate cache when pricing config changes", async () => { const root = await makeSessionCostRoot("cost-cache-pricing"); 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 eaafc0dfea00..0f5541b83913 100644 --- a/src/infra/session-cost-usage.ts +++ b/src/infra/session-cost-usage.ts @@ -21,6 +21,7 @@ import { createSubsystemLogger } from "../logging/subsystem.js"; import { stripEnvelope, stripMessageIdHints } from "../shared/chat-envelope.js"; import { asFiniteNumber } from "../shared/number-coercion.js"; import { normalizeOptionalString } from "../shared/string-coerce.js"; +import { runTasksWithConcurrency } from "../utils/run-with-concurrency.js"; import { countToolResults, extractToolCallNames } from "../utils/transcript-tools.js"; import { estimateUsageCost, @@ -88,6 +89,7 @@ const emptyTotals = (): CostUsageTotals => ({ 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; const logger = createSubsystemLogger("usage-cost-cache"); type UsageCostRefreshState = { @@ -366,24 +368,38 @@ async function writeUsageCostCache(cachePath: string, cache: UsageCostCacheFile) }); } -async function listUsageCountedTranscriptFiles( +async function listUsageCountedTranscriptFileStats( agentId?: string, + params?: { minMtimeMs?: number }, ): Promise { const sessionsDir = resolveSessionTranscriptsDirForAgent(agentId); const entries = await fs.promises.readdir(sessionsDir, { withFileTypes: true }).catch(() => []); - const files = await Promise.all( - entries - .filter((entry) => entry.isFile() && isUsageCountedSessionTranscriptFileName(entry.name)) - .map(async (entry) => { + const tasks = entries + .filter((entry) => entry.isFile() && isUsageCountedSessionTranscriptFileName(entry.name)) + .map( + (entry) => async (): Promise => { const filePath = path.join(sessionsDir, entry.name); const stats = await fs.promises.stat(filePath).catch(() => null); if (!stats) { return undefined; } + if (params?.minMtimeMs !== undefined && stats.mtimeMs < params.minMtimeMs) { + return undefined; + } return { filePath, size: stats.size, mtimeMs: stats.mtimeMs }; - }), - ); - return files.filter((file): file is UsageCostTranscriptFile => Boolean(file)); + }, + ); + const { results } = await runTasksWithConcurrency({ + tasks, + limit: USAGE_COST_TRANSCRIPT_STAT_CONCURRENCY, + }); + return results.filter((file): file is UsageCostTranscriptFile => Boolean(file)); +} + +async function listUsageCountedTranscriptFiles( + agentId?: string, +): Promise { + return await listUsageCountedTranscriptFileStats(agentId); } function isUsageCostCacheEntryFresh(params: { @@ -1283,30 +1299,13 @@ export async function loadCostUsageSummary(params?: { const totals = emptyTotals(); const resolveCost = createUsageCostResolver(params?.config); - const sessionsDir = resolveSessionTranscriptsDirForAgent(params?.agentId); - const entries = await fs.promises.readdir(sessionsDir, { withFileTypes: true }).catch(() => []); - const files = ( - await Promise.all( - entries - .filter((entry) => entry.isFile() && isUsageCountedSessionTranscriptFileName(entry.name)) - .map(async (entry) => { - const filePath = path.join(sessionsDir, entry.name); - const stats = await fs.promises.stat(filePath).catch(() => null); - if (!stats) { - return null; - } - // Include file if it was modified after our start time - if (stats.mtimeMs < sinceTime) { - return null; - } - return filePath; - }), - ) - ).filter((filePath): filePath is string => Boolean(filePath)); + const files = await listUsageCountedTranscriptFileStats(params?.agentId, { + minMtimeMs: sinceTime, + }); - for (const filePath of files) { + for (const file of files) { await scanUsageFile({ - filePath, + filePath: file.filePath, config: params?.config, resolveCost, onEntry: (entry) => { @@ -1870,33 +1869,22 @@ export async function discoverAllSessions(params?: { endMs?: number; includeFirstUserMessage?: boolean; }): Promise { - const sessionsDir = resolveSessionTranscriptsDirForAgent(params?.agentId); - const entries = await fs.promises.readdir(sessionsDir, { withFileTypes: true }).catch(() => []); + const files = await listUsageCountedTranscriptFileStats(params?.agentId, { + minMtimeMs: params?.startMs, + }); const discovered = new Map(); - for (const entry of entries) { - if (!entry.isFile() || !isUsageCountedSessionTranscriptFileName(entry.name)) { - continue; - } - - const filePath = path.join(sessionsDir, entry.name); - const stats = await fs.promises.stat(filePath).catch(() => null); - if (!stats) { - continue; - } - - // Filter by date range if provided - if (params?.startMs && stats.mtimeMs < params.startMs) { - continue; - } + for (const file of files) { // Do not exclude by endMs: a session can have activity in range even if it continued later. + const filePath = file.filePath; + const fileName = path.basename(filePath); - const sessionId = parseUsageCountedSessionIdFromFileName(entry.name); + const sessionId = parseUsageCountedSessionIdFromFileName(fileName); if (!sessionId) { continue; } - const isPrimaryTranscript = isPrimarySessionTranscriptFileName(entry.name); + const isPrimaryTranscript = isPrimarySessionTranscriptFileName(fileName); // Try to read first user message for label extraction let firstUserMessage: string | undefined; @@ -1942,13 +1930,13 @@ export async function discoverAllSessions(params?: { const shouldReplace = !existing || (isPrimaryTranscript && !existingIsPrimary) || - (isPrimaryTranscript === existingIsPrimary && stats.mtimeMs >= existing.mtime); + (isPrimaryTranscript === existingIsPrimary && file.mtimeMs >= existing.mtime); if (shouldReplace) { discovered.set(sessionId, { sessionId, sessionFile: filePath, - mtime: stats.mtimeMs, + mtime: file.mtimeMs, firstUserMessage: firstUserMessage ?? existing?.firstUserMessage, }); continue;