perf(doctor): snapshot SQLite migration validation (#125361)

* perf(doctor): snapshot SQLite migration validation

Amp-Thread-ID: https://ampcode.com/threads/T-01a00a6a-b64e-74a5-8b15-2d3b966a468d

* test(doctor): bound SQLite validation reads

Amp-Thread-ID: https://ampcode.com/threads/T-01a00a6a-b64e-74a5-8b15-2d3b966a468d

---------

Co-authored-by: Amp <amp@ampcode.com>

src/commands/doctor-session-sqlite-readers.ts

@@ -31,13 +31,14 @@ type ReadOnlySqliteSessionEntriesResult =
   | { exists: true; ok: true; summaries: ReadOnlySqliteSessionSummary[] }
   | { error: unknown; exists: true; ok: false };
 
-type ReadOnlySqliteExactSessionEntryResult =
-  | { entry?: ReadOnlySqliteSessionSummary; ok: true }
-  | { error: unknown; ok: false };
+export type ReadOnlySqliteValidationSnapshot = {
+  entriesBySessionKey: ReadonlyMap<string, SessionEntry>;
+  transcriptEventCountsBySessionId: ReadonlyMap<string, number>;
+};
 
-type ReadOnlySqliteTranscriptEventCountResult =
-  | { events: number; exists: boolean; ok: true }
-  | { error: unknown; exists: true; ok: false };
+type ReadOnlySqliteValidationSnapshotResult =
+  | { ok: true; snapshot: ReadOnlySqliteValidationSnapshot }
+  | { error: unknown; ok: false };
 
 type ReadOnlySqliteDbStatsResult =
   | {
@@ -291,18 +292,51 @@ export function readSqliteEntryCount(target: SessionStoreTarget): number {
   return result.ok ? result.summaries.length : 0;
 }
 
-export function readOnlySqliteExactSessionEntry(
+export function readOnlySqliteValidationSnapshot(
   target: SessionStoreTarget,
-  sessionKey: string,
-): ReadOnlySqliteExactSessionEntryResult {
-  const result = readOnlySqliteSessionEntries(target);
-  if (!result.ok) {
-    return { error: result.error, ok: false };
+): ReadOnlySqliteValidationSnapshotResult {
+  const sqlitePath = resolveTargetSqlitePath(target);
+  if (!fs.existsSync(sqlitePath)) {
+    return {
+      ok: true,
+      snapshot: {
+        entriesBySessionKey: new Map(),
+        transcriptEventCountsBySessionId: new Map(),
+      },
+    };
+  }
+  let database: DatabaseSync | undefined;
+  try {
+    database = openNodeSqliteDatabase(sqlitePath, { readOnly: true });
+    const summaries = readOnlySqliteSessionEntriesFromDatabase(database);
+    const hasTranscriptEvents = database
+      .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?")
+      .get("transcript_events");
+    const transcriptRows = hasTranscriptEvents
+      ? (database
+          .prepare(
+            "SELECT session_id, COUNT(*) AS count FROM transcript_events GROUP BY session_id",
+          )
+          .all() as Array<{ count?: unknown; session_id?: unknown }>)
+      : [];
+    return {
+      ok: true,
+      snapshot: {
+        entriesBySessionKey: new Map(summaries.map(({ entry, sessionKey }) => [sessionKey, entry])),
+        transcriptEventCountsBySessionId: new Map(
+          transcriptRows.flatMap((row) =>
+            typeof row.session_id === "string" && typeof row.count === "number"
+              ? [[row.session_id, row.count] as const]
+              : [],
+          ),
+        ),
+      },
+    };
+  } catch (error) {
+    return { error, ok: false };
+  } finally {
+    database?.close();
   }
-  return {
-    entry: result.summaries.find((summary) => summary.sessionKey === sessionKey),
-    ok: true,
-  };
 }
 
 export function readOnlySqliteSessionEntries(
@@ -315,34 +349,10 @@ export function readOnlySqliteSessionEntries(
   let database: DatabaseSync | undefined;
   try {
     database = openNodeSqliteDatabase(sqlitePath, { readOnly: true });
-    const nodeTable = database
-      .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?")
-      .get("session_nodes");
-    const legacyEntryTable = nodeTable
-      ? undefined
-      : database
-          .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?")
-          .get("session_entries");
-    if (!nodeTable && !legacyEntryTable) {
-      return { exists: true, ok: true, summaries: [] };
-    }
-    const rows = database
-      .prepare(
-        nodeTable
-          ? "SELECT session_key, entry_json FROM session_nodes ORDER BY session_key ASC"
-          : "SELECT session_key, entry_json FROM session_entries ORDER BY session_key ASC",
-      )
-      .all() as Array<{ entry_json?: unknown; session_key?: unknown }>;
     return {
       exists: true,
       ok: true,
-      summaries: rows.flatMap((row) => {
-        if (typeof row.session_key !== "string" || typeof row.entry_json !== "string") {
-          return [];
-        }
-        const entry = parseSqliteSessionEntry(row.entry_json);
-        return entry ? [{ entry, sessionKey: row.session_key }] : [];
-      }),
+      summaries: readOnlySqliteSessionEntriesFromDatabase(database),
     };
   } catch (error) {
     return { error, exists: true, ok: false };
@@ -351,37 +361,34 @@ export function readOnlySqliteSessionEntries(
   }
 }
 
-export function readOnlySqliteTranscriptEventCount(
-  target: SessionStoreTarget,
-  sessionId: string,
-): ReadOnlySqliteTranscriptEventCountResult {
-  const sqlitePath = resolveTargetSqlitePath(target);
-  if (!fs.existsSync(sqlitePath)) {
-    return { events: 0, exists: false, ok: true };
+function readOnlySqliteSessionEntriesFromDatabase(
+  database: DatabaseSync,
+): ReadOnlySqliteSessionSummary[] {
+  const nodeTable = database
+    .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?")
+    .get("session_nodes");
+  const legacyEntryTable = nodeTable
+    ? undefined
+    : database
+        .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?")
+        .get("session_entries");
+  if (!nodeTable && !legacyEntryTable) {
+    return [];
   }
-  let database: DatabaseSync | undefined;
-  try {
-    database = openNodeSqliteDatabase(sqlitePath, { readOnly: true });
-    const table = database
-      .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?")
-      .get("transcript_events");
-    if (!table) {
-      return { events: 0, exists: true, ok: true };
+  const rows = database
+    .prepare(
+      nodeTable
+        ? "SELECT session_key, entry_json FROM session_nodes ORDER BY session_key ASC"
+        : "SELECT session_key, entry_json FROM session_entries ORDER BY session_key ASC",
+    )
+    .all() as Array<{ entry_json?: unknown; session_key?: unknown }>;
+  return rows.flatMap((row) => {
+    if (typeof row.session_key !== "string" || typeof row.entry_json !== "string") {
+      return [];
     }
-    const row = database
-      .prepare("SELECT COUNT(*) AS count FROM transcript_events WHERE session_id = ?")
-      .get(sessionId) as { count?: unknown } | undefined;
-    const count = row?.count;
-    return {
-      events: typeof count === "number" && Number.isFinite(count) ? count : 0,
-      exists: true,
-      ok: true,
-    };
-  } catch (error) {
-    return { error, exists: true, ok: false };
-  } finally {
-    database?.close();
-  }
+    const entry = parseSqliteSessionEntry(row.entry_json);
+    return entry ? [{ entry, sessionKey: row.session_key }] : [];
+  });
 }
 
 export function readOnlySqliteDbStats(target: SessionStoreTarget): ReadOnlySqliteDbStatsResult {

src/commands/doctor-session-sqlite.test.ts

@@ -727,6 +727,63 @@ describe("runDoctorSessionSqlite", () => {
     ).toHaveLength(2);
   });
 
+  it("uses target-bounded validation reads for multi-session imports", async () => {
+    const countTargetReads = async (sessionCount: number) => {
+      const store = createLegacyStore();
+      const sessions = Object.fromEntries(
+        Array.from({ length: sessionCount }, (_, offset) => {
+          const index = offset + 1;
+          return [
+            index === 1 ? "agent:main:main" : `agent:main:session-${index}`,
+            {
+              sessionFile: `session-${index}.jsonl`,
+              sessionId: `session-${index}`,
+              updatedAt: 2000 + index,
+            },
+          ];
+        }),
+      );
+      fs.writeFileSync(store.storePath, `${JSON.stringify(sessions)}\n`, { mode: 0o600 });
+      for (let index = 2; index <= sessionCount; index += 1) {
+        fs.writeFileSync(
+          path.join(store.sessionDir, `session-${index}.jsonl`),
+          `{"type":"session","sessionId":"session-${index}"}\n{"type":"event","id":"evt-${index}"}\n`,
+          { mode: 0o600 },
+        );
+      }
+
+      const sqlitePath = path.resolve(
+        resolveTargetSqlitePath({ agentId: "main", storePath: store.storePath }),
+      );
+      const openSqlite = vi.spyOn(nodeSqlite, "openNodeSqliteDatabase");
+      try {
+        const report = await runDoctorSessionSqlite({
+          env: store.env,
+          mode: "import",
+          store: store.storePath,
+        });
+        expect(report.totals).toMatchObject({
+          importedEntries: sessionCount,
+          importedTranscriptEvents: sessionCount * 2,
+          issues: 0,
+          sqliteEntries: sessionCount,
+        });
+        return openSqlite.mock.calls.filter(
+          ([location, options]) =>
+            path.resolve(location) === sqlitePath && options?.readOnly === true,
+        ).length;
+      } finally {
+        openSqlite.mockRestore();
+      }
+    };
+
+    const singleSessionReads = await countTargetReads(1);
+    const multiSessionReads = await countTargetReads(3);
+
+    expect(singleSessionReads).toBeGreaterThan(0);
+    expect(multiSessionReads).toBe(singleSessionReads);
+  });
+
   it("archives legacy stores with valid sessions and invalid cron stubs without failing", async () => {
     const store = createLegacyStore();
     const legacyStore = JSON.parse(fs.readFileSync(store.storePath, "utf-8")) as Record<

src/commands/doctor-session-sqlite.ts

@@ -52,12 +52,12 @@ import {
   countTranscriptEventsForPath,
   createTranscriptEventReader,
   readOnlySqliteDbStats,
-  readOnlySqliteExactSessionEntry,
   readOnlySqliteSessionEntries,
-  readOnlySqliteTranscriptEventCount,
+  readOnlySqliteValidationSnapshot,
   readTranscriptFingerprint,
   readSqliteEntryCount,
   resolveTargetSqlitePath,
+  type ReadOnlySqliteValidationSnapshot,
 } from "./doctor-session-sqlite-readers.js";
 import { recoverDoctorSessionSqliteTargets } from "./doctor-session-sqlite-recover-report.js";
 import { restoreDoctorSessionSqliteTargets } from "./doctor-session-sqlite-restore-report.js";
@@ -349,14 +349,12 @@ async function inspectOrMigrateTarget(params: {
   }
   if (params.mode === "import") {
     await importLegacySessionRecords(params.target, records, report);
-  } else {
+  } else if (params.mode === "dry-run") {
     for (const record of records) {
-      if (params.mode === "dry-run") {
-        countLegacyTranscript(record, report);
-      } else {
-        validateLegacySessionRecord(params.target, record, report);
-      }
+      countLegacyTranscript(record, report);
     }
+  } else {
+    validateLegacySessionRecords(params.target, records, report);
   }
   if (params.mode === "import" && blockingIssueCount(report) === 0) {
     const validationPassed = validateImportedTargetBeforeArchive(params.target, records, report);
@@ -624,6 +622,7 @@ async function importLegacySessionRecords(
   report: DoctorSessionSqliteTargetReport,
 ): Promise<void> {
   const importedTranscriptSources = new Set<string>();
+  const existingSnapshot = readOnlySqliteValidationSnapshot(target);
   for (let offset = 0; offset < records.length; offset += SESSION_IMPORT_BATCH_SIZE) {
     const pending = records.slice(offset, offset + SESSION_IMPORT_BATCH_SIZE).flatMap((record) => {
       const prepared = prepareLegacySessionImport(
@@ -631,6 +630,7 @@ async function importLegacySessionRecords(
         record,
         report,
         importedTranscriptSources,
+        existingSnapshot.ok ? existingSnapshot.snapshot : undefined,
       );
       return prepared ? [prepared] : [];
     });
@@ -649,6 +649,7 @@ function prepareLegacySessionImport(
   record: LegacySessionRecord,
   report: DoctorSessionSqliteTargetReport,
   importedTranscriptSources: Set<string>,
+  existingSnapshot: ReadOnlySqliteValidationSnapshot | undefined,
 ) {
   const transcriptSourceKey = record.transcriptPath
     ? `${record.entry.sessionId}\0${record.transcriptPath}`
@@ -671,7 +672,7 @@ function prepareLegacySessionImport(
     storePath: target.sqlitePath ?? target.storePath,
   };
   if (result.status === "missing") {
-    if (markAlreadyMigratedTranscript(target, record, report)) {
+    if (markAlreadyMigratedTranscript(record, report, existingSnapshot)) {
       return undefined;
     }
     return {
@@ -714,11 +715,11 @@ function prepareLegacySessionImport(
 }
 
 function markAlreadyMigratedTranscript(
-  target: SessionStoreTarget,
   record: LegacySessionRecord,
   report: DoctorSessionSqliteTargetReport,
+  snapshot: ReadOnlySqliteValidationSnapshot | undefined,
 ): boolean {
-  const migratedEvents = countAlreadyMigratedTranscriptEventsForImport(target, record);
+  const migratedEvents = countAlreadyMigratedTranscriptEventsForImport(snapshot, record);
   if (migratedEvents === undefined) {
     return false;
   }
@@ -733,20 +734,28 @@ function validateImportedTargetBeforeArchive(
   report: DoctorSessionSqliteTargetReport,
 ): boolean {
   const issueCountBeforeValidation = report.issues.length;
+  const validation = readOnlySqliteValidationSnapshot(target);
+  if (!validation.ok) {
+    report.issues.push({
+      code: "sqlite_read_failed",
+      message: `SQLite validation read failed: ${String(validation.error)}`,
+    });
+    return false;
+  }
   for (const record of records) {
-    validateImportedRecordBeforeArchive(target, record, report);
+    validateImportedRecordBeforeArchive(record, report, validation.snapshot);
   }
   return report.issues.length === issueCountBeforeValidation;
 }
 
 function validateImportedRecordBeforeArchive(
-  target: SessionStoreTarget,
   record: LegacySessionRecord,
   report: DoctorSessionSqliteTargetReport,
+  snapshot: ReadOnlySqliteValidationSnapshot,
 ): void {
   const normalizedKey = record.sessionKey;
-  const sqliteEntry = readOnlySqliteExactSessionEntry(target, normalizedKey);
-  if (!sqliteEntry.ok || !sqliteEntry.entry) {
+  const sqliteEntry = snapshot.entriesBySessionKey.get(normalizedKey);
+  if (!sqliteEntry) {
     report.issues.push({
       code: "sqlite_entry_missing",
       message: `SQLite entry is missing for ${normalizedKey}.`,
@@ -754,10 +763,10 @@ function validateImportedRecordBeforeArchive(
     });
     return;
   }
-  if (sqliteEntry.entry.entry.sessionId !== record.entry.sessionId) {
+  if (sqliteEntry.sessionId !== record.entry.sessionId) {
     report.issues.push({
       code: "sqlite_entry_mismatch",
-      message: `SQLite sessionId ${sqliteEntry.entry.entry.sessionId} does not match ${record.entry.sessionId}.`,
+      message: `SQLite sessionId ${sqliteEntry.sessionId} does not match ${record.entry.sessionId}.`,
       sessionKey: record.sessionKey,
     });
     return;
@@ -776,19 +785,11 @@ function validateImportedRecordBeforeArchive(
     }
     return;
   }
-  const sqliteEvents = readOnlySqliteTranscriptEventCount(target, record.entry.sessionId);
-  if (!sqliteEvents.ok) {
-    report.issues.push({
-      code: "sqlite_read_failed",
-      message: `SQLite transcript count read failed: ${String(sqliteEvents.error)}`,
-      sessionKey: record.sessionKey,
-    });
-    return;
-  }
-  if (sqliteEvents.events < result.events) {
+  const sqliteEvents = snapshot.transcriptEventCountsBySessionId.get(record.entry.sessionId) ?? 0;
+  if (sqliteEvents < result.events) {
     report.issues.push({
       code: "sqlite_transcript_count_mismatch",
-      message: `SQLite transcript has ${sqliteEvents.events} events; source has ${result.events}.`,
+      message: `SQLite transcript has ${sqliteEvents} events; source has ${result.events}.`,
       sessionKey: record.sessionKey,
     });
   }
@@ -972,22 +973,32 @@ function recordLegacyStoreMoveForTarget(
   recordCompletedMigrationMove(activeRun, createMigrationTargetInput(target), move);
 }
 
-function validateLegacySessionRecord(
+function validateLegacySessionRecords(
   target: SessionStoreTarget,
-  record: LegacySessionRecord,
+  records: readonly LegacySessionRecord[],
   report: DoctorSessionSqliteTargetReport,
 ): void {
-  const normalizedKey = normalizeStoreSessionKey(record.sessionKey);
-  const sqliteEntry = readOnlySqliteExactSessionEntry(target, normalizedKey);
-  if (!sqliteEntry.ok) {
+  const validation = readOnlySqliteValidationSnapshot(target);
+  if (!validation.ok) {
     report.issues.push({
       code: "sqlite_read_failed",
-      message: `SQLite session entry read failed: ${String(sqliteEntry.error)}`,
-      sessionKey: record.sessionKey,
+      message: `SQLite validation read failed: ${String(validation.error)}`,
     });
     return;
   }
-  if (!sqliteEntry.entry) {
+  for (const record of records) {
+    validateLegacySessionRecord(record, report, validation.snapshot);
+  }
+}
+
+function validateLegacySessionRecord(
+  record: LegacySessionRecord,
+  report: DoctorSessionSqliteTargetReport,
+  snapshot: ReadOnlySqliteValidationSnapshot,
+): void {
+  const normalizedKey = normalizeStoreSessionKey(record.sessionKey);
+  const sqliteEntry = snapshot.entriesBySessionKey.get(normalizedKey);
+  if (!sqliteEntry) {
     report.issues.push({
       code: "sqlite_entry_missing",
       message: `SQLite entry is missing for ${normalizedKey}.`,
@@ -995,26 +1006,26 @@ function validateLegacySessionRecord(
     });
     return;
   }
-  if (sqliteEntry.entry.entry.sessionId !== record.entry.sessionId) {
+  if (sqliteEntry.sessionId !== record.entry.sessionId) {
     report.issues.push({
       code: "sqlite_entry_mismatch",
-      message: `SQLite sessionId ${sqliteEntry.entry.entry.sessionId} does not match ${record.entry.sessionId}.`,
+      message: `SQLite sessionId ${sqliteEntry.sessionId} does not match ${record.entry.sessionId}.`,
       sessionKey: record.sessionKey,
     });
     return;
   }
   report.validatedEntries += 1;
-  validateTranscriptEventCount(target, record, report);
+  validateTranscriptEventCount(record, report, snapshot);
 }
 
 function validateTranscriptEventCount(
-  target: SessionStoreTarget,
   record: LegacySessionRecord,
   report: DoctorSessionSqliteTargetReport,
+  snapshot: ReadOnlySqliteValidationSnapshot,
 ): void {
   const result = countTranscriptEvents(record);
   if (result.status === "missing") {
-    const migratedEvents = countAlreadyMigratedTranscriptEventsForValidate(target, record);
+    const migratedEvents = countAlreadyMigratedTranscriptEventsForValidate(snapshot, record);
     if (migratedEvents !== undefined) {
       report.validatedTranscriptEvents += migratedEvents;
     }
@@ -1030,24 +1041,16 @@ function validateTranscriptEventCount(
     }
     return;
   }
-  const sqliteEvents = readOnlySqliteTranscriptEventCount(target, record.entry.sessionId);
-  if (!sqliteEvents.ok) {
-    report.issues.push({
-      code: "sqlite_read_failed",
-      message: `SQLite transcript count read failed: ${String(sqliteEvents.error)}`,
-      sessionKey: record.sessionKey,
-    });
-    return;
-  }
-  if (sqliteEvents.events !== result.events) {
+  const sqliteEvents = snapshot.transcriptEventCountsBySessionId.get(record.entry.sessionId) ?? 0;
+  if (sqliteEvents !== result.events) {
     report.issues.push({
       code: "sqlite_transcript_count_mismatch",
-      message: `SQLite transcript has ${sqliteEvents.events} events; source has ${result.events}.`,
+      message: `SQLite transcript has ${sqliteEvents} events; source has ${result.events}.`,
       sessionKey: record.sessionKey,
     });
     return;
   }
-  report.validatedTranscriptEvents += sqliteEvents.events;
+  report.validatedTranscriptEvents += sqliteEvents;
 }
 
 function hasSessionIssue(
@@ -1059,29 +1062,30 @@ function hasSessionIssue(
 }
 
 function countAlreadyMigratedTranscriptEventsForImport(
-  target: SessionStoreTarget,
+  snapshot: ReadOnlySqliteValidationSnapshot | undefined,
   record: LegacySessionRecord,
 ): number | undefined {
-  const normalizedKey = record.sessionKey;
-  const sqliteEntry = readOnlySqliteExactSessionEntry(target, normalizedKey);
-  if (!sqliteEntry.ok || sqliteEntry.entry?.entry.sessionId !== record.entry.sessionId) {
+  if (!snapshot) {
     return undefined;
   }
-  const eventCount = readOnlySqliteTranscriptEventCount(target, record.entry.sessionId);
-  return eventCount.ok ? eventCount.events : undefined;
+  const normalizedKey = record.sessionKey;
+  const sqliteEntry = snapshot.entriesBySessionKey.get(normalizedKey);
+  if (sqliteEntry?.sessionId !== record.entry.sessionId) {
+    return undefined;
+  }
+  return snapshot.transcriptEventCountsBySessionId.get(record.entry.sessionId) ?? 0;
 }
 
 function countAlreadyMigratedTranscriptEventsForValidate(
-  target: SessionStoreTarget,
+  snapshot: ReadOnlySqliteValidationSnapshot,
   record: LegacySessionRecord,
 ): number | undefined {
   const normalizedKey = normalizeStoreSessionKey(record.sessionKey);
-  const sqliteEntry = readOnlySqliteExactSessionEntry(target, normalizedKey);
-  if (!sqliteEntry.ok || sqliteEntry.entry?.entry.sessionId !== record.entry.sessionId) {
+  const sqliteEntry = snapshot.entriesBySessionKey.get(normalizedKey);
+  if (sqliteEntry?.sessionId !== record.entry.sessionId) {
     return undefined;
   }
-  const eventCount = readOnlySqliteTranscriptEventCount(target, record.entry.sessionId);
-  return eventCount.ok ? eventCount.events : undefined;
+  return snapshot.transcriptEventCountsBySessionId.get(record.entry.sessionId) ?? 0;
 }
 
 function countTranscriptEvents(