agentHost: use memory db for sqlite tests to avoid locking weirdness (#306524)

This commit is contained in:
Connor Peet
2026-03-30 16:38:42 -07:00
committed by GitHub
parent c1f3775929
commit 12f2ac1e27
4 changed files with 69 additions and 79 deletions

View File

@@ -112,7 +112,7 @@ function dbOpen(path: string): Promise<Database> {
* `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<void> {
export async function runMigrations(db: Database, migrations: readonly ISessionDatabaseMigration[]): Promise<void> {
// 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<Database> | undefined;
private _closed: Promise<void> | true | undefined;
protected _dbPromise: Promise<Database> | undefined;
protected _closed: Promise<void> | true | undefined;
private readonly _fileEditSequencer = new SequencerByKey<string>();
constructor(
@@ -174,7 +174,7 @@ export class SessionDatabase implements ISessionDatabase {
return inst;
}
private _ensureDb(): Promise<Database> {
protected _ensureDb(): Promise<Database> {
if (this._closed) {
return Promise.reject(new Error('SessionDatabase has been disposed'));
}

View File

@@ -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();

View File

@@ -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[] = [
{

View File

@@ -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<TestableSessionDatabase> {
const inst = new TestableSessionDatabase(path, migrations);
await inst._ensureDb();
return inst;
}
/** Extract the raw db connection; this instance becomes inert. */
async ejectDb(): Promise<Database> {
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<TestableSessionDatabase> {
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/);
});