fix: preserve workspace-specific usage quotas + harden test isolation (#161)

Dedupe usage accounts by workspace identity (accountId+organizationId) first, refresh token fallback; keep freshest occurrence; skip disabled/identity-less. Harden resolveCodexUsageActiveAccount against sparse/disabled slots. Active-marker match by workspace key. Fix test isolation: debounced-save tests flush+dispose own managers before afterAll.

Author: hawkff

lib/codex-usage.ts

@@ -490,23 +490,92 @@ export async function ensureCodexUsageAccessToken(params: {
 	return { accessToken, refreshed: true, persisted };
 }
 
+/**
+ * Normalize an account identity field to a trimmed string.
+ *
+ * Non-string values collapse to an empty string so callers can treat
+ * "missing" and "blank" identity parts uniformly when building dedupe keys.
+ */
+function normalizeUsageIdentityPart(value: string | undefined): string {
+	return typeof value === "string" ? value.trim() : "";
+}
+
+/**
+ * Derive a stable usage-quota dedupe key for an account.
+ *
+ * A single OAuth credential can authorize multiple ChatGPT workspaces, each
+ * exposing distinct `accountId` / `organizationId` values while sharing one
+ * refresh token. Quota deduplication must therefore key on workspace identity
+ * first so separate workspaces are not collapsed into one row, falling back to
+ * the refresh token only when no workspace identity is available.
+ *
+ * Keys are emitted as `JSON.stringify` arrays (tagged `"workspace"` or
+ * `"refresh"`) so values containing delimiter characters cannot collide.
+ *
+ * @param account - Stored account metadata to derive the key from.
+ * @returns A unique identity key, or `undefined` when the account carries no
+ *   workspace identity and no refresh token.
+ */
+export function getUsageAccountDedupeKey(
+	account: AccountMetadataV3,
+): string | undefined {
+	const accountId = normalizeUsageIdentityPart(account.accountId);
+	const organizationId = normalizeUsageIdentityPart(account.organizationId);
+	if (accountId || organizationId) {
+		return JSON.stringify(["workspace", accountId, organizationId]);
+	}
+
+	const refreshToken = normalizeUsageIdentityPart(account.refreshToken);
+	return refreshToken ? JSON.stringify(["refresh", refreshToken]) : undefined;
+}
+
+/**
+ * Collect the indices of accounts that represent distinct usage quotas.
+ *
+ * Disabled accounts are skipped. Accounts with no usable identity — no
+ * `accountId`, no `organizationId`, and no `refreshToken`, i.e. those for which
+ * {@link getUsageAccountDedupeKey} returns `undefined` — are also dropped, since
+ * they cannot be attributed to a quota and have no token to query. Entries
+ * sharing the same dedupe key are collapsed to a single index.
+ *
+ * When a workspace key appears more than once (e.g. an account re-added after a
+ * token re-issue), the *last* (most recently added) occurrence is kept so the
+ * freshest credential is queried — keeping the first occurrence could surface
+ * an invalidated refresh token after re-auth. First-appearance order is still
+ * used for display stability.
+ *
+ * @param storage - The account storage to scan.
+ * @returns Storage indices of unique, enabled, identifiable usage accounts in
+ *   first-appearance order, each pointing at its freshest occurrence.
+ */
 export function deduplicateUsageAccountIndices(storage: AccountStorageV3): number[] {
-	const seenTokens = new Set<string>();
-	const uniqueIndices: number[] = [];
+	const indexByIdentity = new Map<string, number>();
 	for (let i = 0; i < storage.accounts.length; i += 1) {
 		const account = storage.accounts[i];
 		if (!account) continue;
-		const refreshToken =
-			typeof account.refreshToken === "string"
-				? account.refreshToken.trim()
-				: "";
-		if (refreshToken && seenTokens.has(refreshToken)) continue;
-		if (refreshToken) seenTokens.add(refreshToken);
-		uniqueIndices.push(i);
+		if (account.enabled === false) continue;
+		const key = getUsageAccountDedupeKey(account);
+		if (!key) continue;
+		// Map keeps first-insertion key order (stable display) while overwriting
+		// the value so the latest occurrence's index wins (freshest credential).
+		indexByIdentity.set(key, i);
 	}
-	return uniqueIndices;
+	return [...indexByIdentity.values()];
 }
 
+/**
+ * Resolve which account's usage quota should be shown as active.
+ *
+ * Starts from the persisted active index (preferring the Codex family index)
+ * and then prefers the most-recently-used enabled account by `lastUsed`, so the
+ * displayed quota tracks the credential actually serving requests. Disabled
+ * accounts are ignored, and invalid/missing `lastUsed` values are treated as
+ * oldest.
+ *
+ * @param storage - The account storage to inspect.
+ * @returns The selected account and its index, or `null` when no enabled
+ *   account is available.
+ */
 export function resolveCodexUsageActiveAccount(
 	storage: AccountStorageV3,
 ): UsageAccountSelection | null {
@@ -519,25 +588,41 @@ export function resolveCodexUsageActiveAccount(
 		Math.min(storage.accounts.length - 1, Math.trunc(numericIndex)),
 	);
 	const activeAccount = storage.accounts[index];
-	if (!activeAccount) return null;
+	if (
+		!activeAccount &&
+		storage.accounts.every((account) => !account || account.enabled === false)
+	) {
+		return null;
+	}
 
+	// An enabled active account with a missing/invalid `lastUsed` must fall back
+	// to 0 (same as every other enabled account), not -1. Using -1 would let a
+	// lower-index enabled account with `lastUsed` 0 win the `0 > -1` comparison
+	// and steal the active marker before the active account's own iteration.
+	const activeEnabled = !!activeAccount && activeAccount.enabled !== false;
 	const activeLastUsed =
-		typeof activeAccount.lastUsed === "number" && Number.isFinite(activeAccount.lastUsed)
+		activeEnabled &&
+		typeof activeAccount?.lastUsed === "number" &&
+		Number.isFinite(activeAccount.lastUsed)
 			? activeAccount.lastUsed
-			: 0;
-	let newestIndex = index;
+			: activeEnabled
+				? 0
+				: -1;
+	let newestIndex = activeEnabled ? index : -1;
 	let newestLastUsed = activeLastUsed;
 	for (let i = 0; i < storage.accounts.length; i += 1) {
 		const account = storage.accounts[i];
+		if (!account || account.enabled === false) continue;
 		const lastUsed =
-			typeof account?.lastUsed === "number" && Number.isFinite(account.lastUsed)
+			typeof account.lastUsed === "number" && Number.isFinite(account.lastUsed)
 				? account.lastUsed
 				: 0;
-		if (account && lastUsed > newestLastUsed) {
+		if (lastUsed > newestLastUsed) {
 			newestIndex = i;
 			newestLastUsed = lastUsed;
 		}
 	}
+	if (newestIndex < 0) return null;
 
 	const account = storage.accounts[newestIndex];
 	return account ? { index: newestIndex, account } : null;

lib/tools/codex-limits.ts

@@ -11,6 +11,7 @@ import {
 	fetchCodexUsage,
 	formatUsageLimitSummary,
 	formatUsageLimitTitle,
+	getUsageAccountDedupeKey,
 	hasUsageWindow,
 	parseCodexUsagePayload,
 	resolveCodexUsageAccountId,
@@ -26,6 +27,18 @@ import {
 import { normalizeToolOutputFormat, renderJsonOutput } from "../runtime.js";
 import type { ToolContext } from "./index.js";
 
+/**
+ * Build the `codex-limits` tool.
+ *
+ * The tool fetches and renders Codex usage quotas for each unique account.
+ * Accounts are deduplicated by workspace identity via
+ * {@link getUsageAccountDedupeKey} so distinct workspaces are shown separately,
+ * while the active-account marker is only applied when an account matches both
+ * the active refresh token and the active workspace dedupe key.
+ *
+ * @param ctx - Shared {@link ToolContext} providing UI runtime and account helpers.
+ * @returns The `codex-limits` {@link ToolDefinition}.
+ */
 export function createCodexLimitsTool(ctx: ToolContext): ToolDefinition {
 	const {
 		resolveUiRuntime,
@@ -97,14 +110,45 @@ export function createCodexLimitsTool(ctx: ToolContext): ToolDefinition {
 				activeIndex < storage.accounts.length
 					? storage.accounts[activeIndex]?.refreshToken?.trim() || undefined
 					: undefined;
+			const activeAccount =
+				typeof activeIndex === "number" &&
+				activeIndex >= 0 &&
+				activeIndex < storage.accounts.length
+					? storage.accounts[activeIndex]
+					: undefined;
+			const activeUsageKey = activeAccount
+				? getUsageAccountDedupeKey(activeAccount)
+				: undefined;
+			// If the active account was deduplicated out of uniqueIndices (e.g. it
+			// is a later occurrence of a workspace whose first occurrence carries a
+			// re-issued refresh token), warn so the missing `[active]` marker is
+			// diagnosable. The key-based match below still recovers the marker.
+			if (
+				typeof activeIndex === "number" &&
+				activeIndex >= 0 &&
+				activeIndex < storage.accounts.length &&
+				!uniqueIndices.includes(activeIndex)
+			) {
+				logWarn(
+					`[${PLUGIN_NAME}] active account index ${activeIndex} was deduplicated out of the usage list; matching the active workspace by identity instead.`,
+				);
+			}
 			let storageChanged = false;
 			const jsonAccounts: Array<Record<string, unknown>> = [];
 
 			for (const i of uniqueIndices) {
 				const account = storage.accounts[i];
 				if (!account) continue;
-				const sharesActiveCredential =
-					!!activeRefreshToken && account.refreshToken === activeRefreshToken;
+				const accountUsageKey = getUsageAccountDedupeKey(account);
+				// Match the active account by workspace identity first: two entries
+				// for the same workspace can carry different refresh tokens (e.g. a
+				// re-issued token after re-add), so an exact token match alone would
+				// drop the `[active]` marker. Fall back to refresh-token equality for
+				// accounts that have no workspace identity (token-only dedupe key).
+				const sharesActiveCredential = activeUsageKey
+					? accountUsageKey === activeUsageKey
+					: !!activeRefreshToken &&
+						account.refreshToken === activeRefreshToken;
 				const displayIndex =
 					sharesActiveCredential && typeof activeIndex === "number"
 						? activeIndex

test/codex-usage.test.ts

@@ -87,13 +87,111 @@ describe("codex usage helpers", () => {
 			],
 		};
 
-		expect(deduplicateUsageAccountIndices(storage)).toEqual([0, 2]);
+		expect(deduplicateUsageAccountIndices(storage)).toEqual([1, 2]);
 		expect(resolveCodexUsageActiveAccount(storage)).toMatchObject({
 			index: 2,
 			account: { accountId: "acc-2" },
 		});
 	});
 
+	it("keeps same-token workspace entries distinct, skips disabled, and prefers the freshest duplicate", () => {
+		const storage: AccountStorageV3 = {
+			version: 3,
+			activeIndex: 0,
+			accounts: [
+				{ refreshToken: "r1", accountId: "acc-1", organizationId: "org-1", addedAt: 0, lastUsed: 0 },
+				{ refreshToken: "r1", accountId: "acc-2", organizationId: "org-2", addedAt: 0, lastUsed: 0 },
+				{ refreshToken: "r2", accountId: "acc-3", enabled: false, addedAt: 0, lastUsed: 50 },
+				{ refreshToken: "r3", accountId: "acc-1", organizationId: "org-1", addedAt: 0, lastUsed: 0 },
+			],
+		};
+
+		// org-1 (key W) appears at index 0 and again at index 3 (re-added with a
+		// fresh token r3); org-2 (key X) at index 1; index 2 disabled. Display
+		// order follows first appearance (W then X), but W resolves to its
+		// freshest occurrence (index 3, token r3), not the stale index 0.
+		expect(deduplicateUsageAccountIndices(storage)).toEqual([3, 1]);
+		expect(resolveCodexUsageActiveAccount(storage)).toMatchObject({
+			index: 0,
+			account: { accountId: "acc-1" },
+		});
+	});
+
+	it("deduplicates workspace identities without delimiter collisions", () => {
+		const storage: AccountStorageV3 = {
+			version: 3,
+			activeIndex: 0,
+			accounts: [
+				{ refreshToken: "r1", accountId: "acc:1", organizationId: "org", addedAt: 0, lastUsed: 0 },
+				{ refreshToken: "r2", accountId: "acc", organizationId: "1:org", addedAt: 0, lastUsed: 0 },
+			],
+		};
+
+		expect(deduplicateUsageAccountIndices(storage)).toEqual([0, 1]);
+	});
+
+	it("handles sparse/undefined account slots without throwing", () => {
+		const storage = {
+			version: 3,
+			activeIndex: 5,
+			accounts: [
+				undefined,
+				{ refreshToken: "r1", accountId: "acc-1", organizationId: "org-1", addedAt: 0, lastUsed: 10 },
+			],
+		} as unknown as AccountStorageV3;
+
+		expect(() => resolveCodexUsageActiveAccount(storage)).not.toThrow();
+		expect(resolveCodexUsageActiveAccount(storage)).toMatchObject({
+			index: 1,
+			account: { accountId: "acc-1" },
+		});
+	});
+
+	it("returns null when every account slot is empty or disabled", () => {
+		const storage = {
+			version: 3,
+			activeIndex: 0,
+			accounts: [
+				undefined,
+				{ refreshToken: "r2", accountId: "acc-2", enabled: false, addedAt: 0, lastUsed: 0 },
+			],
+		} as unknown as AccountStorageV3;
+
+		expect(resolveCodexUsageActiveAccount(storage)).toBeNull();
+	});
+
+	it("keeps the active account when its lastUsed is missing", () => {
+		const storage = {
+			version: 3,
+			activeIndex: 1,
+			accounts: [
+				{ refreshToken: "r1", accountId: "acc-1", organizationId: "org-1", addedAt: 0, lastUsed: 0 },
+				{ refreshToken: "r2", accountId: "acc-2", organizationId: "org-2", addedAt: 0 },
+			],
+		} as unknown as AccountStorageV3;
+
+		// The active account (index 1) has no lastUsed. It must not lose the
+		// marker to index 0's lastUsed:0 via a 0 > -1 comparison.
+		expect(resolveCodexUsageActiveAccount(storage)).toMatchObject({
+			index: 1,
+			account: { accountId: "acc-2" },
+		});
+	});
+
+	it("drops accounts that have no workspace identity and no refresh token", () => {
+		const storage: AccountStorageV3 = {
+			version: 3,
+			activeIndex: 0,
+			accounts: [
+				{ addedAt: 0, lastUsed: 0 },
+				{ refreshToken: "r1", accountId: "acc-1", organizationId: "org-1", addedAt: 0, lastUsed: 0 },
+			],
+		};
+
+		// The identity-less entry (index 0) yields no dedupe key and is excluded.
+		expect(deduplicateUsageAccountIndices(storage)).toEqual([1]);
+	});
+
 	it("uses the most recently persisted request account for usage display", () => {
 		const storage: AccountStorageV3 = {
 			version: 3,