From 27cb5d021fd1f7d8dd408871645038318609ce0a Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 17 Aug 2026 15:29:31 -0700 Subject: [PATCH] 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 --- src/commands/doctor-session-sqlite-readers.ts | 147 +++++++++--------- src/commands/doctor-session-sqlite.test.ts | 57 +++++++ src/commands/doctor-session-sqlite.ts | 132 ++++++++-------- 3 files changed, 202 insertions(+), 134 deletions(-) diff --git a/src/commands/doctor-session-sqlite-readers.ts b/src/commands/doctor-session-sqlite-readers.ts index eed645cbc879..e6be2614dfea 100644 --- a/src/commands/doctor-session-sqlite-readers.ts +++ b/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; + transcriptEventCountsBySessionId: ReadonlyMap; +}; -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 { diff --git a/src/commands/doctor-session-sqlite.test.ts b/src/commands/doctor-session-sqlite.test.ts index b0555413bda3..bbd50f9768cd 100644 --- a/src/commands/doctor-session-sqlite.test.ts +++ b/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< diff --git a/src/commands/doctor-session-sqlite.ts b/src/commands/doctor-session-sqlite.ts index 0acb7cae3bf2..d47effc8473a 100644 --- a/src/commands/doctor-session-sqlite.ts +++ b/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 { const importedTranscriptSources = new Set(); + 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, + 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(