Skip to content

Commit 0fd695e

Browse files
ndycodendycode
committed
fix(storage): snapshot exact pre-mutation state
1 parent 105730c commit 0fd695e

4 files changed

Lines changed: 183 additions & 68 deletions

File tree

‎lib/destructive-actions.ts‎

Lines changed: 51 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@ import {
77
clearAccounts,
88
clearFlaggedAccounts,
99
type FlaggedAccountStorageV1,
10+
getStoragePath,
1011
loadFlaggedAccounts,
11-
saveAccounts,
12-
saveFlaggedAccounts,
1312
snapshotAccountStorage,
13+
withAccountAndFlaggedStorageTransaction,
1414
} from "./storage.js";
1515

1616
export const DESTRUCTIVE_ACTION_COPY = {
@@ -99,74 +99,62 @@ export interface DestructiveActionResult {
9999
quotaCacheCleared: boolean;
100100
}
101101

102-
function asError(error: unknown, fallbackMessage: string): Error {
103-
return error instanceof Error
104-
? error
105-
: new Error(`${fallbackMessage}: ${String(error)}`);
106-
}
107-
108102
export async function deleteAccountAtIndex(options: {
109103
storage: AccountStorageV3;
110104
index: number;
111105
}): Promise<DeleteAccountResult | null> {
112-
const target = options.storage.accounts.at(options.index);
113-
if (!target) return null;
114-
const flagged = await loadFlaggedAccounts();
115-
await snapshotAccountStorage({ reason: "delete-account" });
116-
const nextStorage: AccountStorageV3 = {
117-
...options.storage,
118-
accounts: options.storage.accounts.map((account) => ({ ...account })),
119-
activeIndexByFamily: { ...(options.storage.activeIndexByFamily ?? {}) },
120-
};
121-
const previousStorage: AccountStorageV3 = {
122-
...options.storage,
123-
accounts: options.storage.accounts.map((account) => ({ ...account })),
124-
activeIndexByFamily: { ...(options.storage.activeIndexByFamily ?? {}) },
125-
};
106+
const requestedTarget = options.storage.accounts.at(options.index);
107+
if (!requestedTarget) return null;
126108

127-
nextStorage.accounts.splice(options.index, 1);
128-
rebaseActiveIndicesAfterDelete(nextStorage, options.index);
129-
clampActiveIndices(nextStorage);
130-
await saveAccounts(nextStorage);
131-
132-
const remainingFlagged = flagged.accounts.filter(
133-
(account) => account.refreshToken !== target.refreshToken,
134-
);
135-
const removedFlaggedCount = flagged.accounts.length - remainingFlagged.length;
136-
let updatedFlagged = flagged;
137-
if (removedFlaggedCount > 0) {
138-
updatedFlagged = { ...flagged, accounts: remainingFlagged };
139-
try {
140-
await saveFlaggedAccounts(updatedFlagged);
141-
} catch (error) {
142-
const originalError = asError(
143-
error,
144-
"Failed to save flagged account storage after deleting an account",
145-
);
146-
try {
147-
await saveAccounts(previousStorage);
148-
} catch (rollbackError) {
149-
throw new AggregateError(
150-
[
151-
originalError,
152-
asError(
153-
rollbackError,
154-
"Failed to roll back account storage after flagged save failure",
155-
),
156-
],
157-
"Deleting the account partially failed and rollback also failed.",
158-
);
159-
}
160-
throw originalError;
109+
return withAccountAndFlaggedStorageTransaction(async (current, persist) => {
110+
const sourceStorage = current ?? options.storage;
111+
const targetIndex = sourceStorage.accounts.findIndex(
112+
(account) => account.refreshToken === requestedTarget.refreshToken,
113+
);
114+
if (targetIndex < 0) {
115+
return null;
116+
}
117+
const target = sourceStorage.accounts[targetIndex];
118+
if (!target) {
119+
return null;
161120
}
162-
}
163121

164-
return {
165-
storage: nextStorage,
166-
flagged: updatedFlagged,
167-
removedAccount: target,
168-
removedFlaggedCount,
169-
};
122+
const flagged = await loadFlaggedAccounts();
123+
await snapshotAccountStorage({
124+
reason: "delete-account",
125+
storage: sourceStorage,
126+
storagePath: getStoragePath(),
127+
});
128+
129+
const nextStorage: AccountStorageV3 = {
130+
...sourceStorage,
131+
accounts: sourceStorage.accounts.map((account) => ({ ...account })),
132+
activeIndexByFamily: { ...(sourceStorage.activeIndexByFamily ?? {}) },
133+
};
134+
135+
nextStorage.accounts.splice(targetIndex, 1);
136+
rebaseActiveIndicesAfterDelete(nextStorage, targetIndex);
137+
clampActiveIndices(nextStorage);
138+
139+
const remainingFlagged = flagged.accounts.filter(
140+
(account) => account.refreshToken !== target.refreshToken,
141+
);
142+
const removedFlaggedCount =
143+
flagged.accounts.length - remainingFlagged.length;
144+
const updatedFlagged =
145+
removedFlaggedCount > 0
146+
? { ...flagged, accounts: remainingFlagged }
147+
: flagged;
148+
149+
await persist(nextStorage, updatedFlagged);
150+
151+
return {
152+
storage: nextStorage,
153+
flagged: updatedFlagged,
154+
removedAccount: target,
155+
removedFlaggedCount,
156+
};
157+
});
170158
}
171159

172160
/**

‎lib/storage.ts‎

Lines changed: 47 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,8 @@ export interface AccountSnapshotOptions {
219219
force?: boolean;
220220
failurePolicy?: AccountSnapshotFailurePolicy;
221221
createBackup?: typeof createNamedBackup;
222+
storage?: AccountStorageV3 | null;
223+
storagePath?: string;
222224
}
223225

224226
interface LoadedBackupCandidate {
@@ -2171,6 +2173,37 @@ export async function createNamedBackup(
21712173
);
21722174
}
21732175

2176+
async function writeNamedBackupFromStorage(
2177+
name: string,
2178+
storage: AccountStorageV3,
2179+
options: {
2180+
force?: boolean;
2181+
storagePath?: string;
2182+
} = {},
2183+
): Promise<NamedBackupMetadata> {
2184+
const storagePath = options.storagePath ?? getStoragePath();
2185+
const backupPath = resolveNamedBackupPath(name, storagePath);
2186+
if (!options.force && existsSync(backupPath)) {
2187+
throw new Error(`File already exists: ${backupPath}`);
2188+
}
2189+
await fs.mkdir(dirname(backupPath), { recursive: true });
2190+
await fs.writeFile(
2191+
backupPath,
2192+
JSON.stringify(
2193+
{
2194+
version: storage.version,
2195+
accounts: storage.accounts,
2196+
activeIndex: storage.activeIndex,
2197+
activeIndexByFamily: storage.activeIndexByFamily,
2198+
},
2199+
null,
2200+
2,
2201+
),
2202+
{ encoding: "utf-8", mode: 0o600 },
2203+
);
2204+
return buildNamedBackupMetadata(name, backupPath);
2205+
}
2206+
21742207
function formatTimestampForSnapshot(timestamp: number): string {
21752208
const date = new Date(timestamp);
21762209
const pad = (value: number): string => value.toString().padStart(2, "0");
@@ -2226,14 +2259,22 @@ export async function snapshotAccountStorage(
22262259
force = true,
22272260
failurePolicy = "warn",
22282261
createBackup = createNamedBackup,
2262+
storagePath,
22292263
} = options;
2230-
const currentStorage = await loadAccounts();
2264+
const currentStorage =
2265+
options.storage !== undefined ? options.storage : await loadAccounts();
22312266
if (!currentStorage || currentStorage.accounts.length === 0) {
22322267
return null;
22332268
}
22342269

22352270
const backupName = buildAccountSnapshotName(reason, now);
22362271
try {
2272+
if (options.storage !== undefined) {
2273+
return await writeNamedBackupFromStorage(backupName, currentStorage, {
2274+
force,
2275+
storagePath,
2276+
});
2277+
}
22372278
return await createBackup(backupName, { force });
22382279
} catch (error) {
22392280
if (failurePolicy === "error") {
@@ -3631,15 +3672,16 @@ export async function importAccounts(
36313672
throw new Error("Invalid account storage format");
36323673
}
36333674

3634-
// Capture the current pool before merge/limit checks so failed imports still
3635-
// leave a rollback-ready checkpoint for the pre-import state.
3636-
await snapshotAccountStorage({ reason: "import-accounts" });
3637-
36383675
const {
36393676
imported: importedCount,
36403677
total,
36413678
skipped: skippedCount,
36423679
} = await withAccountStorageTransaction(async (existing, persist) => {
3680+
await snapshotAccountStorage({
3681+
reason: "import-accounts",
3682+
storage: existing,
3683+
storagePath: getStoragePath(),
3684+
});
36433685
const existingAccounts = existing?.accounts ?? [];
36443686
const existingActiveIndex = existing?.activeIndex ?? 0;
36453687

‎test/destructive-actions.test.ts‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ const loadFlaggedAccountsMock = vi.fn();
88
const saveAccountsMock = vi.fn();
99
const saveFlaggedAccountsMock = vi.fn();
1010
const snapshotAccountStorageMock = vi.fn();
11+
const withAccountAndFlaggedStorageTransactionMock = vi.fn();
12+
const getStoragePathMock = vi.fn(() => "/mock/openai-codex-accounts.json");
13+
let transactionCurrentStorage: unknown = null;
1114

1215
vi.mock("../lib/codex-cli/state.js", () => ({
1316
clearCodexCliStateCache: clearCodexCliStateCacheMock,
@@ -24,10 +27,13 @@ vi.mock("../lib/quota-cache.js", () => ({
2427
vi.mock("../lib/storage.js", () => ({
2528
clearAccounts: clearAccountsMock,
2629
clearFlaggedAccounts: clearFlaggedAccountsMock,
30+
getStoragePath: getStoragePathMock,
2731
loadFlaggedAccounts: loadFlaggedAccountsMock,
2832
saveAccounts: saveAccountsMock,
2933
saveFlaggedAccounts: saveFlaggedAccountsMock,
3034
snapshotAccountStorage: snapshotAccountStorageMock,
35+
withAccountAndFlaggedStorageTransaction:
36+
withAccountAndFlaggedStorageTransactionMock,
3137
}));
3238

3339
describe("destructive actions", () => {
@@ -41,6 +47,34 @@ describe("destructive actions", () => {
4147
saveAccountsMock.mockResolvedValue(undefined);
4248
saveFlaggedAccountsMock.mockResolvedValue(undefined);
4349
snapshotAccountStorageMock.mockResolvedValue(null);
50+
transactionCurrentStorage = null;
51+
withAccountAndFlaggedStorageTransactionMock.mockImplementation(
52+
async (handler) => {
53+
const previousSnapshot = structuredClone(transactionCurrentStorage);
54+
return handler(
55+
transactionCurrentStorage,
56+
async (accountStorage: unknown, flaggedStorage: unknown) => {
57+
await saveAccountsMock(accountStorage);
58+
try {
59+
await saveFlaggedAccountsMock(flaggedStorage);
60+
transactionCurrentStorage = structuredClone(accountStorage);
61+
} catch (error) {
62+
try {
63+
await saveAccountsMock(previousSnapshot);
64+
transactionCurrentStorage =
65+
structuredClone(previousSnapshot);
66+
} catch (rollbackError) {
67+
throw new AggregateError(
68+
[error, rollbackError],
69+
"Deleting the account partially failed and rollback also failed.",
70+
);
71+
}
72+
throw error;
73+
}
74+
},
75+
);
76+
},
77+
);
4478
});
4579

4680
it("returns delete-only results without pretending kept data was cleared", async () => {
@@ -135,12 +169,16 @@ describe("destructive actions", () => {
135169
},
136170
],
137171
};
172+
transactionCurrentStorage = structuredClone(storage);
138173

139174
const deleted = await deleteAccountAtIndex({ storage, index: 0 });
140175

141176
expect(deleted).not.toBeNull();
177+
expect(withAccountAndFlaggedStorageTransactionMock).toHaveBeenCalledTimes(1);
142178
expect(snapshotAccountStorageMock).toHaveBeenCalledWith({
143179
reason: "delete-account",
180+
storage,
181+
storagePath: "/mock/openai-codex-accounts.json",
144182
});
145183
expect(
146184
snapshotAccountStorageMock.mock.invocationCallOrder[0],
@@ -204,10 +242,12 @@ describe("destructive actions", () => {
204242
},
205243
],
206244
};
245+
transactionCurrentStorage = structuredClone(storage);
207246

208247
const deleted = await deleteAccountAtIndex({ storage, index: 1 });
209248

210249
expect(deleted).not.toBeNull();
250+
expect(withAccountAndFlaggedStorageTransactionMock).toHaveBeenCalledTimes(1);
211251
expect(deleted?.flagged.accounts).toEqual([
212252
expect.objectContaining({ refreshToken: "refresh-newer" }),
213253
]);
@@ -255,6 +295,7 @@ describe("destructive actions", () => {
255295
},
256296
],
257297
};
298+
transactionCurrentStorage = structuredClone(storage);
258299

259300
await expect(deleteAccountAtIndex({ storage, index: 1 })).rejects.toBe(
260301
flaggedSaveError,
@@ -307,6 +348,7 @@ describe("destructive actions", () => {
307348
},
308349
],
309350
};
351+
transactionCurrentStorage = structuredClone(storage);
310352

311353
try {
312354
await deleteAccountAtIndex({ storage, index: 1 });

‎test/storage.test.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2210,6 +2210,49 @@ describe("storage", () => {
22102210
).rejects.toThrow(/snapshot failed/);
22112211
});
22122212

2213+
it("writes the provided storage snapshot instead of re-reading live storage", async () => {
2214+
await saveAccounts({
2215+
version: 3,
2216+
activeIndex: 0,
2217+
accounts: [
2218+
{
2219+
accountId: "live",
2220+
refreshToken: "ref-live",
2221+
addedAt: 1,
2222+
lastUsed: 1,
2223+
},
2224+
],
2225+
});
2226+
2227+
const snapshot = await snapshotAccountStorage({
2228+
reason: "import-accounts",
2229+
storage: {
2230+
version: 3,
2231+
activeIndex: 0,
2232+
accounts: [
2233+
{
2234+
accountId: "provided",
2235+
refreshToken: "ref-provided",
2236+
addedAt: 2,
2237+
lastUsed: 2,
2238+
},
2239+
],
2240+
},
2241+
storagePath: testStoragePath,
2242+
});
2243+
2244+
expect(snapshot?.path && existsSync(snapshot.path)).toBe(true);
2245+
const snapshotContent = JSON.parse(
2246+
await fs.readFile(snapshot!.path, "utf-8"),
2247+
) as { accounts: Array<{ accountId?: string }> };
2248+
expect(snapshotContent.accounts).toEqual([
2249+
expect.objectContaining({ accountId: "provided" }),
2250+
]);
2251+
expect((await loadAccounts())?.accounts).toEqual([
2252+
expect.objectContaining({ accountId: "live" }),
2253+
]);
2254+
});
2255+
22132256
it("creates a named snapshot before deleteSavedAccounts", async () => {
22142257
await saveAccounts({
22152258
version: 3,

0 commit comments

Comments
 (0)