From 12f2ac1e276830afd1b3bd8c6ca2f2e1fc38fc6e Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Mon, 30 Mar 2026 16:38:42 -0700 Subject: [PATCH] agentHost: use memory db for sqlite tests to avoid locking weirdness (#306524) --- .../agentHost/node/sessionDatabase.ts | 8 +- .../test/node/fileEditTracker.test.ts | 11 +- .../test/node/mapSessionEvents.test.ts | 21 +--- .../test/node/sessionDatabase.test.ts | 108 ++++++++++-------- 4 files changed, 69 insertions(+), 79 deletions(-) diff --git a/src/vs/platform/agentHost/node/sessionDatabase.ts b/src/vs/platform/agentHost/node/sessionDatabase.ts index 7f960f89d2a..8c9e51c106d 100644 --- a/src/vs/platform/agentHost/node/sessionDatabase.ts +++ b/src/vs/platform/agentHost/node/sessionDatabase.ts @@ -112,7 +112,7 @@ function dbOpen(path: string): Promise { * `PRAGMA user_version` are run inside a serialized transaction. After all * migrations complete the pragma is updated to the highest applied version. */ -async function runMigrations(db: Database, migrations: readonly ISessionDatabaseMigration[]): Promise { +export async function runMigrations(db: Database, migrations: readonly ISessionDatabaseMigration[]): Promise { // Enable foreign key enforcement — must be set outside a transaction // and every time a connection is opened. await dbExec(db, 'PRAGMA foreign_keys = ON'); @@ -154,8 +154,8 @@ async function runMigrations(db: Database, migrations: readonly ISessionDatabase */ export class SessionDatabase implements ISessionDatabase { - private _dbPromise: Promise | undefined; - private _closed: Promise | true | undefined; + protected _dbPromise: Promise | undefined; + protected _closed: Promise | true | undefined; private readonly _fileEditSequencer = new SequencerByKey(); constructor( @@ -174,7 +174,7 @@ export class SessionDatabase implements ISessionDatabase { return inst; } - private _ensureDb(): Promise { + protected _ensureDb(): Promise { if (this._closed) { return Promise.reject(new Error('SessionDatabase has been disposed')); } diff --git a/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts b/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts index bb62eadc929..dec81f59f3e 100644 --- a/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts +++ b/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts @@ -4,9 +4,6 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; -import { tmpdir } from 'os'; -import { randomUUID } from 'crypto'; -import { mkdirSync, rmSync } from 'fs'; import { VSBuffer } from '../../../../base/common/buffer.js'; import { DisposableStore } from '../../../../base/common/lifecycle.js'; import { URI } from '../../../../base/common/uri.js'; @@ -17,7 +14,6 @@ import { NullLogService } from '../../../log/common/log.js'; import { ToolResultContentType } from '../../common/state/sessionState.js'; import { SessionDatabase } from '../../node/sessionDatabase.js'; import { FileEditTracker, buildSessionDbUri, parseSessionDbUri } from '../../node/copilot/fileEditTracker.js'; -import { join } from '../../../../base/common/path.js'; suite('FileEditTracker', () => { @@ -25,17 +21,13 @@ suite('FileEditTracker', () => { let fileService: FileService; let db: SessionDatabase; let tracker: FileEditTracker; - let testDir: string; setup(async () => { - testDir = join(tmpdir(), `vscode-edit-tracker-test-${randomUUID()}`); - mkdirSync(testDir, { recursive: true }); - fileService = disposables.add(new FileService(new NullLogService())); const sourceFs = disposables.add(new InMemoryFileSystemProvider()); disposables.add(fileService.registerProvider('file', sourceFs)); - db = disposables.add(await SessionDatabase.open(join(testDir, 'session.db'))); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); tracker = new FileEditTracker('copilot:/test-session', db, fileService, new NullLogService()); @@ -44,7 +36,6 @@ suite('FileEditTracker', () => { teardown(async () => { disposables.clear(); await db.close(); - rmSync(testDir, { recursive: true, force: true }); }); ensureNoDisposablesAreLeakedInTestSuite(); diff --git a/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts b/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts index 9ebeb2fec2c..034dbb49435 100644 --- a/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts +++ b/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts @@ -4,9 +4,6 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; -import { tmpdir } from 'os'; -import { randomUUID } from 'crypto'; -import { mkdirSync, rmSync } from 'fs'; import { DisposableStore } from '../../../../base/common/lifecycle.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { AgentSession } from '../../common/agentService.js'; @@ -14,31 +11,19 @@ import { ToolResultContentType } from '../../common/state/sessionState.js'; import { SessionDatabase } from '../../node/sessionDatabase.js'; import { parseSessionDbUri } from '../../node/copilot/fileEditTracker.js'; import { mapSessionEvents, type ISessionEvent } from '../../node/copilot/mapSessionEvents.js'; -import { join } from '../../../../base/common/path.js'; suite('mapSessionEvents', () => { const disposables = new DisposableStore(); - let testDir: string; let db: SessionDatabase | undefined; const session = AgentSession.uri('copilot', 'test-session'); - setup(() => { - testDir = join(tmpdir(), `vscode-map-events-test-${randomUUID()}`); - mkdirSync(testDir, { recursive: true }); - }); - teardown(async () => { disposables.clear(); await db?.close(); - rmSync(testDir, { recursive: true, force: true }); }); ensureNoDisposablesAreLeakedInTestSuite(); - function dbPath(): string { - return join(testDir, 'session.db'); - } - // ---- Basic event mapping -------------------------------------------- test('maps user and assistant messages', async () => { @@ -111,7 +96,7 @@ suite('mapSessionEvents', () => { suite('file edit restoration', () => { test('restores file edits from database for edit tools', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.storeFileEdit({ turnId: 'turn-1', @@ -156,7 +141,7 @@ suite('mapSessionEvents', () => { }); test('handles multiple file edits for one tool call', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.storeFileEdit({ turnId: 'turn-1', @@ -217,7 +202,7 @@ suite('mapSessionEvents', () => { }); test('non-edit tools do not get file edits even if db has data', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); const events: ISessionEvent[] = [ { diff --git a/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts b/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts index 662743c3fb7..f71a9f2ee86 100644 --- a/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts +++ b/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts @@ -4,35 +4,52 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; -import { tmpdir } from 'os'; -import { randomUUID } from 'crypto'; -import { mkdirSync, rmSync } from 'fs'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { DisposableStore } from '../../../../base/common/lifecycle.js'; -import { SessionDatabase, type ISessionDatabaseMigration } from '../../node/sessionDatabase.js'; -import { join } from '../../../../base/common/path.js'; +import { SessionDatabase, runMigrations, sessionDatabaseMigrations, type ISessionDatabaseMigration } from '../../node/sessionDatabase.js'; +import type { Database } from '@vscode/sqlite3'; suite('SessionDatabase', () => { const disposables = new DisposableStore(); - let testDir: string; let db: SessionDatabase | undefined; let db2: SessionDatabase | undefined; - setup(() => { - testDir = join(tmpdir(), `vscode-session-db-test-${randomUUID()}`); - mkdirSync(testDir, { recursive: true }); - }); - teardown(async () => { disposables.clear(); await Promise.all([db?.close(), db2?.close()]); - rmSync(testDir, { recursive: true, force: true }); }); ensureNoDisposablesAreLeakedInTestSuite(); - function dbPath(name = 'session.db'): string { - return join(testDir, name); + /** + * Extends SessionDatabase to allow ejecting/injecting the raw sqlite3 + * Database instance, enabling reopen tests with :memory: databases. + */ + class TestableSessionDatabase extends SessionDatabase { + static override async open(path: string, migrations: readonly ISessionDatabaseMigration[] = sessionDatabaseMigrations): Promise { + const inst = new TestableSessionDatabase(path, migrations); + await inst._ensureDb(); + return inst; + } + + /** Extract the raw db connection; this instance becomes inert. */ + async ejectDb(): Promise { + const rawDb = await this._ensureDb(); + this._dbPromise = undefined; + this._closed = true; + return rawDb; + } + + /** Create a TestableSessionDatabase wrapping an existing raw db. */ + static async fromDb( + rawDb: Database, + migrations: readonly ISessionDatabaseMigration[] = sessionDatabaseMigrations, + ): Promise { + await runMigrations(rawDb, migrations); + const inst = new TestableSessionDatabase(':memory:', migrations); + inst._dbPromise = Promise.resolve(rawDb); + return inst; + } } // ---- Migration system ----------------------------------------------- @@ -45,7 +62,7 @@ suite('SessionDatabase', () => { { version: 2, sql: 'CREATE TABLE t2 (id INTEGER PRIMARY KEY)' }, ]; - db = disposables.add(await SessionDatabase.open(dbPath(), migrations)); + db = disposables.add(await SessionDatabase.open(':memory:', migrations)); const tables = (await db.getAllTables()).sort(); assert.deepStrictEqual(tables, ['t1', 't2']); @@ -56,11 +73,11 @@ suite('SessionDatabase', () => { { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, ]; - const db1 = await SessionDatabase.open(dbPath(), migrations); - await db1.close(); + const db1 = await TestableSessionDatabase.open(':memory:', migrations); + const rawDb = await db1.ejectDb(); // Reopen — should not throw (table already exists, migration skipped) - db2 = disposables.add(await SessionDatabase.open(dbPath(), migrations)); + db2 = disposables.add(await TestableSessionDatabase.fromDb(rawDb, migrations)); assert.deepStrictEqual(await db2.getAllTables(), ['t1']); }); @@ -68,14 +85,14 @@ suite('SessionDatabase', () => { const v1: ISessionDatabaseMigration[] = [ { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, ]; - const db1 = await SessionDatabase.open(dbPath(), v1); - await db1.close(); + const db1 = await TestableSessionDatabase.open(':memory:', v1); + const rawDb = await db1.ejectDb(); const v2: ISessionDatabaseMigration[] = [ ...v1, { version: 2, sql: 'CREATE TABLE t2 (id INTEGER PRIMARY KEY)' }, ]; - db2 = disposables.add(await SessionDatabase.open(dbPath(), v2)); + db2 = disposables.add(await TestableSessionDatabase.fromDb(rawDb, v2)); const tables = (await db2.getAllTables()).sort(); assert.deepStrictEqual(tables, ['t1', 't2']); @@ -87,11 +104,10 @@ suite('SessionDatabase', () => { { version: 2, sql: 'THIS IS INVALID SQL' }, ]; - await assert.rejects(() => SessionDatabase.open(dbPath(), migrations)); + await assert.rejects(() => SessionDatabase.open(':memory:', migrations)); - // Reopen with only v1 — t1 should not exist because the whole - // transaction was rolled back - db = disposables.add(await SessionDatabase.open(dbPath(), [ + // A fresh :memory: open with valid migrations succeeds + db = disposables.add(await SessionDatabase.open(':memory:', [ { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, ])); assert.deepStrictEqual(await db.getAllTables(), ['t1']); @@ -103,7 +119,7 @@ suite('SessionDatabase', () => { suite('file edits', () => { test('store and retrieve a file edit', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -127,7 +143,7 @@ suite('SessionDatabase', () => { }); test('retrieve multiple edits for a single tool call', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -156,7 +172,7 @@ suite('SessionDatabase', () => { }); test('retrieve edits across multiple tool calls', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -188,19 +204,19 @@ suite('SessionDatabase', () => { }); test('returns empty array for unknown tool call IDs', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); const edits = await db.getFileEdits(['nonexistent']); assert.deepStrictEqual(edits, []); }); test.skip('returns empty array when given empty array' /* Flaky https://github.com/microsoft/vscode/issues/306057 */, async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); const edits = await db.getFileEdits([]); assert.deepStrictEqual(edits, []); }); test('replace on conflict (same toolCallId + filePath)', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -232,7 +248,7 @@ suite('SessionDatabase', () => { }); test('readFileEditContent returns content on demand', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -252,13 +268,13 @@ suite('SessionDatabase', () => { }); test('readFileEditContent returns undefined for missing edit', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); const content = await db.readFileEditContent('tc-missing', '/no/such/file'); assert.strictEqual(content, undefined); }); test('persists binary content correctly', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); const binary = new Uint8Array([0, 1, 2, 255, 128, 64]); await db.createTurn('turn-1'); @@ -278,7 +294,7 @@ suite('SessionDatabase', () => { }); test('auto-creates turn if it does not exist', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); // storeFileEdit should succeed even without a prior createTurn call await db.storeFileEdit({ @@ -302,13 +318,13 @@ suite('SessionDatabase', () => { suite('turns', () => { test('createTurn is idempotent', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.createTurn('turn-1'); // should not throw }); test('deleteTurn cascades to file edits', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -330,7 +346,7 @@ suite('SessionDatabase', () => { }); test('deleteTurn only removes its own edits', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.createTurn('turn-1'); await db.createTurn('turn-2'); @@ -360,7 +376,7 @@ suite('SessionDatabase', () => { }); test('deleteTurn is a no-op for unknown turn', async () => { - db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(':memory:')); await db.deleteTurn('nonexistent'); // should not throw }); }); @@ -370,7 +386,7 @@ suite('SessionDatabase', () => { suite('dispose', () => { test('methods throw after dispose', async () => { - db = await SessionDatabase.open(dbPath()); + db = await SessionDatabase.open(':memory:'); db.close(); await assert.rejects( @@ -380,7 +396,7 @@ suite('SessionDatabase', () => { }); test('double dispose is safe', async () => { - db = await SessionDatabase.open(dbPath()); + db = await SessionDatabase.open(':memory:'); await db.close(); await db.close(); // should not throw }); @@ -391,22 +407,20 @@ suite('SessionDatabase', () => { suite('lazy open', () => { test('constructor does not open the database', () => { - // Should not throw even if path does not exist yet - db = new SessionDatabase(join(testDir, 'lazy', 'session.db')); + db = new SessionDatabase(':memory:'); disposables.add(db); // No error — the database is not opened until first use }); test('first async call opens and migrates the database', async () => { - db = disposables.add(new SessionDatabase(dbPath())); - // Database file may not exist yet — first call triggers open + db = disposables.add(new SessionDatabase(':memory:')); await db.createTurn('turn-1'); const edits = await db.getFileEdits(['nonexistent']); assert.deepStrictEqual(edits, []); }); test('multiple concurrent calls share the same open promise', async () => { - db = disposables.add(new SessionDatabase(dbPath())); + db = disposables.add(new SessionDatabase(':memory:')); // Fire multiple calls concurrently — all should succeed await Promise.all([ db.createTurn('turn-1'), @@ -416,7 +430,7 @@ suite('SessionDatabase', () => { }); test('dispose during open rejects subsequent calls', async () => { - db = new SessionDatabase(dbPath()); + db = new SessionDatabase(':memory:'); await db.close(); await assert.rejects(() => db!.createTurn('turn-1'), /disposed/); });