diff --git a/sdk/typescript/src/cost.ts b/sdk/typescript/src/cost.ts index 321b975..fdb66e0 100644 --- a/sdk/typescript/src/cost.ts +++ b/sdk/typescript/src/cost.ts @@ -77,10 +77,25 @@ export class ScanCostTracker { 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(); diff --git a/sdk/typescript/tests-ts/cost.test.ts b/sdk/typescript/tests-ts/cost.test.ts index 16d6fc5..107adb7 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 new Promise((resolve) => setTimeout(resolve, 5)); + } + throw new Error("Timed out waiting for the cost tracker."); +} + afterEach(async () => { await Promise.all( temporaryDirectories @@ -142,6 +150,78 @@ describe("scan cost", () => { }); describe("live scan cost tracking", () => { + 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, + }); + 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]!(); + releases[1]!(); + + expect((await stopped).cost?.inputTokens).toBe(100); + expect(releases).toHaveLength(2); + }); + + 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; + 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 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 () => { const home = await codexHome(); await writeSession(home, "scan-thread", {