Skip to content

Commit 99016e4

Browse files
ndycodendycode
committed
fix: roll back flagged recovery partial writes
1 parent 066a145 commit 99016e4

4 files changed

Lines changed: 325 additions & 54 deletions

File tree

‎lib/codex-manager.ts‎

Lines changed: 22 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ import {
5858
saveFlaggedAccounts,
5959
saveAccounts,
6060
setStoragePath,
61+
withAccountAndFlaggedStorageTransaction,
6162
withAccountStorageTransaction,
6263
type AccountMetadataV3,
6364
type AccountStorageV3,
@@ -2553,10 +2554,9 @@ async function runVerifyFlagged(args: string[]): Promise<number> {
25532554
});
25542555
}
25552556

2556-
const applyRefreshChecks = async (
2557+
const applyRefreshChecks = (
25572558
storage: AccountStorageV3,
2558-
persist?: (nextStorage: AccountStorageV3) => Promise<void>,
2559-
): Promise<void> => {
2559+
): void => {
25602560
for (const check of refreshChecks) {
25612561
const { index: i, flagged, label, result } = check;
25622562
if (result.type === "success") {
@@ -2638,28 +2638,33 @@ async function runVerifyFlagged(args: string[]): Promise<number> {
26382638
message: detail,
26392639
});
26402640
}
2641-
2642-
if (persist && storageChanged) {
2643-
normalizeDoctorIndexes(storage);
2644-
await persist(storage);
2645-
}
26462641
};
26472642

26482643
if (options.restore) {
26492644
if (options.dryRun) {
2650-
await applyRefreshChecks(
2645+
applyRefreshChecks(
26512646
(await loadAccounts()) ?? createEmptyAccountStorage(),
26522647
);
26532648
} else {
2654-
await withAccountStorageTransaction(async (loadedStorage, persist) => {
2655-
await applyRefreshChecks(
2656-
loadedStorage ? structuredClone(loadedStorage) : createEmptyAccountStorage(),
2657-
persist,
2658-
);
2659-
});
2649+
await withAccountAndFlaggedStorageTransaction(
2650+
async (loadedStorage, persist) => {
2651+
const nextStorage = loadedStorage
2652+
? structuredClone(loadedStorage)
2653+
: createEmptyAccountStorage();
2654+
applyRefreshChecks(nextStorage);
2655+
if (!storageChanged) {
2656+
return;
2657+
}
2658+
normalizeDoctorIndexes(nextStorage);
2659+
await persist(nextStorage, {
2660+
version: 1,
2661+
accounts: nextFlaggedAccounts,
2662+
});
2663+
},
2664+
);
26602665
}
26612666
} else {
2662-
await applyRefreshChecks(createEmptyAccountStorage());
2667+
applyRefreshChecks(createEmptyAccountStorage());
26632668
}
26642669

26652670
const remainingFlagged = nextFlaggedAccounts.length;
@@ -2668,7 +2673,7 @@ async function runVerifyFlagged(args: string[]): Promise<number> {
26682673
const stillFlagged = reports.filter((report) => report.outcome === "still-flagged").length;
26692674
const changed = storageChanged || flaggedChanged;
26702675

2671-
if (!options.dryRun && flaggedChanged) {
2676+
if (!options.dryRun && flaggedChanged && (!options.restore || !storageChanged)) {
26722677
await saveFlaggedAccounts({
26732678
version: 1,
26742679
accounts: nextFlaggedAccounts,

‎lib/storage.ts‎

Lines changed: 103 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1845,6 +1845,21 @@ async function saveAccountsUnlocked(storage: AccountStorageV3): Promise<void> {
18451845
}
18461846
}
18471847

1848+
function cloneAccountStorageForPersistence(
1849+
storage: AccountStorageV3 | null | undefined,
1850+
): AccountStorageV3 {
1851+
return {
1852+
version: 3,
1853+
accounts: structuredClone(storage?.accounts ?? []),
1854+
activeIndex:
1855+
typeof storage?.activeIndex === "number" &&
1856+
Number.isFinite(storage.activeIndex)
1857+
? storage.activeIndex
1858+
: 0,
1859+
activeIndexByFamily: structuredClone(storage?.activeIndexByFamily ?? {}),
1860+
};
1861+
}
1862+
18481863
export async function withAccountStorageTransaction<T>(
18491864
handler: (
18501865
current: AccountStorageV3 | null,
@@ -1867,6 +1882,53 @@ export async function withAccountStorageTransaction<T>(
18671882
});
18681883
}
18691884

1885+
export async function withAccountAndFlaggedStorageTransaction<T>(
1886+
handler: (
1887+
current: AccountStorageV3 | null,
1888+
persist: (
1889+
accountStorage: AccountStorageV3,
1890+
flaggedStorage: FlaggedAccountStorageV1,
1891+
) => Promise<void>,
1892+
) => Promise<T>,
1893+
): Promise<T> {
1894+
return withStorageLock(async () => {
1895+
const state = {
1896+
snapshot: await loadAccountsInternal(saveAccountsUnlocked),
1897+
active: true,
1898+
};
1899+
const current = state.snapshot;
1900+
const persist = async (
1901+
accountStorage: AccountStorageV3,
1902+
flaggedStorage: FlaggedAccountStorageV1,
1903+
): Promise<void> => {
1904+
const previousAccounts = cloneAccountStorageForPersistence(state.snapshot);
1905+
const nextAccounts = cloneAccountStorageForPersistence(accountStorage);
1906+
await saveAccountsUnlocked(nextAccounts);
1907+
try {
1908+
await saveFlaggedAccountsUnlocked(flaggedStorage);
1909+
state.snapshot = nextAccounts;
1910+
} catch (error) {
1911+
try {
1912+
await saveAccountsUnlocked(previousAccounts);
1913+
state.snapshot = previousAccounts;
1914+
} catch (rollbackError) {
1915+
log.error(
1916+
"Failed to rollback account storage after flagged save failure",
1917+
{
1918+
error: String(error),
1919+
rollbackError: String(rollbackError),
1920+
},
1921+
);
1922+
}
1923+
throw error;
1924+
}
1925+
};
1926+
return transactionSnapshotContext.run(state, () =>
1927+
handler(current, persist),
1928+
);
1929+
});
1930+
}
1931+
18701932
/**
18711933
* Persists account storage to disk using atomic write (temp file + rename).
18721934
* Creates the Codex multi-auth storage directory if it doesn't exist.
@@ -2091,47 +2153,53 @@ export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
20912153
}
20922154
}
20932155

2094-
export async function saveFlaggedAccounts(
2156+
async function saveFlaggedAccountsUnlocked(
20952157
storage: FlaggedAccountStorageV1,
20962158
): Promise<void> {
2097-
return withStorageLock(async () => {
2098-
const path = getFlaggedAccountsPath();
2099-
const markerPath = getIntentionalResetMarkerPath(path);
2100-
const uniqueSuffix = `${Date.now()}.${Math.random().toString(36).slice(2, 8)}`;
2101-
const tempPath = `${path}.${uniqueSuffix}.tmp`;
2159+
const path = getFlaggedAccountsPath();
2160+
const markerPath = getIntentionalResetMarkerPath(path);
2161+
const uniqueSuffix = `${Date.now()}.${Math.random().toString(36).slice(2, 8)}`;
2162+
const tempPath = `${path}.${uniqueSuffix}.tmp`;
21022163

2103-
try {
2104-
await fs.mkdir(dirname(path), { recursive: true });
2105-
if (existsSync(path)) {
2106-
try {
2107-
await fs.copyFile(path, `${path}.bak`);
2108-
} catch (backupError) {
2109-
log.warn("Failed to create flagged backup snapshot", {
2110-
path,
2111-
error: String(backupError),
2112-
});
2113-
}
2114-
}
2115-
const content = JSON.stringify(normalizeFlaggedStorage(storage), null, 2);
2116-
await fs.writeFile(tempPath, content, { encoding: "utf-8", mode: 0o600 });
2117-
await fs.rename(tempPath, path);
2118-
try {
2119-
await fs.unlink(markerPath);
2120-
} catch {
2121-
// Best effort cleanup.
2122-
}
2123-
} catch (error) {
2164+
try {
2165+
await fs.mkdir(dirname(path), { recursive: true });
2166+
if (existsSync(path)) {
21242167
try {
2125-
await fs.unlink(tempPath);
2126-
} catch {
2127-
// Ignore cleanup failures.
2168+
await fs.copyFile(path, `${path}.bak`);
2169+
} catch (backupError) {
2170+
log.warn("Failed to create flagged backup snapshot", {
2171+
path,
2172+
error: String(backupError),
2173+
});
21282174
}
2129-
log.error("Failed to save flagged account storage", {
2130-
path,
2131-
error: String(error),
2132-
});
2133-
throw error;
21342175
}
2176+
const content = JSON.stringify(normalizeFlaggedStorage(storage), null, 2);
2177+
await fs.writeFile(tempPath, content, { encoding: "utf-8", mode: 0o600 });
2178+
await fs.rename(tempPath, path);
2179+
try {
2180+
await fs.unlink(markerPath);
2181+
} catch {
2182+
// Best effort cleanup.
2183+
}
2184+
} catch (error) {
2185+
try {
2186+
await fs.unlink(tempPath);
2187+
} catch {
2188+
// Ignore cleanup failures.
2189+
}
2190+
log.error("Failed to save flagged account storage", {
2191+
path,
2192+
error: String(error),
2193+
});
2194+
throw error;
2195+
}
2196+
}
2197+
2198+
export async function saveFlaggedAccounts(
2199+
storage: FlaggedAccountStorageV1,
2200+
): Promise<void> {
2201+
return withStorageLock(async () => {
2202+
await saveFlaggedAccountsUnlocked(storage);
21352203
});
21362204
}
21372205

‎test/codex-manager-cli.test.ts‎

Lines changed: 98 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ const promptQuestionMock = vi.fn();
2626
const detectOcChatgptMultiAuthTargetMock = vi.fn();
2727
const normalizeAccountStorageMock = vi.fn((value) => value);
2828
const withAccountStorageTransactionMock = vi.fn();
29+
const withAccountAndFlaggedStorageTransactionMock = vi.fn();
2930

3031
vi.mock("../lib/logger.js", () => ({
3132
createLogger: vi.fn(() => ({
@@ -94,6 +95,8 @@ vi.mock("../lib/storage.js", async () => {
9495
loadFlaggedAccounts: loadFlaggedAccountsMock,
9596
saveAccounts: saveAccountsMock,
9697
saveFlaggedAccounts: saveFlaggedAccountsMock,
98+
withAccountAndFlaggedStorageTransaction:
99+
withAccountAndFlaggedStorageTransactionMock,
97100
withAccountStorageTransaction: withAccountStorageTransactionMock,
98101
setStoragePath: setStoragePathMock,
99102
getStoragePath: getStoragePathMock,
@@ -412,6 +415,7 @@ describe("codex manager cli commands", () => {
412415
loadFlaggedAccountsMock.mockReset();
413416
saveAccountsMock.mockReset();
414417
saveFlaggedAccountsMock.mockReset();
418+
withAccountAndFlaggedStorageTransactionMock.mockReset();
415419
withAccountStorageTransactionMock.mockReset();
416420
queuedRefreshMock.mockReset();
417421
setCodexCliActiveSelectionMock.mockReset();
@@ -455,6 +459,34 @@ describe("codex manager cli commands", () => {
455459
);
456460
},
457461
);
462+
withAccountAndFlaggedStorageTransactionMock.mockImplementation(
463+
async (handler) => {
464+
const current = await loadAccountsMock();
465+
let snapshot =
466+
current == null
467+
? {
468+
version: 3,
469+
accounts: [],
470+
activeIndex: 0,
471+
activeIndexByFamily: {},
472+
}
473+
: structuredClone(current);
474+
return handler(
475+
structuredClone(snapshot),
476+
async (storage: unknown, flaggedStorage: unknown) => {
477+
const previousSnapshot = structuredClone(snapshot);
478+
await saveAccountsMock(storage);
479+
try {
480+
await saveFlaggedAccountsMock(flaggedStorage);
481+
snapshot = structuredClone(storage);
482+
} catch (error) {
483+
await saveAccountsMock(previousSnapshot);
484+
throw error;
485+
}
486+
},
487+
);
488+
},
489+
);
458490
loadDashboardDisplaySettingsMock.mockResolvedValue({
459491
showPerAccountRows: true,
460492
showQuotaDetails: true,
@@ -642,7 +674,7 @@ describe("codex manager cli commands", () => {
642674
]);
643675

644676
expect(exitCode).toBe(0);
645-
expect(withAccountStorageTransactionMock).toHaveBeenCalledTimes(1);
677+
expect(withAccountAndFlaggedStorageTransactionMock).toHaveBeenCalledTimes(1);
646678
expect(saveAccountsMock).toHaveBeenCalledWith(
647679
expect.objectContaining({
648680
accounts: expect.arrayContaining([
@@ -699,7 +731,7 @@ describe("codex manager cli commands", () => {
699731
]);
700732

701733
expect(exitCode).toBe(0);
702-
expect(withAccountStorageTransactionMock).toHaveBeenCalledTimes(1);
734+
expect(withAccountAndFlaggedStorageTransactionMock).toHaveBeenCalledTimes(1);
703735
const savedStorage = saveAccountsMock.mock.calls.at(-1)?.[0];
704736
expect(savedStorage).toEqual(
705737
expect.objectContaining({
@@ -715,6 +747,70 @@ describe("codex manager cli commands", () => {
715747
extractAccountIdMock.mockImplementation(() => "acc_test");
716748
});
717749

750+
it("rolls back active storage when flagged persistence fails during recovery", async () => {
751+
const now = Date.now();
752+
loadFlaggedAccountsMock.mockResolvedValueOnce({
753+
version: 1,
754+
accounts: [
755+
{
756+
refreshToken: "flagged-refresh",
757+
accountId: "acc_flagged",
758+
email: "flagged@example.com",
759+
addedAt: now - 1_000,
760+
lastUsed: now - 1_000,
761+
flaggedAt: now - 5_000,
762+
},
763+
],
764+
});
765+
loadAccountsMock.mockResolvedValueOnce({
766+
version: 3,
767+
activeIndex: 0,
768+
activeIndexByFamily: { codex: 0 },
769+
accounts: [
770+
{
771+
refreshToken: "refresh-existing",
772+
accountId: "acc_existing",
773+
email: "existing@example.com",
774+
addedAt: now - 10_000,
775+
lastUsed: now - 10_000,
776+
},
777+
],
778+
});
779+
queuedRefreshMock.mockResolvedValueOnce({
780+
type: "success",
781+
access: "access-restored",
782+
refresh: "refresh-restored",
783+
expires: now + 3_600_000,
784+
});
785+
saveFlaggedAccountsMock.mockRejectedValueOnce(
786+
new Error("flagged write failed"),
787+
);
788+
const { runCodexMultiAuthCli } = await import("../lib/codex-manager.js");
789+
790+
await expect(
791+
runCodexMultiAuthCli(["auth", "verify-flagged", "--json"]),
792+
).rejects.toThrow("flagged write failed");
793+
expect(withAccountAndFlaggedStorageTransactionMock).toHaveBeenCalledTimes(
794+
1,
795+
);
796+
expect(saveAccountsMock).toHaveBeenCalledTimes(2);
797+
expect(saveAccountsMock.mock.calls[0]?.[0]).toEqual(
798+
expect.objectContaining({
799+
accounts: expect.arrayContaining([
800+
expect.objectContaining({ refreshToken: "refresh-existing" }),
801+
expect.objectContaining({ refreshToken: "refresh-restored" }),
802+
]),
803+
}),
804+
);
805+
expect(saveAccountsMock.mock.calls[1]?.[0]).toEqual(
806+
expect.objectContaining({
807+
accounts: [
808+
expect.objectContaining({ refreshToken: "refresh-existing" }),
809+
],
810+
}),
811+
);
812+
});
813+
718814
it("keeps flagged account when verification still fails", async () => {
719815
const now = Date.now();
720816
loadFlaggedAccountsMock.mockResolvedValueOnce({

0 commit comments

Comments
 (0)