Coalesce overlapping cost refreshes

sdk/typescript/src/cost.ts

@@ -49,6 +49,8 @@ interface ScanCostSnapshot {
   cost: ScanCost | null;
 }
 
+type ScanCostRefreshTask = () => Promise<void>;
+
 const MODEL_PRICING_NANODOLLARS: Readonly<Record<string, ModelPricing>> = {
   "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<string, SessionUsage>();
   #threadId: string | null = null;
   #timer: NodeJS.Timeout | null = null;
-  #pending: Promise<void> = Promise.resolve();
+  #activeRefresh: Promise<void> | null = null;
+  #queuedRefresh: Promise<void> | 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<ScanCostSnapshot> {
-    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<void> {
+    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>): 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<void> {
     if (this.#threadId === null) return;
     const unreadable: Array<{ session: SessionUsage; error: unknown }> = [];

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<void> {
+  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<void>((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", {