fix workspace identity follow-ups in quota cleanup

index.ts

@@ -119,6 +119,7 @@ import {
 	saveFlaggedAccounts,
 	withFlaggedAccountStorageTransaction,
 	clearFlaggedAccounts,
+	getWorkspaceIdentityKey,
 	StorageError,
 	formatStorageErrorHint,
 	type AccountStorageV3,
@@ -185,23 +186,6 @@ import {
 	getRecoveryToastContent,
 } from "./lib/recovery.js";
 
-function getWorkspaceIdentityKey(account: {
-	organizationId?: string;
-	accountId?: string;
-	refreshToken: string;
-}): string {
-	const organizationId = account.organizationId?.trim();
-	const accountId = account.accountId?.trim();
-	const refreshToken = account.refreshToken.trim();
-	if (organizationId) {
-		return accountId
-			? `organizationId:${organizationId}|accountId:${accountId}`
-			: `organizationId:${organizationId}`;
-	}
-	if (accountId) return `accountId:${accountId}`;
-	return `refreshToken:${refreshToken}`;
-}
-
 function matchesWorkspaceIdentity(
 	account: {
 		organizationId?: string;

lib/storage.ts

@@ -34,6 +34,28 @@ export interface FlaggedAccountStorageV1 {
 	accounts: FlaggedAccountMetadataV1[];
 }
 
+const normalizeWorkspaceIdentityPart = (value: unknown): string | undefined =>
+	typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
+
+export function getWorkspaceIdentityKey(account: {
+	organizationId?: string;
+	accountId?: string;
+	refreshToken: string;
+}): string {
+	const organizationId = normalizeWorkspaceIdentityPart(account.organizationId);
+	const accountId = normalizeWorkspaceIdentityPart(account.accountId);
+	const refreshToken = normalizeWorkspaceIdentityPart(account.refreshToken) ?? "";
+	if (organizationId) {
+		return accountId
+			? `organizationId:${organizationId}|accountId:${accountId}`
+			: `organizationId:${organizationId}`;
+	}
+	if (accountId) {
+		return `accountId:${accountId}`;
+	}
+	return `refreshToken:${refreshToken}`;
+}
+
 export type ImportBackupMode = "none" | "best-effort" | "required";
 
 export interface ImportAccountsOptions {
@@ -921,30 +943,6 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 {
 		return { version: 1, accounts: [] };
 	}
 
-	const normalizeFlaggedIdentityPart = (value: unknown): string | undefined =>
-		typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
-	// Flagged storage must keep sibling workspace entries separate when they share an
-	// organization but have different accountIds, so this key is more specific than
-	// the normal account identity collapse used in active storage.
-	const getFlaggedIdentityKey = (account: {
-		organizationId?: string;
-		accountId?: string;
-		refreshToken: string;
-	}): string => {
-		const organizationId = normalizeFlaggedIdentityPart(account.organizationId);
-		const accountId = normalizeFlaggedIdentityPart(account.accountId);
-		const refreshToken = normalizeFlaggedIdentityPart(account.refreshToken) ?? "";
-		if (organizationId) {
-			return accountId
-				? `organizationId:${organizationId}|accountId:${accountId}`
-				: `organizationId:${organizationId}`;
-		}
-		if (accountId) {
-			return `accountId:${accountId}`;
-		}
-		return `refreshToken:${refreshToken}`;
-	};
-
 	const byIdentityKey = new Map<string, FlaggedAccountMetadataV1>();
 	for (const rawAccount of data.accounts) {
 		if (!isRecord(rawAccount)) continue;
@@ -1024,7 +1022,9 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 {
 			flaggedReason: typeof rawAccount.flaggedReason === "string" ? rawAccount.flaggedReason : undefined,
 			lastError: typeof rawAccount.lastError === "string" ? rawAccount.lastError : undefined,
 		};
-		byIdentityKey.set(getFlaggedIdentityKey(normalized), normalized);
+		// Keep flagged dedup aligned with active cleanup so sibling workspaces only
+		// collapse when they resolve to the same shared workspace identity.
+		byIdentityKey.set(getWorkspaceIdentityKey(normalized), normalized);
 	}
 
 	return {

test/index.test.ts

@@ -272,7 +272,10 @@ const mockSaveFlaggedAccounts = vi.fn(async (nextStorage: typeof mockFlaggedStor
 	await persistMockFlaggedStorage(nextStorage);
 });
 
-vi.mock("../lib/storage.js", () => ({
+vi.mock("../lib/storage.js", async () => {
+	const actual = await vi.importActual<typeof import("../lib/storage.js")>("../lib/storage.js");
+
+	return {
 	getStoragePath: () => "/mock/path/accounts.json",
 	loadAccounts: vi.fn(async () => cloneMockStorage()),
 	saveAccounts: vi.fn(async (nextStorage: typeof mockStorage) => {
@@ -311,6 +314,7 @@ vi.mock("../lib/storage.js", () => ({
 	previewImportAccounts: vi.fn(async () => ({ imported: 2, skipped: 1, total: 5 })),
 	createTimestampedBackupPath: vi.fn((prefix?: string) => `/tmp/${prefix ?? "codex-backup"}-20260101-000000.json`),
 	loadFlaggedAccounts: vi.fn(async () => cloneMockFlaggedStorage()),
+	getWorkspaceIdentityKey: vi.fn(actual.getWorkspaceIdentityKey),
 	saveFlaggedAccounts: mockSaveFlaggedAccounts,
 	withFlaggedAccountStorageTransaction: vi.fn(
 		async <T>(
@@ -335,7 +339,8 @@ vi.mock("../lib/storage.js", () => ({
 		}
 	},
 	formatStorageErrorHint: () => "Check file permissions",
-}));
+	};
+});
 
 vi.mock("../lib/accounts.js", () => {
 	class MockAccountManager {
@@ -438,12 +443,15 @@ vi.mock("../lib/accounts.js", () => {
 		),
 		extractAccountEmail: vi.fn(() => "user@example.com"),
 		extractAccountId: vi.fn(() => "account-1"),
-		resolveRequestAccountId: (_storedId: string | undefined, _source: string | undefined, tokenId: string | undefined) => tokenId,
+		resolveRequestAccountId: vi.fn(
+			(_storedId: string | undefined, _source: string | undefined, tokenId: string | undefined) =>
+				tokenId,
+		),
 		formatAccountLabel: (_account: unknown, index: number) => `Account ${index + 1}`,
 		formatCooldown: () => null,
 		formatWaitTime: (ms: number) => `${Math.round(ms / 1000)}s`,
 		sanitizeEmail: (email: string) => email,
-		shouldUpdateAccountIdFromToken: () => true,
+		shouldUpdateAccountIdFromToken: vi.fn(() => true),
 		parseRateLimitReason: () => "unknown",
 		lookupCodexCliTokensByEmail: vi.fn(async () => null),
 	};
@@ -2639,18 +2647,16 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
 			expect(removeAccount).toHaveBeenCalledTimes(1);
 			expect(removeAccountsWithSameRefreshToken).not.toHaveBeenCalled();
 			expect(accounts.map((account) => account.accountId)).toEqual(["org-live"]);
-			expect(vi.mocked(storageModule.saveFlaggedAccounts)).toHaveBeenCalledTimes(1);
-			expect(vi.mocked(storageModule.saveFlaggedAccounts).mock.calls[0]?.[0]).toEqual(
-				expect.objectContaining({
-					accounts: expect.arrayContaining([
-						expect.objectContaining({
-							accountId: "org-dead",
-							organizationId: "org-dead",
-							flaggedReason: "workspace-deactivated",
-							lastError: "deactivated_workspace",
-						}),
-					]),
-				}),
+			expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalledTimes(1);
+			expect(mockFlaggedStorage.accounts).toEqual(
+				expect.arrayContaining([
+					expect.objectContaining({
+						accountId: "org-dead",
+						organizationId: "org-dead",
+						flaggedReason: "workspace-deactivated",
+						lastError: "deactivated_workspace",
+					}),
+				]),
 			);
 		});
 
@@ -3618,16 +3624,137 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
 		expect(mockStorage.accounts.some((account) => account.accountId === "workspace-dead")).toBe(false);
 		expect(mockStorage.accounts[0]?.organizationId).toBe("org-shared");
 		expect(mockStorage.accounts[0]?.refreshToken).toBe("shared-refresh");
-		expect(vi.mocked(storageModule.saveFlaggedAccounts)).toHaveBeenCalledWith(
-			expect.objectContaining({
-				accounts: expect.arrayContaining([
-					expect.objectContaining({
-						accountId: "workspace-dead",
-						organizationId: "org-shared",
-						flaggedReason: "token-invalid",
-					}),
-				]),
-			}),
+		expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalled();
+		expect(mockFlaggedStorage.accounts).toEqual(
+			expect.arrayContaining([
+				expect.objectContaining({
+					accountId: "workspace-dead",
+					organizationId: "org-shared",
+					flaggedReason: "token-invalid",
+				}),
+			]),
+		);
+	});
+
+	it("removes only the deactivated org-scoped workspace during quota-check cleanup", async () => {
+		const accountsModule = await import("../lib/accounts.js");
+		const cliModule = await import("../lib/cli.js");
+		const refreshQueueModule = await import("../lib/refresh-queue.js");
+		const fetchHelpersModule = await import("../lib/request/fetch-helpers.js");
+		const storageModule = await import("../lib/storage.js");
+
+		mockStorage.accounts = [
+			{
+				refreshToken: "shared-refresh",
+				organizationId: "org-shared",
+				accountId: "workspace-dead",
+				accountIdSource: "manual",
+				email: "dead@example.com",
+				addedAt: 1,
+				lastUsed: 1,
+			},
+			{
+				refreshToken: "shared-refresh",
+				organizationId: "org-shared",
+				accountId: "workspace-live",
+				accountIdSource: "manual",
+				email: "live@example.com",
+				addedAt: 2,
+				lastUsed: 2,
+			},
+		];
+
+		vi.mocked(cliModule.promptLoginMode)
+			.mockResolvedValueOnce({ mode: "check" })
+			.mockResolvedValueOnce({ mode: "cancel" });
+		vi.mocked(accountsModule.shouldUpdateAccountIdFromToken).mockImplementation(
+			(source: string | undefined, currentAccountId?: string) => {
+				if (!currentAccountId) return true;
+				if (!source) return true;
+				return source === "token" || source === "id_token";
+			},
+		);
+		vi.mocked(accountsModule.resolveRequestAccountId).mockImplementation(
+			(storedId: string | undefined, source: string | undefined, tokenId: string | undefined) => {
+				if (!storedId) return tokenId;
+				if (!accountsModule.shouldUpdateAccountIdFromToken(source, storedId)) {
+					return storedId;
+				}
+				return tokenId ?? storedId;
+			},
+		);
+		vi.mocked(refreshQueueModule.queuedRefresh)
+			.mockResolvedValueOnce({
+				type: "success",
+				access: "access-dead",
+				refresh: "shared-refresh",
+				expires: Date.now() + 300_000,
+			})
+			.mockResolvedValueOnce({
+				type: "success",
+				access: "access-live",
+				refresh: "shared-refresh",
+				expires: Date.now() + 300_000,
+			});
+		vi.mocked(fetchHelpersModule.createCodexHeaders).mockImplementation(
+			(_requestId, _accountId, accessToken) => {
+				const headers = new Headers();
+				if (typeof accessToken === "string" && accessToken.length > 0) {
+					headers.set("x-test-access-token", accessToken);
+				}
+				return headers;
+			},
+		);
+		globalThis.fetch = vi.fn(async (_url, init) => {
+			const headers = new Headers(init?.headers);
+			const accessToken = headers.get("x-test-access-token");
+			if (accessToken === "access-dead") {
+				return new Response(JSON.stringify({ error: { code: "deactivated_workspace", message: "workspace dead" } }), {
+					status: 402,
+					headers: { "content-type": "application/json" },
+				});
+			}
+			return new Response("", {
+				status: 200,
+				headers: {
+					"x-codex-primary-used-percent": "20",
+					"x-codex-primary-window-minutes": "180",
+					"x-codex-primary-reset-after-seconds": "900",
+					"x-codex-secondary-used-percent": "10",
+					"x-codex-secondary-window-minutes": "10080",
+					"x-codex-secondary-reset-after-seconds": "86400",
+					"x-codex-plan-type": "plus",
+					"x-codex-active-limit": "40",
+				},
+			});
+		});
+
+		const mockClient = createMockClient();
+		const { OpenAIOAuthPlugin } = await import("../index.js");
+		const plugin = (await OpenAIOAuthPlugin({
+			client: mockClient,
+		} as never)) as unknown as PluginType;
+		const autoMethod = plugin.auth.methods[0] as unknown as {
+			authorize: (inputs?: Record<string, string>) => Promise<{ instructions: string }>;
+		};
+
+		const authResult = await autoMethod.authorize();
+		expect(authResult.instructions).toBe("Authentication cancelled");
+
+		expect(globalThis.fetch).toHaveBeenCalledTimes(2);
+		expect(mockStorage.accounts).toHaveLength(1);
+		expect(mockStorage.accounts.some((account) => account.accountId === "workspace-dead")).toBe(false);
+		expect(mockStorage.accounts[0]?.accountId).toBe("workspace-live");
+		expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalled();
+		expect(mockFlaggedStorage.accounts).toEqual(
+			expect.arrayContaining([
+				expect.objectContaining({
+					accountId: "workspace-dead",
+					organizationId: "org-shared",
+					flaggedReason: "workspace-deactivated",
+					lastError: "deactivated_workspace",
+				}),
+			]),
 		);
 	});
 });

test/storage.test.ts

@@ -20,6 +20,7 @@ import {
   importAccounts,
   previewImportAccounts,
   createTimestampedBackupPath,
+  getWorkspaceIdentityKey,
   withAccountStorageTransaction,
   withFlaggedAccountStorageTransaction,
 } from "../lib/storage.js";
@@ -30,6 +31,37 @@ import {
 // But Task 0 says: "Tests should fail initially (RED phase)"
 
 describe("storage", () => {
+  describe("getWorkspaceIdentityKey", () => {
+    it("uses organizationId and accountId when both are present", () => {
+      expect(
+        getWorkspaceIdentityKey({
+          organizationId: " org-shared ",
+          accountId: " workspace-a ",
+          refreshToken: " refresh-a ",
+        }),
+      ).toBe("organizationId:org-shared|accountId:workspace-a");
+    });
+
+    it("falls back to accountId when organizationId is missing", () => {
+      expect(
+        getWorkspaceIdentityKey({
+          accountId: " workspace-only ",
+          refreshToken: " refresh-b ",
+        }),
+      ).toBe("accountId:workspace-only");
+    });
+
+    it("falls back to refreshToken when workspace ids are missing", () => {
+      expect(
+        getWorkspaceIdentityKey({
+          organizationId: "   ",
+          accountId: "",
+          refreshToken: " refresh-c ",
+        }),
+      ).toBe("refreshToken:refresh-c");
+    });
+  });
+
   describe("deduplication", () => {
     it("remaps activeIndex after deduplication using active account key", () => {
       const now = Date.now();