diff --git a/sdk/typescript/src/cost.ts b/sdk/typescript/src/cost.ts index 321b975..2df163e 100644 --- a/sdk/typescript/src/cost.ts +++ b/sdk/typescript/src/cost.ts @@ -49,6 +49,8 @@ interface ScanCostSnapshot { cost: ScanCost | null; } +type ScanCostRefreshTask = () => Promise; + const MODEL_PRICING_NANODOLLARS: Readonly> = { "gpt-5.6": [5_000, 500, 6_250, 30_000], "gpt-5.6-sol": [5_000, 500, 6_250, 30_000], @@ -62,15 +64,21 @@ const MAX_SESSION_EVENT_BYTES = 1 * 1_024 * 1_024; export class ScanCostTracker { readonly #options: ScanCostTrackerOptions; + readonly #refreshTask: ScanCostRefreshTask; readonly #sessions = new Map(); #threadId: string | null = null; #timer: NodeJS.Timeout | null = null; - #pending: Promise = Promise.resolve(); + #activeRefresh: Promise | null = null; + #queuedRefresh: Promise | null = null; #snapshot: ScanCostSnapshot = { usage: null, cost: null }; #lastCost: number | null = null; - public constructor(options: ScanCostTrackerOptions) { + public constructor( + options: ScanCostTrackerOptions, + refreshTask?: ScanCostRefreshTask, + ) { this.#options = options; + this.#refreshTask = refreshTask ?? (() => this.#readSessions()); } public start(threadId: string): void { @@ -88,10 +96,12 @@ export class ScanCostTracker { } public async refresh(): Promise { - const update = this.#pending.then(async () => { - await this.#readSessions(); - }); - this.#pending = update.catch(() => {}); + const update = + this.#activeRefresh === null + ? this.#startRefresh() + : (this.#queuedRefresh ??= this.#activeRefresh + .catch(() => {}) + .then(this.#refreshTask)); await update; return this.#snapshot; } @@ -109,6 +119,31 @@ export class ScanCostTracker { return this.#snapshot; } + #startRefresh(): Promise { + const update = Promise.resolve().then(this.#refreshTask); + this.#activeRefresh = update; + void update.then( + () => this.#finishRefresh(update), + () => this.#finishRefresh(update), + ); + return update; + } + + #finishRefresh(finished: Promise): void { + if (this.#activeRefresh !== finished) return; + const queued = this.#queuedRefresh; + if (queued === null) { + this.#activeRefresh = null; + return; + } + this.#activeRefresh = queued; + this.#queuedRefresh = null; + void queued.then( + () => this.#finishRefresh(queued), + () => this.#finishRefresh(queued), + ); + } + async #readSessions(): Promise { if (this.#threadId === null) return; const unreadable: Array<{ session: SessionUsage; error: unknown }> = []; diff --git a/sdk/typescript/tests-ts/cost.test.ts b/sdk/typescript/tests-ts/cost.test.ts index 16d6fc5..a1bad53 100644 --- a/sdk/typescript/tests-ts/cost.test.ts +++ b/sdk/typescript/tests-ts/cost.test.ts @@ -13,6 +13,14 @@ import { estimateScanCost, ScanCostTracker } from "../src/cost.js"; const temporaryDirectories: string[] = []; +async function waitFor(check: () => boolean): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (check()) return; + await Promise.resolve(); + } + throw new Error("Timed out waiting for the cost tracker."); +} + afterEach(async () => { await Promise.all( temporaryDirectories @@ -142,6 +150,70 @@ describe("scan cost", () => { }); describe("live scan cost tracking", () => { + test("coalesces overlapping refreshes and bounds final work", async () => { + const releases: Array<() => void> = []; + const tracker = new ScanCostTracker( + { + codexHome: await codexHome(), + model: "gpt-5.6-sol", + }, + () => { + return new Promise((resolve) => { + releases.push(() => resolve()); + }); + }, + ); + tracker.start("scan-thread"); + + const first = tracker.refresh(); + await waitFor(() => releases.length === 1); + const overlapping = [ + tracker.refresh(), + tracker.refresh(), + tracker.refresh(), + ]; + const stopped = tracker.stop({ + input_tokens: 100, + output_tokens: 10, + }); + await Promise.resolve(); + expect(releases).toHaveLength(1); + + releases[0]!(); + await waitFor(() => releases.length >= 2); + releases[1]!(); + for (let attempt = 0; attempt < 10; attempt += 1) { + await Promise.resolve(); + } + expect(releases).toHaveLength(2); + + const snapshots = await Promise.all([first, ...overlapping]); + expect(snapshots.every((snapshot) => snapshot === snapshots[0])).toBe(true); + expect((await stopped).cost?.inputTokens).toBe(100); + }); + + test("recovers after a failed refresh", async () => { + let traversals = 0; + const tracker = new ScanCostTracker( + { + codexHome: await codexHome(), + model: "gpt-5.6-sol", + }, + async () => { + traversals += 1; + if (traversals === 1) throw new Error("session read failed"); + }, + ); + tracker.start("scan-thread"); + + await expect(tracker.refresh()).rejects.toThrow("session read failed"); + await expect(tracker.refresh()).resolves.toEqual({ + usage: null, + cost: null, + }); + expect(traversals).toBe(2); + }); + test("counts the scan and delegated workers without including other scans", async () => { const home = await codexHome(); await writeSession(home, "scan-thread", {