From 7134cc8604c03838b43aedfe6a09f40b0849e8f6 Mon Sep 17 00:00:00 2001 From: SiavashShams Date: Fri, 31 Jul 2026 02:25:48 -0700 Subject: [PATCH 1/2] fix: coalesce overlapping cost refreshes --- sdk/typescript/src/cost.ts | 47 +++++++++++++++--- sdk/typescript/tests-ts/cost.test.ts | 72 ++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 6 deletions(-) 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", { From 9c6a5889387d5a8490dae25a481a8271eabb8727 Mon Sep 17 00:00:00 2001 From: Michael D'Angelo Date: Sat, 1 Aug 2026 06:49:54 -0700 Subject: [PATCH 2/2] fix: coalesce cost polling without changing tracker API --- sdk/typescript/src/cost.ts | 68 +++++++----------- sdk/typescript/tests-ts/cost.test.ts | 102 +++++++++++++++------------ 2 files changed, 79 insertions(+), 91 deletions(-) diff --git a/sdk/typescript/src/cost.ts b/sdk/typescript/src/cost.ts index 2df163e..fdb66e0 100644 --- a/sdk/typescript/src/cost.ts +++ b/sdk/typescript/src/cost.ts @@ -49,8 +49,6 @@ 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], @@ -64,31 +62,40 @@ 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; - #activeRefresh: Promise | null = null; - #queuedRefresh: Promise | null = null; + #pending: Promise = Promise.resolve(); #snapshot: ScanCostSnapshot = { usage: null, cost: null }; #lastCost: number | null = null; - public constructor( - options: ScanCostTrackerOptions, - refreshTask?: ScanCostRefreshTask, - ) { + public constructor(options: ScanCostTrackerOptions) { this.#options = options; - this.#refreshTask = refreshTask ?? (() => this.#readSessions()); } public start(threadId: string): void { if (this.#threadId !== null) return; this.#threadId = threadId; if (this.#options.maxCostUsd === undefined) return; + let polling = false; + let rerun = false; const poll = () => { - void this.refresh().catch((error: unknown) => { - this.#options.onError?.(error); - }); + if (polling) { + rerun = true; + return; + } + polling = true; + void this.refresh() + .catch((error: unknown) => { + this.#options.onError?.(error); + }) + .finally(() => { + polling = false; + if (rerun && this.#timer !== null) { + rerun = false; + poll(); + } + }); }; this.#timer = setInterval(poll, COST_POLL_INTERVAL_MS); this.#timer.unref(); @@ -96,12 +103,10 @@ export class ScanCostTracker { } public async refresh(): Promise { - const update = - this.#activeRefresh === null - ? this.#startRefresh() - : (this.#queuedRefresh ??= this.#activeRefresh - .catch(() => {}) - .then(this.#refreshTask)); + const update = this.#pending.then(async () => { + await this.#readSessions(); + }); + this.#pending = update.catch(() => {}); await update; return this.#snapshot; } @@ -119,31 +124,6 @@ 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 a1bad53..107adb7 100644 --- a/sdk/typescript/tests-ts/cost.test.ts +++ b/sdk/typescript/tests-ts/cost.test.ts @@ -16,7 +16,7 @@ const temporaryDirectories: string[] = []; async function waitFor(check: () => boolean): Promise { for (let attempt = 0; attempt < 100; attempt += 1) { if (check()) return; - await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 5)); } throw new Error("Timed out waiting for the cost tracker."); } @@ -150,68 +150,76 @@ 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({ + test("coalesces overlapping polling ticks and bounds final work", async () => { + const home = await codexHome(); + await writeSession(home, "scan-thread", { input_tokens: 100, output_tokens: 10, }); - await Promise.resolve(); + const releases: Array<() => void> = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-sol", + maxCostUsd: 1, + }); + const refresh = tracker.refresh.bind(tracker); + tracker.refresh = async () => { + await new Promise((resolve) => releases.push(resolve)); + return refresh(); + }; + tracker.start("scan-thread"); + + await new Promise((resolve) => setTimeout(resolve, 350)); expect(releases).toHaveLength(1); + const stopped = tracker.stop(); + expect(releases).toHaveLength(2); 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); + expect(releases).toHaveLength(2); }); - test("recovers after a failed refresh", async () => { + test("retries one coalesced poll after a failed refresh", async () => { + const home = await codexHome(); + await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const errors: string[] = []; 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"); + let release: (() => void) | undefined; + const blocked = new Promise((resolve) => { + release = resolve; + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-sol", + maxCostUsd: 1, + onError: (error) => { + if (error instanceof Error) errors.push(error.message); }, - ); + }); + const refresh = tracker.refresh.bind(tracker); + tracker.refresh = async () => { + traversals += 1; + if (traversals === 1) { + await blocked; + throw new Error("session read failed"); + } + return refresh(); + }; tracker.start("scan-thread"); - await expect(tracker.refresh()).rejects.toThrow("session read failed"); - await expect(tracker.refresh()).resolves.toEqual({ - usage: null, - cost: null, - }); + await new Promise((resolve) => setTimeout(resolve, 250)); + expect(traversals).toBe(1); + release!(); + await waitFor(() => traversals === 2); + + expect(errors).toEqual(["session read failed"]); expect(traversals).toBe(2); + expect((await tracker.stop()).cost?.inputTokens).toBe(100); }); test("counts the scan and delegated workers without including other scans", async () => {