diff --git a/src/infra/session-cost-usage.test.ts b/src/infra/session-cost-usage.test.ts index b8efd3d3b074..0a31451efcc3 100644 --- a/src/infra/session-cost-usage.test.ts +++ b/src/infra/session-cost-usage.test.ts @@ -33,6 +33,7 @@ import { loadSessionUsageTimeSeries as loadSessionUsageTimeSeriesForAgent, resolveExistingUsageSessionFile as resolveExistingUsageSessionFileForAgent, } from "./session-cost-usage.js"; +import { testing as sessionCostUsageTestApi } from "./session-cost-usage.test-support.js"; type WithOptionalAgentId = T extends (params: infer P) => unknown ? Omit & { agentId?: string } @@ -77,6 +78,13 @@ function waitForFast( return vi.waitFor(callback, { interval: 1, ...options }); } +async function refreshSessionCostUsageForTest(sessionFile: string): Promise { + await sessionCostUsageTestApi.usageCostRefreshRuntime.refreshCostUsageCacheForAgent({ + agentId: "main", + sessionFiles: [sessionFile], + }); +} + function clearGatewayModelPricingState(): void { replaceGatewayModelPricingCache(new Map(), 0); clearGatewayModelPricingFailures(); @@ -931,21 +939,20 @@ describe("session cost usage", () => { [ JSON.stringify(assistantEntry(undefined, 1_000)), JSON.stringify(assistantEntry("2026-02-05T12:00:00.000Z", 20)), + "", ].join("\n"), "utf-8", ); await withStateDir(root, async () => { const session = { sessionId: "sess-v8-upgrade", sessionFile }; - await loadSessionCostSummariesFromCache({ sessions: [session], agentId: "main" }); - await waitForFast(async () => { - const current = await loadSessionCostSummariesFromCache({ - sessions: [session], - agentId: "main", - requestRefresh: false, - }); - expect(current.cacheStatus.status).toBe("fresh"); + await refreshSessionCostUsageForTest(sessionFile); + const current = await loadSessionCostSummariesFromCache({ + sessions: [session], + agentId: "main", + requestRefresh: false, }); + expect(current.cacheStatus.status).toBe("fresh"); const currentRow = requireValue( readSessionCostUsageRollupRows("main").find((row) => row.key === sessionFile), @@ -968,41 +975,32 @@ describe("session cost usage", () => { ).toBe(true); const rangeEndMs = Date.UTC(2026, 1, 5) + 24 * 60 * 60 * 1000 - 1; - await loadSessionCostSummariesFromCache({ + await refreshSessionCostUsageForTest(sessionFile); + const rebuilt = await loadSessionCostSummariesFromCache({ sessions: [session], agentId: "main", startMs: Date.UTC(2026, 1, 5), endMs: rangeEndMs, + requestRefresh: false, }); - await waitForFast(async () => { - const rebuilt = await loadSessionCostSummariesFromCache({ - sessions: [session], - agentId: "main", - startMs: Date.UTC(2026, 1, 5), - endMs: rangeEndMs, - requestRefresh: false, - }); - expect(rebuilt.cacheStatus.status).toBe("fresh"); - expect(rebuilt.summaries[0]?.totalTokens).toBe(20); - }); + expect(rebuilt.cacheStatus.status).toBe("fresh"); + expect(rebuilt.summaries[0]?.totalTokens).toBe(20); await fs.appendFile( sessionFile, - `\n${JSON.stringify(assistantEntry("2026-02-05T13:00:00.000Z", 5))}`, + `${JSON.stringify(assistantEntry("2026-02-05T13:00:00.000Z", 5))}\n`, "utf-8", ); - await loadSessionCostSummariesFromCache({ sessions: [session], agentId: "main" }); - await waitForFast(async () => { - const appended = await loadSessionCostSummariesFromCache({ - sessions: [session], - agentId: "main", - startMs: Date.UTC(2026, 1, 5), - endMs: rangeEndMs, - requestRefresh: false, - }); - expect(appended.cacheStatus.status).toBe("fresh"); - expect(appended.summaries[0]?.totalTokens).toBe(25); + await refreshSessionCostUsageForTest(sessionFile); + const appended = await loadSessionCostSummariesFromCache({ + sessions: [session], + agentId: "main", + startMs: Date.UTC(2026, 1, 5), + endMs: rangeEndMs, + requestRefresh: false, }); + expect(appended.cacheStatus.status).toBe("fresh"); + expect(appended.summaries[0]?.totalTokens).toBe(25); const appendedRow = requireValue( readSessionCostUsageRollupRows("main").find((row) => row.key === sessionFile),