feat: workspace rotation within accounts

Author: Microck

index.ts

@@ -116,6 +116,7 @@ import {
         parseRateLimitReason,
 	lookupCodexCliTokensByEmail,
 	isCodexCliSyncEnabled,
+	type Workspace,
 } from "./lib/accounts.js";
 import {
 	getStoragePath,
@@ -376,6 +377,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
                 accountIdOverride?: string;
                 accountIdSource?: AccountIdSource;
                 accountLabel?: string;
+		workspaces?: Workspace[];
         };
 
         const resolveAccountSelection = (
@@ -398,6 +400,14 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
                         return tokens;
                 }
 
+                // Convert candidates to workspaces
+                const workspaces: Workspace[] = candidates.map((c) => ({
+                        id: c.accountId,
+                        name: c.label,
+                        enabled: true,
+                        isDefault: c.isDefault,
+                }));
+
                 if (candidates.length === 1) {
 				const [candidate] = candidates;
 				if (candidate) {
@@ -406,6 +416,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 						accountIdOverride: candidate.accountId,
 						accountIdSource: candidate.source,
 						accountLabel: candidate.label,
+						workspaces,
 					};
 				}
 			}
@@ -420,6 +431,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
                         accountIdOverride: choice.accountId,
                         accountIdSource: choice.source ?? "token",
                         accountLabel: choice.label,
+			workspaces,
                 };
         };
 
@@ -588,6 +600,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 								expiresAt: result.expires,
 								addedAt: now,
 								lastUsed: now,
+								workspaces: result.workspaces,
 							});
 							indexByRefreshToken.set(result.refresh, newIndex);
 							if (accountId) {
@@ -609,6 +622,11 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 						const nextAccountIdSource =
 							accountId ? accountIdSource ?? existing.accountIdSource : existing.accountIdSource;
 						const nextAccountLabel = accountLabel ?? existing.accountLabel;
+						// Merge workspaces: preserve enabled status from existing, add new ones
+						const mergedWorkspaces = result.workspaces?.map((newWs) => {
+							const existingWs = existing.workspaces?.find((w) => w.id === newWs.id);
+							return existingWs ? { ...newWs, enabled: existingWs.enabled } : newWs;
+						}) ?? result.workspaces;
 						accounts[existingIndex] = {
 							...existing,
 							accountId: nextAccountId,
@@ -619,6 +637,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 							accessToken: result.access,
 							expiresAt: result.expires,
 							lastUsed: now,
+							workspaces: mergedWorkspaces,
 						};
 						if (oldToken !== result.refresh) {
 							indexByRefreshToken.delete(oldToken);
@@ -1554,8 +1573,13 @@ while (attempted.size < Math.max(1, accountCount)) {
 
 				const hadAccountId = !!account.accountId;
 					const tokenAccountId = extractAccountId(accountAuth.access);
+					
+					// Use workspace ID if workspaces are configured, otherwise use account ID
+					const currentWorkspace = accountManager.getCurrentWorkspace(account);
+					const workspaceId = currentWorkspace?.id;
+					
 					const accountId = resolveRequestAccountId(
-						account.accountId,
+						workspaceId ?? account.accountId,
 						account.accountIdSource,
 						tokenAccountId,
 					);
@@ -1955,35 +1979,62 @@ while (attempted.size < Math.max(1, accountCount)) {
 			);
 		}
 
-		// Handle workspace disabled/expired errors by marking account as disabled and rotating
+		// Handle workspace disabled/expired errors by rotating to next workspace within account
 		if (errorResponse.status === 403 && errorBody) {
 			const errorCode = (errorBody as { error?: { code?: string } })?.error?.code ?? "";
 			const errorMessage = (errorBody as { error?: { message?: string } })?.error?.message ?? "";
 
 			if (isWorkspaceDisabledError(errorResponse.status, errorCode, errorMessage)) {
 				runtimeMetrics.failedRequests++;
-				runtimeMetrics.accountRotations++;
 				runtimeMetrics.lastError = `Workspace disabled for account ${account.index + 1}`;
 
+				// Get current workspace info for logging
+				const currentWorkspace = accountManager.getCurrentWorkspace(account);
+				const workspaceName = currentWorkspace?.name ?? currentWorkspace?.id ?? "unknown";
+				
 				logWarn(
-					`Workspace disabled/expired for account ${account.index + 1} (${account.email ?? "unknown"}). Disabling account and rotating.`,
+					`Workspace disabled/expired for account ${account.index + 1} (${account.email ?? "unknown"}) - workspace: ${workspaceName}. Rotating to next workspace.`,
 					{ errorCode, errorMessage },
 				);
 
-				// Mark the account as disabled
-				accountManager.setAccountEnabled(account.index, false);
-				accountManager.saveToDiskDebounced();
+				// Disable the current workspace
+				accountManager.disableCurrentWorkspace(account);
 
-				// Show a toast notification
-				await showToast(
-					`Workspace disabled for account ${account.index + 1}. Switching to another account.`,
-					"warning",
-					{ duration: toastDurationMs },
-				);
+				// Try to rotate to next enabled workspace
+				const nextWorkspace = accountManager.rotateToNextWorkspace(account);
 
-				// Forget session affinity and rotate to next account
-				sessionAffinityStore?.forgetSession(sessionAffinityKey);
-				break;
+				if (nextWorkspace) {
+					// Found another enabled workspace, retry with it
+					accountManager.saveToDiskDebounced();
+					
+					const newWorkspaceName = nextWorkspace.name ?? nextWorkspace.id;
+					await showToast(
+						`Workspace ${workspaceName} disabled. Switched to ${newWorkspaceName}.`,
+						"warning",
+						{ duration: toastDurationMs },
+					);
+					
+					logInfo(`Rotated to workspace ${newWorkspaceName} for account ${account.index + 1}`);
+					
+					// Retry the request with the new workspace (continue the loop)
+					continue;
+				} else {
+					// No more enabled workspaces, disable the entire account
+					logWarn(`All workspaces disabled for account ${account.index + 1}. Disabling account.`);
+					
+					accountManager.setAccountEnabled(account.index, false);
+					accountManager.saveToDiskDebounced();
+					
+					await showToast(
+						`All workspaces disabled for account ${account.index + 1}. Switching to another account.`,
+						"warning",
+						{ duration: toastDurationMs },
+					);
+					
+					// Forget session affinity and rotate to next account
+					sessionAffinityStore?.forgetSession(sessionAffinityKey);
+					break;
+				}
 			}
 		}
 

lib/accounts.ts

@@ -79,6 +79,14 @@ function initFamilyState(defaultValue: number): Record<ModelFamily, number> {
 	) as Record<ModelFamily, number>;
 }
 
+export interface Workspace {
+	id: string;
+	name?: string;
+	enabled: boolean;
+	disabledAt?: number;
+	isDefault?: boolean;
+}
+
 export interface ManagedAccount {
 	index: number;
 	accountId?: string;
@@ -97,6 +105,8 @@ export interface ManagedAccount {
 	coolingDownUntil?: number;
 	cooldownReason?: CooldownReason;
 	consecutiveAuthFailures?: number;
+	workspaces?: Workspace[];
+	currentWorkspaceIndex?: number;
 }
 
 export class AccountManager {
@@ -245,6 +255,8 @@ export class AccountManager {
 						rateLimitResetTimes: account.rateLimitResetTimes ?? {},
 						coolingDownUntil: account.coolingDownUntil,
 						cooldownReason: account.cooldownReason,
+						workspaces: account.workspaces,
+						currentWorkspaceIndex: account.currentWorkspaceIndex,
 					};
 				})
 				.filter((account): account is ManagedAccount => account !== null);
@@ -764,6 +776,8 @@ export class AccountManager {
 					Object.keys(account.rateLimitResetTimes).length > 0 ? account.rateLimitResetTimes : undefined,
 				coolingDownUntil: account.coolingDownUntil,
 				cooldownReason: account.cooldownReason,
+				workspaces: account.workspaces,
+				currentWorkspaceIndex: account.currentWorkspaceIndex,
 			})),
 			activeIndex,
 			activeIndexByFamily,
@@ -805,6 +819,75 @@ export class AccountManager {
 			await this.pendingSave;
 		}
 	}
+
+	// Workspace management methods
+	setWorkspaces(account: ManagedAccount, workspaces: Workspace[]): void {
+		account.workspaces = workspaces;
+		if (workspaces.length > 0 && account.currentWorkspaceIndex === undefined) {
+			// Find first enabled workspace or default to 0
+			const firstEnabled = workspaces.findIndex((w) => w.enabled !== false);
+			account.currentWorkspaceIndex = firstEnabled >= 0 ? firstEnabled : 0;
+		}
+	}
+
+	getCurrentWorkspace(account: ManagedAccount): Workspace | null {
+		if (!account.workspaces || account.workspaces.length === 0) {
+			return null;
+		}
+		const idx = account.currentWorkspaceIndex ?? 0;
+		return account.workspaces[idx] ?? null;
+	}
+
+	disableCurrentWorkspace(account: ManagedAccount): boolean {
+		if (!account.workspaces || account.workspaces.length === 0) {
+			return false;
+		}
+		const idx = account.currentWorkspaceIndex ?? 0;
+		if (idx < 0 || idx >= account.workspaces.length) {
+			return false;
+		}
+		const workspace = account.workspaces[idx];
+		if (!workspace) return false;
+		workspace.enabled = false;
+		workspace.disabledAt = Date.now();
+		return true;
+	}
+
+	rotateToNextWorkspace(account: ManagedAccount): Workspace | null {
+		if (!account.workspaces || account.workspaces.length === 0) {
+			return null;
+		}
+		const currentIdx = account.currentWorkspaceIndex ?? 0;
+		const totalWorkspaces = account.workspaces.length;
+		
+		// Find next enabled workspace
+		for (let i = 1; i <= totalWorkspaces; i++) {
+			const nextIdx = (currentIdx + i) % totalWorkspaces;
+			const workspace = account.workspaces[nextIdx];
+			if (workspace && workspace.enabled !== false) {
+				account.currentWorkspaceIndex = nextIdx;
+				return workspace;
+			}
+		}
+		
+		return null; // No enabled workspaces found
+	}
+
+	hasEnabledWorkspaces(account: ManagedAccount): boolean {
+		if (!account.workspaces || account.workspaces.length === 0) {
+			return true; // No workspaces tracked yet, assume single workspace
+		}
+		return account.workspaces.some((w) => w.enabled !== false);
+	}
+
+	getWorkspaceCount(account: ManagedAccount): number {
+		return account.workspaces?.length ?? 0;
+	}
+
+	getEnabledWorkspaceCount(account: ManagedAccount): number {
+		if (!account.workspaces) return 0;
+		return account.workspaces.filter((w) => w.enabled !== false).length;
+	}
 }
 
 export function formatAccountLabel(

lib/storage/migrations.ts

@@ -8,6 +8,14 @@ import type { AccountIdSource } from "../types.js";
 
 export type CooldownReason = "auth-failure" | "network-error" | "rate-limit";
 
+export interface Workspace {
+	id: string;
+	name?: string;
+	enabled: boolean;
+	disabledAt?: number;
+	isDefault?: boolean;
+}
+
 export interface RateLimitStateV3 {
 	[key: string]: number | undefined;
 }
@@ -54,6 +62,8 @@ export interface AccountMetadataV3 {
 	rateLimitResetTimes?: RateLimitStateV3;
 	coolingDownUntil?: number;
 	cooldownReason?: CooldownReason;
+	workspaces?: Workspace[];
+	currentWorkspaceIndex?: number;
 }
 
 export interface AccountStorageV3 {

test/index-retry.test.ts

@@ -107,6 +107,14 @@ vi.mock("../lib/accounts.js", () => {
 		}
 
 		markToastShown() {}
+
+		getCurrentWorkspace() { return null; }
+
+		disableCurrentWorkspace() { return false; }
+
+		rotateToNextWorkspace() { return null; }
+
+		hasEnabledWorkspaces() { return true; }
 	}
 
 	return {