From 94de5a0d6b982f1401580886085f26c2107a2e27 Mon Sep 17 00:00:00 2001 From: Hieu Le <78847613+hieu-lee@users.noreply.github.com> Date: Tue, 27 Oct 2026 01:50:46 +0100 Subject: [PATCH 1/2] perf(usage): show saved totals before scanning history --- apps/server/src/usage/UsageService.test.ts | 128 ++++++++++++++++++ apps/server/src/usage/UsageService.ts | 55 ++++++-- .../server/src/usage/usageTranscriptReader.ts | 64 ++++++--- .../usage/usageTranscriptStreaming.test.ts | 29 ++++ packages/client-runtime/src/state/server.ts | 2 +- .../src/state/serverUsage.test.ts | 4 + .../client-runtime/src/state/usage.test.ts | 59 ++++---- packages/client-runtime/src/state/usage.ts | 35 +++-- packages/contracts/src/usage.ts | 2 + 9 files changed, 309 insertions(+), 69 deletions(-) diff --git a/apps/server/src/usage/UsageService.test.ts b/apps/server/src/usage/UsageService.test.ts index 4cd1143ac3787..783da76762d8b 100644 --- a/apps/server/src/usage/UsageService.test.ts +++ b/apps/server/src/usage/UsageService.test.ts @@ -465,6 +465,134 @@ describe("Codex usage history", () => { }); describe("UsageService", () => { + it.effect("keeps pricing metadata paired with the rates used during a concurrent refresh", () => + Effect.gen(function* () { + const { home, settings } = yield* setup; + const sessions = NodePath.join(home, "codex", "sessions"); + yield* Effect.promise(async () => { + await NodeFSP.mkdir(sessions, { recursive: true }); + await NodeFSP.writeFile( + NodePath.join(sessions, "session.jsonl"), + [ + { type: "session_meta", payload: { id: "session", model_provider: "openai" } }, + { type: "turn_context", payload: { model: "gpt-6-astra", turn_id: "turn" } }, + { + type: "event_msg", + timestamp: "2026-08-01T10:00:00Z", + payload: { + type: "token_count", + info: { last_token_usage: { input_tokens: 0, output_tokens: 1 } }, + }, + }, + ] + .map((record) => encodeUnknownJsonString(record)) + .join("\n") + "\n", + ); + }); + let rate = 0; + yield* Effect.gen(function* () { + const resolving = yield* Deferred.make(); + const release = yield* Deferred.make(); + const service = yield* UsageService.make.pipe( + Effect.provideService(CodexUsageHistory.CodexUsageHistory, { + resolve: () => + Deferred.succeed(resolving, undefined).pipe( + Effect.andThen(Deferred.await(release)), + Effect.as(new Map()), + ), + }), + ); + const initialPricing = yield* service.refreshRates; + yield* TestClock.adjust(61_000); + const pending = yield* service.readSummary(WINDOW).pipe(Effect.forkChild); + yield* Deferred.await(resolving); + yield* service.refreshRates; + yield* Deferred.succeed(release, undefined); + const original = yield* Fiber.join(pending); + assert.strictEqual(original.buckets[0]?.costUsd, 1); + assert.strictEqual(original.pricing.fetchedAt, initialPricing.fetchedAt); + const updated = yield* service.readSummary(WINDOW); + assert.strictEqual(updated.buckets[0]?.costUsd, 2); + assert.notStrictEqual(updated.pricing.fetchedAt, initialPricing.fetchedAt); + }).pipe( + Effect.provide( + layerService({ + prefix: "concurrent-rates", + home, + settings, + ratesResponse: () => + Effect.succeed( + Response.json({ + "gpt-6-astra": { input_cost_per_token: 0, output_cost_per_token: ++rate }, + }), + ), + }), + ), + ); + }).pipe(Effect.scoped), + ); + + it.live("shows saved usage before refreshing changed transcripts after restart", () => + Effect.gen(function* () { + const { home, settings, transcript } = yield* setup; + let ratesFetches = 0; + const ratesDocument = { + "claude-fable-5": { input_cost_per_token: 0.000003, output_cost_per_token: 0.000015 }, + }; + yield* Effect.promise(() => NodeFSP.writeFile(transcript, claudeLine(1, 5))); + yield* Effect.gen(function* () { + const first = yield* UsageService.make; + const cold = yield* first.readSummary({ ...WINDOW, preferCached: true }); + assert.strictEqual(totalOutputTokens(cold), 0); + assert.strictEqual( + cold.sources.some((source) => source.refreshing), + true, + ); + assert.strictEqual(totalOutputTokens(yield* first.readSummary(WINDOW)), 5); + yield* first.awaitPersisted; + yield* Effect.promise(() => NodeFSP.appendFile(transcript, claudeLine(2, 7))); + const config = yield* ServerConfig.ServerConfig; + yield* Effect.promise(() => + NodeFSP.writeFile( + NodePath.join(config.stateDir, "usage-model-rates.json"), + encodeUnknownJsonString({ fetchedAtMs: 0, document: ratesDocument }), + ), + ); + + const restarted = yield* UsageService.make; + const savedInput = { ...WINDOW, preferCached: true }; + const saved = yield* restarted.readSummary(savedInput); + assert.strictEqual(totalOutputTokens(saved), 5); + assert.strictEqual(saved.pricing.status, "cached"); + assert.strictEqual(ratesFetches, 1); + assert.strictEqual( + saved.sources.find((source) => source.fingerprint.provider === "claude")?.refreshing, + true, + ); + const fresh = yield* restarted.readSummary({ ...savedInput, awaitRefresh: true }); + assert.strictEqual(totalOutputTokens(fresh), 12); + assert.strictEqual(fresh.pricing.status, "fresh"); + assert.strictEqual(ratesFetches, 2); + assert.strictEqual( + fresh.sources.some((source) => source.refreshing), + false, + ); + }).pipe( + Effect.provide( + layerService({ + prefix: "saved-first", + home, + settings, + ratesDocument, + onRatesFetch: () => { + ratesFetches += 1; + }, + }), + ), + ); + }).pipe(Effect.scoped), + ); + it.live( "repairs initial Codex prices after unchanged-file retries, persists evidence and corrections, and honors later native tiers", () => diff --git a/apps/server/src/usage/UsageService.ts b/apps/server/src/usage/UsageService.ts index 78dd32139dc13..61d28f044d610 100644 --- a/apps/server/src/usage/UsageService.ts +++ b/apps/server/src/usage/UsageService.ts @@ -207,12 +207,8 @@ export const make = Effect.gen(function* () { Effect.catchCause(() => Effect.succeed(null)), ); - const loadRates = Effect.fn("UsageService.loadRates")(function* (force: boolean) { - const now = yield* Clock.currentTimeMillis; - const maxAgeMs = force ? RATES_REFRESH_FLOOR_MS : RATES_TTL_MS; - if (ratesFetchedAtMs !== null && now - ratesFetchedAtMs < maxAgeMs) return; - - if (ratesFetchedAtMs === null) { + const ensureSavedRates = yield* Effect.cached( + Effect.gen(function* () { const fromDisk = yield* fileSystem.readFileString(ratesCachePath).pipe( Effect.flatMap((raw) => decodeRatesCache(raw)), Effect.catchCause(() => Effect.succeed(null)), @@ -223,10 +219,16 @@ export const make = Effect.gen(function* () { rates = parsed; ratesFetchedAtMs = fromDisk.fetchedAtMs; ratesStatus = "cached"; - if (now - fromDisk.fetchedAtMs < maxAgeMs) return; } } - } + }), + ); + + const loadRates = Effect.fn("UsageService.loadRates")(function* (force: boolean) { + yield* ensureSavedRates; + const now = yield* Clock.currentTimeMillis; + const maxAgeMs = force ? RATES_REFRESH_FLOOR_MS : RATES_TTL_MS; + if (ratesFetchedAtMs !== null && now - ratesFetchedAtMs < maxAgeMs) return; if (ratesLastFailedAtMs !== null && now - ratesLastFailedAtMs < RATES_REFRESH_FLOOR_MS) return; @@ -532,8 +534,25 @@ export const make = Effect.gen(function* () { const scanTranscriptDir = Effect.fn("UsageService.scanTranscriptDir")(function* ( source: TranscriptSource, windowStartMs: number, + preferCached: boolean, ) { const { provider, format, dir, volumeId, fileName } = source; + if (preferCached) { + const files = [...fileCache].flatMap(([filePath, entry]) => + entry.provider === provider && + entry.mtimeMs >= windowStartMs && + isWithinDirectory(filePath, dir) + ? [ + { + path: filePath, + records: [...entry.records, ...entry.tailRecords], + malformedRecords: entry.malformedRecords, + }, + ] + : [], + ); + return { provider, dir, volumeId, files, refreshing: true } satisfies ScannedDir; + } const exists = yield* fileSystem .exists(dir) .pipe(Effect.catchCause(() => Effect.succeed(false))); @@ -589,6 +608,7 @@ export const make = Effect.gen(function* () { settings: ServerSettingsValue, retentionCutoffMs: number, awaitRefresh: boolean, + preferCached: boolean, ) { const dirs = yield* resolveTranscriptDirs(settings, retentionCutoffMs).pipe( Effect.provideContext(readerContext), @@ -623,7 +643,7 @@ export const make = Effect.gen(function* () { // Scan independent sources together, directories sequentially to bound open files; preserve order for first-copy dedupe. const [transcripts, scanDirs] = yield* Effect.all( - [Effect.forEach(dirs, (dir) => scanTranscriptDir(dir, windowStartMs)), scans], + [Effect.forEach(dirs, (dir) => scanTranscriptDir(dir, windowStartMs, preferCached)), scans], { concurrency: "unbounded" }, ); const scanned: readonly ScannedDir[] = [...transcripts, ...scanDirs.flat()]; @@ -680,15 +700,23 @@ export const make = Effect.gen(function* () { (hourlyWindow?.sinceTimeMs ?? DateTime.toEpochMillis(windowStart.value)) - MTIME_SLACK_MS; const retentionCutoffMs = startedAtMs - CACHE_RETENTION_DAYS * 24 * 60 * 60 * 1000; + const preferCached = input.preferCached === true && input.awaitRefresh !== true; const [, scannedDirs] = yield* Effect.all( [ - ensureRates(false), - collectDirs(windowStartMs, settings, retentionCutoffMs, input.awaitRefresh === true), + preferCached ? ensureSavedRates : ensureRates(false), + collectDirs( + windowStartMs, + settings, + retentionCutoffMs, + input.awaitRefresh === true, + preferCached, + ), ], { concurrency: 2 }, ); + const summaryPricing = pricing(); const aggregator = new UsageAggregator({ timeZone: input.timeZone, sinceDay: input.sinceDay, @@ -728,7 +756,7 @@ export const make = Effect.gen(function* () { 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 (preferCached || reader.resolveServiceTiers === undefined) continue; const pending = new Set(); for (const [index, files] of filesByDir.entries()) { if (scannedDirs[index]?.provider !== reader.provider) continue; @@ -857,7 +885,7 @@ export const make = Effect.gen(function* () { untilDay: input.untilDay, buckets: aggregated.buckets, sources, - pricing: pricing(), + pricing: summaryPricing, scanDurationMs: Math.max(0, finishedAtMs - startedAtMs), } satisfies UsageSummary; }); @@ -874,6 +902,7 @@ export const make = Effect.gen(function* () { input.untilTime ?? null, settings, input.awaitRefresh === true, + input.preferCached === true && input.awaitRefresh !== true, ]); const readSummary = Effect.fn("UsageService.readSummary")(function* (input: UsageSummaryInput) { diff --git a/apps/server/src/usage/usageTranscriptReader.ts b/apps/server/src/usage/usageTranscriptReader.ts index 4f131f310f418..00f88e44836d7 100644 --- a/apps/server/src/usage/usageTranscriptReader.ts +++ b/apps/server/src/usage/usageTranscriptReader.ts @@ -210,32 +210,53 @@ export async function readTranscriptRecords( let scanOffset = start; let pendingChunks: Buffer[] = []; let pendingBytes = 0; - let streaming: ReturnType | undefined; - let decoder: NodeStringDecoder.StringDecoder | undefined; + let oversized = false; + let carriesUsage = false; + let markerTail = Buffer.alloc(0); + const markerOverlap = + Math.max(...format.usageMarkers.map((marker) => Buffer.byteLength(marker))) - 1; const selectPath = selectUsageFields(format.selectFields); + const checkMarkers = (segment: Buffer) => { + if (carriesUsage) return; + const boundary = Buffer.concat([markerTail, segment.subarray(0, markerOverlap)]); + carriesUsage = mightCarryUsage(segment) || mightCarryUsage(boundary); + markerTail = Buffer.from( + segment.length >= markerOverlap + ? segment.subarray(segment.length - markerOverlap) + : Buffer.concat([markerTail, segment]).subarray(-markerOverlap), + ); + }; const append = (segment: Buffer) => { - if (!streaming && pendingBytes + segment.length <= streamingThresholdBytes) { + if (!oversized && pendingBytes + segment.length <= streamingThresholdBytes) { if (segment.length > 0) pendingChunks.push(segment); pendingBytes += segment.length; return; } - if (!streaming) { - // Usage has no import-history budget; retain selected fields regardless of surrounding tool payload size. - streaming = createTranscriptJsonReader(() => {}, selectPath, { maxDepth: Infinity }); - decoder = new NodeStringDecoder.StringDecoder("utf8"); - for (const pending of pendingChunks) streaming.write(decoder.write(pending)); + if (!oversized) { + oversized = true; + for (const pending of pendingChunks) checkMarkers(pending); pendingChunks = []; pendingBytes = 0; } - streaming.write(decoder!.write(segment)); + checkMarkers(segment); }; - const finish = (lineState: State, out: UsageRecord[]) => { - if (streaming) { - streaming.write(decoder!.end()); + const finish = async (lineState: State, out: UsageRecord[], end: number) => { + if (oversized && carriesUsage) { + // Replay only usage-bearing oversized lines; tool output can otherwise cost minutes of JSON tokenization. + const streaming = createTranscriptJsonReader(() => {}, selectPath, { maxDepth: Infinity }); + const decoder = new NodeStringDecoder.StringDecoder("utf8"); + const line = handle.createReadStream({ + start: resumeOffset, + end: end - 1, + autoClose: false, + highWaterMark: 1024 * 1024, + }) as AsyncIterable; + for await (const chunk of line) streaming.write(decoder.write(chunk)); + streaming.write(decoder.end()); const projected = streaming.finish(); for (const record of format.parseProjected(projected, lineState)) out.push(record); - } else if (pendingBytes > 0) { + } else if (!oversized && pendingBytes > 0) { const line = pendingChunks.length === 1 ? pendingChunks[0]! @@ -244,13 +265,14 @@ export async function readTranscriptRecords( } pendingChunks = []; pendingBytes = 0; - streaming = undefined; - decoder = undefined; + oversized = false; + carriesUsage = false; + markerTail = Buffer.alloc(0); }; const stream = handle.createReadStream({ start, autoClose: false, - highWaterMark: 256 * 1024, + highWaterMark: 1024 * 1024, }) as AsyncIterable; for await (const chunk of stream) { let lineStart = 0; @@ -260,11 +282,15 @@ export async function readTranscriptRecords( append(chunk.subarray(lineStart)); break; } - if (!streaming && pendingBytes === 0) { + if ( + !oversized && + pendingBytes === 0 && + newlineIndex - lineStart <= streamingThresholdBytes + ) { parseLine(chunk.subarray(lineStart, newlineIndex), state, records); } else { append(chunk.subarray(lineStart, newlineIndex)); - finish(state, records); + await finish(state, records, scanOffset + newlineIndex); } lineStart = newlineIndex + 1; resumeOffset = scanOffset + lineStart; @@ -277,7 +303,7 @@ export async function readTranscriptRecords( const tailRecords: UsageRecord[] = []; const tailState = decodeState === undefined ? state : Option.getOrThrow(decodeState(encodedState)); - finish(tailState, tailRecords); + await finish(tailState, tailRecords, scanOffset); const guardLength = Math.min(GUARD_LENGTH, resumeOffset); let guardHash = 0; diff --git a/apps/server/src/usage/usageTranscriptStreaming.test.ts b/apps/server/src/usage/usageTranscriptStreaming.test.ts index 856ba2ac5b9d7..841155e935c96 100644 --- a/apps/server/src/usage/usageTranscriptStreaming.test.ts +++ b/apps/server/src/usage/usageTranscriptStreaming.test.ts @@ -4,6 +4,7 @@ import * as NodeOS from "node:os"; import * as NodePath from "node:path"; import { afterEach, beforeEach, describe, expect, it } from "@effect/vitest"; +import { vi } from "vite-plus/test"; import { TEST_FORMATS } from "./usageTestFormats.ts"; import { @@ -108,6 +109,34 @@ async function scan( } describe("large usage records", () => { + it("skips JSON projection for oversized lines without usage markers", async () => { + const path = NodePath.join(dir, "unrelated.jsonl"); + await NodeFSP.writeFile( + path, + JSON.stringify({ type: "response_item", payload: { content } }) + + "\n" + + codex.map((line) => JSON.stringify(line)).join("\n") + + "\n", + ); + const parseProjected = vi.fn(TEST_FORMATS.codex.parseProjected); + const result = await readWithDefaultThreshold( + path, + { ...TEST_FORMATS.codex, parseProjected }, + undefined, + { streamingThresholdBytes: 256 * 1024 }, + ); + expect(result?.records.map((record) => record.totals)).toEqual([ + { + uncachedInputTokens: 75, + cachedInputTokens: 20, + cacheCreationTokens: 5, + outputTokens: 99, + reasoningTokens: 10, + }, + ]); + expect(parseProjected).not.toHaveBeenCalled(); + }); + it("streams Muse completion envelopes and resumes a partial appended event", async () => { const entry = { schema_version: 1, diff --git a/packages/client-runtime/src/state/server.ts b/packages/client-runtime/src/state/server.ts index f1c8d3c19e83d..de224b650bd1e 100644 --- a/packages/client-runtime/src/state/server.ts +++ b/packages/client-runtime/src/state/server.ts @@ -1274,7 +1274,7 @@ export function createServerEnvironmentAtoms( staleTimeMs: 60_000, refreshTrigger: ({ environmentId }) => usageScanSettingsAtom(environmentId), execute: (input, emit) => - request(WS_METHODS.serverGetUsageSummary, input).pipe( + request(WS_METHODS.serverGetUsageSummary, { ...input, preferCached: true }).pipe( Effect.filterOrElse( (summary) => !summary.sources.some((source) => source.refreshing === true), (summary) => diff --git a/packages/client-runtime/src/state/serverUsage.test.ts b/packages/client-runtime/src/state/serverUsage.test.ts index 21ad521badce0..1a90e743664f9 100644 --- a/packages/client-runtime/src/state/serverUsage.test.ts +++ b/packages/client-runtime/src/state/serverUsage.test.ts @@ -295,6 +295,10 @@ it.effect("keeps the cached usage summary when the awaited refresh fails", () => ]); expect(AsyncResult.isSuccess(harness.registry.get(summary))).toBe(true); expect(harness.requests()).toBe(2); + expect(harness.inputs()).toEqual([ + { ...INPUT, preferCached: true }, + { ...INPUT, awaitRefresh: true }, + ]); unmount(); }), ), diff --git a/packages/client-runtime/src/state/usage.test.ts b/packages/client-runtime/src/state/usage.test.ts index 29aac2ae98ecf..85caab5e7fc80 100644 --- a/packages/client-runtime/src/state/usage.test.ts +++ b/packages/client-runtime/src/state/usage.test.ts @@ -5,6 +5,7 @@ import { UsageDay, USAGE_CONTRACT_VERSION, type ServerProvider, + type UsagePricing, type UsageSummary, } from "@t3tools/contracts"; import * as Effect from "effect/Effect"; @@ -24,7 +25,12 @@ const input = { untilDay: UsageDay.make("2026-09-05"), timeZone: "UTC", }; -const pricing = { status: "fresh" as const, source: "test", fetchedAt: null, knownModels: 1 }; +const pricing: UsagePricing = { + status: "fresh" as const, + source: "test", + fetchedAt: null, + knownModels: 1, +}; const summary: UsageSummary = { ...input, contractVersion: USAGE_CONTRACT_VERSION, @@ -90,28 +96,31 @@ function harness(ids = ["a"]) { } describe("manual usage refresh", () => { - it.each(["success", "failure"])("waits for the rescan after a pricing %s", async (result) => { - const { - environments: [environment], - refresh, - } = harness(); - const entry = environment!; - let finished = false; - const refreshing = refresh().then(() => { - finished = true; - }); - expect(finished).toBe(false); - entry.rates.resolve( - result === "success" - ? AsyncResult.success(pricing) - : AsyncResult.fail(new Error("Pricing offline")), - ); - await entry.scanStarted.promise; - expect(finished).toBe(false); - entry.scan.resolve(summary); - await refreshing; - expect(finished).toBe(true); - }); + it.each(["success", "failure"])( + "reads usage while pricing is pending and waits for both after a pricing %s", + async (result) => { + const { + environments: [environment], + refresh, + } = harness(); + const entry = environment!; + let finished = false; + const refreshing = refresh().then(() => { + finished = true; + }); + expect(finished).toBe(false); + await entry.scanStarted.promise; + entry.rates.resolve( + result === "success" + ? AsyncResult.success(pricing) + : AsyncResult.fail(new Error("Pricing offline")), + ); + expect(finished).toBe(false); + entry.scan.resolve(summary); + await refreshing; + expect(finished).toBe(true); + }, + ); it("settles when an environment disconnects during the rescan", async () => { const { @@ -166,10 +175,10 @@ describe("manual usage refresh", () => { const unmount = registry.mount(query); expect(reads).toBe(1); const refreshing = refresh(); - entry.rates.resolve(AsyncResult.success(pricing)); + entry.rates.resolve(AsyncResult.success({ ...pricing, fetchedAt: "2026-09-05T13:00:00Z" })); await rescanned.promise; await refreshing; - expect(reads).toBe(2); + expect(reads).toBe(3); unmount(); }); }); diff --git a/packages/client-runtime/src/state/usage.ts b/packages/client-runtime/src/state/usage.ts index 7bc582eb7da39..acc794a37c1fd 100644 --- a/packages/client-runtime/src/state/usage.ts +++ b/packages/client-runtime/src/state/usage.ts @@ -103,21 +103,34 @@ export async function refreshUsage({ const unsubscribe = registry.subscribe(presentation, abortWhenDisconnected); abortWhenDisconnected(); try { - const ratesResult = await runAtomCommand( - registry, - server.refreshUsageRates, - { environmentId, input: {} }, - { reportFailure: false }, - ); + registry.refresh(query); + const [usageResult, ratesResult] = await Promise.all([ + executeAtomQuery(registry, query, { + reportFailure: false, + signal: controller.signal, + }), + runAtomCommand( + registry, + server.refreshUsageRates, + { environmentId, input: {} }, + { reportFailure: false }, + ), + ]); const sessionUnavailable = ratesResult._tag === "Failure" && isEnvironmentRpcUnavailable(squashAtomCommandFailure(ratesResult)); - registry.refresh(query); if (sessionUnavailable || controller.signal.aborted) return; - await executeAtomQuery(registry, query, { - reportFailure: false, - signal: controller.signal, - }); + if ( + usageResult._tag === "Success" && + ratesResult._tag === "Success" && + usageResult.value.pricing.fetchedAt !== ratesResult.value.fetchedAt + ) { + registry.refresh(query); + await executeAtomQuery(registry, query, { + reportFailure: false, + signal: controller.signal, + }); + } } finally { unsubscribe(); } diff --git a/packages/contracts/src/usage.ts b/packages/contracts/src/usage.ts index 8e34c10eb885e..5a83fd5d1bf65 100644 --- a/packages/contracts/src/usage.ts +++ b/packages/contracts/src/usage.ts @@ -116,6 +116,8 @@ export const UsageSummaryInput = Schema.Struct({ resolution: Schema.optional(UsageResolution), sinceTime: Schema.optional(TrimmedNonEmptyString), untilTime: Schema.optional(TrimmedNonEmptyString), + /** Return saved transcript usage first; refreshing sources require an awaitRefresh read. */ + preferCached: Schema.optional(Schema.Boolean), awaitRefresh: Schema.optional(Schema.Boolean), }); export type UsageSummaryInput = typeof UsageSummaryInput.Type; From 7b6307bb51aa975976adc0fe5172c05fa9bf11d2 Mon Sep 17 00:00:00 2001 From: maria-rcks Date: Fri, 30 Oct 2026 05:00:49 +0000 Subject: [PATCH 2/2] perf(usage): keep saved-first reads focused on transcript cache --- apps/server/src/usage/UsageService.ts | 15 +---- .../server/src/usage/usageTranscriptReader.ts | 64 ++++++------------- .../usage/usageTranscriptStreaming.test.ts | 29 --------- .../client-runtime/src/state/usage.test.ts | 59 ++++++++--------- packages/client-runtime/src/state/usage.ts | 35 ++++------ 5 files changed, 56 insertions(+), 146 deletions(-) diff --git a/apps/server/src/usage/UsageService.ts b/apps/server/src/usage/UsageService.ts index e540b7e5b3276..79a000fd66416 100644 --- a/apps/server/src/usage/UsageService.ts +++ b/apps/server/src/usage/UsageService.ts @@ -564,20 +564,7 @@ export const make = Effect.gen(function* () { ) { const { provider, format, dir, volumeId, fileName } = source; if (preferCached) { - const files = [...fileCache].flatMap(([filePath, entry]) => - entry.provider === provider && - entry.mtimeMs >= windowStartMs && - isWithinDirectory(filePath, dir) - ? [ - { - path: filePath, - records: [...entry.records, ...entry.tailRecords], - malformedRecords: entry.malformedRecords, - }, - ] - : [], - ); - return { provider, dir, volumeId, files, refreshing: true } satisfies ScannedDir; + return { provider, dir, volumeId, files: [], refreshing: true } satisfies ScannedDir; } const exists = yield* fileSystem .exists(dir) diff --git a/apps/server/src/usage/usageTranscriptReader.ts b/apps/server/src/usage/usageTranscriptReader.ts index 00f88e44836d7..4f131f310f418 100644 --- a/apps/server/src/usage/usageTranscriptReader.ts +++ b/apps/server/src/usage/usageTranscriptReader.ts @@ -210,53 +210,32 @@ export async function readTranscriptRecords( let scanOffset = start; let pendingChunks: Buffer[] = []; let pendingBytes = 0; - let oversized = false; - let carriesUsage = false; - let markerTail = Buffer.alloc(0); - const markerOverlap = - Math.max(...format.usageMarkers.map((marker) => Buffer.byteLength(marker))) - 1; + let streaming: ReturnType | undefined; + let decoder: NodeStringDecoder.StringDecoder | undefined; const selectPath = selectUsageFields(format.selectFields); - const checkMarkers = (segment: Buffer) => { - if (carriesUsage) return; - const boundary = Buffer.concat([markerTail, segment.subarray(0, markerOverlap)]); - carriesUsage = mightCarryUsage(segment) || mightCarryUsage(boundary); - markerTail = Buffer.from( - segment.length >= markerOverlap - ? segment.subarray(segment.length - markerOverlap) - : Buffer.concat([markerTail, segment]).subarray(-markerOverlap), - ); - }; const append = (segment: Buffer) => { - if (!oversized && pendingBytes + segment.length <= streamingThresholdBytes) { + if (!streaming && pendingBytes + segment.length <= streamingThresholdBytes) { if (segment.length > 0) pendingChunks.push(segment); pendingBytes += segment.length; return; } - if (!oversized) { - oversized = true; - for (const pending of pendingChunks) checkMarkers(pending); + if (!streaming) { + // Usage has no import-history budget; retain selected fields regardless of surrounding tool payload size. + streaming = createTranscriptJsonReader(() => {}, selectPath, { maxDepth: Infinity }); + decoder = new NodeStringDecoder.StringDecoder("utf8"); + for (const pending of pendingChunks) streaming.write(decoder.write(pending)); pendingChunks = []; pendingBytes = 0; } - checkMarkers(segment); + streaming.write(decoder!.write(segment)); }; - const finish = async (lineState: State, out: UsageRecord[], end: number) => { - if (oversized && carriesUsage) { - // Replay only usage-bearing oversized lines; tool output can otherwise cost minutes of JSON tokenization. - const streaming = createTranscriptJsonReader(() => {}, selectPath, { maxDepth: Infinity }); - const decoder = new NodeStringDecoder.StringDecoder("utf8"); - const line = handle.createReadStream({ - start: resumeOffset, - end: end - 1, - autoClose: false, - highWaterMark: 1024 * 1024, - }) as AsyncIterable; - for await (const chunk of line) streaming.write(decoder.write(chunk)); - streaming.write(decoder.end()); + const finish = (lineState: State, out: UsageRecord[]) => { + if (streaming) { + streaming.write(decoder!.end()); const projected = streaming.finish(); for (const record of format.parseProjected(projected, lineState)) out.push(record); - } else if (!oversized && pendingBytes > 0) { + } else if (pendingBytes > 0) { const line = pendingChunks.length === 1 ? pendingChunks[0]! @@ -265,14 +244,13 @@ export async function readTranscriptRecords( } pendingChunks = []; pendingBytes = 0; - oversized = false; - carriesUsage = false; - markerTail = Buffer.alloc(0); + streaming = undefined; + decoder = undefined; }; const stream = handle.createReadStream({ start, autoClose: false, - highWaterMark: 1024 * 1024, + highWaterMark: 256 * 1024, }) as AsyncIterable; for await (const chunk of stream) { let lineStart = 0; @@ -282,15 +260,11 @@ export async function readTranscriptRecords( append(chunk.subarray(lineStart)); break; } - if ( - !oversized && - pendingBytes === 0 && - newlineIndex - lineStart <= streamingThresholdBytes - ) { + if (!streaming && pendingBytes === 0) { parseLine(chunk.subarray(lineStart, newlineIndex), state, records); } else { append(chunk.subarray(lineStart, newlineIndex)); - await finish(state, records, scanOffset + newlineIndex); + finish(state, records); } lineStart = newlineIndex + 1; resumeOffset = scanOffset + lineStart; @@ -303,7 +277,7 @@ export async function readTranscriptRecords( const tailRecords: UsageRecord[] = []; const tailState = decodeState === undefined ? state : Option.getOrThrow(decodeState(encodedState)); - await finish(tailState, tailRecords, scanOffset); + finish(tailState, tailRecords); const guardLength = Math.min(GUARD_LENGTH, resumeOffset); let guardHash = 0; diff --git a/apps/server/src/usage/usageTranscriptStreaming.test.ts b/apps/server/src/usage/usageTranscriptStreaming.test.ts index 841155e935c96..856ba2ac5b9d7 100644 --- a/apps/server/src/usage/usageTranscriptStreaming.test.ts +++ b/apps/server/src/usage/usageTranscriptStreaming.test.ts @@ -4,7 +4,6 @@ import * as NodeOS from "node:os"; import * as NodePath from "node:path"; import { afterEach, beforeEach, describe, expect, it } from "@effect/vitest"; -import { vi } from "vite-plus/test"; import { TEST_FORMATS } from "./usageTestFormats.ts"; import { @@ -109,34 +108,6 @@ async function scan( } describe("large usage records", () => { - it("skips JSON projection for oversized lines without usage markers", async () => { - const path = NodePath.join(dir, "unrelated.jsonl"); - await NodeFSP.writeFile( - path, - JSON.stringify({ type: "response_item", payload: { content } }) + - "\n" + - codex.map((line) => JSON.stringify(line)).join("\n") + - "\n", - ); - const parseProjected = vi.fn(TEST_FORMATS.codex.parseProjected); - const result = await readWithDefaultThreshold( - path, - { ...TEST_FORMATS.codex, parseProjected }, - undefined, - { streamingThresholdBytes: 256 * 1024 }, - ); - expect(result?.records.map((record) => record.totals)).toEqual([ - { - uncachedInputTokens: 75, - cachedInputTokens: 20, - cacheCreationTokens: 5, - outputTokens: 99, - reasoningTokens: 10, - }, - ]); - expect(parseProjected).not.toHaveBeenCalled(); - }); - it("streams Muse completion envelopes and resumes a partial appended event", async () => { const entry = { schema_version: 1, diff --git a/packages/client-runtime/src/state/usage.test.ts b/packages/client-runtime/src/state/usage.test.ts index 85caab5e7fc80..29aac2ae98ecf 100644 --- a/packages/client-runtime/src/state/usage.test.ts +++ b/packages/client-runtime/src/state/usage.test.ts @@ -5,7 +5,6 @@ import { UsageDay, USAGE_CONTRACT_VERSION, type ServerProvider, - type UsagePricing, type UsageSummary, } from "@t3tools/contracts"; import * as Effect from "effect/Effect"; @@ -25,12 +24,7 @@ const input = { untilDay: UsageDay.make("2026-09-05"), timeZone: "UTC", }; -const pricing: UsagePricing = { - status: "fresh" as const, - source: "test", - fetchedAt: null, - knownModels: 1, -}; +const pricing = { status: "fresh" as const, source: "test", fetchedAt: null, knownModels: 1 }; const summary: UsageSummary = { ...input, contractVersion: USAGE_CONTRACT_VERSION, @@ -96,31 +90,28 @@ function harness(ids = ["a"]) { } describe("manual usage refresh", () => { - it.each(["success", "failure"])( - "reads usage while pricing is pending and waits for both after a pricing %s", - async (result) => { - const { - environments: [environment], - refresh, - } = harness(); - const entry = environment!; - let finished = false; - const refreshing = refresh().then(() => { - finished = true; - }); - expect(finished).toBe(false); - await entry.scanStarted.promise; - entry.rates.resolve( - result === "success" - ? AsyncResult.success(pricing) - : AsyncResult.fail(new Error("Pricing offline")), - ); - expect(finished).toBe(false); - entry.scan.resolve(summary); - await refreshing; - expect(finished).toBe(true); - }, - ); + it.each(["success", "failure"])("waits for the rescan after a pricing %s", async (result) => { + const { + environments: [environment], + refresh, + } = harness(); + const entry = environment!; + let finished = false; + const refreshing = refresh().then(() => { + finished = true; + }); + expect(finished).toBe(false); + entry.rates.resolve( + result === "success" + ? AsyncResult.success(pricing) + : AsyncResult.fail(new Error("Pricing offline")), + ); + await entry.scanStarted.promise; + expect(finished).toBe(false); + entry.scan.resolve(summary); + await refreshing; + expect(finished).toBe(true); + }); it("settles when an environment disconnects during the rescan", async () => { const { @@ -175,10 +166,10 @@ describe("manual usage refresh", () => { const unmount = registry.mount(query); expect(reads).toBe(1); const refreshing = refresh(); - entry.rates.resolve(AsyncResult.success({ ...pricing, fetchedAt: "2026-09-05T13:00:00Z" })); + entry.rates.resolve(AsyncResult.success(pricing)); await rescanned.promise; await refreshing; - expect(reads).toBe(3); + expect(reads).toBe(2); unmount(); }); }); diff --git a/packages/client-runtime/src/state/usage.ts b/packages/client-runtime/src/state/usage.ts index acc794a37c1fd..7bc582eb7da39 100644 --- a/packages/client-runtime/src/state/usage.ts +++ b/packages/client-runtime/src/state/usage.ts @@ -103,34 +103,21 @@ export async function refreshUsage({ const unsubscribe = registry.subscribe(presentation, abortWhenDisconnected); abortWhenDisconnected(); try { - registry.refresh(query); - const [usageResult, ratesResult] = await Promise.all([ - executeAtomQuery(registry, query, { - reportFailure: false, - signal: controller.signal, - }), - runAtomCommand( - registry, - server.refreshUsageRates, - { environmentId, input: {} }, - { reportFailure: false }, - ), - ]); + const ratesResult = await runAtomCommand( + registry, + server.refreshUsageRates, + { environmentId, input: {} }, + { reportFailure: false }, + ); const sessionUnavailable = ratesResult._tag === "Failure" && isEnvironmentRpcUnavailable(squashAtomCommandFailure(ratesResult)); + registry.refresh(query); if (sessionUnavailable || controller.signal.aborted) return; - if ( - usageResult._tag === "Success" && - ratesResult._tag === "Success" && - usageResult.value.pricing.fetchedAt !== ratesResult.value.fetchedAt - ) { - registry.refresh(query); - await executeAtomQuery(registry, query, { - reportFailure: false, - signal: controller.signal, - }); - } + await executeAtomQuery(registry, query, { + reportFailure: false, + signal: controller.signal, + }); } finally { unsubscribe(); }