diff --git a/packages/shared/src/usageLimits.test.ts b/packages/shared/src/usageLimits.test.ts index bed48428a93c4..8f7fa26d7e728 100644 --- a/packages/shared/src/usageLimits.test.ts +++ b/packages/shared/src/usageLimits.test.ts @@ -804,6 +804,52 @@ describe("pools", () => { ).toEqual(["weekly", "weekly:secondary"]); }); + it("weights pooled Codex quota and resets by each account's plan capacity", () => { + const [pool] = collectLimitPools( + collectLimitAccounts( + new Map([ + [ + EnvironmentId.make("env-a"), + { + ...laptop, + serverConfig: { + providers: [], + usageLimitSources: [ + { + ...source, + accounts: [ + ["pro", "ChatGPT Pro 200 Subscription", 100], + ["plus", "ChatGPT Plus Subscription", 0], + ].map(([id, plan, usedPercent]) => ({ + id: String(id), + email: `${id}@example.com`, + driver: ProviderDriverKind.make("codex"), + plan: String(plan), + usageLimits: { + checkedAt, + windows: [{ ...weekly, id: "secondary", usedPercent: Number(usedPercent) }], + }, + })), + }, + ], + }, + }, + ], + ]), + ), + now, + ); + const [pooledWeekly] = pool?.windows ?? []; + // Pro 20x holds 20 of the pool's 21 capacity parts. + expect(pooledWeekly?.remainingPercent).toBe(5); + expect( + pooledWeekly?.resets.map(({ member, restoresPercent }) => [ + member.account.email, + restoresPercent, + ]), + ).toContainEqual(["pro@example.com", 95]); + }); + it("pools windows by id across accounts and orders resets by when they land", () => { const input = new Map([ [ diff --git a/packages/shared/src/usageLimits.ts b/packages/shared/src/usageLimits.ts index 716103ac2fc9b..ed38fdeb45b4e 100644 --- a/packages/shared/src/usageLimits.ts +++ b/packages/shared/src/usageLimits.ts @@ -539,6 +539,25 @@ function accountSortName(account: LimitAccount): string { return (account.displayName ?? account.email ?? account.key).toLowerCase(); } +const CODEX_PLAN_CAPACITY: Record = { + "ChatGPT Plus Subscription": 1, + "ChatGPT Pro 100 Subscription": 5, + "ChatGPT Pro 200 Subscription": 20, +}; +// Claude's published plan multipliers describe the session allowance only. +const CLAUDE_SESSION_CAPACITY: Record = { Pro: 1, Max: 5, Ultra: 20 }; + +function planCapacity(account: LimitAccount, kind: ServerProviderUsageWindow["kind"]) { + if (!account.plan) return undefined; + if (account.driver === "codex" && (kind === "session" || kind === "weekly")) { + return CODEX_PLAN_CAPACITY[account.plan]; + } + if (account.driver === "claudeAgent" && kind === "session") { + return CLAUDE_SESSION_CAPACITY[account.plan]; + } + return undefined; +} + // Codex ids name slots, so pool by kind; other drivers' ids identify distinct allowances. function poolKey(driver: ServerProvider["driver"], window: ServerProviderUsageWindow): string { return driver === "codex" ? window.kind : `${window.kind}:${window.id}`; @@ -607,15 +626,28 @@ function poolWindows(accounts: readonly LimitAccount[], now: number): readonly L const memberByAccount = new Map(members.map((member) => [member.account.key, member])); const first = members[0]!.window; const id = members[0]!.account.driver === "codex" ? key : first.id; - const usedPercent = members.reduce((sum, m) => sum + m.window.usedPercent, 0) / members.length; + // One unknown plan makes the pool incomparable by plan, so every member then counts once. + const capacities = members.map((m) => planCapacity(m.account, first.kind)); + const weights = capacities.every((capacity) => capacity !== undefined) + ? capacities + : members.map(() => 1); + const weightOf = new Map(members.map((m, index) => [m, weights[index]!])); + const totalWeight = weights.reduce((sum, weight) => sum + weight, 0); + const usedPercent = + members.reduce((sum, m) => sum + m.window.usedPercent * weightOf.get(m)!, 0) / totalWeight; // Judge pace only for windows with reset clocks; others add spend without elapsed time. const timed = members.flatMap((m) => { const share = elapsedShare(m.window, now); - return share === null ? [] : [{ used: m.window.usedPercent, elapsed: share }]; + return share === null + ? [] + : [{ used: m.window.usedPercent, elapsed: share, weight: weightOf.get(m)! }]; }); - const timedUsed = timed.reduce((sum, t) => sum + t.used, 0) / timed.length; + const timedWeight = timed.reduce((sum, t) => sum + t.weight, 0); + const timedUsed = timed.reduce((sum, t) => sum + t.used * t.weight, 0) / timedWeight; const meanElapsed = - timed.length > 0 ? timed.reduce((sum, t) => sum + t.elapsed, 0) / timed.length : null; + timed.length > 0 + ? timed.reduce((sum, t) => sum + t.elapsed * t.weight, 0) / timedWeight + : null; const resets = members .flatMap((member) => { const at = resetMillis(member.window); @@ -625,7 +657,9 @@ function poolWindows(accounts: readonly LimitAccount[], now: number): readonly L { member, at, - restoresPercent: Math.round(member.refillPercent / members.length), + restoresPercent: Math.round( + (member.refillPercent * weightOf.get(member)!) / totalWeight, + ), }, ]; })