fix-runtime-recovery-token-restore-safety

ndycode · ndycode · commit bafd2c8cce1a · 2026-04-01T23:29:45.000+08:00
diff --git a/lib/runtime/account-check.ts b/lib/runtime/account-check.ts
@@ -65,6 +65,13 @@ export async function runRuntimeAccountCheck(
 			version: 1;
 			accounts: FlaggedAccountMetadataV1[];
 		}) => Promise<void>;
+		persistAccountAndFlaggedStorage?: (
+			accountStorage: AccountStorageV3,
+			flaggedStorage: {
+				version: 1;
+				accounts: FlaggedAccountMetadataV1[];
+			},
+		) => Promise<void>;
 		now?: () => number;
 		showLine: (message: string) => void;
 	},
@@ -145,14 +152,6 @@ export async function runRuntimeAccountCheck(
 				) {
 					accessToken = cached.accessToken;
 					authDetail = "OK (Codex CLI cache)";
-
-					if (
-						cached.refreshToken &&
-						cached.refreshToken !== account.refreshToken
-					) {
-						account.refreshToken = cached.refreshToken;
-						state.storageChanged = true;
-					}
 					if (
 						cached.accessToken &&
 						cached.accessToken !== account.accessToken
@@ -318,12 +317,23 @@ export async function runRuntimeAccountCheck(
 		state.storageChanged = true;
 	}
 
-	if (state.flaggedChanged) {
-		await deps.saveFlaggedAccounts(state.flaggedStorage);
-	}
-	if (state.storageChanged) {
-		await deps.saveAccounts(workingStorage);
+	const persistAccountAndFlaggedStorage =
+		deps.persistAccountAndFlaggedStorage;
+	const canPersistTogether =
+		state.flaggedChanged &&
+		state.storageChanged &&
+		typeof persistAccountAndFlaggedStorage === "function";
+	if (canPersistTogether) {
+		await persistAccountAndFlaggedStorage(workingStorage, state.flaggedStorage);
 		deps.invalidateAccountManagerCache();
+	} else {
+		if (state.flaggedChanged) {
+		await deps.saveFlaggedAccounts(state.flaggedStorage);
+		}
+		if (state.storageChanged) {
+			await deps.saveAccounts(workingStorage);
+			deps.invalidateAccountManagerCache();
+		}
 	}
 
 	deps.showLine("");
diff --git a/lib/runtime/verify-flagged.ts b/lib/runtime/verify-flagged.ts
@@ -43,6 +43,14 @@ export async function verifyRuntimeFlaggedAccounts(deps: {
 		results: Array<TokenSuccessWithAccount<SuccessfulAccountTokens>>,
 		replaceAll?: boolean,
 	) => Promise<void>;
+	persistAccountsAndFlagged?: (
+		results: Array<TokenSuccessWithAccount<SuccessfulAccountTokens>>,
+		flaggedStorage: {
+			version: 1;
+			accounts: FlaggedAccountMetadataV1[];
+		},
+		replaceAll?: boolean,
+	) => Promise<void>;
 	invalidateAccountManagerCache: () => void;
 	saveFlaggedAccounts: (storage: {
 		version: 1;
@@ -78,7 +86,8 @@ export async function verifyRuntimeFlaggedAccounts(deps: {
 				const refreshToken =
 					typeof cached.refreshToken === "string" && cached.refreshToken.trim()
 						? cached.refreshToken.trim()
-						: flagged.refreshToken;
+						: undefined;
+				if (refreshToken) {
 				const resolved = deps.resolveTokenSuccessAccount({
 					type: "success",
 					access: cached.accessToken,
@@ -98,6 +107,7 @@ export async function verifyRuntimeFlaggedAccounts(deps: {
 					`[${i + 1}/${flaggedStorage.accounts.length}] ${label}: RESTORED (Codex CLI cache)`,
 				);
 				continue;
+				}
 			}
 
 			const refreshResult = await deps.queuedRefresh(flagged.refreshToken);
@@ -134,16 +144,21 @@ export async function verifyRuntimeFlaggedAccounts(deps: {
 		}
 	}
 
-	if (state.restored.length > 0) {
-		await deps.persistAccounts(state.restored, false);
+	const nextFlaggedStorage = {
+		version: 1 as const,
+		accounts: state.remaining,
+	};
+	if (state.restored.length > 0 && deps.persistAccountsAndFlagged) {
+		await deps.persistAccountsAndFlagged(state.restored, nextFlaggedStorage, false);
 		deps.invalidateAccountManagerCache();
+	} else {
+		if (state.restored.length > 0) {
+			await deps.persistAccounts(state.restored, false);
+			deps.invalidateAccountManagerCache();
+		}
+		await deps.saveFlaggedAccounts(nextFlaggedStorage);
 	}
 
-	await deps.saveFlaggedAccounts({
-		version: 1,
-		accounts: state.remaining,
-	});
-
 	deps.showLine("");
 	deps.showLine(
 		`Results: ${state.restored.length} restored, ${state.remaining.length} still flagged`,
diff --git a/test/runtime-account-check.test.ts b/test/runtime-account-check.test.ts
@@ -207,14 +207,62 @@ describe("runRuntimeAccountCheck", () => {
 		const saved = saveAccounts.mock.calls[0]?.[0];
 		expect(saved.accounts[0]).toMatchObject({
 			email: "fresh@example.com",
-			refreshToken: "fresh-refresh",
+			refreshToken: "stale-refresh",
 			accessToken: "cached-access",
 			expiresAt: 70_000,
 			accountId: "new-account",
 			accountIdSource: "token",
 		});
 	});
 
+	it("uses combined persistence when flagging an invalid account removes it from active storage", async () => {
+		const persistAccountAndFlaggedStorage = vi.fn(async () => {});
+		const saveAccounts = vi.fn(async () => {});
+		const saveFlaggedAccounts = vi.fn(async () => {});
+
+		await runRuntimeAccountCheck(true, {
+			hydrateEmails: async (storage) => storage,
+			loadAccounts: async () => ({
+				version: 3,
+				accounts: [{ email: "one@example.com", refreshToken: "refresh-1", accessToken: undefined, addedAt: 1, lastUsed: 1 }],
+				activeIndex: 0,
+				activeIndexByFamily: { codex: 0 },
+			}),
+			createEmptyStorage: () => ({ version: 3, accounts: [], activeIndex: 0, activeIndexByFamily: {} }),
+			loadFlaggedAccounts: async () => ({ version: 1, accounts: [] }),
+			createAccountCheckWorkingState: (flaggedStorage) => ({ flaggedStorage, removeFromActive: new Set(), storageChanged: false, flaggedChanged: false, ok: 0, errors: 0, disabled: 0 }),
+			lookupCodexCliTokensByEmail: async () => null,
+			extractAccountId: () => undefined,
+			shouldUpdateAccountIdFromToken: () => false,
+			sanitizeEmail: (email) => email,
+			extractAccountEmail: () => undefined,
+			queuedRefresh: async () => ({ type: "failed", reason: "invalid_grant", message: "refresh failed" }),
+			isRuntimeFlaggableFailure: () => true,
+			fetchCodexQuotaSnapshot: async () => { throw new Error("should not probe quota in deep mode"); },
+			resolveRequestAccountId: () => undefined,
+			formatCodexQuotaLine: () => "quota",
+			clampRuntimeActiveIndices: vi.fn(),
+			MODEL_FAMILIES: ["codex"],
+			saveAccounts,
+			invalidateAccountManagerCache: vi.fn(),
+			saveFlaggedAccounts,
+			persistAccountAndFlaggedStorage,
+			showLine: vi.fn(),
+		});
+
+		expect(persistAccountAndFlaggedStorage).toHaveBeenCalledWith(
+			expect.objectContaining({
+				accounts: [],
+			}),
+			expect.objectContaining({
+				version: 1,
+				accounts: [expect.objectContaining({ refreshToken: "refresh-1" })],
+			}),
+		);
+		expect(saveAccounts).not.toHaveBeenCalled();
+		expect(saveFlaggedAccounts).not.toHaveBeenCalled();
+	});
+
 	it("treats cache lookup failures as a cache miss and still refreshes", async () => {
 		const saveAccounts = vi.fn(async () => {});
 		await runRuntimeAccountCheck(false, {
diff --git a/test/runtime-verify-flagged.test.ts b/test/runtime-verify-flagged.test.ts
@@ -94,6 +94,53 @@ describe("verifyRuntimeFlaggedAccounts", () => {
 		);
 	});
 
+	it("falls back to a real refresh when cached access has no refresh token", async () => {
+		const now = 10_000;
+		const queuedRefresh = vi.fn(async () => ({
+			type: "success" as const,
+			access: "new-access",
+			refresh: "new-refresh",
+			expires: now + 60_000,
+		}));
+		const persistAccounts = vi.fn(async () => {});
+
+		await verifyRuntimeFlaggedAccounts({
+			loadFlaggedAccounts: async () => ({
+				version: 1,
+				accounts: [
+					{
+						email: "refresh@example.com",
+						refreshToken: "flagged-refresh",
+						addedAt: 1,
+						lastUsed: 1,
+					},
+				],
+			}),
+			lookupCodexCliTokensByEmail: async () => ({
+				accessToken: "cached-access",
+				expiresAt: now + 60_000,
+			}),
+			queuedRefresh,
+			resolveTokenSuccessAccount: (tokens) => ({ ...tokens }) as never,
+			persistAccounts,
+			invalidateAccountManagerCache: vi.fn(),
+			saveFlaggedAccounts: vi.fn(async () => {}),
+			now: () => now,
+			showLine: vi.fn(),
+		});
+
+		expect(queuedRefresh).toHaveBeenCalledWith("flagged-refresh");
+		expect(persistAccounts).toHaveBeenCalledWith(
+			[
+				expect.objectContaining({
+					access: "new-access",
+					refresh: "new-refresh",
+				}),
+			],
+			false,
+		);
+	});
+
 	it("restores accounts after a successful refresh when cache misses", async () => {
 		const now = 10_000;
 		const persistAccounts = vi.fn(async () => {});
@@ -245,6 +292,59 @@ describe("verifyRuntimeFlaggedAccounts", () => {
 		expect(calls).toEqual(["persistAccounts", "saveFlaggedAccounts"]);
 	});
 
+	it("uses combined persistence when restored accounts and flagged cleanup must commit together", async () => {
+		const now = 10_000;
+		const persistAccountsAndFlagged = vi.fn(async () => {});
+		const persistAccounts = vi.fn(async () => {});
+		const saveFlaggedAccounts = vi.fn(async () => {});
+
+		await verifyRuntimeFlaggedAccounts({
+			loadFlaggedAccounts: async () => ({
+				version: 1,
+				accounts: [
+					{
+						email: "refresh@example.com",
+						refreshToken: "refresh-token",
+						addedAt: 1,
+						lastUsed: 1,
+					},
+				],
+			}),
+			lookupCodexCliTokensByEmail: async () => {
+				throw new Error("busy");
+			},
+			queuedRefresh: async () => ({
+				type: "success" as const,
+				access: "new-access",
+				refresh: "new-refresh",
+				expires: now + 60_000,
+			}),
+			resolveTokenSuccessAccount: (tokens) => ({ ...tokens }) as never,
+			persistAccounts,
+			persistAccountsAndFlagged,
+			invalidateAccountManagerCache: vi.fn(),
+			saveFlaggedAccounts,
+			now: () => now,
+			showLine: vi.fn(),
+		});
+
+		expect(persistAccountsAndFlagged).toHaveBeenCalledWith(
+			[
+				expect.objectContaining({
+					access: "new-access",
+					refresh: "new-refresh",
+				}),
+			],
+			{
+				version: 1,
+				accounts: [],
+			},
+			false,
+		);
+		expect(persistAccounts).not.toHaveBeenCalled();
+		expect(saveFlaggedAccounts).not.toHaveBeenCalled();
+	});
+
 	it("leaves flagged state untouched when persistAccounts throws EBUSY", async () => {
 		const now = 10_000;
 		const saveFlaggedAccounts = vi.fn(async () => {});
