diff --git a/apps/server/src/usage/UsageService.test.ts b/apps/server/src/usage/UsageService.test.ts index bff9bdb1f57fd..281884dcc3bc2 100644 --- a/apps/server/src/usage/UsageService.test.ts +++ b/apps/server/src/usage/UsageService.test.ts @@ -249,6 +249,7 @@ const layerService = (input: { ); const BEFORE_NARROW_WINDOW = Date.parse("2026-08-01T10:00:00Z") / 1000; +const NARROW_WINDOW_MTIME = Date.parse("2026-09-01T10:00:00Z") / 1000; const NARROW_WINDOW: UsageSummaryInput = { timeZone: "UTC", sinceDay: UsageDay.make("2026-09-01"), @@ -2449,6 +2450,75 @@ describe("UsageService", () => { }).pipe(Effect.scoped), ); + it.live( + "answers a provisional read from the last scan, then the exact read adds what changed", + () => + Effect.gen(function* () { + const { transcript, settings, home } = yield* setup; + yield* Effect.promise(() => NodeFSP.writeFile(transcript, claudeLine(1, 5))); + const provisional = { ...WINDOW, provisional: true }; + const claude = (summary: UsageSummary) => + summary.sources.find((source) => source.fingerprint.provider === "claude"); + yield* Effect.gen(function* () { + // With nothing persisted, the first answer is a full scan. + const service = yield* UsageService.make; + const loop = NodePath.join(NodePath.dirname(transcript), "loop.jsonl"); + yield* Effect.promise(() => NodeFSP.symlink(loop, loop, "junction")); + assert.strictEqual(claude(yield* service.readSummary(provisional))?.status, "partial"); + yield* Effect.promise(() => NodeFSP.rm(loop)); + // A scan that missed a path covers nothing, so this is a full scan too. + assert.isUndefined(claude(yield* service.readSummary(provisional))?.refreshing); + yield* service.awaitPersisted; + + yield* Effect.promise(() => NodeFSP.appendFile(transcript, claudeLine(2, 7))); + const restarted = yield* UsageService.make; + const saved = yield* restarted.readSummary(provisional); + assert.strictEqual(totalOutputTokens(saved), 5); + assert.isTrue(claude(saved)?.refreshing); + // A window that starts before the saved scan could only show part of its usage. + const wider = yield* restarted.readSummary({ + ...provisional, + sinceDay: UsageDay.make("2026-07-01"), + }); + assert.strictEqual(totalOutputTokens(wider), 12); + assert.isUndefined(claude(wider)?.refreshing); + const refreshed = yield* restarted.readSummary({ ...WINDOW, awaitRefresh: true }); + assert.strictEqual(totalOutputTokens(refreshed), 12); + assert.isUndefined(claude(refreshed)?.refreshing); + // Saved usage with nothing to show marks nothing refreshing, so it must not be the answer. + yield* Effect.promise(() => + NodeFSP.appendFile(transcript, claudeLine(3, 11, "claude-other")), + ); + const unseen = yield* restarted.readSummary({ + ...provisional, + modelFilter: { include: [{ model: "claude-other" }] }, + }); + assert.strictEqual(totalOutputTokens(unseen), 11); + yield* restarted.awaitPersisted; + + // A saved scan with an entry this server cannot read covers nothing. + const { stateDir } = yield* ServerConfig.ServerConfig; + const cachePath = NodePath.join(stateDir, SCAN_CACHE_FILE_NAME); + yield* Effect.promise(async () => { + const document = decodeUnknownJsonString(await NodeFSP.readFile(cachePath, "utf8")) as { + files: Record>; + }; + const [entry] = Object.values(document.files); + document.files[NodePath.join(NodePath.dirname(transcript), "broken.jsonl")] = { + ...entry, + r: "broken", + }; + await NodeFSP.writeFile(cachePath, JSON.stringify(document)); + }); + const reloaded = yield* UsageService.make; + assert.isUndefined(claude(yield* reloaded.readSummary(provisional))?.refreshing); + }).pipe( + Effect.scoped, + Effect.provide(layerService({ prefix: "usage-service-provisional", home, settings })), + ); + }).pipe(Effect.scoped), + ); + it.live("filters original model IDs before aliases, pricing, buckets and session counts", () => Effect.gen(function* () { const { transcript, settings, home } = yield* setup; @@ -3204,6 +3274,64 @@ describe("UsageService", () => { }).pipe(Effect.scoped), ); + it.live("credits the same copy of a duplicate from the last scan as an exact read does", () => + Effect.gen(function* () { + yield* TestClock.setTime(Date.parse("2026-08-15T00:00:00Z")); + const { transcript, settings, home } = yield* setup; + const dir = NodePath.dirname(transcript); + const [first = "", second = ""] = yield* Effect.promise(async () => { + await NodeFSP.writeFile(NodePath.join(dir, "a.jsonl"), ""); + await NodeFSP.writeFile(NodePath.join(dir, "b.jsonl"), ""); + return (await NodeFSP.readdir(dir)).map((name) => NodePath.join(dir, name)); + }); + const forked = (line: string) => line.replace('"session-1"', '"session-2"'); + yield* Effect.promise(async () => { + await NodeFSP.writeFile(first, claudeLine(1, 5)); + await NodeFSP.utimes(first, BEFORE_NARROW_WINDOW, BEFORE_NARROW_WINDOW); + await NodeFSP.writeFile(second, forked(claudeLine(1, 5)) + forked(claudeLine(2, 7))); + }); + yield* Effect.gen(function* () { + const service = yield* UsageService.make; + // The narrow window saves the second copy before the wide one reads the first. + yield* service.readSummary(NARROW_WINDOW); + const exact = yield* service.readSummary(WINDOW); + yield* service.awaitPersisted; + const restarted = yield* UsageService.make; + const saved = yield* restarted.readSummary({ ...WINDOW, provisional: true }); + assert.isTrue(saved.sources.some((source) => source.refreshing === true)); + assert.deepStrictEqual(saved.buckets, exact.buckets); + }).pipe( + Effect.provide(layerService({ prefix: "usage-service-saved-order-test", home, settings })), + ); + + // A narrow scan that first saves the earlier copy still puts it before a later file it skipped. + yield* Effect.promise(async () => { + await NodeFSP.writeFile(first, forked(claudeLine(1, 5)) + forked(claudeLine(2, 7))); + await NodeFSP.utimes( + first, + Date.parse("2026-07-01T00:00:00Z") / 1000, + Date.parse("2026-07-01T00:00:00Z") / 1000, + ); + await NodeFSP.writeFile(second, claudeLine(1, 5)); + await NodeFSP.utimes(second, BEFORE_NARROW_WINDOW, BEFORE_NARROW_WINDOW); + }); + yield* Effect.gen(function* () { + const service = yield* UsageService.make; + yield* service.readSummary(WINDOW); + yield* Effect.promise(() => + NodeFSP.utimes(first, NARROW_WINDOW_MTIME, NARROW_WINDOW_MTIME), + ); + yield* service.readSummary(NARROW_WINDOW); + const saved = yield* service.readSummary({ ...WINDOW, provisional: true }); + assert.isTrue(saved.sources.some((source) => source.refreshing === true)); + const exact = yield* service.readSummary(WINDOW); + assert.deepStrictEqual(saved.buckets, exact.buckets); + }).pipe( + Effect.provide(layerService({ prefix: "usage-service-narrow-order-test", home, settings })), + ); + }).pipe(Effect.scoped, Effect.provide(TestClock.layer())), + ); + it.live.skipIf(HostProcess.Platform.defaultValue() === "win32")( "keeps a newer cached read when a slower scan of another window finishes later", () => @@ -3345,7 +3473,14 @@ describe("UsageService", () => { HttpClient.HttpClient, HttpClient.make((request) => Deferred.await(releaseRates).pipe( - Effect.as(HttpClientResponse.fromWeb(request, Response.json({}))), + Effect.as( + HttpClientResponse.fromWeb( + request, + Response.json({ + "example-model": { input_cost_per_token: 1e-6, output_cost_per_token: 1e-5 }, + }), + ), + ), ), ), ), @@ -3364,7 +3499,8 @@ describe("UsageService", () => { const original = yield* Fiber.join(first); const updated = yield* Fiber.join(second); - assert.strictEqual(original.buckets[0]?.costUsd, 0); + // Fetched rates price the first read, so scans must start while the fetch is pending. + assert.isAbove(original.buckets[0]?.costUsd ?? 0, 0); assert.closeTo(updated.buckets[0]?.costUsd ?? -1, 0.00006, 1e-12); }).pipe( Effect.provide(layerService({ prefix: "usage-service-price-race-test", home, settings })), diff --git a/apps/server/src/usage/UsageService.ts b/apps/server/src/usage/UsageService.ts index 912ed51c96b3e..cd04e4bc4c1c9 100644 --- a/apps/server/src/usage/UsageService.ts +++ b/apps/server/src/usage/UsageService.ts @@ -137,7 +137,10 @@ function sharedSessions( return shared; } const decodeCachedSources = Schema.decodeUnknownOption( - Schema.Struct({ sources: Schema.Record(Schema.String, CachedSource) }), + Schema.Struct({ + sources: Schema.Record(Schema.String, CachedSource), + coveredSince: Schema.optional(Schema.Record(Schema.String, Schema.Number)), + }), ); export class UsageService extends Context.Service< @@ -165,7 +168,15 @@ export const make = Effect.gen(function* () { const correctedBatch = (records: readonly UsageRecord[]) => records.some((record) => resolvedRecords.has(record)) ? records.map(correctedRecord) : records; const sourceCache = new Map(); + // Earliest window start a complete scan read each transcript directory from. + const coveredSince = new Map(); + const coverageKey = (source: TranscriptSource) => `${source.provider}\0${source.dir}`; let cacheDirty = false; + const cover = (source: TranscriptSource, windowStartMs: number) => { + if ((coveredSince.get(coverageKey(source)) ?? Infinity) <= windowStartMs) return; + coveredSince.set(coverageKey(source), windowStartMs); + cacheDirty = true; + }; const isWithinDirectory = (filePath: string, dir: string) => { const relative = path.relative(dir, filePath); return relative !== ".." && !relative.startsWith(".." + path.sep) && !path.isAbsolute(relative); @@ -235,27 +246,26 @@ export const make = Effect.gen(function* () { Effect.catchCause(() => Effect.succeed(null)), ); + // Outside ratesLock, which a slow fetch holds: a provisional answer prices with these. + const loadSavedRates = yield* Effect.cached( + fileSystem.readFileString(ratesCachePath).pipe( + Effect.flatMap((raw) => decodeRatesCache(raw)), + Effect.map((saved) => { + const parsed = parseRateTable(saved.document); + if (parsed.size === 0) return; + rates = parsed; + ratesFetchedAtMs = saved.fetchedAtMs; + ratesStatus = "cached"; + }), + Effect.ignoreCause, + ), + ); + const loadRates = Effect.fn("UsageService.loadRates")(function* (force: boolean) { const now = yield* Clock.currentTimeMillis; const maxAgeMs = force ? RATES_REFRESH_FLOOR_MS : RATES_TTL_MS; + yield* loadSavedRates; if (ratesFetchedAtMs !== null && now - ratesFetchedAtMs < maxAgeMs) return; - - if (ratesFetchedAtMs === null) { - const fromDisk = yield* fileSystem.readFileString(ratesCachePath).pipe( - Effect.flatMap((raw) => decodeRatesCache(raw)), - Effect.catchCause(() => Effect.succeed(null)), - ); - if (fromDisk !== null) { - const parsed = parseRateTable(fromDisk.document); - if (parsed.size > 0) { - rates = parsed; - ratesFetchedAtMs = fromDisk.fetchedAtMs; - ratesStatus = "cached"; - if (now - fromDisk.fetchedAtMs < maxAgeMs) return; - } - } - } - if (ratesLastFailedAtMs !== null && now - ratesLastFailedAtMs < RATES_REFRESH_FLOOR_MS) return; let fetched = yield* fetchRates(LITELLM_RATES_URL); @@ -396,19 +406,27 @@ export const make = Effect.gen(function* () { Effect.catchCause(() => Effect.succeed(null)), ); let document = yield* readDocument(scanCachePath); + const current = document !== null; for (const fileName of LEGACY_SCAN_CACHE_FILE_NAMES) { if (document !== null) break; document = yield* readDocument(path.join(config.stateDir, fileName)); cacheDirty = document !== null; } if (document === null) return; - for (const [path, entry] of decodeScanCache(document, transcriptFormats)) { - fileCache.set(path, entry); - } + const decoded = decodeScanCache(document, transcriptFormats); + for (const [path, entry] of decoded) fileCache.set(path, entry); + const saved = (document as { readonly files?: unknown }).files; + const complete = + typeof saved === "object" && saved !== null && Object.keys(saved).length === decoded.size; const sources = decodeCachedSources(document); if (Option.isSome(sources)) { for (const [key, source] of Object.entries(sources.value.sources)) sourceCache.set(key, source); + // Coverage describes the saved files, so it needs every one, in this version's format. + if (current && complete && fileCache.size > 0) { + for (const [key, since] of Object.entries(sources.value.coveredSince ?? {})) + coveredSince.set(key, since); + } } }), ); @@ -448,7 +466,10 @@ export const make = Effect.gen(function* () { if (cacheDirty) { cacheDirty = false; yield* Effect.sync(() => - writeScanCache(fileCache, { sources: Object.fromEntries(sourceCache) }), + writeScanCache(fileCache, { + sources: Object.fromEntries(sourceCache), + coveredSince: Object.fromEntries(coveredSince), + }), ).pipe( Effect.flatMap((contents) => writeCacheFile(scanCachePath, contents)), Effect.andThen( @@ -595,19 +616,27 @@ export const make = Effect.gen(function* () { readonly malformedRecords: number; }[] | null; + /** Every transcript path the walk listed, in walk order, whatever its mtime. */ + readonly walked?: readonly string[]; readonly refreshing?: true; } const scanTranscriptDir = Effect.fn("UsageService.scanTranscriptDir")(function* ( source: TranscriptSource, windowStartMs: number, + fromMemory: boolean, ) { const { provider, format, dir, volumeId, fileName } = source; const exists = yield* fileSystem .exists(dir) .pipe(Effect.catchCause(() => Effect.succeed(false))); - if (!exists) return { provider, dir, volumeId, files: null } satisfies ScannedDir; - const { files, failedPaths } = yield* Effect.promise(() => + if (!exists) { + cover(source, windowStartMs); + return { provider, dir, volumeId, files: null } satisfies ScannedDir; + } + // scanSummary aggregates the saved files, as it does for a deleted directory. + if (fromMemory) return { provider, dir, volumeId, files: [] } satisfies ScannedDir; + const { files, walked, failedPaths } = yield* Effect.promise(() => listTranscriptFiles(dir, windowStartMs, fileName === undefined ? undefined : { fileName }), ); const read = yield* Effect.forEach( @@ -646,6 +675,7 @@ export const make = Effect.gen(function* () { if (!livePaths.has(originalPath) && fileCache.delete(originalPath)) cacheDirty = true; } const unread = failedPaths + read.filter((file) => file.failed).length; + if (unread === 0) cover(source, windowStartMs); yield* Effect.annotateCurrentSpan({ "usage.files": files.length, "usage.files_parsed": read.filter((file) => file.update !== undefined).length, @@ -655,6 +685,7 @@ export const make = Effect.gen(function* () { dir, volumeId, files: parsedFiles, + walked, ...(unread > 0 ? { status: "partial", @@ -665,17 +696,13 @@ export const make = Effect.gen(function* () { }); const collectDirs = Effect.fn("UsageService.collectDirs")(function* ( + { dirs, pending }: Effect.Success>, windowStartMs: number, settings: ServerSettingsValue, retentionCutoffMs: number, awaitRefresh: boolean, + fromMemory: boolean, ) { - const { dirs, pending } = yield* resolveTranscriptDirs( - settings, - retentionCutoffMs, - awaitRefresh, - ).pipe(Effect.provideContext(readerContext)); - const scans = Effect.forEach( scanReaders, ({ driver, reader }) => @@ -705,7 +732,7 @@ export const make = Effect.gen(function* () { // Scan independent sources together, directories sequentially to bound open files; keep order, since the first copy decides a duplicate's bucket. const [transcripts, scanDirs] = yield* Effect.all( - [Effect.forEach(dirs, (dir) => scanTranscriptDir(dir, windowStartMs)), scans], + [Effect.forEach(dirs, (dir) => scanTranscriptDir(dir, windowStartMs, fromMemory)), scans], { concurrency: "unbounded" }, ); const scanned: readonly ScannedDir[] = [...transcripts, ...pending, ...scanDirs.flat()]; @@ -715,6 +742,7 @@ export const make = Effect.gen(function* () { const scanSummary = Effect.fn("UsageService.scanSummary")(function* ( input: UsageSummaryInput, settings: ServerSettingsValue, + provisional: boolean, ) { if (input.sinceDay > input.untilDay) { return yield* new UsageReadError({ @@ -762,11 +790,38 @@ export const make = Effect.gen(function* () { const retentionCutoffMs = startedAtMs - CACHE_RETENTION_DAYS * 24 * 60 * 60 * 1000; - const [, scannedDirs] = yield* Effect.all( - [ - ensureRates(false), - collectDirs(windowStartMs, settings, retentionCutoffMs, input.awaitRefresh === true), - ], + const awaitRefresh = input.awaitRefresh === true; + const collect = Effect.gen(function* () { + const resolved = yield* resolveTranscriptDirs(settings, retentionCutoffMs, awaitRefresh).pipe( + Effect.provideContext(readerContext), + ); + // A window no complete scan has read would show partial totals, so every directory must cover it. + // Saved files are keyed by path, so two providers reading one directory keep only one parse. + const fromMemory = + provisional && + resolved.dirs.length > 0 && + new Set(resolved.dirs.map((source) => source.dir)).size === resolved.dirs.length && + resolved.dirs.every( + (source) => (coveredSince.get(coverageKey(source)) ?? Infinity) <= windowStartMs, + ); + const [, scanned] = yield* Effect.all( + [ + fromMemory ? loadSavedRates : provisional ? ensureRates(false) : Effect.void, + collectDirs( + resolved, + windowStartMs, + settings, + retentionCutoffMs, + awaitRefresh, + fromMemory, + ), + ], + { concurrency: 2 }, + ); + return { scannedDirs: scanned, fromMemory }; + }); + const [{ scannedDirs, fromMemory }] = yield* Effect.all( + [collect, provisional ? Effect.void : ensureRates(false)], { concurrency: 2 }, ); @@ -786,31 +841,73 @@ export const make = Effect.gen(function* () { // Deleted transcripts' saved usage still aggregates and deduplicates normally. const retainedSinceMs = Math.max(windowStartMs, retentionCutoffMs); + // Saved paths are files below a scanned root, so their folder decides; folders repeat thousands of times. + const withinDir = (dir: string) => { + const foldersWithin = new Map(); + return (filePath: string) => { + const folder = path.dirname(filePath); + let within = foldersWithin.get(folder); + if (within === undefined) { + within = isWithinDirectory(folder, dir); + foldersWithin.set(folder, within); + } + return within; + }; + }; const filesByDir = scannedDirs.map(({ provider, dir, files }) => { const retainedFiles = [...(files ?? [])]; const livePaths = new Set(retainedFiles.map((file) => file.path)); + const isRetained = withinDir(dir); for (const [filePath, entry] of fileCache) { if ( entry.provider !== provider || entry.mtimeMs < retainedSinceMs || livePaths.has(filePath) || - !isWithinDirectory(filePath, dir) + !isRetained(filePath) ) continue; retainedFiles.push({ path: filePath, - records: [...entry.records, ...entry.tailRecords], + records: + entry.tailRecords.length === 0 + ? entry.records + : [...entry.records, ...entry.tailRecords], malformedRecords: entry.malformedRecords, }); } return retainedFiles; }); + // Usage two transcripts share counts in the first one, so saved files take the order an exact + // scan reads them in: walked files in walk order, then the directory's deleted files. + for (const { provider, dir, walked } of fromMemory ? [] : scannedDirs) { + if (walked === undefined) continue; + const within = withinDir(dir); + const saved = [...fileCache.keys()]; + const members = new Set( + saved.filter( + (filePath) => fileCache.get(filePath)!.provider === provider && within(filePath), + ), + ); + if (members.size < 2) continue; + const listed = new Set(walked); + const target = [ + ...walked.filter((filePath) => members.has(filePath)), + ...saved.filter((filePath) => members.has(filePath) && !listed.has(filePath)), + ]; + let next = 0; + const ordered = saved.map((filePath) => (members.has(filePath) ? target[next++]! : filePath)); + if (ordered.every((filePath, index) => filePath === saved[index])) continue; + const entries = new Map(fileCache); + fileCache.clear(); + for (const filePath of ordered) fileCache.set(filePath, entries.get(filePath)!); + cacheDirty = true; + } const windowEndMs = hourlyWindow?.untilTimeMs ?? Date.parse(`${input.untilDay}T00:00:00Z`) + 86400000 + MTIME_SLACK_MS; let resolvedAny = false; for (const { reader } of transcriptReaders) { - if (reader.resolveServiceTiers === undefined) continue; + if (reader.resolveServiceTiers === undefined || fromMemory) continue; const pending = new Set(); for (const [index, files] of filesByDir.entries()) { if (scannedDirs[index]?.provider !== reader.provider) continue; @@ -860,6 +957,7 @@ export const make = Effect.gen(function* () { // Distinct per directory. Buckets carry per-cell session counts, but a // session spans days and models, so clients total this figure instead. const sessionIds = new Set(); + let counted = false; for (const file of filesByDir[index] ?? []) { malformedRecords += file.malformedRecords; @@ -888,8 +986,9 @@ export const make = Effect.gen(function* () { eventOccurrences.set(key, occurrence); usageRecord = { ...record, dedupeKey: key + ":" + occurrence }; } - if (aggregator.add(usageRecord, dir) && record.sessionId.length > 0) { - sessionIds.add(record.sessionId); + if (aggregator.add(usageRecord, dir)) { + counted = true; + if (record.sessionId.length > 0) sessionIds.add(record.sessionId); } } } @@ -916,18 +1015,28 @@ export const make = Effect.gen(function* () { .filter(Boolean) .join(" ") : sourceMessage, - ...(refreshing ? { refreshing } : {}), + // A saved answer misses what changed since its scan, so the client must ask again. An + // empty source stays unmarked: it would only flash a loading row. + ...(refreshing || (fromMemory && counted && transcriptFormats.has(provider)) + ? { refreshing: true } + : {}), }); } const pruned = pruneScanCache(fileCache, retentionCutoffMs); - if (pruned > 0) cacheDirty = true; + if (pruned > 0) { + cacheDirty = true; + // Pruned files are no longer saved, so nothing before the cutoff stays covered. + for (const [key, since] of coveredSince) { + if (since < retentionCutoffMs) coveredSince.set(key, retentionCutoffMs); + } + } const aggregated = aggregator.finish(); const readAt = yield* DateTime.now; const finishedAtMs = yield* Clock.currentTimeMillis; - return { + const summary = { contractVersion: USAGE_CONTRACT_VERSION, readAt: DateTime.formatIso(readAt), timeZone: input.timeZone, @@ -939,11 +1048,28 @@ export const make = Effect.gen(function* () { scanDurationMs: Math.max(0, finishedAtMs - startedAtMs), ...(input.modelFilter === undefined ? {} : { modelFilter: input.modelFilter }), } satisfies UsageSummary; + return { summary, fromMemory }; + }); + + const readWindow = Effect.fn("UsageService.readWindow")(function* ( + input: UsageSummaryInput, + settings: ServerSettingsValue, + provisional: boolean, + ) { + const { summary, fromMemory } = yield* scanSummary(input, settings, provisional); + // With nothing marked refreshing the client would keep a saved answer as final. + return fromMemory && !summary.sources.some((source) => source.refreshing === true) + ? (yield* scanSummary(input, settings, false)).summary + : summary; }); const inflightScans = new Map>(); - const scanKey = (input: UsageSummaryInput, settings: ServerSettingsValue): string => + const scanKey = ( + input: UsageSummaryInput, + settings: ServerSettingsValue, + provisional: boolean, + ): string => JSON.stringify([ input.timeZone, input.sinceDay, @@ -953,6 +1079,7 @@ export const make = Effect.gen(function* () { input.untilTime ?? null, settings, input.awaitRefresh === true, + provisional, input.modelFilter ?? null, ]); @@ -963,7 +1090,8 @@ export const make = Effect.gen(function* () { const modelFilter = canonicalModelFilter(requestedFilter); const input = modelFilter === undefined ? unfiltered : { ...unfiltered, modelFilter }; const settings = yield* readSettings; - const key = scanKey(input, settings); + const provisional = input.provisional === true && input.awaitRefresh !== true; + const key = scanKey(input, settings, provisional); const deferred = yield* Effect.uninterruptible( Effect.gen(function* () { const existing = inflightScans.get(key); @@ -974,7 +1102,7 @@ export const make = Effect.gen(function* () { inflightScans.set(key, created); // Detach so caller departure cannot cancel shared scans; finished scans warm the cache. // Register writes before releasing waiters, but start writing after delivering their summary. - yield* scanSummary(input, settings).pipe( + yield* readWindow(input, settings, provisional).pipe( Effect.onExit((exit) => Effect.sync(() => inflightScans.delete(key)).pipe( Effect.andThen(Exit.isSuccess(exit) ? schedulePersist : Effect.void), diff --git a/apps/server/src/usage/usageTranscriptReader.test.ts b/apps/server/src/usage/usageTranscriptReader.test.ts index 105884d81ccfb..f4cb72ec1e25c 100644 --- a/apps/server/src/usage/usageTranscriptReader.test.ts +++ b/apps/server/src/usage/usageTranscriptReader.test.ts @@ -99,6 +99,30 @@ describe("compressed transcripts", () => { assert.isAbove(listed.files.find((file) => file.path === compressed)?.size ?? 0, 0); }); + it("lists files in directory order when more directories wait than it reads at once", async () => { + for (let index = 0; index < 40; index += 1) { + const nested = NodePath.join(dir, `project-${index}`, "sessions"); + await NodeFSP.mkdir(nested, { recursive: true }); + await NodeFSP.writeFile(NodePath.join(nested, "a.jsonl"), ""); + await NodeFSP.writeFile(NodePath.join(nested, "..", "b.jsonl"), ""); + } + // Usage that two transcripts share counts in the first one listed, so the order is part of the result. + const inOrder = async (folder: string): Promise => { + const paths: string[] = []; + for (const entry of await NodeFSP.readdir(folder, { withFileTypes: true })) { + const child = NodePath.join(folder, entry.name); + if (entry.isDirectory()) paths.push(...(await inOrder(child))); + else paths.push(child); + } + return paths; + }; + const listed = await listTranscriptFiles(dir, 0); + assert.deepStrictEqual( + listed.files.map((file) => file.path), + await inOrder(dir), + ); + }); + it.each([ { parser: "native", streamingThresholdBytes: undefined }, { parser: "projected", streamingThresholdBytes: 1024 }, diff --git a/apps/server/src/usage/usageTranscriptReader.ts b/apps/server/src/usage/usageTranscriptReader.ts index 38b80d36c8665..d80c51964a20c 100644 --- a/apps/server/src/usage/usageTranscriptReader.ts +++ b/apps/server/src/usage/usageTranscriptReader.ts @@ -1,8 +1,10 @@ // @effect-diagnostics nodeBuiltinImport:off // Cold windows span ~1.4 GB/1,500 files; node:fs buffer streaming is ~10x cheaper than materializing files. +import * as NodeFS from "node:fs"; import * as NodeFSP from "node:fs/promises"; import * as NodePath from "node:path"; import * as NodeStringDecoder from "node:string_decoder"; +import * as NodeUtil from "node:util"; import * as NodeZlib from "node:zlib"; import type { @@ -43,6 +45,9 @@ export const GUARD_LENGTH = 64; const STREAMING_THRESHOLD_BYTES = 8 * 1024 * 1024; const NEWLINE = 0x0a; const STAT_CONCURRENCY = 32; +const WALK_CONCURRENCY = 16; +// Measured about 3x faster than the promise API's stat across 20,000 transcripts. +const statFile = NodeUtil.promisify(NodeFS.stat); const CARRIAGE_RETURN = 0x0d; function selectUsageFields(fields: SelectedFields) { @@ -71,38 +76,54 @@ export async function listTranscriptFiles( root: string, sinceMs: number, options?: { readonly fileName?: string }, -): Promise<{ readonly files: readonly TranscriptFile[]; readonly failedPaths: number }> { +): Promise<{ + readonly files: readonly TranscriptFile[]; + readonly walked: readonly string[]; + readonly failedPaths: number; +}> { const fileName = options?.fileName; - const candidates: string[] = []; let failedPaths = 0; - const walk = async (dir: string): Promise => { + let reading = 0; + const waiting: Array<() => void> = []; + // Sibling directories read together; each returns its own list so the result keeps readdir order. + const walk = async (dir: string): Promise => { + if (reading < WALK_CONCURRENCY) reading += 1; + else await new Promise((resume) => waiting.push(resume)); let entries; try { entries = await NodeFSP.readdir(dir, { withFileTypes: true }); } catch { failedPaths += 1; - return; - } - for (const entry of entries) { - const child = NodePath.join(dir, entry.name); - if (entry.isDirectory()) await walk(child); - else if ( - fileName !== undefined - ? entry.name === fileName - : entry.name.endsWith(".jsonl") || entry.name.endsWith(".jsonl.zst") - ) { - candidates.push(child); - } + return []; + } finally { + // A waiting directory takes over this slot, so the count only drops when none waits. + const next = waiting.shift(); + if (next === undefined) reading -= 1; + else next(); } + const found = await Promise.all( + entries.map((entry) => { + const child = NodePath.join(dir, entry.name); + if (entry.isDirectory()) return walk(child); + return ( + fileName !== undefined + ? entry.name === fileName + : entry.name.endsWith(".jsonl") || entry.name.endsWith(".jsonl.zst") + ) + ? [child] + : []; + }), + ); + return found.flat(); }; - await walk(root); + const candidates = await walk(root); const found: Array = Array.from({ length: candidates.length }); const queue = candidates.entries(); const statQueued = async (): Promise => { for (const [index, path] of queue) { try { - const stats = await NodeFSP.stat(path); + const stats = await statFile(path); if (stats.mtimeMs >= sinceMs) { found[index] = { path, size: stats.size, mtimeMs: stats.mtimeMs }; } @@ -116,7 +137,7 @@ export async function listTranscriptFiles( await Promise.all( Array.from({ length: Math.min(STAT_CONCURRENCY, candidates.length) }, statQueued), ); - return { files: found.filter((file) => file !== undefined), failedPaths }; + return { files: found.filter((file) => file !== undefined), walked: candidates, failedPaths }; } export async function readDirectoryVolumeId(path: string): Promise { diff --git a/packages/client-runtime/src/state/server.ts b/packages/client-runtime/src/state/server.ts index 5a42d07c9b34f..6162df0809c83 100644 --- a/packages/client-runtime/src/state/server.ts +++ b/packages/client-runtime/src/state/server.ts @@ -9,6 +9,7 @@ import { type ServerSelfUpdateProgressEvent, type ServerSelfUpdateResult, shouldRecommendDatabaseCompaction, + type UsageSummary, WS_METHODS, } from "@t3tools/contracts"; import * as Cause from "effect/Cause"; @@ -1497,30 +1498,42 @@ export function createServerEnvironmentAtoms( tag: WS_METHODS.scheduledTasksSubscribe, }), // A cold transcript scan is measured in seconds, so keep the result around - // long enough that switching windows or re-rendering does not rescan. + // long enough that switching windows or re-rendering does not rescan. With + // nothing to show yet, the server's last scan answers first, then a plain + // read, then a read that waits for slow sources. usageSummary: createEnvironmentRpcQueryAtomFamily(runtime, { label: "environment-data:server:usage-summary", tag: WS_METHODS.serverGetUsageSummary, staleTimeMs: 60_000, refreshTrigger: ({ environmentId }) => usageScanSettingsAtom(environmentId), - execute: (input, emit) => - request(WS_METHODS.serverGetUsageSummary, input).pipe( - Effect.filterOrElse( - (summary) => !summary.sources.some((source) => source.refreshing === true), - (summary) => - emit(summary).pipe( - Effect.andThen( - request(WS_METHODS.serverGetUsageSummary, { ...input, awaitRefresh: true }), - ), - Effect.catch((error) => - Effect.logWarning("Could not refresh slow usage sources.").pipe( - Effect.annotateLogs({ ...safeErrorLogAttributes(error) }), - Effect.as(summary), - ), + execute: (input, emit, previous) => { + const settled = (summary: UsageSummary) => + !summary.sources.some((source) => source.refreshing === true); + const read = request(WS_METHODS.serverGetUsageSummary, input).pipe( + Effect.filterOrElse(settled, (summary) => + emit(summary).pipe( + Effect.andThen( + request(WS_METHODS.serverGetUsageSummary, { ...input, awaitRefresh: true }), + ), + Effect.catch((error) => + Effect.logWarning("Could not refresh slow usage sources.").pipe( + Effect.annotateLogs({ ...safeErrorLogAttributes(error) }), + Effect.as(summary), ), ), + ), ), - ), + ); + if (Option.isSome(previous)) return read; + return request(WS_METHODS.serverGetUsageSummary, { ...input, provisional: true }).pipe( + Effect.matchEffect({ + // The saved answer only speeds things up; without it the plain read still answers. + onFailure: () => read, + onSuccess: (saved) => + settled(saved) ? Effect.succeed(saved) : emit(saved).pipe(Effect.andThen(read)), + }), + ); + }, }), resourceTelemetry: createEnvironmentRpcSubscriptionAtomFamily(runtime, { label: "environment-data:server:resource-telemetry", diff --git a/packages/client-runtime/src/state/serverUsage.test.ts b/packages/client-runtime/src/state/serverUsage.test.ts index b16550c7c487c..9d12d1114f896 100644 --- a/packages/client-runtime/src/state/serverUsage.test.ts +++ b/packages/client-runtime/src/state/serverUsage.test.ts @@ -232,7 +232,13 @@ const makeRefreshingHarness = Effect.fn("ServerUsageTest.makeRefreshingHarness") input.awaitRefresh === true ? Deferred.succeed(followUpStarted, undefined).pipe(Effect.andThen(followUp)) : Effect.void, - (input) => [cursorSource(input.awaitRefresh !== true)], + // The saved answer differs from the plain one, so a test can tell which was kept. + (input) => [ + { + ...cursorSource(input.awaitRefresh !== true), + ...(input.provisional === true ? { scannedFiles: 0 } : {}), + }, + ], ); return { ...harness, followUpStarted }; }); @@ -295,7 +301,42 @@ it.effect("keeps the cached usage summary when the awaited refresh fails", () => cursorSource(true), ]); expect(AsyncResult.isSuccess(harness.registry.get(summary))).toBe(true); - expect(harness.requests()).toBe(2); + // The saved answer is followed by a plain read before the one that waits for Cursor. + expect( + harness.inputs().map(({ provisional, awaitRefresh }) => ({ provisional, awaitRefresh })), + ).toEqual([ + { provisional: true, awaitRefresh: undefined }, + { provisional: undefined, awaitRefresh: undefined }, + { provisional: undefined, awaitRefresh: true }, + ]); + // With figures on screen, a refresh reads them exactly instead of asking for saved ones. + yield* harness.updateSettings({ + ...DEFAULT_SERVER_SETTINGS, + usagePriceOverrides: { + "custom-model": { inputCostPerMillionTokens: 5, outputCostPerMillionTokens: 10 }, + }, + }); + yield* waitForCost(harness.registry, summary, 5); + expect(harness.inputs()[3]?.provisional).toBeUndefined(); + unmount(); + }), + ), +); + +it.effect("reads usage in full when the saved answer cannot be read", () => + Effect.scoped( + Effect.gen(function* () { + const harness = yield* makeHarness((_request, input) => + input.provisional === true + ? Effect.fail( + new UsageReadError({ reason: "scanFailed", detail: "saved scan unreadable" }), + ) + : Effect.void, + ); + const summary = harness.summary(); + const unmount = harness.registry.mount(summary); + yield* settledSummary(harness.registry, summary); + expect(harness.inputs().map(({ provisional }) => provisional)).toEqual([true, undefined]); unmount(); }), ), diff --git a/packages/contracts/src/usage.ts b/packages/contracts/src/usage.ts index 3663687811729..76dd69ff91699 100644 --- a/packages/contracts/src/usage.ts +++ b/packages/contracts/src/usage.ts @@ -135,6 +135,8 @@ export const UsageSummaryInput = Schema.Struct({ sinceTime: Schema.optional(TrimmedNonEmptyString), untilTime: Schema.optional(TrimmedNonEmptyString), awaitRefresh: Schema.optional(Schema.Boolean), + /** Accept an immediate answer from the last scan; its sources are marked refreshing. */ + provisional: Schema.optional(Schema.Boolean), modelFilter: Schema.optional(UsageModelFilter), }); export type UsageSummaryInput = typeof UsageSummaryInput.Type;