diff --git a/src/index.ts b/src/index.ts index da5e004..2f58473 100644 --- a/src/index.ts +++ b/src/index.ts @@ -59,7 +59,7 @@ export default Plugin.define({ await quotaRegistration?.events.emit("updated", {}); }); quotaRegistration = await context.rpc.register(QuotaRpc, { - status: async ({ sessionID }) => (await quota.status(sessionID)) ?? null, + status: async () => quota.status() ?? null, }); await quota.init(); // Subscribe before the first read: configuration can publish its models while diff --git a/src/quota-rpc.ts b/src/quota-rpc.ts index 5e43e6f..c9b78bd 100644 --- a/src/quota-rpc.ts +++ b/src/quota-rpc.ts @@ -8,14 +8,13 @@ const window = z.object({ used: z.number(), resetsAt: z.number(), remaining: z.number(), - sessionConsumed: z.number(), }); export const QuotaRpc = Rpc.define({ id: "stanp-opencode.quota", methods: { status: { - input: z.object({ sessionID: z.string() }), + input: z.object({}), output: z .object({ plan: z.string(), fetchedAt: z.number(), windows: z.array(window) }) .nullable(), diff --git a/src/quota.ts b/src/quota.ts index 9eeea3b..3420249 100644 --- a/src/quota.ts +++ b/src/quota.ts @@ -54,7 +54,7 @@ export interface QuotaSnapshot { export interface QuotaStatus { plan: string; fetchedAt: number; - windows: Array; + windows: Array; } const storedState = z.object({ @@ -74,14 +74,8 @@ const storedState = z.object({ ), }) .optional(), - sessions: z.record( - z.string(), - z.record( - z.string(), - z.object({ resetsAt: z.number(), lastUsed: z.number(), consumed: z.number() }), - ), - ), }); +const legacyStoredState = storedState.extend({ sessions: z.unknown() }); type StoredState = z.infer; function withoutExcludedState(state: StoredState): StoredState { @@ -89,20 +83,10 @@ function withoutExcludedState(state: StoredState): StoredState { return { ...state, latest: state.latest && windows?.length ? { ...state.latest, windows } : undefined, - sessions: Object.fromEntries( - Object.entries(state.sessions).map(([session, tracked]) => [ - session, - Object.fromEntries( - Object.entries(tracked).filter( - ([id]) => !isExcludedLimit(id.split(":", 1)[0], undefined), - ), - ), - ]), - ), }; } -type QuotaContext = Pick; +type QuotaContext = Pick; type QuotaFetch = (input: string | URL | Request, init?: RequestInit) => Promise; const quotaWindow = z.object({ @@ -239,7 +223,8 @@ export class QuotaTracker { private timer: ReturnType | undefined; private request: Promise | undefined; private readonly requests = new Set>(); - private readonly active = new Map>(); + private readonly migrations = new Set>(); + private readonly active = new Set(); private controller: AbortController | undefined; private generation = 0; private disposed = false; @@ -276,14 +261,22 @@ export class QuotaTracker { : (this.state?.account ?? (await fingerprint(`credential:${credential.refresh}`))); if (generation !== this.generation) return; if (this.state?.account !== account) { - const parsed = storedState.safeParse(await this.context.storage.get(STORAGE_KEY)); + const value = await this.context.storage.get(STORAGE_KEY); + const parsed = storedState.safeParse(value); if (generation !== this.generation) return; const stored = parsed.success ? parsed.data : undefined; if (generation !== this.generation) return; - this.state = - stored?.account === account - ? withoutExcludedState(stored) - : { account, sessions: {} }; + this.state = stored?.account === account ? withoutExcludedState(stored) : { account }; + const legacy = legacyStoredState.safeParse(value); + if (stored?.account === account && legacy.success) { + const migration = this.persist(); + this.migrations.add(migration); + try { + await migration; + } finally { + this.migrations.delete(migration); + } + } } if (generation !== this.generation) return; if (!this.timer) this.timer = setInterval(() => void this.refresh(), this.pollMs); @@ -345,7 +338,7 @@ export class QuotaTracker { if (parsed.success && parsed.data.account_id) { const account = await fingerprint(`account:${parsed.data.account_id}`); if (generation !== this.generation || !this.state) return; - if (account !== this.state.account) this.state = { account, sessions: {} }; + if (account !== this.state.account) this.state = { account }; } const snapshot = normalizeQuota(payload); if (!snapshot) return; @@ -360,102 +353,30 @@ export class QuotaTracker { private apply(snapshot: QuotaSnapshot) { if (!this.state) return; - for (const root of this.active.keys()) { - const tracked = (this.state.sessions[root] ??= {}); - for (const window of snapshot.windows) { - const previous = tracked[window.id]; - if (!previous) { - tracked[window.id] = { - resetsAt: window.resetsAt, - lastUsed: window.used, - consumed: 0, - }; - continue; - } - if (previous.resetsAt !== window.resetsAt) { - previous.resetsAt = window.resetsAt; - previous.lastUsed = window.used; - previous.consumed += window.used; - } else if (window.used >= previous.lastUsed) { - previous.consumed += window.used - previous.lastUsed; - previous.lastUsed = window.used; - } - } - } this.state.latest = snapshot; } - private async root(sessionID: string) { - let current = sessionID; - const seen = new Set(); - while (!seen.has(current)) { - seen.add(current); - const session = await this.context.session.get({ sessionID: current }); - if (!session.parentID) return current; - current = session.parentID; - } - return sessionID; - } - async start(sessionID: string) { + if (this.active.has(sessionID)) return; const generation = this.generation; - try { - const root = await this.root(sessionID); - if (generation !== this.generation) return; - const family = this.active.get(root); - if (family?.has(sessionID)) return; - const snapshot = await this.refresh(); - if (!snapshot || !this.state || generation !== this.generation) return; - if (!family) { - const previous = this.state.sessions[root] ?? {}; - this.state.sessions[root] = Object.fromEntries( - snapshot.windows.map((window) => [ - window.id, - { - resetsAt: window.resetsAt, - lastUsed: window.used, - consumed: previous[window.id]?.consumed ?? 0, - }, - ]), - ); - } - await this.persist(); - if (generation !== this.generation) return; - await this.changed(); - if (generation !== this.generation) return; - const active = family ?? new Set(); - active.add(sessionID); - this.active.set(root, active); - } catch { - return; - } + const snapshot = await this.refresh(); + if (!snapshot || generation !== this.generation) return; + this.active.add(sessionID); } async complete(sessionID: string) { - try { - const root = await this.root(sessionID); - const family = this.active.get(root); - if (!family?.has(sessionID)) return; - await this.refresh(); - family.delete(sessionID); - if (family.size === 0) this.active.delete(root); - await this.persist(); - await this.changed(); - } catch { - return; - } + if (!this.active.delete(sessionID)) return; + await this.refresh(); } - async status(sessionID: string): Promise { + status(): QuotaStatus | undefined { if (!this.state?.latest) return; - const tracked = this.state.sessions[await this.root(sessionID)]; return { plan: this.state.latest.plan, fetchedAt: this.state.latest.fetchedAt, windows: this.state.latest.windows.map((window) => ({ ...window, remaining: Math.max(0, 100 - window.used), - sessionConsumed: tracked?.[window.id]?.consumed ?? 0, })), }; } @@ -465,10 +386,9 @@ export class QuotaTracker { const value: Schema.Json = this.state.latest ? { account: this.state.account, - sessions: this.state.sessions, latest: this.state.latest, } - : { account: this.state.account, sessions: this.state.sessions }; + : { account: this.state.account }; await this.context.storage.set(STORAGE_KEY, value); } @@ -478,6 +398,8 @@ export class QuotaTracker { this.controller?.abort(); this.controller = undefined; this.request = undefined; + this.active.clear(); + await Promise.all(this.migrations); if (this.state) await this.clear(); } await this.refresh(); @@ -489,6 +411,6 @@ export class QuotaTracker { this.controller?.abort(); if (this.timer) clearInterval(this.timer); this.timer = undefined; - await Promise.all(this.requests); + await Promise.all([...this.requests, ...this.migrations]); } } diff --git a/src/tui.tsx b/src/tui.tsx index 344c4f3..479788b 100644 --- a/src/tui.tsx +++ b/src/tui.tsx @@ -33,14 +33,11 @@ function resetTiming(resetsAt: number, now: number) { return `${minutes}m`; } -export function formatQuotaWindow( - window: QuotaWindow & { remaining: number; sessionConsumed: number }, - now = Date.now(), -) { +export function formatQuotaWindow(window: QuotaWindow & { remaining: number }, now = Date.now()) { const label = window.label ? `${window.label} ${duration(window.seconds)}` : duration(window.seconds); - return `${label}: ${percent(window.remaining)}% · ~${percent(window.sessionConsumed)}pp · ${resetTiming(window.resetsAt, now)}`; + return `${label}: ${percent(window.remaining)}% · ${resetTiming(window.resetsAt, now)}`; } export async function loadQuotaStatus(load: () => Promise) { @@ -66,7 +63,7 @@ export default Plugin.define({ const Slot = (props: { sessionID: string }) => { const [status] = createResource( () => `${props.sessionID}:${revision()}`, - () => loadQuotaStatus(() => quota.status({ sessionID: props.sessionID })), + () => loadQuotaStatus(() => quota.status({})), ); return ( diff --git a/tests/quota.test.ts b/tests/quota.test.ts index cd20e7d..ae536d8 100644 --- a/tests/quota.test.ts +++ b/tests/quota.test.ts @@ -173,10 +173,6 @@ function fixture( storageReads: Promise[] = [], ) { const stored = new Map(); - const sessions = new Map([ - ["root", {}], - ["child", { parentID: "root" }], - ]); let calls = 0; let storageCalls = 0; let changes = 0; @@ -198,12 +194,6 @@ function fixture( }, }, } as never, - session: { - async get({ sessionID }: { sessionID: string }) { - // SAFETY: QuotaTracker reads only parentID from this focused fixture. - return sessions.get(sessionID) as never; - }, - } as never, storage: { async get(key: string) { storageCalls++; @@ -242,6 +232,7 @@ function fixture( storageCalls: () => storageCalls, changes: () => changes, headers, + stored, switch(credential: typeof current) { current = credential; }, @@ -256,7 +247,7 @@ async function accountFingerprint(id: string) { return Buffer.from(bytes).toString("hex"); } -test("removes cached Spark and representative reserve windows and consumption while offline", async () => { +test("reads old cached snapshots without session estimates and removes excluded windows", async () => { const cached = { account: await accountFingerprint("account-one"), latest: { @@ -331,22 +322,23 @@ test("removes cached Spark and representative reserve windows and consumption wh ); await result.tracker.init(); - const status = await result.tracker.status("root"); + const status = result.tracker.status(); expect(status?.windows.map((window) => window.id)).toEqual([ "codex:primary", "codex_luna:primary", ]); - expect(status?.windows[0]?.sessionConsumed).toBe(2); - expect(status?.windows[1]?.sessionConsumed).toBe(3); + expect(status?.windows[0]).not.toHaveProperty("sessionConsumed"); + expect(result.stored.get("chatgpt-quota")).not.toHaveProperty("sessions"); + expect(result.calls()).toBe(0); await result.tracker.dispose(); }); -test("activates only for OAuth and establishes an awaited root baseline shared by children", async () => { +test("activates only for OAuth and refreshes quota at request start", async () => { const key = fixture({ type: "key", key: "secret" }, [plus]); await key.tracker.init(); await key.tracker.start("root"); expect(key.calls()).toBe(0); - expect(await key.tracker.status("root")).toBeUndefined(); + expect(key.tracker.status()).toBeUndefined(); await key.tracker.dispose(); const oauth = fixture( @@ -372,13 +364,14 @@ test("activates only for OAuth and establishes an awaited root baseline shared b expect(oauth.calls()).toBe(1); expect(oauth.headers[0]?.get("chatgpt-account-id")).toBe("account-one"); await oauth.tracker.refresh(); - const child = await oauth.tracker.status("child"); - expect(child?.windows[0]?.sessionConsumed).toBe(5); - expect(child?.windows[0]?.remaining).toBe(82.5); + const status = oauth.tracker.status(); + expect(status?.windows[0]?.remaining).toBe(82.5); + expect(status?.windows[0]).not.toHaveProperty("sessionConsumed"); + expect(oauth.stored.get("chatgpt-quota")).not.toHaveProperty("sessions"); await oauth.tracker.dispose(); }); -test("polls, carries consumption across quota resets, tolerates errors, and clears on logout", async () => { +test("polls, tolerates errors, and clears on logout", async () => { const reset = { ...plus, rate_limit: { @@ -403,12 +396,12 @@ test("polls, carries consumption across quota resets, tolerates errors, and clea await result.tracker.init(); await result.tracker.start("root"); await Bun.sleep(35); - const status = await result.tracker.status("root"); + const status = result.tracker.status(); expect(result.calls()).toBeGreaterThanOrEqual(3); - expect(status?.windows[0]?.sessionConsumed).toBe(2); + expect(status?.windows[0]?.remaining).toBe(98); result.logout(); await result.tracker.authChanged(true); - expect(await result.tracker.status("root")).toBeUndefined(); + expect(result.tracker.status()).toBeUndefined(); const calls = result.calls(); await result.tracker.dispose(); await Bun.sleep(20); @@ -424,12 +417,11 @@ test("sidebar text places reset timing beside latest remaining quota", () => { label: null, used: 25, remaining: 75, - sessionConsumed: 1.5, resetsAt: 100_000 + 2 * 86_400_000 + 5 * 3_600_000, }, 100_000, ), - ).toBe("7d: 75% · ~1.5pp · 2d 5h"); + ).toBe("7d: 75% · 2d 5h"); }); test("sidebar reset timing uses compact future units and handles boundaries", () => { @@ -442,17 +434,16 @@ test("sidebar reset timing uses compact future units and handles boundaries", () label: null, used: 25, remaining: 75, - sessionConsumed: 1.5, resetsAt: now + resetIn, }, now, ); - expect(format(2 * 3_600_000 + 15 * 60_000)).toBe("5h: 75% · ~1.5pp · 2h 15m"); - expect(format(3_600_000)).toBe("5h: 75% · ~1.5pp · 1h"); - expect(format(15 * 60_000)).toBe("5h: 75% · ~1.5pp · 15m"); - expect(format(60_000)).toBe("5h: 75% · ~1.5pp · 1m"); - expect(format(59_999)).toBe("5h: 75% · ~1.5pp · <1m"); + expect(format(2 * 3_600_000 + 15 * 60_000)).toBe("5h: 75% · 2h 15m"); + expect(format(3_600_000)).toBe("5h: 75% · 1h"); + expect(format(15 * 60_000)).toBe("5h: 75% · 15m"); + expect(format(60_000)).toBe("5h: 75% · 1m"); + expect(format(59_999)).toBe("5h: 75% · <1m"); }); test("sidebar reset timing preserves labels and distinguishes stale and unavailable timestamps", () => { @@ -462,17 +453,16 @@ test("sidebar reset timing preserves labels and distinguishes stale and unavaila label: "GPT-5.6 Luna", used: 99.95, remaining: 0.05, - sessionConsumed: 1.5, }; expect(formatQuotaWindow({ ...window, resetsAt: 999 }, 1_000)).toBe( - "GPT-5.6 Luna 5h: <0.1% · ~1.5pp · due", + "GPT-5.6 Luna 5h: <0.1% · due", ); expect(formatQuotaWindow({ ...window, resetsAt: Number.NaN }, 1_000)).toBe( - "GPT-5.6 Luna 5h: <0.1% · ~1.5pp · unavailable", + "GPT-5.6 Luna 5h: <0.1% · unavailable", ); }); -test("inactive roots freeze, resume from a cached baseline on failure, and deactivate by family", async () => { +test("request completion refreshes once across terminal and idle events", async () => { const used = (percent: number) => ({ ...plus, rate_limit: { @@ -492,16 +482,10 @@ test("inactive roots freeze, resume from a cached baseline on failure, and deact await result.tracker.init(); await result.tracker.start("root"); await result.tracker.complete("root"); - expect((await result.tracker.status("root"))?.windows[0]?.sessionConsumed).toBe(1); - await result.tracker.refresh(); - expect((await result.tracker.status("root"))?.windows[0]?.sessionConsumed).toBe(1); - - await result.tracker.start("root"); - await result.tracker.start("child"); + expect(result.tracker.status()?.windows[0]?.remaining).toBe(89); await result.tracker.complete("root"); - await result.tracker.refresh(); - await result.tracker.complete("child"); - expect((await result.tracker.status("root"))?.windows[0]?.sessionConsumed).toBe(6); + expect(result.calls()).toBe(2); + expect(result.tracker.status()?.windows[0]?.remaining).toBe(89); await result.tracker.dispose(); }); @@ -542,10 +526,10 @@ test("an auth switch detaches an in-flight old-account response and immediately expect(result.calls()).toBe(2); pending.resolve(plus); await old; - expect((await result.tracker.status("root"))?.windows[0]?.used).toBe(70); + expect(result.tracker.status()?.windows[0]?.used).toBe(70); expect(result.headers[1]?.get("chatgpt-account-id")).toBe("account-two"); await result.tracker.refresh(); - expect((await result.tracker.status("root"))?.windows[0]?.sessionConsumed).toBe(0); + expect(result.tracker.status()?.windows[0]).not.toHaveProperty("sessionConsumed"); await result.tracker.dispose(); }); @@ -577,7 +561,7 @@ test("an auth switch during an old account storage read cannot restore stale sta expect(result.calls()).toBe(1); pending.resolve(undefined); await old; - expect((await result.tracker.status("root"))?.windows[0]?.used).toBe(64); + expect(result.tracker.status()?.windows[0]?.used).toBe(64); expect(result.headers[0]?.get("chatgpt-account-id")).toBe("account-two"); await result.tracker.dispose(); }); diff --git a/tests/tui.test.ts b/tests/tui.test.ts index 10552b8..9d3a26f 100644 --- a/tests/tui.test.ts +++ b/tests/tui.test.ts @@ -18,7 +18,6 @@ const status = (plan: string, remaining: number): QuotaStatus => ({ used: 100 - remaining, resetsAt: 2, remaining, - sessionConsumed: 1.5, }, ], }); @@ -32,13 +31,13 @@ test("built quota slot stays reactive across pending, empty, failure, update, an const { default: plugin } = tui; const pending = Promise.withResolvers(); const responses: Array<() => Promise> = [() => pending.promise]; - const sessions: string[] = []; + let statusCalls = 0; let slotRender: SlotRender | undefined; let notifyUpdated: (() => void) | undefined; const quota = { - status(input: { sessionID: string }) { - sessions.push(input.sessionID); + status() { + statusCalls++; const response = responses.shift(); if (!response) throw new Error("Missing test response"); return response(); @@ -84,7 +83,8 @@ test("built quota slot stays reactive across pending, empty, failure, update, an pending.resolve(status("plus", 75)); await setup.waitForFrame((frame) => frame.includes("ChatGPT quota (plus)")); - expect(setup.captureCharFrame()).toContain("5h: 75% · ~1.5pp · due"); + expect(setup.captureCharFrame()).toContain("5h: 75% · due"); + expect(setup.captureCharFrame()).not.toContain("~"); responses.push(async () => null); notifyUpdated(); @@ -94,13 +94,13 @@ test("built quota slot stays reactive across pending, empty, failure, update, an throw new Error("disconnected"); }); notifyUpdated(); - await setup.waitFor(() => sessions.length === 3); + await setup.waitFor(() => statusCalls === 3); expect(setup.captureCharFrame()).not.toContain("ChatGPT quota"); responses.push(async () => status("pro", 42)); setSessionID("session-two"); await setup.waitForFrame((frame) => frame.includes("ChatGPT quota (pro)")); - expect(sessions).toEqual(["session-one", "session-one", "session-one", "session-two"]); + expect(statusCalls).toBe(4); } finally { setup.renderer.destroy(); if (cleanup) await cleanup();