Follow-up: address PR #46 review findings

lib/accounts.ts

@@ -74,7 +74,7 @@ import {
 
 const log = createLogger("accounts");
 type StoredAccount = AccountStorageV3["accounts"][number];
-const DISK_PREFERRED_MERGE_KEYS = new Set(["refreshToken", "accessToken", "expiresAt"]);
+const CREDENTIAL_MERGE_KEYS = new Set(["refreshToken", "accessToken", "expiresAt"]);
 
 function initFamilyState(defaultValue: number): Record<ModelFamily, number> {
 	return Object.fromEntries(
@@ -854,13 +854,18 @@ export class AccountManager {
 	private mergeStoredAccountRecords(current: StoredAccount, incoming: StoredAccount): StoredAccount {
 		const next: StoredAccount = { ...current };
 		const nextRecord = next as unknown as Record<string, unknown>;
+		const preferredCredentialSource = this.selectCredentialMergeSource(current, incoming);
 		for (const [rawKey, rawValue] of Object.entries(incoming)) {
 			const value = rawValue as unknown;
 			if (value === undefined) {
 				continue;
 			}
 			const currentValue = nextRecord[rawKey];
-			if (DISK_PREFERRED_MERGE_KEYS.has(rawKey) && currentValue !== undefined) {
+			if (
+				CREDENTIAL_MERGE_KEYS.has(rawKey) &&
+				currentValue !== undefined &&
+				preferredCredentialSource === "current"
+			) {
 				continue;
 			}
 			if (
@@ -899,6 +904,27 @@ export class AccountManager {
 		return next;
 	}
 
+	private selectCredentialMergeSource(
+		current: StoredAccount,
+		incoming: StoredAccount,
+	): "current" | "incoming" {
+		const currentExpiresAt =
+			typeof current.expiresAt === "number" ? current.expiresAt : undefined;
+		const incomingExpiresAt =
+			typeof incoming.expiresAt === "number" ? incoming.expiresAt : undefined;
+
+		if (incomingExpiresAt === undefined && currentExpiresAt === undefined) {
+			return "incoming";
+		}
+		if (incomingExpiresAt === undefined) {
+			return "current";
+		}
+		if (currentExpiresAt === undefined) {
+			return "incoming";
+		}
+		return incomingExpiresAt > currentExpiresAt ? "incoming" : "current";
+	}
+
 	private applyPersistedStorageSnapshot(storage: AccountStorageV3): void {
 		const previousByRefreshToken = new Map(
 			this.accounts.map((account) => [account.refreshToken, account] as const),

lib/storage.ts

@@ -2092,7 +2092,7 @@ function cloneFlaggedStorageForPersist(storage: FlaggedAccountStorageV1): Flagge
 	};
 }
 
-export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
+async function loadFlaggedAccountsUnlocked(): Promise<FlaggedAccountStorageV1> {
 	const path = getFlaggedAccountsPath();
 	const empty: FlaggedAccountStorageV1 = { version: 1, accounts: [] };
 
@@ -2121,7 +2121,7 @@ export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
 		const legacyData = JSON.parse(legacyContent) as unknown;
 		const migrated = normalizeFlaggedStorage(legacyData);
 		if (migrated.accounts.length > 0) {
-			await saveFlaggedAccounts(migrated);
+			await saveFlaggedAccountsUnlocked(migrated);
 		}
 		try {
 			await fs.unlink(legacyPath);
@@ -2148,28 +2148,34 @@ export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
 	}
 }
 
-export async function saveFlaggedAccounts(storage: FlaggedAccountStorageV1): Promise<void> {
-	return withStorageLock(async () => {
-		const path = getFlaggedAccountsPath();
-		const uniqueSuffix = `${Date.now()}.${Math.random().toString(36).slice(2, 8)}`;
-		const tempPath = `${path}.${uniqueSuffix}.tmp`;
+async function saveFlaggedAccountsUnlocked(storage: FlaggedAccountStorageV1): Promise<void> {
+	const path = getFlaggedAccountsPath();
+	const uniqueSuffix = `${Date.now()}.${Math.random().toString(36).slice(2, 8)}`;
+	const tempPath = `${path}.${uniqueSuffix}.tmp`;
 
+	try {
+		await fs.mkdir(dirname(path), { recursive: true });
+		const normalized = normalizeFlaggedStorage(storage);
+		const content = JSON.stringify(cloneFlaggedStorageForPersist(normalized), null, 2);
+		await fs.writeFile(tempPath, content, { encoding: "utf-8", mode: 0o600 });
+		await fs.rename(tempPath, path);
+	} catch (error) {
 		try {
-			await fs.mkdir(dirname(path), { recursive: true });
-			const normalized = normalizeFlaggedStorage(storage);
-			const content = JSON.stringify(cloneFlaggedStorageForPersist(normalized), null, 2);
-			await fs.writeFile(tempPath, content, { encoding: "utf-8", mode: 0o600 });
-			await fs.rename(tempPath, path);
-		} catch (error) {
-			try {
-				await fs.unlink(tempPath);
-			} catch {
-				// Ignore cleanup failures.
-			}
-			log.error("Failed to save flagged account storage", { path, error: String(error) });
-			throw error;
+			await fs.unlink(tempPath);
+		} catch {
+			// Ignore cleanup failures.
 		}
-	});
+		log.error("Failed to save flagged account storage", { path, error: String(error) });
+		throw error;
+	}
+}
+
+export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
+	return loadFlaggedAccountsUnlocked();
+}
+
+export async function saveFlaggedAccounts(storage: FlaggedAccountStorageV1): Promise<void> {
+	return withStorageLock(() => saveFlaggedAccountsUnlocked(storage));
 }
 
 export async function clearFlaggedAccounts(): Promise<void> {
@@ -2292,11 +2298,14 @@ export async function rotateStoredSecretEncryption(): Promise<{
 		return current.accounts.length;
 	});
 
-	const flagged = await loadFlaggedAccounts();
-	const flaggedCount = flagged.accounts.length;
-	if (flaggedCount > 0) {
-		await saveFlaggedAccounts(flagged);
-	}
+	const flaggedCount = await withStorageLock(async () => {
+		const flagged = await loadFlaggedAccountsUnlocked();
+		const count = flagged.accounts.length;
+		if (count > 0) {
+			await saveFlaggedAccountsUnlocked(flagged);
+		}
+		return count;
+	});
 
 	return {
 		accounts: accountCount,

lib/telemetry.ts

@@ -1,11 +1,4 @@
-import {
-	existsSync,
-	promises as fs,
-	readdirSync,
-	renameSync,
-	statSync,
-	unlinkSync,
-} from "node:fs";
+import { existsSync, promises as fs, readdirSync } from "node:fs";
 import { join } from "node:path";
 import { getCorrelationId, maskEmail } from "./logger.js";
 import { getCodexLogDir } from "./runtime-paths.js";
@@ -144,21 +137,42 @@ function parseArchiveSuffix(fileName: string): number | null {
 	return Number.isFinite(parsed) ? parsed : null;
 }
 
-function rotateLogsIfNeeded(): void {
+function isMissingFsError(error: unknown): boolean {
+	return (error as NodeJS.ErrnoException | undefined)?.code === "ENOENT";
+}
+
+async function rotateLogsIfNeeded(): Promise<void> {
 	const logPath = getTelemetryPath();
-	if (!existsSync(logPath)) return;
+	let size: number;
+	try {
+		size = (await fs.stat(logPath)).size;
+	} catch (error) {
+		if (isMissingFsError(error)) {
+			return;
+		}
+		throw error;
+	}
 
-	const size = statSync(logPath).size;
 	if (size < telemetryConfig.maxFileSizeBytes) return;
 
 	for (let i = telemetryConfig.maxFiles - 1; i >= 1; i -= 1) {
 		const target = `${logPath}.${i}`;
 		const source = i === 1 ? logPath : `${logPath}.${i - 1}`;
-		if (i === telemetryConfig.maxFiles - 1 && existsSync(target)) {
-			unlinkSync(target);
+		if (i === telemetryConfig.maxFiles - 1) {
+			try {
+				await fs.unlink(target);
+			} catch (error) {
+				if (!isMissingFsError(error)) {
+					throw error;
+				}
+			}
 		}
-		if (existsSync(source)) {
-			renameSync(source, target);
+		try {
+			await fs.rename(source, target);
+		} catch (error) {
+			if (!isMissingFsError(error)) {
+				throw error;
+			}
 		}
 	}
 }
@@ -253,7 +267,7 @@ export async function recordTelemetryEvent(input: TelemetryEventInput): Promise<
 	try {
 		await queueAppend(async () => {
 			await ensureLogDir();
-			rotateLogsIfNeeded();
+			await rotateLogsIfNeeded();
 			const line = `${JSON.stringify(entry)}\n`;
 			await fs.appendFile(getTelemetryPath(), line, "utf8");
 		});

test/accounts-edge.test.ts

@@ -564,6 +564,58 @@ describe("accounts edge branches", () => {
     );
   });
 
+  it("prefers fresher local credentials over stale disk tokens during conflict merge", async () => {
+    const staleExpiresAt = Date.now() + 1_000;
+    const freshExpiresAt = Date.now() + 120_000;
+    const stored = buildStored([
+      buildStoredAccount({
+        refreshToken: "refresh-fresh",
+        accessToken: "access-fresh",
+        expiresAt: freshExpiresAt,
+        email: "identity@example.com",
+        accountId: "account-identity-1",
+      }),
+    ]);
+
+    const latestDisk = buildStored([
+      buildStoredAccount({
+        refreshToken: "refresh-stale",
+        accessToken: "access-stale",
+        expiresAt: staleExpiresAt,
+        email: "identity@example.com",
+        accountId: "account-identity-1",
+      }),
+    ]);
+
+    const conflictError = Object.assign(new Error("conflict"), {
+      code: "ECONFLICT",
+    });
+    mockSaveAccounts
+      .mockRejectedValueOnce(conflictError)
+      .mockResolvedValueOnce(undefined);
+    mockLoadAccounts.mockResolvedValueOnce(latestDisk);
+
+    const { AccountManager } = await importAccountsModule();
+    const manager = new AccountManager(undefined, stored as never);
+
+    await manager.saveToDisk();
+
+    const retriedPayload = mockSaveAccounts.mock.calls[1]?.[0] as {
+      accounts: Array<{
+        accountId?: string;
+        refreshToken: string;
+        accessToken?: string;
+        expiresAt?: number;
+      }>;
+    };
+    const mergedAccount = retriedPayload.accounts.find(
+      (account) => account.accountId === "account-identity-1",
+    );
+    expect(mergedAccount?.refreshToken).toBe("refresh-fresh");
+    expect(mergedAccount?.accessToken).toBe("access-fresh");
+    expect(mergedAccount?.expiresAt).toBe(freshExpiresAt);
+  });
+
   it("prefers fresher timestamp and rate-limit reset values during conflict merge", async () => {
     const localNow = Date.now();
     const localAccount = buildStoredAccount({

test/storage-flagged.test.ts

@@ -7,6 +7,7 @@ import {
   getFlaggedAccountsPath,
   getStoragePath,
   loadFlaggedAccounts,
+  rotateStoredSecretEncryption,
   saveFlaggedAccounts,
   setStoragePathDirect,
 } from "../lib/storage.js";
@@ -239,4 +240,82 @@ describe("flagged account storage", () => {
 
     renameSpy.mockRestore();
   });
+
+  it("preserves concurrent flagged-account updates during secret rotation", async () => {
+    const previousEncryptionKey = process.env.CODEX_AUTH_ENCRYPTION_KEY;
+    const previousPreviousKey = process.env.CODEX_AUTH_PREVIOUS_ENCRYPTION_KEY;
+    process.env.CODEX_AUTH_ENCRYPTION_KEY = "0123456789abcdef0123456789abcdef";
+    delete process.env.CODEX_AUTH_PREVIOUS_ENCRYPTION_KEY;
+
+    await saveFlaggedAccounts({
+      version: 1,
+      accounts: [
+        {
+          refreshToken: "flagged-alpha",
+          flaggedAt: 1,
+          addedAt: 1,
+          lastUsed: 1,
+        },
+      ],
+    });
+
+    let concurrentSavePromise: Promise<void> | null = null;
+    const flaggedPath = getFlaggedAccountsPath();
+    const originalReadFile = fs.readFile.bind(fs);
+    const readFileSpy = vi.spyOn(fs, "readFile").mockImplementation(async (path, options) => {
+      if (
+        concurrentSavePromise === null &&
+        typeof path === "string" &&
+        path === flaggedPath &&
+        options === "utf-8"
+      ) {
+        queueMicrotask(() => {
+          concurrentSavePromise = saveFlaggedAccounts({
+            version: 1,
+            accounts: [
+              {
+                refreshToken: "flagged-alpha",
+                flaggedAt: 1,
+                addedAt: 1,
+                lastUsed: 1,
+              },
+              {
+                refreshToken: "flagged-beta",
+                flaggedAt: 2,
+                addedAt: 2,
+                lastUsed: 2,
+              },
+            ],
+          });
+        });
+      }
+      return originalReadFile(path, options);
+    });
+
+    try {
+      const result = await rotateStoredSecretEncryption();
+      if (concurrentSavePromise) {
+        await concurrentSavePromise;
+      }
+
+      const flagged = await loadFlaggedAccounts();
+      expect(result.flaggedAccounts).toBe(1);
+      expect(flagged.accounts.map((account) => account.refreshToken).sort()).toEqual([
+        "flagged-alpha",
+        "flagged-beta",
+      ]);
+    } finally {
+      readFileSpy.mockRestore();
+      if (previousEncryptionKey === undefined) {
+        delete process.env.CODEX_AUTH_ENCRYPTION_KEY;
+      } else {
+        process.env.CODEX_AUTH_ENCRYPTION_KEY = previousEncryptionKey;
+      }
+      if (previousPreviousKey === undefined) {
+        delete process.env.CODEX_AUTH_PREVIOUS_ENCRYPTION_KEY;
+      } else {
+        process.env.CODEX_AUTH_PREVIOUS_ENCRYPTION_KEY = previousPreviousKey;
+      }
+    }
+  });
 });