Skip to content

Commit f1caab6

Browse files
ndycodendycode
committed
Harden workspace rotation edge cases
1 parent 9a916a6 commit f1caab6

5 files changed

Lines changed: 215 additions & 42 deletions

File tree

‎index.ts‎

Lines changed: 68 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -593,6 +593,28 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
593593
: newWs;
594594
})
595595
: existing.workspaces;
596+
const nextCurrentWorkspaceIndex =
597+
mergedWorkspaces && mergedWorkspaces.length > 0
598+
? (() => {
599+
const currentIndex =
600+
typeof existing.currentWorkspaceIndex === "number"
601+
? existing.currentWorkspaceIndex
602+
: 0;
603+
if (currentIndex >= 0 && currentIndex < mergedWorkspaces.length) {
604+
return currentIndex;
605+
}
606+
const defaultWorkspaceIndex = mergedWorkspaces.findIndex(
607+
(workspace) => workspace.isDefault === true,
608+
);
609+
if (defaultWorkspaceIndex >= 0) {
610+
return defaultWorkspaceIndex;
611+
}
612+
const firstEnabledWorkspaceIndex = mergedWorkspaces.findIndex(
613+
(workspace) => workspace.enabled !== false,
614+
);
615+
return firstEnabledWorkspaceIndex >= 0 ? firstEnabledWorkspaceIndex : 0;
616+
})()
617+
: existing.currentWorkspaceIndex;
596618
accounts[existingIndex] = {
597619
...existing,
598620
accountId: nextAccountId,
@@ -604,6 +626,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
604626
expiresAt: result.expires,
605627
lastUsed: now,
606628
workspaces: mergedWorkspaces,
629+
currentWorkspaceIndex: nextCurrentWorkspaceIndex,
607630
};
608631
}
609632

@@ -1940,56 +1963,63 @@ accountAttemptLoop: while (attempted.size < Math.max(1, accountCount)) {
19401963
runtimeMetrics.failedRequests++;
19411964
runtimeMetrics.lastError = `Workspace disabled for account ${account.index + 1}`;
19421965

1943-
const currentWorkspace = accountManager.getCurrentWorkspace(account);
1944-
const workspaceName = currentWorkspace?.name ?? currentWorkspace?.id ?? "unknown";
1966+
if (!account.workspaces || account.workspaces.length === 0) {
1967+
logWarn(
1968+
`Workspace disabled/expired for account ${account.index + 1} without tracked workspaces. Leaving account enabled.`,
1969+
{ errorCode: workspaceErrorCode },
1970+
);
1971+
} else {
1972+
const currentWorkspace = accountManager.getCurrentWorkspace(account);
1973+
const workspaceName = currentWorkspace?.name ?? currentWorkspace?.id ?? "unknown";
19451974

1946-
logWarn(
1947-
`Workspace disabled/expired for account ${account.index + 1} - workspace: ${workspaceName}. Rotating to next workspace.`,
1948-
{ errorCode: workspaceErrorCode },
1949-
);
1975+
logWarn(
1976+
`Workspace disabled/expired for account ${account.index + 1} - workspace: ${workspaceName}. Rotating to next workspace.`,
1977+
{ errorCode: workspaceErrorCode },
1978+
);
19501979

1951-
const disabledWorkspace = currentWorkspace
1952-
? accountManager.disableCurrentWorkspace(account, currentWorkspace.id)
1953-
: false;
1954-
let nextWorkspace = disabledWorkspace
1955-
? accountManager.rotateToNextWorkspace(account)
1956-
: accountManager.getCurrentWorkspace(account);
1957-
if (!disabledWorkspace && (!nextWorkspace || nextWorkspace.enabled === false)) {
1958-
nextWorkspace = accountManager.rotateToNextWorkspace(account);
1959-
}
1980+
const disabledWorkspace = currentWorkspace
1981+
? accountManager.disableCurrentWorkspace(account, currentWorkspace.id)
1982+
: false;
1983+
let nextWorkspace = disabledWorkspace
1984+
? accountManager.rotateToNextWorkspace(account)
1985+
: accountManager.getCurrentWorkspace(account);
1986+
if (!disabledWorkspace && (!nextWorkspace || nextWorkspace.enabled === false)) {
1987+
nextWorkspace = accountManager.rotateToNextWorkspace(account);
1988+
}
19601989

1961-
if (nextWorkspace) {
1990+
if (nextWorkspace) {
1991+
accountManager.saveToDiskDebounced();
1992+
1993+
const newWorkspaceName = nextWorkspace.name ?? nextWorkspace.id;
1994+
await showToast(
1995+
`Workspace ${workspaceName} disabled. Switched to ${newWorkspaceName}.`,
1996+
"warning",
1997+
{ duration: toastDurationMs },
1998+
);
1999+
2000+
logInfo(`Rotated to workspace ${newWorkspaceName} for account ${account.index + 1}`);
2001+
2002+
// Allow the same account to be selected again with fresh request state.
2003+
attempted.delete(account.index);
2004+
continue accountAttemptLoop;
2005+
}
2006+
2007+
logWarn(`All workspaces disabled for account ${account.index + 1}. Disabling account.`);
2008+
2009+
accountManager.setAccountEnabled(account.index, false);
19622010
accountManager.saveToDiskDebounced();
19632011

1964-
const newWorkspaceName = nextWorkspace.name ?? nextWorkspace.id;
19652012
await showToast(
1966-
`Workspace ${workspaceName} disabled. Switched to ${newWorkspaceName}.`,
2013+
`All workspaces disabled for account ${account.index + 1}. Switching to another account.`,
19672014
"warning",
19682015
{ duration: toastDurationMs },
19692016
);
19702017

1971-
logInfo(`Rotated to workspace ${newWorkspaceName} for account ${account.index + 1}`);
1972-
1973-
// Allow the same account to be selected again with fresh request state.
1974-
attempted.delete(account.index);
2018+
// Forget session affinity and continue the outer loop so another
2019+
// enabled account can service the request.
2020+
sessionAffinityStore?.forgetSession(sessionAffinityKey);
19752021
continue accountAttemptLoop;
19762022
}
1977-
1978-
logWarn(`All workspaces disabled for account ${account.index + 1}. Disabling account.`);
1979-
1980-
accountManager.setAccountEnabled(account.index, false);
1981-
accountManager.saveToDiskDebounced();
1982-
1983-
await showToast(
1984-
`All workspaces disabled for account ${account.index + 1}. Switching to another account.`,
1985-
"warning",
1986-
{ duration: toastDurationMs },
1987-
);
1988-
1989-
// Forget session affinity and continue the outer loop so another
1990-
// enabled account can service the request.
1991-
sessionAffinityStore?.forgetSession(sessionAffinityKey);
1992-
continue accountAttemptLoop;
19932023
}
19942024

19952025
if (errorResponse.status === 403 && !unsupportedModelInfo.isUnsupported && !isDisabledWorkspaceError) {

‎lib/accounts.ts‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -859,10 +859,7 @@ export class AccountManager {
859859
return;
860860
}
861861

862-
let resetIndex = account.workspaces.findIndex((workspace) => workspace.isDefault === true);
863-
if (resetIndex < 0) {
864-
resetIndex = account.workspaces.findIndex((workspace) => workspace.enabled !== false);
865-
}
862+
const resetIndex = account.workspaces.findIndex((workspace) => workspace.isDefault === true);
866863

867864
for (const workspace of account.workspaces) {
868865
workspace.enabled = true;

‎test/accounts.test.ts‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,41 @@ describe("AccountManager", () => {
294294
]);
295295
});
296296

297+
it("re-enabling an exhausted account without a default workspace resets to the first workspace", () => {
298+
const now = Date.now();
299+
const stored = {
300+
version: 3 as const,
301+
activeIndex: 0,
302+
accounts: [
303+
{
304+
refreshToken: "token-1",
305+
addedAt: now,
306+
lastUsed: now,
307+
enabled: false,
308+
workspaces: [
309+
{ id: "workspace-1", name: "Workspace 1", enabled: false, disabledAt: now - 2_000 },
310+
{ id: "workspace-2", name: "Workspace 2", enabled: false, disabledAt: now - 1_000 },
311+
],
312+
currentWorkspaceIndex: 1,
313+
},
314+
],
315+
};
316+
317+
const manager = new AccountManager(undefined, stored);
318+
const account = manager.setAccountEnabled(0, true);
319+
expect(account).not.toBeNull();
320+
if (!account) {
321+
throw new Error("account should exist");
322+
}
323+
324+
expect(account.currentWorkspaceIndex).toBe(0);
325+
expect(manager.getCurrentWorkspace(account)?.id).toBe("workspace-1");
326+
expect(account.workspaces).toEqual([
327+
{ id: "workspace-1", name: "Workspace 1", enabled: true },
328+
{ id: "workspace-2", name: "Workspace 2", enabled: true },
329+
]);
330+
});
331+
297332
it("checks account availability by enabled/rate-limit/cooldown state", () => {
298333
const now = Date.now();
299334
const stored = {

‎test/index-retry.test.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -486,5 +486,52 @@ describe("OpenAIAuthPlugin rate-limit retry", () => {
486486
expect(firstHeaders.get("x-account-id")).toBe("workspace-1");
487487
expect(secondHeaders.get("x-account-id")).toBe("account-2");
488488
});
489+
490+
it("does not disable workspace-less accounts on workspace-disabled responses", async () => {
491+
const account = createMockAccount({
492+
workspaces: undefined,
493+
currentWorkspaceIndex: undefined,
494+
});
495+
accountManagerState.accounts = [account];
496+
accountManagerState.accountSelections = [account];
497+
498+
globalThis.fetch = vi.fn().mockResolvedValue(
499+
new Response(
500+
JSON.stringify({
501+
error: {
502+
code: "workspace_disabled",
503+
message: "Workspace expired",
504+
},
505+
}),
506+
{
507+
status: 403,
508+
headers: { "content-type": "application/json" },
509+
},
510+
),
511+
) as any;
512+
513+
const { OpenAIAuthPlugin } = await import("../index.js");
514+
const client = {
515+
tui: { showToast: vi.fn() },
516+
auth: { set: vi.fn() },
517+
} as any;
518+
const plugin = await OpenAIAuthPlugin({ client });
519+
520+
const getAuth = async () => ({
521+
type: "oauth" as const,
522+
access: "a",
523+
refresh: "r",
524+
expires: Date.now() + 60_000,
525+
multiAccount: true,
526+
});
527+
528+
const sdk = (await plugin.auth.loader(getAuth, { options: {}, models: {} })) as any;
529+
const response = await sdk.fetch("https://example.com", {});
530+
531+
expect(response.status).toBe(403);
532+
expect(accountManagerState.disableCurrentWorkspaceCalls).toBe(0);
533+
expect(accountManagerState.rotateToNextWorkspaceCalls).toBe(0);
534+
expect(accountManagerState.setAccountEnabledCalls).toEqual([]);
535+
});
489536
});
490537

‎test/index.test.ts‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2243,6 +2243,70 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
22432243
]);
22442244
});
22452245

2246+
it("clamps currentWorkspaceIndex when refreshed workspace metadata removes the active workspace", async () => {
2247+
process.env.CODEX_AUTH_ACCOUNT_ID = "shared-workspace";
2248+
mockStorage.accounts = [
2249+
{
2250+
accountId: "shared-workspace",
2251+
accountIdSource: "manual",
2252+
accountLabel: "Override [id:space]",
2253+
email: "user@example.com",
2254+
refreshToken: "refresh-a",
2255+
addedAt: Date.now() - 200000,
2256+
lastUsed: Date.now() - 200000,
2257+
workspaces: [
2258+
{ id: "workspace-a", name: "Workspace A", enabled: true, isDefault: true },
2259+
{ id: "workspace-b", name: "Workspace B", enabled: true },
2260+
],
2261+
currentWorkspaceIndex: 1,
2262+
},
2263+
];
2264+
2265+
const authModule = await import("../lib/auth/auth.js");
2266+
const accountsModule = await import("../lib/accounts.js");
2267+
vi.mocked(authModule.createAuthorizationFlow).mockResolvedValueOnce({
2268+
pkce: { verifier: "persist-clamp-workspace-index", challenge: "persist-clamp-workspace-index" },
2269+
state: "persist-clamp-workspace-index",
2270+
url: "https://auth.openai.com/test?state=persist-clamp-workspace-index",
2271+
});
2272+
vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({
2273+
type: "success",
2274+
access: "access-token",
2275+
refresh: "refresh-updated",
2276+
expires: Date.now() + 3600_000,
2277+
idToken: undefined,
2278+
workspaces: [
2279+
{ id: "workspace-a", name: "Workspace A", enabled: true, isDefault: true },
2280+
],
2281+
});
2282+
vi.mocked(accountsModule.extractAccountEmail).mockReturnValueOnce("user@example.com");
2283+
vi.mocked(accountsModule.extractAccountId).mockReturnValueOnce("shared-workspace");
2284+
2285+
const mockClient = createMockClient();
2286+
const { OpenAIOAuthPlugin } = await import("../index.js");
2287+
const plugin =
2288+
(await OpenAIOAuthPlugin({
2289+
client: mockClient,
2290+
} as never)) as unknown as PluginType;
2291+
const manualMethod = plugin.auth.methods[1] as unknown as {
2292+
authorize: () => Promise<{
2293+
callback: (input: string) => Promise<{ type: string }>;
2294+
}>;
2295+
};
2296+
2297+
const flow = await manualMethod.authorize();
2298+
const result = await flow.callback(
2299+
"http://127.0.0.1:1455/auth/callback?code=abc123&state=persist-clamp-workspace-index",
2300+
);
2301+
2302+
expect(result.type).toBe("success");
2303+
expect(mockStorage.accounts).toHaveLength(1);
2304+
expect(mockStorage.accounts[0]?.currentWorkspaceIndex).toBe(0);
2305+
expect(mockStorage.accounts[0]?.workspaces).toEqual([
2306+
{ id: "workspace-a", name: "Workspace A", enabled: true, isDefault: true },
2307+
]);
2308+
});
2309+
22462310
it("preserves duplicate shared accountId entries when a login has no email claim", async () => {
22472311
process.env.CODEX_AUTH_ACCOUNT_ID = "shared-workspace";
22482312
mockStorage.accounts = [

0 commit comments

Comments
 (0)