Skip to content

Commit ea3b559

Browse files
ndycodendycode
committed
fix: tighten workspace identity follow-ups
1 parent 1d3752d commit ea3b559

3 files changed

Lines changed: 182 additions & 57 deletions

File tree

‎index.ts‎

Lines changed: 1 addition & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ import {
119119
saveFlaggedAccounts,
120120
withFlaggedAccountStorageTransaction,
121121
clearFlaggedAccounts,
122+
getWorkspaceIdentityKey,
122123
StorageError,
123124
formatStorageErrorHint,
124125
type AccountStorageV3,
@@ -185,23 +186,6 @@ import {
185186
getRecoveryToastContent,
186187
} from "./lib/recovery.js";
187188

188-
function getWorkspaceIdentityKey(account: {
189-
organizationId?: string;
190-
accountId?: string;
191-
refreshToken: string;
192-
}): string {
193-
const organizationId = account.organizationId?.trim();
194-
const accountId = account.accountId?.trim();
195-
const refreshToken = account.refreshToken.trim();
196-
if (organizationId) {
197-
return accountId
198-
? `organizationId:${organizationId}|accountId:${accountId}`
199-
: `organizationId:${organizationId}`;
200-
}
201-
if (accountId) return `accountId:${accountId}`;
202-
return `refreshToken:${refreshToken}`;
203-
}
204-
205189
function matchesWorkspaceIdentity(
206190
account: {
207191
organizationId?: string;

‎lib/storage.ts‎

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,28 @@ export interface FlaggedAccountStorageV1 {
3434
accounts: FlaggedAccountMetadataV1[];
3535
}
3636

37+
const normalizeWorkspaceIdentityPart = (value: unknown): string | undefined =>
38+
typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
39+
40+
export function getWorkspaceIdentityKey(account: {
41+
organizationId?: string;
42+
accountId?: string;
43+
refreshToken: string;
44+
}): string {
45+
const organizationId = normalizeWorkspaceIdentityPart(account.organizationId);
46+
const accountId = normalizeWorkspaceIdentityPart(account.accountId);
47+
const refreshToken = normalizeWorkspaceIdentityPart(account.refreshToken) ?? "";
48+
if (organizationId) {
49+
return accountId
50+
? `organizationId:${organizationId}|accountId:${accountId}`
51+
: `organizationId:${organizationId}`;
52+
}
53+
if (accountId) {
54+
return `accountId:${accountId}`;
55+
}
56+
return `refreshToken:${refreshToken}`;
57+
}
58+
3759
export type ImportBackupMode = "none" | "best-effort" | "required";
3860

3961
export interface ImportAccountsOptions {
@@ -921,29 +943,14 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 {
921943
return { version: 1, accounts: [] };
922944
}
923945

924-
const normalizeFlaggedIdentityPart = (value: unknown): string | undefined =>
925-
typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
926946
// Flagged storage must keep sibling workspace entries separate when they share an
927947
// organization but have different accountIds, so this key is more specific than
928948
// the normal account identity collapse used in active storage.
929949
const getFlaggedIdentityKey = (account: {
930950
organizationId?: string;
931951
accountId?: string;
932952
refreshToken: string;
933-
}): string => {
934-
const organizationId = normalizeFlaggedIdentityPart(account.organizationId);
935-
const accountId = normalizeFlaggedIdentityPart(account.accountId);
936-
const refreshToken = normalizeFlaggedIdentityPart(account.refreshToken) ?? "";
937-
if (organizationId) {
938-
return accountId
939-
? `organizationId:${organizationId}|accountId:${accountId}`
940-
: `organizationId:${organizationId}`;
941-
}
942-
if (accountId) {
943-
return `accountId:${accountId}`;
944-
}
945-
return `refreshToken:${refreshToken}`;
946-
};
953+
}): string => getWorkspaceIdentityKey(account);
947954

948955
const byIdentityKey = new Map<string, FlaggedAccountMetadataV1>();
949956
for (const rawAccount of data.accounts) {

‎test/index.test.ts‎

Lines changed: 158 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -311,6 +311,18 @@ vi.mock("../lib/storage.js", () => ({
311311
previewImportAccounts: vi.fn(async () => ({ imported: 2, skipped: 1, total: 5 })),
312312
createTimestampedBackupPath: vi.fn((prefix?: string) => `/tmp/${prefix ?? "codex-backup"}-20260101-000000.json`),
313313
loadFlaggedAccounts: vi.fn(async () => cloneMockFlaggedStorage()),
314+
getWorkspaceIdentityKey: vi.fn((account: { organizationId?: string; accountId?: string; refreshToken: string }) => {
315+
const organizationId = account.organizationId?.trim();
316+
const accountId = account.accountId?.trim();
317+
const refreshToken = account.refreshToken.trim();
318+
if (organizationId) {
319+
return accountId
320+
? `organizationId:${organizationId}|accountId:${accountId}`
321+
: `organizationId:${organizationId}`;
322+
}
323+
if (accountId) return `accountId:${accountId}`;
324+
return `refreshToken:${refreshToken}`;
325+
}),
314326
saveFlaggedAccounts: mockSaveFlaggedAccounts,
315327
withFlaggedAccountStorageTransaction: vi.fn(
316328
async <T>(
@@ -438,12 +450,15 @@ vi.mock("../lib/accounts.js", () => {
438450
),
439451
extractAccountEmail: vi.fn(() => "user@example.com"),
440452
extractAccountId: vi.fn(() => "account-1"),
441-
resolveRequestAccountId: (_storedId: string | undefined, _source: string | undefined, tokenId: string | undefined) => tokenId,
453+
resolveRequestAccountId: vi.fn(
454+
(_storedId: string | undefined, _source: string | undefined, tokenId: string | undefined) =>
455+
tokenId,
456+
),
442457
formatAccountLabel: (_account: unknown, index: number) => `Account ${index + 1}`,
443458
formatCooldown: () => null,
444459
formatWaitTime: (ms: number) => `${Math.round(ms / 1000)}s`,
445460
sanitizeEmail: (email: string) => email,
446-
shouldUpdateAccountIdFromToken: () => true,
461+
shouldUpdateAccountIdFromToken: vi.fn(() => true),
447462
parseRateLimitReason: () => "unknown",
448463
lookupCodexCliTokensByEmail: vi.fn(async () => null),
449464
};
@@ -2639,18 +2654,16 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
26392654
expect(removeAccount).toHaveBeenCalledTimes(1);
26402655
expect(removeAccountsWithSameRefreshToken).not.toHaveBeenCalled();
26412656
expect(accounts.map((account) => account.accountId)).toEqual(["org-live"]);
2642-
expect(vi.mocked(storageModule.saveFlaggedAccounts)).toHaveBeenCalledTimes(1);
2643-
expect(vi.mocked(storageModule.saveFlaggedAccounts).mock.calls[0]?.[0]).toEqual(
2644-
expect.objectContaining({
2645-
accounts: expect.arrayContaining([
2646-
expect.objectContaining({
2647-
accountId: "org-dead",
2648-
organizationId: "org-dead",
2649-
flaggedReason: "workspace-deactivated",
2650-
lastError: "deactivated_workspace",
2651-
}),
2652-
]),
2653-
}),
2657+
expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalledTimes(1);
2658+
expect(mockFlaggedStorage.accounts).toEqual(
2659+
expect.arrayContaining([
2660+
expect.objectContaining({
2661+
accountId: "org-dead",
2662+
organizationId: "org-dead",
2663+
flaggedReason: "workspace-deactivated",
2664+
lastError: "deactivated_workspace",
2665+
}),
2666+
]),
26542667
);
26552668
});
26562669

@@ -3618,16 +3631,137 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
36183631
expect(mockStorage.accounts.some((account) => account.accountId === "workspace-dead")).toBe(false);
36193632
expect(mockStorage.accounts[0]?.organizationId).toBe("org-shared");
36203633
expect(mockStorage.accounts[0]?.refreshToken).toBe("shared-refresh");
3621-
expect(vi.mocked(storageModule.saveFlaggedAccounts)).toHaveBeenCalledWith(
3622-
expect.objectContaining({
3623-
accounts: expect.arrayContaining([
3624-
expect.objectContaining({
3625-
accountId: "workspace-dead",
3626-
organizationId: "org-shared",
3627-
flaggedReason: "token-invalid",
3628-
}),
3629-
]),
3630-
}),
3634+
expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalled();
3635+
expect(mockFlaggedStorage.accounts).toEqual(
3636+
expect.arrayContaining([
3637+
expect.objectContaining({
3638+
accountId: "workspace-dead",
3639+
organizationId: "org-shared",
3640+
flaggedReason: "token-invalid",
3641+
}),
3642+
]),
3643+
);
3644+
});
3645+
3646+
it("removes only the deactivated org-scoped workspace during quota-check cleanup", async () => {
3647+
const accountsModule = await import("../lib/accounts.js");
3648+
const cliModule = await import("../lib/cli.js");
3649+
const refreshQueueModule = await import("../lib/refresh-queue.js");
3650+
const fetchHelpersModule = await import("../lib/request/fetch-helpers.js");
3651+
const storageModule = await import("../lib/storage.js");
3652+
3653+
mockStorage.accounts = [
3654+
{
3655+
refreshToken: "shared-refresh",
3656+
organizationId: "org-shared",
3657+
accountId: "workspace-dead",
3658+
accountIdSource: "manual",
3659+
email: "dead@example.com",
3660+
addedAt: 1,
3661+
lastUsed: 1,
3662+
},
3663+
{
3664+
refreshToken: "shared-refresh",
3665+
organizationId: "org-shared",
3666+
accountId: "workspace-live",
3667+
accountIdSource: "manual",
3668+
email: "live@example.com",
3669+
addedAt: 2,
3670+
lastUsed: 2,
3671+
},
3672+
];
3673+
3674+
vi.mocked(cliModule.promptLoginMode)
3675+
.mockResolvedValueOnce({ mode: "check" })
3676+
.mockResolvedValueOnce({ mode: "cancel" });
3677+
vi.mocked(accountsModule.shouldUpdateAccountIdFromToken).mockImplementation(
3678+
(source: string | undefined, currentAccountId?: string) => {
3679+
if (!currentAccountId) return true;
3680+
if (!source) return true;
3681+
return source === "token" || source === "id_token";
3682+
},
3683+
);
3684+
vi.mocked(accountsModule.resolveRequestAccountId).mockImplementation(
3685+
(storedId: string | undefined, source: string | undefined, tokenId: string | undefined) => {
3686+
if (!storedId) return tokenId;
3687+
if (!accountsModule.shouldUpdateAccountIdFromToken(source, storedId)) {
3688+
return storedId;
3689+
}
3690+
return tokenId ?? storedId;
3691+
},
3692+
);
3693+
vi.mocked(refreshQueueModule.queuedRefresh)
3694+
.mockResolvedValueOnce({
3695+
type: "success",
3696+
access: "access-dead",
3697+
refresh: "shared-refresh",
3698+
expires: Date.now() + 300_000,
3699+
})
3700+
.mockResolvedValueOnce({
3701+
type: "success",
3702+
access: "access-live",
3703+
refresh: "shared-refresh",
3704+
expires: Date.now() + 300_000,
3705+
});
3706+
vi.mocked(fetchHelpersModule.createCodexHeaders).mockImplementation(
3707+
(_requestId, _accountId, accessToken) => {
3708+
const headers = new Headers();
3709+
if (typeof accessToken === "string" && accessToken.length > 0) {
3710+
headers.set("x-test-access-token", accessToken);
3711+
}
3712+
return headers;
3713+
},
3714+
);
3715+
globalThis.fetch = vi.fn(async (_url, init) => {
3716+
const headers = new Headers(init?.headers);
3717+
const accessToken = headers.get("x-test-access-token");
3718+
if (accessToken === "access-dead") {
3719+
return new Response(JSON.stringify({ error: { code: "deactivated_workspace", message: "workspace dead" } }), {
3720+
status: 402,
3721+
headers: { "content-type": "application/json" },
3722+
});
3723+
}
3724+
return new Response("", {
3725+
status: 200,
3726+
headers: {
3727+
"x-codex-primary-used-percent": "20",
3728+
"x-codex-primary-window-minutes": "180",
3729+
"x-codex-primary-reset-after-seconds": "900",
3730+
"x-codex-secondary-used-percent": "10",
3731+
"x-codex-secondary-window-minutes": "10080",
3732+
"x-codex-secondary-reset-after-seconds": "86400",
3733+
"x-codex-plan-type": "plus",
3734+
"x-codex-active-limit": "40",
3735+
},
3736+
});
3737+
});
3738+
3739+
const mockClient = createMockClient();
3740+
const { OpenAIOAuthPlugin } = await import("../index.js");
3741+
const plugin = (await OpenAIOAuthPlugin({
3742+
client: mockClient,
3743+
} as never)) as unknown as PluginType;
3744+
const autoMethod = plugin.auth.methods[0] as unknown as {
3745+
authorize: (inputs?: Record<string, string>) => Promise<{ instructions: string }>;
3746+
};
3747+
3748+
const authResult = await autoMethod.authorize();
3749+
expect(authResult.instructions).toBe("Authentication cancelled");
3750+
3751+
expect(globalThis.fetch).toHaveBeenCalledTimes(2);
3752+
expect(mockStorage.accounts).toHaveLength(1);
3753+
expect(mockStorage.accounts.some((account) => account.accountId === "workspace-dead")).toBe(false);
3754+
expect(mockStorage.accounts[0]?.accountId).toBe("workspace-live");
3755+
expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalled();
3756+
expect(mockFlaggedStorage.accounts).toEqual(
3757+
expect.arrayContaining([
3758+
expect.objectContaining({
3759+
accountId: "workspace-dead",
3760+
organizationId: "org-shared",
3761+
flaggedReason: "workspace-deactivated",
3762+
lastError: "deactivated_workspace",
3763+
}),
3764+
]),
36313765
);
36323766
});
36333767
});

0 commit comments

Comments
 (0)