fix: remap stale transaction selections

Author: ndycode

lib/storage.ts

@@ -1942,7 +1942,7 @@ function setTransactionSnapshotRevision(
 		value: revision,
 		writable: true,
 		configurable: true,
-		enumerable: true,
+		enumerable: false,
 	});
 	return storage;
 }
@@ -1965,11 +1965,44 @@ function mergeStaleTransactionSnapshot(
 	latest: AccountStorageV3,
 	stale: AccountStorageV3,
 ): AccountStorageV3 {
+	const mergedAccounts = deduplicateAccounts([...latest.accounts, ...stale.accounts]);
+	const staleActiveIndex = clampIndex(stale.activeIndex, stale.accounts.length);
+	const staleActiveRef = extractActiveAccountRef(stale.accounts, staleActiveIndex);
+	const activeIndex = resolveAccountSelectionIndex(
+		mergedAccounts,
+		{
+			accountId: staleActiveRef.accountId,
+			email: staleActiveRef.emailKey,
+			refreshToken: staleActiveRef.refreshToken,
+		},
+		staleActiveIndex,
+	);
+	const activeIndexByFamily: Partial<Record<ModelFamily, number>> = {};
+	for (const family of MODEL_FAMILIES) {
+		const rawFamilyIndex = stale.activeIndexByFamily?.[family];
+		const staleFamilyIndex =
+			typeof rawFamilyIndex === "number" && Number.isFinite(rawFamilyIndex)
+				? clampIndex(rawFamilyIndex, stale.accounts.length)
+				: staleActiveIndex;
+		const staleFamilyRef = extractActiveAccountRef(
+			stale.accounts,
+			staleFamilyIndex,
+		);
+		activeIndexByFamily[family] = resolveAccountSelectionIndex(
+			mergedAccounts,
+			{
+				accountId: staleFamilyRef.accountId,
+				email: staleFamilyRef.emailKey,
+				refreshToken: staleFamilyRef.refreshToken,
+			},
+			staleFamilyIndex,
+		);
+	}
 	const merged = normalizeAccountStorage({
 		version: 3,
-		accounts: deduplicateAccounts([...latest.accounts, ...stale.accounts]),
-		activeIndex: stale.activeIndex,
-		activeIndexByFamily: stale.activeIndexByFamily,
+		accounts: mergedAccounts,
+		activeIndex,
+		activeIndexByFamily,
 	});
 	return cloneAccountStorageForPersistence(merged ?? stale);
 }
@@ -2542,8 +2575,8 @@ export async function importAccounts(
 			// Nested imports write immediately to avoid re-entering the storage
 			// mutex. Later persist() calls compare the payload's snapshot revision
 			// against the live transaction state so stale payloads keep these
-			// imported accounts, while payloads built from getCurrent() can still
-			// intentionally remove them.
+			// imported accounts, while handlers that mutate the live getCurrent()
+			// snapshot in place can still intentionally remove them.
 			return applyImport(transactionState.snapshot, async (storage) => {
 				const nextStorage = cloneAccountStorageForPersistence(storage);
 				transactionState.revision += 1;

test/storage.test.ts

@@ -718,6 +718,8 @@ describe("storage", () => {
 
 				const stalePersistPayload = {
 					...current,
+					activeIndex: 1,
+					activeIndexByFamily: { codex: 1 },
 					accounts: [
 						...current.accounts,
 						{
@@ -751,6 +753,18 @@ describe("storage", () => {
 					}),
 				]),
 			);
+			expect(loaded?.accounts[loaded.activeIndex]).toEqual(
+				expect.objectContaining({
+					accountId: "acct-appended",
+					refreshToken: "refresh-appended",
+				}),
+			);
+			expect(loaded?.accounts[loaded.activeIndexByFamily?.codex ?? -1]).toEqual(
+				expect.objectContaining({
+					accountId: "acct-appended",
+					refreshToken: "refresh-appended",
+				}),
+			);
 		});
 
 		it("allows a live transaction payload to intentionally remove imported accounts", async () => {
@@ -798,13 +812,10 @@ describe("storage", () => {
 				if (!liveCurrent) {
 					throw new Error("expected live transaction snapshot");
 				}
-
-				await persist({
-					...liveCurrent,
-					accounts: liveCurrent.accounts.filter(
-						(account) => account.accountId !== "acct-imported",
-					),
-				});
+				liveCurrent.accounts = liveCurrent.accounts.filter(
+					(account) => account.accountId !== "acct-imported",
+				);
+				await persist(liveCurrent);
 			});
 
 			const loaded = await loadAccounts();
@@ -859,13 +870,10 @@ describe("storage", () => {
 				if (!liveCurrent) {
 					throw new Error("expected live transaction snapshot");
 				}
-
-				await persist({
-					...liveCurrent,
-					accounts: liveCurrent.accounts.filter(
-						(account) => account.accountId !== "acct-imported",
-					),
-				});
+				liveCurrent.accounts = liveCurrent.accounts.filter(
+					(account) => account.accountId !== "acct-imported",
+				);
+				await persist(liveCurrent);
 			});
 
 			const loaded = await loadAccounts();
@@ -999,6 +1007,118 @@ describe("storage", () => {
 			);
 		});
 
+		it("getCurrent exposes the live snapshot after persist inside a flagged transaction", async () => {
+			const now = Date.now();
+			let liveSnapshot: Awaited<ReturnType<typeof loadAccounts>> = null;
+
+			await saveAccounts({
+				version: 3,
+				activeIndex: 0,
+				activeIndexByFamily: { codex: 0 },
+				accounts: [
+					{
+						accountId: "acct-existing",
+						email: "existing@example.com",
+						refreshToken: "refresh-existing",
+						addedAt: now - 1_000,
+						lastUsed: now - 1_000,
+					},
+				],
+			});
+
+			await withAccountAndFlaggedStorageTransaction(
+				async (current, persist, getCurrent) => {
+					if (!current) {
+						throw new Error("expected existing account storage");
+					}
+					expect(getCurrent()).toBe(current);
+
+					await persist(
+						{
+							...current,
+							accounts: [
+								...current.accounts,
+								{
+									accountId: "acct-added",
+									email: "added@example.com",
+									refreshToken: "refresh-added",
+									addedAt: now,
+									lastUsed: now,
+								},
+							],
+						},
+						{
+							version: 1,
+							accounts: [],
+						},
+					);
+
+					liveSnapshot = getCurrent();
+				},
+			);
+
+			expect(liveSnapshot?.accounts).toEqual(
+				expect.arrayContaining([
+					expect.objectContaining({
+						accountId: "acct-existing",
+						refreshToken: "refresh-existing",
+					}),
+					expect.objectContaining({
+						accountId: "acct-added",
+						refreshToken: "refresh-added",
+					}),
+				]),
+			);
+		});
+
+		it("getCurrent starts null and becomes live after persist inside a flagged transaction", async () => {
+			const storagePath = getStoragePath();
+			const now = Date.now();
+			let liveSnapshot: Awaited<ReturnType<typeof loadAccounts>> = null;
+
+			await fs.mkdir(dirname(storagePath), { recursive: true });
+			await fs.writeFile(storagePath, "{invalid-json", "utf-8");
+
+			await withAccountAndFlaggedStorageTransaction(
+				async (current, persist, getCurrent) => {
+					expect(current).toBeNull();
+					expect(getCurrent()).toBeNull();
+
+					await persist(
+						{
+							version: 3,
+							activeIndex: 0,
+							activeIndexByFamily: { codex: 0 },
+							accounts: [
+								{
+									accountId: "acct-created",
+									email: "created@example.com",
+									refreshToken: "refresh-created",
+									addedAt: now,
+									lastUsed: now,
+								},
+							],
+						},
+						{
+							version: 1,
+							accounts: [],
+						},
+					);
+
+					liveSnapshot = getCurrent();
+				},
+			);
+
+			expect(liveSnapshot?.accounts).toEqual(
+				expect.arrayContaining([
+					expect.objectContaining({
+						accountId: "acct-created",
+						refreshToken: "refresh-created",
+					}),
+				]),
+			);
+		});
+
 		it("should preserve distinct shared-accountId imports when the imported row has no email", async () => {
 			const existing = {
 				version: 3,