Skip to content

Commit ba92bd9

Browse files
ndycodendycode
committed
fix: address remaining sync review findings
1 parent d9a0ea3 commit ba92bd9

4 files changed

Lines changed: 239 additions & 28 deletions

File tree

‎index.ts‎

Lines changed: 39 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3374,6 +3374,10 @@ while (attempted.size < Math.max(1, accountCount)) {
33743374
return `${target.organizationId ?? ""}|${target.accountId ?? ""}|${target.refreshToken}`;
33753375
};
33763376

3377+
const getSyncRemovalRefreshTokenKey = (refreshToken: string | undefined): string => {
3378+
return refreshToken?.trim() ?? "";
3379+
};
3380+
33773381
const findAccountIndexByExactIdentity = (
33783382
accounts: AccountStorageV3["accounts"],
33793383
target: SyncRemovalTarget | null | undefined,
@@ -3445,11 +3449,38 @@ while (attempted.size < Math.max(1, accountCount)) {
34453449
if (!normalizedAccounts) {
34463450
throw new Error("Prune backup account snapshot failed validation.");
34473451
}
3448-
await withAccountAndFlaggedStorageTransaction(async (_current, persist) => {
3452+
await withAccountAndFlaggedStorageTransaction(async (current, persist) => {
3453+
const rollbackAccounts =
3454+
current.accounts ??
3455+
({
3456+
version: 3,
3457+
accounts: [],
3458+
activeIndex: 0,
3459+
activeIndexByFamily: {},
3460+
} satisfies AccountStorageV3);
34493461
await persist.accounts(normalizedAccounts);
3450-
await persist.flagged(
3451-
restoreFlaggedSnapshot as { version: 1; accounts: FlaggedAccountMetadataV1[] },
3452-
);
3462+
try {
3463+
await persist.flagged(
3464+
restoreFlaggedSnapshot as {
3465+
version: 1;
3466+
accounts: FlaggedAccountMetadataV1[];
3467+
},
3468+
);
3469+
} catch (flaggedError) {
3470+
try {
3471+
await persist.accounts(rollbackAccounts);
3472+
} catch (rollbackError) {
3473+
const flaggedMessage =
3474+
flaggedError instanceof Error ? flaggedError.message : String(flaggedError);
3475+
const rollbackMessage =
3476+
rollbackError instanceof Error ? rollbackError.message : String(rollbackError);
3477+
throw new Error(
3478+
`Failed to restore sync prune flagged storage: ${flaggedMessage}; ` +
3479+
`failed to roll back account restore: ${rollbackMessage}`,
3480+
);
3481+
}
3482+
throw flaggedError;
3483+
}
34533484
});
34543485
invalidateAccountManagerCache();
34553486
},
@@ -3491,25 +3522,17 @@ while (attempted.size < Math.max(1, accountCount)) {
34913522
return;
34923523
}
34933524

3494-
const removedFlaggedKeys = new Set(
3525+
const removedRefreshTokens = new Set(
34953526
removedTargets.map((entry) =>
3496-
getSyncRemovalTargetKey({
3497-
refreshToken: entry.account.refreshToken,
3498-
organizationId: entry.account.organizationId,
3499-
accountId: entry.account.accountId,
3500-
}),
3527+
getSyncRemovalRefreshTokenKey(entry.account.refreshToken),
35013528
),
35023529
);
35033530
const nextFlaggedStorage = {
35043531
version: 1 as const,
35053532
accounts: currentFlaggedStorage.accounts.filter(
35063533
(flagged) =>
3507-
!removedFlaggedKeys.has(
3508-
getSyncRemovalTargetKey({
3509-
refreshToken: flagged.refreshToken,
3510-
organizationId: flagged.organizationId,
3511-
accountId: flagged.accountId,
3512-
}),
3534+
!removedRefreshTokens.has(
3535+
getSyncRemovalRefreshTokenKey(flagged.refreshToken),
35133536
),
35143537
),
35153538
};

‎lib/codex-multi-auth-sync.ts‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -570,12 +570,24 @@ function normalizeIdentity(value: string | undefined): string | undefined {
570570
return trimmed && trimmed.length > 0 ? trimmed.toLowerCase() : undefined;
571571
}
572572

573+
function toCleanupExactIdentityKey(account: {
574+
organizationId?: string;
575+
accountId?: string;
576+
refreshToken: string;
577+
}): string | undefined {
578+
const refreshToken = normalizeTrimmedIdentity(account.refreshToken);
579+
if (!refreshToken) return undefined;
580+
return `exact:${normalizeIdentity(account.organizationId) ?? ""}|${normalizeIdentity(account.accountId) ?? ""}|${refreshToken}`;
581+
}
582+
573583
function toCleanupIdentityKeys(account: {
574584
organizationId?: string;
575585
accountId?: string;
576586
refreshToken: string;
577587
}): string[] {
578588
const keys: string[] = [];
589+
const exactIdentity = toCleanupExactIdentityKey(account);
590+
if (exactIdentity) keys.push(exactIdentity);
579591
const organizationId = normalizeIdentity(account.organizationId);
580592
if (organizationId) keys.push(`org:${organizationId}`);
581593
const accountId = normalizeIdentity(account.accountId);
@@ -1128,13 +1140,13 @@ function buildCodexMultiAuthOverlapCleanupPlan(existing: AccountStorageV3): {
11281140
const removed = Math.max(0, before - after);
11291141
const originalAccountsByKey = new Map<string, AccountStorageV3["accounts"][number]>();
11301142
for (const account of existing.accounts) {
1131-
const key = toCleanupIdentityKeys(account)[0];
1143+
const key = toCleanupExactIdentityKey(account);
11321144
if (key) {
11331145
originalAccountsByKey.set(key, account);
11341146
}
11351147
}
11361148
const updated = normalized.accounts.reduce((count, account) => {
1137-
const key = toCleanupIdentityKeys(account)[0];
1149+
const key = toCleanupExactIdentityKey(account);
11381150
if (!key) return count;
11391151
const original = originalAccountsByKey.get(key);
11401152
if (!original) return count;

‎test/codex-multi-auth-sync.test.ts‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1873,6 +1873,9 @@ describe("codex-multi-auth sync", () => {
18731873
it("remaps active indices when synced overlap cleanup reorders accounts", async () => {
18741874
const storageModule = await import("../lib/storage.js");
18751875
const persist = vi.fn(async (_next: AccountStorageV3) => {});
1876+
vi.mocked(storageModule.deduplicateAccounts).mockImplementationOnce(
1877+
(accounts: AccountStorageV3["accounts"]) => [accounts[1], accounts[0]].filter(Boolean),
1878+
);
18761879
vi.mocked(storageModule.withAccountStorageTransaction).mockImplementationOnce(async (handler) =>
18771880
handler(
18781881
{
@@ -1881,21 +1884,22 @@ describe("codex-multi-auth sync", () => {
18811884
activeIndexByFamily: { codex: 0 },
18821885
accounts: [
18831886
{
1884-
accountId: "org-sync",
1885-
organizationId: "org-sync",
1887+
accountId: "sync-primary",
1888+
organizationId: "shared-org",
18861889
accountIdSource: "org",
18871890
accountTags: ["codex-multi-auth-sync"],
1888-
email: "sync@example.com",
1889-
refreshToken: "sync-token",
1891+
email: "primary@example.com",
1892+
refreshToken: "sync-token-primary",
18901893
addedAt: 3,
18911894
lastUsed: 3,
18921895
},
18931896
{
1894-
accountId: "org-local",
1895-
organizationId: "org-local",
1897+
accountId: "sync-sibling",
1898+
organizationId: "shared-org",
18961899
accountIdSource: "org",
1897-
email: "local@example.com",
1898-
refreshToken: "local-token",
1900+
accountTags: ["codex-multi-auth-sync"],
1901+
email: "sibling@example.com",
1902+
refreshToken: "sync-token-sibling",
18991903
addedAt: 4,
19001904
lastUsed: 4,
19011905
},
@@ -1912,7 +1916,7 @@ describe("codex-multi-auth sync", () => {
19121916
if (!saved) {
19131917
throw new Error("Expected persisted overlap cleanup result");
19141918
}
1915-
expect(saved.accounts.map((account) => account.accountId)).toEqual(["org-local", "org-sync"]);
1919+
expect(saved.accounts.map((account) => account.accountId)).toEqual(["sync-sibling", "sync-primary"]);
19161920
expect(saved.activeIndex).toBe(1);
19171921
expect(saved.activeIndexByFamily?.codex).toBe(1);
19181922
});

‎test/index.test.ts‎

Lines changed: 173 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2198,7 +2198,8 @@ describe("OpenAIOAuthPlugin", () => {
21982198
mockStorage.activeIndexByFamily = { codex: 1 };
21992199
mockFlaggedStorage.accounts = [
22002200
{
2201-
accountId: "remove-me",
2201+
accountId: "flagged-remove-me",
2202+
organizationId: "flagged-org",
22022203
email: "remove@example.com",
22032204
refreshToken: "refresh-remove",
22042205
accessToken: "flagged-access-remove",
@@ -2277,6 +2278,7 @@ describe("OpenAIOAuthPlugin", () => {
22772278
expect(backupContent).toContain("\"idToken\": \"id-remove\"");
22782279
expect(backupContent).toContain("\"accessToken\": \"flagged-access-remove\"");
22792280
expect(backupContent).toContain("\"idToken\": \"flagged-id-remove\"");
2281+
expect(mockFlaggedStorage.accounts).toHaveLength(0);
22802282
expect(renameSpy).toHaveBeenCalled();
22812283
expect(mkdirSpy).toHaveBeenCalled();
22822284
} finally {
@@ -2506,6 +2508,176 @@ describe("OpenAIOAuthPlugin", () => {
25062508
expect(vi.mocked(storageModule.withAccountAndFlaggedStorageTransaction)).toHaveBeenCalledTimes(2);
25072509
expect(vi.mocked(syncModule.syncFromCodexMultiAuth)).toHaveBeenCalledTimes(1);
25082510
});
2511+
2512+
it("rolls back account restore when flagged restore fails", async () => {
2513+
const cliModule = await import("../lib/cli.js");
2514+
const confirmModule = await import("../lib/ui/confirm.js");
2515+
const configModule = await import("../lib/config.js");
2516+
const loggerModule = await import("../lib/logger.js");
2517+
const storageModule = await import("../lib/storage.js");
2518+
const syncModule = await import("../lib/codex-multi-auth-sync.js");
2519+
2520+
mockStorage.accounts = [
2521+
{
2522+
accountId: "remove-me",
2523+
email: "remove@example.com",
2524+
refreshToken: "refresh-remove",
2525+
accessToken: "access-remove",
2526+
idToken: "id-remove",
2527+
addedAt: 1,
2528+
lastUsed: 1,
2529+
},
2530+
{
2531+
accountId: "keep-me",
2532+
email: "keep@example.com",
2533+
refreshToken: "refresh-keep",
2534+
accessToken: "access-keep",
2535+
idToken: "id-keep",
2536+
addedAt: 2,
2537+
lastUsed: 2,
2538+
},
2539+
];
2540+
mockStorage.activeIndex = 1;
2541+
mockStorage.activeIndexByFamily = { codex: 1 };
2542+
mockFlaggedStorage.accounts = [
2543+
{
2544+
accountId: "flagged-remove-me",
2545+
organizationId: "flagged-org",
2546+
email: "remove@example.com",
2547+
refreshToken: "refresh-remove",
2548+
accessToken: "flagged-access-remove",
2549+
idToken: "flagged-id-remove",
2550+
flaggedAt: 123,
2551+
addedAt: 1,
2552+
lastUsed: 1,
2553+
},
2554+
];
2555+
2556+
const { CodexMultiAuthSyncCapacityError } = syncModule;
2557+
vi.mocked(configModule.getSyncFromCodexMultiAuthEnabled).mockReturnValue(true);
2558+
vi.mocked(cliModule.promptLoginMode)
2559+
.mockResolvedValueOnce({ mode: "experimental-sync-now" })
2560+
.mockResolvedValueOnce({ mode: "cancel" })
2561+
.mockResolvedValue({ mode: "cancel" });
2562+
vi.mocked(cliModule.promptCodexMultiAuthSyncPrune).mockResolvedValueOnce([0]);
2563+
vi.mocked(confirmModule.confirm).mockResolvedValue(true);
2564+
vi.mocked(syncModule.previewSyncFromCodexMultiAuth)
2565+
.mockRejectedValueOnce(
2566+
new CodexMultiAuthSyncCapacityError({
2567+
rootDir: "/tmp/codex-source",
2568+
accountsPath: "/tmp/codex-source/openai-codex-accounts.json",
2569+
scope: "global",
2570+
currentCount: 2,
2571+
sourceCount: 1,
2572+
sourceDedupedTotal: 1,
2573+
dedupedTotal: 3,
2574+
maxAccounts: 2,
2575+
needToRemove: 1,
2576+
importableNewAccounts: 1,
2577+
skippedOverlaps: 0,
2578+
suggestedRemovals: [
2579+
{
2580+
index: 0,
2581+
email: "remove@example.com",
2582+
accountLabel: "Remove Me",
2583+
refreshToken: "refresh-remove",
2584+
organizationId: undefined,
2585+
accountId: "remove-me",
2586+
isCurrentAccount: false,
2587+
score: 100,
2588+
reason: "test removal",
2589+
},
2590+
],
2591+
}),
2592+
)
2593+
.mockResolvedValueOnce({
2594+
rootDir: "/tmp/codex-source",
2595+
accountsPath: "/tmp/codex-source/openai-codex-accounts.json",
2596+
scope: "global",
2597+
imported: 1,
2598+
skipped: 0,
2599+
total: 2,
2600+
});
2601+
vi.mocked(syncModule.syncFromCodexMultiAuth).mockRejectedValueOnce(
2602+
new Error("sync failed after prune"),
2603+
);
2604+
2605+
const prunedAccounts = {
2606+
version: 3 as const,
2607+
accounts: [
2608+
{
2609+
accountId: "keep-me",
2610+
email: "keep@example.com",
2611+
refreshToken: "refresh-keep",
2612+
accessToken: "access-keep",
2613+
idToken: "id-keep",
2614+
addedAt: 2,
2615+
lastUsed: 2,
2616+
},
2617+
],
2618+
activeIndex: 0,
2619+
activeIndexByFamily: { codex: 0 },
2620+
};
2621+
const prunedFlagged = {
2622+
version: 1 as const,
2623+
accounts: [] as Array<(typeof mockFlaggedStorage.accounts)[number]>,
2624+
};
2625+
let transactionCalls = 0;
2626+
vi.mocked(storageModule.withAccountAndFlaggedStorageTransaction).mockImplementation(
2627+
async (callback) => {
2628+
transactionCalls += 1;
2629+
if (transactionCalls === 2) {
2630+
return await callback(
2631+
{
2632+
accounts: cloneMockStorage(),
2633+
flagged: cloneMockFlaggedStorage(),
2634+
},
2635+
{
2636+
accounts: async (nextStorage: typeof mockStorage) => {
2637+
mockStorage.version = nextStorage.version;
2638+
mockStorage.accounts = nextStorage.accounts.map(cloneAccount);
2639+
mockStorage.activeIndex = nextStorage.activeIndex;
2640+
mockStorage.activeIndexByFamily = { ...nextStorage.activeIndexByFamily };
2641+
},
2642+
flagged: async (_nextStorage: typeof mockFlaggedStorage) => {
2643+
throw new Error("flagged restore failed");
2644+
},
2645+
},
2646+
);
2647+
}
2648+
return await callback(
2649+
{
2650+
accounts: cloneMockStorage(),
2651+
flagged: cloneMockFlaggedStorage(),
2652+
},
2653+
{
2654+
accounts: async (nextStorage: typeof mockStorage) => {
2655+
mockStorage.version = nextStorage.version;
2656+
mockStorage.accounts = nextStorage.accounts.map(cloneAccount);
2657+
mockStorage.activeIndex = nextStorage.activeIndex;
2658+
mockStorage.activeIndexByFamily = { ...nextStorage.activeIndexByFamily };
2659+
},
2660+
flagged: async (nextStorage: typeof mockFlaggedStorage) => {
2661+
mockFlaggedStorage.version = nextStorage.version;
2662+
mockFlaggedStorage.accounts = nextStorage.accounts.map(cloneFlaggedAccount);
2663+
},
2664+
},
2665+
);
2666+
},
2667+
);
2668+
2669+
const autoMethod = plugin.auth.methods[0] as unknown as {
2670+
authorize: (inputs?: Record<string, string>) => Promise<{ instructions: string }>;
2671+
};
2672+
2673+
const result = await autoMethod.authorize();
2674+
expect(result.instructions).toBe("Authentication cancelled");
2675+
expect(mockStorage).toMatchObject(prunedAccounts);
2676+
expect(mockFlaggedStorage).toEqual(prunedFlagged);
2677+
expect(vi.mocked(loggerModule.logWarn)).toHaveBeenCalledWith(
2678+
expect.stringContaining("Failed to restore sync prune backup: flagged restore failed"),
2679+
);
2680+
});
25092681
});
25102682
});
25112683

0 commit comments

Comments
 (0)