Skip to content

Commit 9a916a6

Browse files
ndycodendycode
committed
Preserve rotated workspace state on refresh
1 parent 8d877ab commit 9a916a6

4 files changed

Lines changed: 187 additions & 12 deletions

File tree

‎index.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -580,11 +580,19 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
580580
const nextAccountIdSource =
581581
accountId ? accountIdSource ?? existing.accountIdSource : existing.accountIdSource;
582582
const nextAccountLabel = accountLabel ?? existing.accountLabel;
583-
// Merge workspaces: preserve enabled status from existing, add new ones
584-
const mergedWorkspaces = result.workspaces?.map((newWs) => {
585-
const existingWs = existing.workspaces?.find((w) => w.id === newWs.id);
586-
return existingWs ? { ...newWs, enabled: existingWs.enabled } : newWs;
587-
}) ?? result.workspaces;
583+
// Preserve tracked workspace state when auth refreshes do not return workspace metadata.
584+
const mergedWorkspaces = result.workspaces
585+
? result.workspaces.map((newWs) => {
586+
const existingWs = existing.workspaces?.find((w) => w.id === newWs.id);
587+
return existingWs
588+
? {
589+
...newWs,
590+
enabled: existingWs.enabled,
591+
disabledAt: existingWs.disabledAt,
592+
}
593+
: newWs;
594+
})
595+
: existing.workspaces;
588596
accounts[existingIndex] = {
589597
...existing,
590598
accountId: nextAccountId,

‎lib/accounts.ts‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -859,15 +859,17 @@ 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+
}
866+
862867
for (const workspace of account.workspaces) {
863868
workspace.enabled = true;
864869
delete workspace.disabledAt;
865870
}
866871

867-
const currentIdx = account.currentWorkspaceIndex ?? 0;
868-
if (currentIdx < 0 || currentIdx >= account.workspaces.length) {
869-
account.currentWorkspaceIndex = 0;
870-
}
872+
account.currentWorkspaceIndex = resetIndex >= 0 ? resetIndex : 0;
871873
}
872874

873875
getCurrentWorkspace(account: ManagedAccount): Workspace | null {

‎test/accounts.test.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -268,7 +268,7 @@ describe("AccountManager", () => {
268268
lastUsed: now,
269269
enabled: false,
270270
workspaces: [
271-
{ id: "workspace-1", name: "Workspace 1", enabled: false, disabledAt: now - 2_000 },
271+
{ id: "workspace-1", name: "Workspace 1", enabled: false, disabledAt: now - 2_000, isDefault: true },
272272
{ id: "workspace-2", name: "Workspace 2", enabled: false, disabledAt: now - 1_000 },
273273
],
274274
currentWorkspaceIndex: 1,
@@ -286,9 +286,10 @@ describe("AccountManager", () => {
286286
expect(account.enabled).toBe(true);
287287
expect(manager.hasEnabledWorkspaces(account)).toBe(true);
288288
expect(manager.getEnabledWorkspaceCount(account)).toBe(2);
289-
expect(manager.getCurrentWorkspace(account)?.id).toBe("workspace-2");
289+
expect(account.currentWorkspaceIndex).toBe(0);
290+
expect(manager.getCurrentWorkspace(account)?.id).toBe("workspace-1");
290291
expect(account.workspaces).toEqual([
291-
{ id: "workspace-1", name: "Workspace 1", enabled: true },
292+
{ id: "workspace-1", name: "Workspace 1", enabled: true, isDefault: true },
292293
{ id: "workspace-2", name: "Workspace 2", enabled: true },
293294
]);
294295
});

‎test/index.test.ts‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,14 @@ const mockStorage = {
227227
coolingDownUntil?: number;
228228
rateLimitResetTimes?: Record<string, number>;
229229
lastSwitchReason?: string;
230+
workspaces?: Array<{
231+
id: string;
232+
name?: string;
233+
enabled: boolean;
234+
disabledAt?: number;
235+
isDefault?: boolean;
236+
}>;
237+
currentWorkspaceIndex?: number;
230238
}>,
231239
activeIndex: 0,
232240
activeIndexByFamily: {} as Record<string, number>,
@@ -394,6 +402,22 @@ vi.mock("../lib/accounts.js", () => {
394402

395403
markToastShown() {}
396404

405+
getCurrentWorkspace() {
406+
return null;
407+
}
408+
409+
disableCurrentWorkspace() {
410+
return false;
411+
}
412+
413+
rotateToNextWorkspace() {
414+
return null;
415+
}
416+
417+
hasEnabledWorkspaces() {
418+
return true;
419+
}
420+
397421
setActiveIndex(index: number) {
398422
return this.accounts[index] ?? null;
399423
}
@@ -1674,6 +1698,10 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
16741698
getMinWaitTimeForFamily: () => 0,
16751699
shouldShowAccountToast: () => false,
16761700
markToastShown: () => {},
1701+
getCurrentWorkspace: () => null,
1702+
disableCurrentWorkspace: () => false,
1703+
rotateToNextWorkspace: () => null,
1704+
hasEnabledWorkspaces: () => true,
16771705
setActiveIndex: () => accountOne,
16781706
getAccountsSnapshot: () => [accountOne, accountTwo],
16791707
};
@@ -2079,6 +2107,142 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
20792107
]);
20802108
});
20812109

2110+
it("preserves tracked workspaces when manual login updates an account without workspace metadata", async () => {
2111+
process.env.CODEX_AUTH_ACCOUNT_ID = "shared-workspace";
2112+
mockStorage.accounts = [
2113+
{
2114+
accountId: "shared-workspace",
2115+
accountIdSource: "manual",
2116+
accountLabel: "Override [id:space]",
2117+
email: "user@example.com",
2118+
refreshToken: "refresh-a",
2119+
addedAt: Date.now() - 200000,
2120+
lastUsed: Date.now() - 200000,
2121+
workspaces: [
2122+
{ id: "workspace-a", name: "Workspace A", enabled: false, disabledAt: 111 },
2123+
{ id: "workspace-b", name: "Workspace B", enabled: true, isDefault: true },
2124+
],
2125+
currentWorkspaceIndex: 1,
2126+
},
2127+
];
2128+
2129+
const authModule = await import("../lib/auth/auth.js");
2130+
const accountsModule = await import("../lib/accounts.js");
2131+
vi.mocked(authModule.createAuthorizationFlow).mockResolvedValueOnce({
2132+
pkce: { verifier: "persist-preserve-workspaces", challenge: "persist-preserve-workspaces" },
2133+
state: "persist-preserve-workspaces",
2134+
url: "https://auth.openai.com/test?state=persist-preserve-workspaces",
2135+
});
2136+
vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({
2137+
type: "success",
2138+
access: "access-token",
2139+
refresh: "refresh-updated",
2140+
expires: Date.now() + 3600_000,
2141+
idToken: undefined,
2142+
});
2143+
vi.mocked(accountsModule.extractAccountEmail).mockReturnValueOnce("user@example.com");
2144+
vi.mocked(accountsModule.extractAccountId).mockReturnValueOnce("shared-workspace");
2145+
2146+
const mockClient = createMockClient();
2147+
const { OpenAIOAuthPlugin } = await import("../index.js");
2148+
const plugin =
2149+
(await OpenAIOAuthPlugin({
2150+
client: mockClient,
2151+
} as never)) as unknown as PluginType;
2152+
const manualMethod = plugin.auth.methods[1] as unknown as {
2153+
authorize: () => Promise<{
2154+
callback: (input: string) => Promise<{ type: string }>;
2155+
}>;
2156+
};
2157+
2158+
const flow = await manualMethod.authorize();
2159+
const result = await flow.callback(
2160+
"http://127.0.0.1:1455/auth/callback?code=abc123&state=persist-preserve-workspaces",
2161+
);
2162+
2163+
expect(result.type).toBe("success");
2164+
expect(mockStorage.accounts).toHaveLength(1);
2165+
expect(mockStorage.accounts[0]).toEqual(
2166+
expect.objectContaining({
2167+
accountId: "shared-workspace",
2168+
refreshToken: "refresh-updated",
2169+
currentWorkspaceIndex: 1,
2170+
workspaces: [
2171+
{ id: "workspace-a", name: "Workspace A", enabled: false, disabledAt: 111 },
2172+
{ id: "workspace-b", name: "Workspace B", enabled: true, isDefault: true },
2173+
],
2174+
}),
2175+
);
2176+
});
2177+
2178+
it("preserves disabledAt when refreshed workspace metadata is merged into an existing account", async () => {
2179+
process.env.CODEX_AUTH_ACCOUNT_ID = "shared-workspace";
2180+
mockStorage.accounts = [
2181+
{
2182+
accountId: "shared-workspace",
2183+
accountIdSource: "manual",
2184+
accountLabel: "Override [id:space]",
2185+
email: "user@example.com",
2186+
refreshToken: "refresh-a",
2187+
addedAt: Date.now() - 200000,
2188+
lastUsed: Date.now() - 200000,
2189+
workspaces: [
2190+
{ id: "workspace-a", name: "Workspace A", enabled: false, disabledAt: 222, isDefault: true },
2191+
],
2192+
currentWorkspaceIndex: 0,
2193+
},
2194+
];
2195+
2196+
const authModule = await import("../lib/auth/auth.js");
2197+
const accountsModule = await import("../lib/accounts.js");
2198+
vi.mocked(authModule.createAuthorizationFlow).mockResolvedValueOnce({
2199+
pkce: { verifier: "persist-merge-workspaces", challenge: "persist-merge-workspaces" },
2200+
state: "persist-merge-workspaces",
2201+
url: "https://auth.openai.com/test?state=persist-merge-workspaces",
2202+
});
2203+
vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({
2204+
type: "success",
2205+
access: "access-token",
2206+
refresh: "refresh-updated",
2207+
expires: Date.now() + 3600_000,
2208+
idToken: undefined,
2209+
workspaces: [
2210+
{ id: "workspace-a", name: "Workspace A Renamed", enabled: true, isDefault: true },
2211+
],
2212+
});
2213+
vi.mocked(accountsModule.extractAccountEmail).mockReturnValueOnce("user@example.com");
2214+
vi.mocked(accountsModule.extractAccountId).mockReturnValueOnce("shared-workspace");
2215+
2216+
const mockClient = createMockClient();
2217+
const { OpenAIOAuthPlugin } = await import("../index.js");
2218+
const plugin =
2219+
(await OpenAIOAuthPlugin({
2220+
client: mockClient,
2221+
} as never)) as unknown as PluginType;
2222+
const manualMethod = plugin.auth.methods[1] as unknown as {
2223+
authorize: () => Promise<{
2224+
callback: (input: string) => Promise<{ type: string }>;
2225+
}>;
2226+
};
2227+
2228+
const flow = await manualMethod.authorize();
2229+
const result = await flow.callback(
2230+
"http://127.0.0.1:1455/auth/callback?code=abc123&state=persist-merge-workspaces",
2231+
);
2232+
2233+
expect(result.type).toBe("success");
2234+
expect(mockStorage.accounts).toHaveLength(1);
2235+
expect(mockStorage.accounts[0]?.workspaces).toEqual([
2236+
{
2237+
id: "workspace-a",
2238+
name: "Workspace A Renamed",
2239+
enabled: false,
2240+
disabledAt: 222,
2241+
isDefault: true,
2242+
},
2243+
]);
2244+
});
2245+
20822246
it("preserves duplicate shared accountId entries when a login has no email claim", async () => {
20832247
process.env.CODEX_AUTH_ACCOUNT_ID = "shared-workspace";
20842248
mockStorage.accounts = [

0 commit comments

Comments
 (0)