Skip to content

Commit 8165b93

Browse files
ndycodendycode
committed
Fix workspace rotation review regressions
1 parent 4803bec commit 8165b93

6 files changed

Lines changed: 185 additions & 16 deletions

File tree

‎index.ts‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,23 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
557557
});
558558

559559
if (existingIndex === undefined) {
560+
const initialWorkspaceIndex =
561+
result.workspaces && result.workspaces.length > 0
562+
? (() => {
563+
if (accountId) {
564+
const matchingWorkspaceIndex = result.workspaces.findIndex(
565+
(workspace) => workspace.id === accountId,
566+
);
567+
if (matchingWorkspaceIndex >= 0) {
568+
return matchingWorkspaceIndex;
569+
}
570+
}
571+
const firstEnabledWorkspaceIndex = result.workspaces.findIndex(
572+
(workspace) => workspace.enabled !== false,
573+
);
574+
return firstEnabledWorkspaceIndex >= 0 ? firstEnabledWorkspaceIndex : 0;
575+
})()
576+
: undefined;
560577
accounts.push({
561578
accountId,
562579
accountIdSource,
@@ -568,6 +585,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
568585
addedAt: now,
569586
lastUsed: now,
570587
workspaces: result.workspaces,
588+
currentWorkspaceIndex: initialWorkspaceIndex,
571589
});
572590
continue;
573591
}
@@ -1957,7 +1975,6 @@ accountAttemptLoop: while (attempted.size < Math.max(1, accountCount)) {
19571975
const workspaceErrorMessage =
19581976
(errorBody as { error?: { message?: string } } | undefined)?.error?.message ?? "";
19591977
const isDisabledWorkspaceError =
1960-
errorResponse.status === 403 &&
19611978
isWorkspaceDisabledError(
19621979
errorResponse.status,
19631980
workspaceErrorCode,
@@ -1975,6 +1992,9 @@ accountAttemptLoop: while (attempted.size < Math.max(1, accountCount)) {
19751992
`Workspace disabled/expired for account ${account.index + 1} without tracked workspaces. Leaving account enabled.`,
19761993
{ errorCode: workspaceErrorCode },
19771994
);
1995+
if (hasRemainingAccounts) {
1996+
continue accountAttemptLoop;
1997+
}
19781998
} else {
19791999
const currentWorkspace = accountManager.getCurrentWorkspace(account);
19802000
const workspaceName = currentWorkspace?.name ?? currentWorkspace?.id ?? "unknown";

‎lib/accounts.ts‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -845,15 +845,6 @@ export class AccountManager {
845845
}
846846

847847
// Workspace management methods
848-
setWorkspaces(account: ManagedAccount, workspaces: Workspace[]): void {
849-
account.workspaces = workspaces;
850-
if (workspaces.length > 0 && account.currentWorkspaceIndex === undefined) {
851-
// Find first enabled workspace or default to 0
852-
const firstEnabled = workspaces.findIndex((w) => w.enabled !== false);
853-
account.currentWorkspaceIndex = firstEnabled >= 0 ? firstEnabled : 0;
854-
}
855-
}
856-
857848
private resetWorkspaces(account: ManagedAccount): void {
858849
if (!account.workspaces || account.workspaces.length === 0) {
859850
return;

‎test/accounts.test.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,9 @@ describe("AccountManager", () => {
252252
expect(manager.rotateToNextWorkspace(account)?.id).toBe("workspace-2");
253253

254254
expect(manager.disableCurrentWorkspace(account, "workspace-1")).toBe(false);
255+
expect(account.enabled).not.toBe(false);
256+
expect(manager.hasEnabledWorkspaces(account)).toBe(true);
257+
expect(manager.getEnabledWorkspaceCount(account)).toBe(1);
255258
expect(manager.getCurrentWorkspace(account)?.id).toBe("workspace-2");
256259
expect(account.workspaces?.map((workspace) => workspace.enabled)).toEqual([false, true]);
257260
});

‎test/fetch-helpers.test.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,12 @@ describe('isWorkspaceDisabledError', () => {
342342
expect(isWorkspaceDisabledError(403, 'organization_disabled', '')).toBe(true);
343343
});
344344

345+
it('matches wrapped workspace error tokens but not partial token text', () => {
346+
expect(isWorkspaceDisabledError(403, 'error.workspace_disabled', '')).toBe(true);
347+
expect(isWorkspaceDisabledError(403, 'workspace_expired:error', '')).toBe(true);
348+
expect(isWorkspaceDisabledError(403, 'error.usage_not_included', '')).toBe(false);
349+
});
350+
345351
it('returns false for non-403 status even with disabled message', () => {
346352
expect(isWorkspaceDisabledError(400, '', 'Your workspace has been disabled')).toBe(false);
347353
expect(isWorkspaceDisabledError(401, '', 'Your workspace has been disabled')).toBe(false);

‎test/index-retry.test.ts‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -533,5 +533,67 @@ describe("OpenAIAuthPlugin rate-limit retry", () => {
533533
expect(accountManagerState.rotateToNextWorkspaceCalls).toBe(0);
534534
expect(accountManagerState.setAccountEnabledCalls).toEqual([]);
535535
});
536+
537+
it("retries with a fallback account after a workspace-less account gets a workspace-disabled response", async () => {
538+
const firstAccount = createMockAccount({
539+
index: 0,
540+
accountId: "account-1",
541+
workspaces: undefined,
542+
currentWorkspaceIndex: undefined,
543+
});
544+
const secondAccount = createMockAccount({
545+
index: 1,
546+
accountId: "account-2",
547+
email: "fallback@example.com",
548+
refreshToken: "refresh-token-2",
549+
workspaces: undefined,
550+
currentWorkspaceIndex: undefined,
551+
});
552+
accountManagerState.accounts = [firstAccount, secondAccount];
553+
accountManagerState.accountSelections = [firstAccount, secondAccount];
554+
555+
const fetchMock = vi
556+
.fn()
557+
.mockResolvedValueOnce(
558+
new Response(
559+
JSON.stringify({
560+
error: {
561+
code: "workspace_disabled",
562+
message: "Workspace expired",
563+
},
564+
}),
565+
{
566+
status: 403,
567+
headers: { "content-type": "application/json" },
568+
},
569+
),
570+
)
571+
.mockResolvedValueOnce(new Response("ok", { status: 200 }));
572+
globalThis.fetch = fetchMock as any;
573+
574+
const { OpenAIAuthPlugin } = await import("../index.js");
575+
const client = {
576+
tui: { showToast: vi.fn() },
577+
auth: { set: vi.fn() },
578+
} as any;
579+
const plugin = await OpenAIAuthPlugin({ client });
580+
581+
const getAuth = async () => ({
582+
type: "oauth" as const,
583+
access: "a",
584+
refresh: "r",
585+
expires: Date.now() + 60_000,
586+
multiAccount: true,
587+
});
588+
589+
const sdk = (await plugin.auth.loader(getAuth, { options: {}, models: {} })) as any;
590+
const response = await sdk.fetch("https://example.com", {});
591+
592+
expect(response.status).toBe(200);
593+
expect(fetchMock).toHaveBeenCalledTimes(2);
594+
expect(accountManagerState.disableCurrentWorkspaceCalls).toBe(0);
595+
expect(accountManagerState.rotateToNextWorkspaceCalls).toBe(0);
596+
expect(accountManagerState.setAccountEnabledCalls).toEqual([]);
597+
});
536598
});
537599

‎test/index.test.ts‎

Lines changed: 93 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,12 @@ vi.mock("../lib/request/rate-limit-backoff.js", () => ({
208208
refreshAndUpdateToken: vi.fn(async (auth: unknown) => auth),
209209
createCodexHeaders: vi.fn(() => new Headers()),
210210
handleErrorResponse: vi.fn(async (response: Response) => ({ response })),
211+
isWorkspaceDisabledError: (status: number, code: string, bodyText: string) =>
212+
status === 403 &&
213+
(code.toLowerCase().includes("workspace_disabled") ||
214+
code.toLowerCase().includes("workspace_expired") ||
215+
bodyText.toLowerCase().includes("workspace disabled") ||
216+
bodyText.toLowerCase().includes("workspace expired")),
211217
getUnsupportedCodexModelInfo: vi.fn(() => ({ isUnsupported: false })),
212218
resolveUnsupportedCodexFallbackModel: vi.fn(() => undefined),
213219
shouldFallbackToGpt52OnUnsupportedGpt53: vi.fn(() => false),
@@ -240,6 +246,11 @@ const mockStorage = {
240246
activeIndexByFamily: {} as Record<string, number>,
241247
};
242248

249+
const cloneMockAccount = (account: (typeof mockStorage.accounts)[number]) => ({
250+
...account,
251+
workspaces: account.workspaces?.map((workspace) => ({ ...workspace })),
252+
});
253+
243254
const loadAccountsMock = vi.fn(async () => mockStorage);
244255
const saveAccountsMock = vi.fn(
245256
async (storage: {
@@ -249,7 +260,7 @@ const saveAccountsMock = vi.fn(
249260
activeIndexByFamily?: Record<string, number>;
250261
}) => {
251262
mockStorage.version = storage.version;
252-
mockStorage.accounts = storage.accounts.map((account) => ({ ...account }));
263+
mockStorage.accounts = storage.accounts.map((account) => cloneMockAccount(account));
253264
mockStorage.activeIndex = storage.activeIndex;
254265
mockStorage.activeIndexByFamily = {
255266
...(storage.activeIndexByFamily ?? {}),
@@ -283,7 +294,7 @@ const withAccountStorageTransactionMock = vi.fn(
283294
handler(
284295
{
285296
version: 3,
286-
accounts: mockStorage.accounts.map((account) => ({ ...account })),
297+
accounts: mockStorage.accounts.map((account) => cloneMockAccount(account)),
287298
activeIndex: mockStorage.activeIndex,
288299
activeIndexByFamily: { ...mockStorage.activeIndexByFamily },
289300
},
@@ -331,6 +342,10 @@ vi.mock("../lib/storage.js", async () => {
331342

332343
const extractAccountEmailMock = vi.fn(() => "user@example.com");
333344
const extractAccountIdMock = vi.fn(() => "account-1");
345+
const getAccountIdCandidatesMock = vi.fn(() => [{ accountId: "acc-1", source: "token", label: "Test" }]);
346+
const selectBestAccountCandidateMock = vi.fn(
347+
(candidates: Array<{ accountId: string }>) => candidates[0] ?? null,
348+
);
334349

335350
vi.mock("../lib/accounts.js", () => {
336351
class MockAccountManager {
@@ -429,8 +444,8 @@ vi.mock("../lib/accounts.js", () => {
429444

430445
return {
431446
AccountManager: MockAccountManager,
432-
getAccountIdCandidates: () => [{ accountId: "acc-1", source: "token", label: "Test" }],
433-
selectBestAccountCandidate: (candidates: Array<{ accountId: string }>) => candidates[0] ?? null,
447+
getAccountIdCandidates: getAccountIdCandidatesMock,
448+
selectBestAccountCandidate: selectBestAccountCandidateMock,
434449
extractAccountEmail: extractAccountEmailMock,
435450
extractAccountId: extractAccountIdMock,
436451
resolveRequestAccountId: (_storedId: string | undefined, _source: string | undefined, tokenId: string | undefined) => tokenId,
@@ -513,6 +528,14 @@ describe("OpenAIOAuthPlugin", () => {
513528
extractAccountEmailMock.mockImplementation(() => "user@example.com");
514529
extractAccountIdMock.mockReset();
515530
extractAccountIdMock.mockImplementation(() => "account-1");
531+
getAccountIdCandidatesMock.mockReset();
532+
getAccountIdCandidatesMock.mockImplementation(() => [
533+
{ accountId: "acc-1", source: "token", label: "Test" },
534+
]);
535+
selectBestAccountCandidateMock.mockReset();
536+
selectBestAccountCandidateMock.mockImplementation(
537+
(candidates: Array<{ accountId: string }>) => candidates[0] ?? null,
538+
);
516539

517540
mockStorage.accounts = [];
518541
mockStorage.activeIndex = 0;
@@ -1924,6 +1947,14 @@ describe("OpenAIOAuthPlugin resolveAccountSelection", () => {
19241947
describe("OpenAIOAuthPlugin persistAccountPool", () => {
19251948
beforeEach(() => {
19261949
vi.clearAllMocks();
1950+
getAccountIdCandidatesMock.mockReset();
1951+
getAccountIdCandidatesMock.mockImplementation(() => [
1952+
{ accountId: "acc-1", source: "token", label: "Test" },
1953+
]);
1954+
selectBestAccountCandidateMock.mockReset();
1955+
selectBestAccountCandidateMock.mockImplementation(
1956+
(candidates: Array<{ accountId: string }>) => candidates[0] ?? null,
1957+
);
19271958
mockStorage.accounts = [];
19281959
mockStorage.activeIndex = 0;
19291960
mockStorage.activeIndexByFamily = {};
@@ -1968,6 +1999,62 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
19681999
expect(mockStorage.accounts).toHaveLength(1);
19692000
});
19702001

2002+
it("initializes currentWorkspaceIndex for a newly persisted selected workspace", async () => {
2003+
const authModule = await import("../lib/auth/auth.js");
2004+
const accountsModule = await import("../lib/accounts.js");
2005+
vi.mocked(authModule.createAuthorizationFlow).mockResolvedValueOnce({
2006+
pkce: { verifier: "persist-new-workspace-index", challenge: "persist-new-workspace-index" },
2007+
state: "persist-new-workspace-index",
2008+
url: "https://auth.openai.com/test?state=persist-new-workspace-index",
2009+
});
2010+
vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({
2011+
type: "success",
2012+
access: "access-token",
2013+
refresh: "refresh-token",
2014+
expires: Date.now() + 3600_000,
2015+
idToken: "id-token",
2016+
});
2017+
vi.mocked(accountsModule.getAccountIdCandidates).mockReturnValueOnce([
2018+
{ accountId: "workspace-a", source: "org", label: "Workspace A" },
2019+
{ accountId: "workspace-b", source: "org", label: "Workspace B", isDefault: true },
2020+
{ accountId: "workspace-c", source: "org", label: "Workspace C" },
2021+
]);
2022+
vi.mocked(accountsModule.selectBestAccountCandidate).mockImplementationOnce(
2023+
(candidates) => candidates[1] ?? null,
2024+
);
2025+
2026+
const mockClient = createMockClient();
2027+
const { OpenAIOAuthPlugin } = await import("../index.js");
2028+
const plugin =
2029+
(await OpenAIOAuthPlugin({
2030+
client: mockClient,
2031+
} as never)) as unknown as PluginType;
2032+
const manualMethod = plugin.auth.methods[1] as unknown as {
2033+
authorize: () => Promise<{
2034+
callback: (input: string) => Promise<{ type: string }>;
2035+
}>;
2036+
};
2037+
2038+
const flow = await manualMethod.authorize();
2039+
const result = await flow.callback(
2040+
"http://127.0.0.1:1455/auth/callback?code=abc123&state=persist-new-workspace-index",
2041+
);
2042+
2043+
expect(result.type).toBe("success");
2044+
expect(mockStorage.accounts).toHaveLength(1);
2045+
expect(mockStorage.accounts[0]).toEqual(
2046+
expect.objectContaining({
2047+
accountId: "workspace-b",
2048+
currentWorkspaceIndex: 1,
2049+
workspaces: [
2050+
{ id: "workspace-a", name: "Workspace A", enabled: true, isDefault: undefined },
2051+
{ id: "workspace-b", name: "Workspace B", enabled: true, isDefault: true },
2052+
{ id: "workspace-c", name: "Workspace C", enabled: true, isDefault: undefined },
2053+
],
2054+
}),
2055+
);
2056+
});
2057+
19712058
it("preserves distinct accountId plus email pairs during manual login", async () => {
19722059
process.env.CODEX_AUTH_ACCOUNT_ID = "shared-workspace";
19732060
mockStorage.accounts = [
@@ -2459,15 +2546,15 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
24592546
.mockImplementationOnce(async (storage) => {
24602547
await firstPersist.promise;
24612548
mockStorage.version = storage.version;
2462-
mockStorage.accounts = storage.accounts.map((account) => ({ ...account }));
2549+
mockStorage.accounts = storage.accounts.map((account) => cloneMockAccount(account));
24632550
mockStorage.activeIndex = storage.activeIndex;
24642551
mockStorage.activeIndexByFamily = {
24652552
...(storage.activeIndexByFamily ?? {}),
24662553
};
24672554
})
24682555
.mockImplementation(async (storage) => {
24692556
mockStorage.version = storage.version;
2470-
mockStorage.accounts = storage.accounts.map((account) => ({ ...account }));
2557+
mockStorage.accounts = storage.accounts.map((account) => cloneMockAccount(account));
24712558
mockStorage.activeIndex = storage.activeIndex;
24722559
mockStorage.activeIndexByFamily = {
24732560
...(storage.activeIndexByFamily ?? {}),

0 commit comments

Comments
 (0)