agentHost: add session-specific metadata

Adds a SQLite DB for session-specific metadata. Stores edits in there.
It can _almost_ restore edits, but I still need to make undoStops be
similarly persisted. That is a project for later this evening.

src/vs/platform/agentHost/test/node/agentService.test.ts

@@ -26,6 +26,7 @@ suite('AgentService (node dispatcher)', () => {
 			_serviceBrand: undefined,
 			getSessionDataDir: () => URI.parse('inmemory:/session-data'),
 			getSessionDataDirById: () => URI.parse('inmemory:/session-data'),
+			openDatabase: () => { throw new Error('not implemented'); },
 			deleteSessionData: async () => { },
 			cleanupOrphanedData: async () => { },
 		};

src/vs/platform/agentHost/test/node/agentSideEffects.test.ts

@@ -74,6 +74,7 @@ suite('AgentSideEffects', () => {
 				_serviceBrand: undefined,
 				getSessionDataDir: () => URI.from({ scheme: Schemas.inMemory, path: '/session-data' }),
 				getSessionDataDirById: () => URI.from({ scheme: Schemas.inMemory, path: '/session-data' }),
+				openDatabase: () => { throw new Error('not implemented'); },
 				deleteSessionData: async () => { },
 				cleanupOrphanedData: async () => { },
 			} satisfies ISessionDataService,

src/vs/platform/agentHost/test/node/fileEditTracker.test.ts

@@ -4,89 +4,161 @@
  *--------------------------------------------------------------------------------------------*/
 
 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 { Schemas } from '../../../../base/common/network.js';
 import { URI } from '../../../../base/common/uri.js';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js';
 import { FileService } from '../../../files/common/fileService.js';
 import { InMemoryFileSystemProvider } from '../../../files/common/inMemoryFilesystemProvider.js';
 import { NullLogService } from '../../../log/common/log.js';
-import { ISessionDataService } from '../../common/sessionDataService.js';
 import { ToolResultContentType } from '../../common/state/sessionState.js';
-import { SessionDataService } from '../../node/sessionDataService.js';
-import { FileEditTracker } from '../../node/copilot/fileEditTracker.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', () => {
 
 	const disposables = new DisposableStore();
 	let fileService: FileService;
-	let sessionDataService: ISessionDataService;
+	let db: SessionDatabase;
 	let tracker: FileEditTracker;
+	let testDir: string;
 
-	const basePath = URI.from({ scheme: Schemas.inMemory, path: '/userData' });
+	setup(async () => {
+		testDir = join(tmpdir(), `vscode-edit-tracker-test-${randomUUID()}`);
+		mkdirSync(testDir, { recursive: true });
 
-	setup(() => {
 		fileService = disposables.add(new FileService(new NullLogService()));
-		disposables.add(fileService.registerProvider(Schemas.inMemory, disposables.add(new InMemoryFileSystemProvider())));
-		sessionDataService = new SessionDataService(basePath, fileService, new NullLogService());
-		tracker = new FileEditTracker('test-session', sessionDataService, 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')));
+		await db.createTurn('turn-1');
+
+		tracker = new FileEditTracker('copilot:/test-session', db, fileService, new NullLogService());
 	});
 
-	teardown(() => disposables.clear());
+	teardown(() => {
+		disposables.clear();
+		rmSync(testDir, { recursive: true, force: true });
+	});
 	ensureNoDisposablesAreLeakedInTestSuite();
 
 	test('tracks edit start and complete for existing file', async () => {
-		const sourceFs = disposables.add(new InMemoryFileSystemProvider());
-		disposables.add(fileService.registerProvider(Schemas.file, sourceFs));
 		await fileService.writeFile(URI.file('/workspace/test.txt'), VSBuffer.fromString('original content\nline 2'));
 
 		await tracker.trackEditStart('/workspace/test.txt');
 		await fileService.writeFile(URI.file('/workspace/test.txt'), VSBuffer.fromString('modified content\nline 2\nline 3'));
 		await tracker.completeEdit('/workspace/test.txt');
 
-		const fileEdit = tracker.takeCompletedEdit('/workspace/test.txt');
+		const fileEdit = await tracker.takeCompletedEdit('turn-1', 'tc-1', '/workspace/test.txt');
 		assert.ok(fileEdit);
 		assert.strictEqual(fileEdit.type, ToolResultContentType.FileEdit);
-		// Both URIs point to snapshots in the session data directory
-		const sessionDir = sessionDataService.getSessionDataDirById('test-session');
-		assert.ok(fileEdit.beforeURI.startsWith(sessionDir.toString()));
-		assert.ok(fileEdit.afterURI.startsWith(sessionDir.toString()));
+
+		// URIs are parseable session-db: URIs
+		const beforeFields = parseSessionDbUri(fileEdit.beforeURI);
+		assert.ok(beforeFields);
+		assert.strictEqual(beforeFields.sessionUri, 'copilot:/test-session');
+		assert.strictEqual(beforeFields.toolCallId, 'tc-1');
+		assert.strictEqual(beforeFields.filePath, '/workspace/test.txt');
+		assert.strictEqual(beforeFields.part, 'before');
+
+		const afterFields = parseSessionDbUri(fileEdit.afterURI);
+		assert.ok(afterFields);
+		assert.strictEqual(afterFields.part, 'after');
+
+		// Content is persisted in the database (wait for fire-and-forget write)
+		await new Promise(r => setTimeout(r, 50));
+
+		const content = await db.readFileEditContent('tc-1', '/workspace/test.txt');
+		assert.ok(content);
+		assert.strictEqual(new TextDecoder().decode(content.beforeContent), 'original content\nline 2');
+		assert.strictEqual(new TextDecoder().decode(content.afterContent), 'modified content\nline 2\nline 3');
 	});
 
 	test('tracks edit for newly created file (no before content)', async () => {
-		const sourceFs = disposables.add(new InMemoryFileSystemProvider());
-		disposables.add(fileService.registerProvider(Schemas.file, sourceFs));
-
 		await tracker.trackEditStart('/workspace/new-file.txt');
 		await fileService.writeFile(URI.file('/workspace/new-file.txt'), VSBuffer.fromString('new file\ncontent'));
 		await tracker.completeEdit('/workspace/new-file.txt');
 
-		const fileEdit = tracker.takeCompletedEdit('/workspace/new-file.txt');
+		const fileEdit = await tracker.takeCompletedEdit('turn-1', 'tc-2', '/workspace/new-file.txt');
 		assert.ok(fileEdit);
-		const sessionDir = sessionDataService.getSessionDataDirById('test-session');
-		assert.ok(fileEdit.afterURI.startsWith(sessionDir.toString()));
+
+		// Wait for the fire-and-forget DB write to complete
+		await new Promise(r => setTimeout(r, 50));
+
+		const content = await db.readFileEditContent('tc-2', '/workspace/new-file.txt');
+		assert.ok(content);
+		assert.strictEqual(new TextDecoder().decode(content.beforeContent), '');
+		assert.strictEqual(new TextDecoder().decode(content.afterContent), 'new file\ncontent');
 	});
 
-	test('takeCompletedEdit returns undefined for unknown file path', () => {
-		const result = tracker.takeCompletedEdit('/nonexistent');
+	test('takeCompletedEdit returns undefined for unknown file path', async () => {
+		const result = await tracker.takeCompletedEdit('turn-1', 'tc-x', '/nonexistent');
 		assert.strictEqual(result, undefined);
 	});
 
-	test('before and after snapshot content can be read back', async () => {
-		const sourceFs = disposables.add(new InMemoryFileSystemProvider());
-		disposables.add(fileService.registerProvider(Schemas.file, sourceFs));
+	test('before and after content can be read from database', async () => {
 		await fileService.writeFile(URI.file('/workspace/file.ts'), VSBuffer.fromString('original'));
 
 		await tracker.trackEditStart('/workspace/file.ts');
 		await fileService.writeFile(URI.file('/workspace/file.ts'), VSBuffer.fromString('modified'));
 		await tracker.completeEdit('/workspace/file.ts');
 
-		const fileEdit = tracker.takeCompletedEdit('/workspace/file.ts');
-		assert.ok(fileEdit);
-		const beforeContent = await fileService.readFile(URI.parse(fileEdit.beforeURI));
-		assert.strictEqual(beforeContent.value.toString(), 'original');
-		const afterContent = await fileService.readFile(URI.parse(fileEdit.afterURI));
-		assert.strictEqual(afterContent.value.toString(), 'modified');
+		await tracker.takeCompletedEdit('turn-1', 'tc-3', '/workspace/file.ts');
+
+		// Wait for the fire-and-forget DB write to complete
+		await new Promise(r => setTimeout(r, 50));
+
+		const content = await db.readFileEditContent('tc-3', '/workspace/file.ts');
+		assert.ok(content);
+		assert.strictEqual(new TextDecoder().decode(content.beforeContent), 'original');
+		assert.strictEqual(new TextDecoder().decode(content.afterContent), 'modified');
+	});
+});
+
+suite('buildSessionDbUri / parseSessionDbUri', () => {
+
+	ensureNoDisposablesAreLeakedInTestSuite();
+
+	test('round-trips a simple URI', () => {
+		const uri = buildSessionDbUri('copilot:/abc-123', 'tc-1', '/workspace/file.ts', 'before');
+		const parsed = parseSessionDbUri(uri);
+		assert.ok(parsed);
+		assert.deepStrictEqual(parsed, {
+			sessionUri: 'copilot:/abc-123',
+			toolCallId: 'tc-1',
+			filePath: '/workspace/file.ts',
+			part: 'before',
+		});
+	});
+
+	test('round-trips with special characters in filePath', () => {
+		const uri = buildSessionDbUri('copilot:/s1', 'tc-2', '/work space/file (1).ts', 'after');
+		const parsed = parseSessionDbUri(uri);
+		assert.ok(parsed);
+		assert.strictEqual(parsed.filePath, '/work space/file (1).ts');
+		assert.strictEqual(parsed.part, 'after');
+	});
+
+	test('round-trips with special characters in toolCallId', () => {
+		const uri = buildSessionDbUri('copilot:/s1', 'call_abc=123&x', '/file.ts', 'before');
+		const parsed = parseSessionDbUri(uri);
+		assert.ok(parsed);
+		assert.strictEqual(parsed.toolCallId, 'call_abc=123&x');
+	});
+
+	test('parseSessionDbUri returns undefined for non-session-db URIs', () => {
+		assert.strictEqual(parseSessionDbUri('file:///foo/bar'), undefined);
+		assert.strictEqual(parseSessionDbUri('https://example.com'), undefined);
+	});
+
+	test('parseSessionDbUri returns undefined for malformed session-db URIs', () => {
+		assert.strictEqual(parseSessionDbUri('session-db:copilot:/s1'), undefined);
+		assert.strictEqual(parseSessionDbUri('session-db:copilot:/s1?toolCallId=tc-1'), undefined);
+		assert.strictEqual(parseSessionDbUri('session-db:copilot:/s1?toolCallId=tc-1&filePath=/f&part=middle'), undefined);
 	});
 });

src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts

@@ -0,0 +1,238 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+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';
+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;
+	const session = AgentSession.uri('copilot', 'test-session');
+
+	setup(() => {
+		testDir = join(tmpdir(), `vscode-map-events-test-${randomUUID()}`);
+		mkdirSync(testDir, { recursive: true });
+	});
+
+	teardown(() => {
+		disposables.clear();
+		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 () => {
+		const events: ISessionEvent[] = [
+			{ type: 'user.message', data: { messageId: 'msg-1', content: 'hello' } },
+			{ type: 'assistant.message', data: { messageId: 'msg-2', content: 'world' } },
+		];
+
+		const result = await mapSessionEvents(session, undefined, events);
+		assert.strictEqual(result.length, 2);
+		assert.deepStrictEqual(result[0], {
+			session,
+			type: 'message',
+			role: 'user',
+			messageId: 'msg-1',
+			content: 'hello',
+			toolRequests: undefined,
+			reasoningOpaque: undefined,
+			reasoningText: undefined,
+			encryptedContent: undefined,
+			parentToolCallId: undefined,
+		});
+		assert.strictEqual(result[1].type, 'message');
+		assert.strictEqual((result[1] as { role: string }).role, 'assistant');
+	});
+
+	test('maps tool start and complete events', async () => {
+		const events: ISessionEvent[] = [
+			{
+				type: 'tool.execution_start',
+				data: { toolCallId: 'tc-1', toolName: 'shell', arguments: { command: 'echo hi' } },
+			},
+			{
+				type: 'tool.execution_complete',
+				data: { toolCallId: 'tc-1', success: true, result: { content: 'hi\n' } },
+			},
+		];
+
+		const result = await mapSessionEvents(session, undefined, events);
+		assert.strictEqual(result.length, 2);
+		assert.strictEqual(result[0].type, 'tool_start');
+		assert.strictEqual(result[1].type, 'tool_complete');
+
+		const complete = result[1] as { result: { content?: readonly { type: string; text?: string }[] } };
+		assert.ok(complete.result.content);
+		assert.strictEqual(complete.result.content[0].type, ToolResultContentType.Text);
+	});
+
+	test('skips tool_complete without matching tool_start', async () => {
+		const events: ISessionEvent[] = [
+			{ type: 'tool.execution_complete', data: { toolCallId: 'orphan', success: true } },
+		];
+
+		const result = await mapSessionEvents(session, undefined, events);
+		assert.strictEqual(result.length, 0);
+	});
+
+	test('ignores unknown event types', async () => {
+		const events: ISessionEvent[] = [
+			{ type: 'some.unknown.event', data: {} },
+			{ type: 'user.message', data: { messageId: 'msg-1', content: 'test' } },
+		];
+
+		const result = await mapSessionEvents(session, undefined, events);
+		assert.strictEqual(result.length, 1);
+	});
+
+	// ---- File edit restoration ------------------------------------------
+
+	suite('file edit restoration', () => {
+
+		test('restores file edits from database for edit tools', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-edit',
+				filePath: '/workspace/file.ts',
+				beforeContent: new TextEncoder().encode('before'),
+				afterContent: new TextEncoder().encode('after'),
+				addedLines: 3,
+				removedLines: 1,
+			});
+
+			const events: ISessionEvent[] = [
+				{
+					type: 'tool.execution_start',
+					data: { toolCallId: 'tc-edit', toolName: 'edit', arguments: { filePath: '/workspace/file.ts' } },
+				},
+				{
+					type: 'tool.execution_complete',
+					data: { toolCallId: 'tc-edit', success: true, result: { content: 'Edited file.ts' } },
+				},
+			];
+
+			const result = await mapSessionEvents(session, db, events);
+			const complete = result[1];
+			assert.strictEqual(complete.type, 'tool_complete');
+
+			const content = (complete as { result: { content?: readonly Record<string, unknown>[] } }).result.content;
+			assert.ok(content);
+			// Should have text content + file edit
+			assert.strictEqual(content.length, 2);
+			assert.strictEqual(content[0].type, ToolResultContentType.Text);
+			assert.strictEqual(content[1].type, ToolResultContentType.FileEdit);
+
+			// File edit URIs should be parseable
+			const fileEdit = content[1] as { beforeURI: string; afterURI: string; diff?: { added?: number; removed?: number } };
+			const beforeFields = parseSessionDbUri(fileEdit.beforeURI);
+			assert.ok(beforeFields);
+			assert.strictEqual(beforeFields.toolCallId, 'tc-edit');
+			assert.strictEqual(beforeFields.filePath, '/workspace/file.ts');
+			assert.strictEqual(beforeFields.part, 'before');
+			assert.deepStrictEqual(fileEdit.diff, { added: 3, removed: 1 });
+		});
+
+		test('handles multiple file edits for one tool call', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-multi',
+				filePath: '/workspace/a.ts',
+				beforeContent: new Uint8Array(0),
+				afterContent: new TextEncoder().encode('a'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-multi',
+				filePath: '/workspace/b.ts',
+				beforeContent: new Uint8Array(0),
+				afterContent: new TextEncoder().encode('b'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+
+			const events: ISessionEvent[] = [
+				{
+					type: 'tool.execution_start',
+					data: { toolCallId: 'tc-multi', toolName: 'write' },
+				},
+				{
+					type: 'tool.execution_complete',
+					data: { toolCallId: 'tc-multi', success: true },
+				},
+			];
+
+			const result = await mapSessionEvents(session, db, events);
+			const content = (result[1] as { result: { content?: readonly Record<string, unknown>[] } }).result.content;
+			assert.ok(content);
+			// Two file edits (no text since result had no content)
+			const fileEdits = content.filter(c => c.type === ToolResultContentType.FileEdit);
+			assert.strictEqual(fileEdits.length, 2);
+		});
+
+		test('works without database (no file edits restored)', async () => {
+			const events: ISessionEvent[] = [
+				{
+					type: 'tool.execution_start',
+					data: { toolCallId: 'tc-1', toolName: 'edit', arguments: { filePath: '/workspace/file.ts' } },
+				},
+				{
+					type: 'tool.execution_complete',
+					data: { toolCallId: 'tc-1', success: true, result: { content: 'done' } },
+				},
+			];
+
+			const result = await mapSessionEvents(session, undefined, events);
+			const content = (result[1] as { result: { content?: readonly Record<string, unknown>[] } }).result.content;
+			assert.ok(content);
+			// Only text content, no file edits
+			assert.strictEqual(content.length, 1);
+			assert.strictEqual(content[0].type, ToolResultContentType.Text);
+		});
+
+		test('non-edit tools do not get file edits even if db has data', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			const events: ISessionEvent[] = [
+				{
+					type: 'tool.execution_start',
+					data: { toolCallId: 'tc-1', toolName: 'shell', arguments: { command: 'ls' } },
+				},
+				{
+					type: 'tool.execution_complete',
+					data: { toolCallId: 'tc-1', success: true, result: { content: 'files' } },
+				},
+			];
+
+			const result = await mapSessionEvents(session, db, events);
+			const content = (result[1] as { result: { content?: readonly Record<string, unknown>[] } }).result.content;
+			assert.ok(content);
+			assert.strictEqual(content.length, 1);
+			assert.strictEqual(content[0].type, ToolResultContentType.Text);
+		});
+	});
+});

src/vs/platform/agentHost/test/node/sessionDataService.test.ts

@@ -4,16 +4,21 @@
  *--------------------------------------------------------------------------------------------*/
 
 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 { Schemas } from '../../../../base/common/network.js';
 import { URI } from '../../../../base/common/uri.js';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js';
 import { FileService } from '../../../files/common/fileService.js';
+import { DiskFileSystemProvider } from '../../../files/node/diskFileSystemProvider.js';
 import { InMemoryFileSystemProvider } from '../../../files/common/inMemoryFilesystemProvider.js';
 import { NullLogService } from '../../../log/common/log.js';
 import { AgentSession } from '../../common/agentService.js';
 import { SessionDataService } from '../../node/sessionDataService.js';
+import { join } from '../../../../base/common/path.js';
 
 suite('SessionDataService', () => {
 
@@ -80,3 +85,72 @@ suite('SessionDataService', () => {
 		await service.cleanupOrphanedData(new Set());
 	});
 });
+
+suite('SessionDataService — openDatabase ref-counting', () => {
+
+	const disposables = new DisposableStore();
+	let service: SessionDataService;
+	let testDir: string;
+
+	setup(() => {
+		testDir = join(tmpdir(), `vscode-session-data-test-${randomUUID()}`);
+		mkdirSync(testDir, { recursive: true });
+
+		const fileService = disposables.add(new FileService(new NullLogService()));
+		disposables.add(fileService.registerProvider(Schemas.file, disposables.add(new DiskFileSystemProvider(new NullLogService()))));
+		service = new SessionDataService(URI.file(testDir), fileService, new NullLogService());
+	});
+
+	teardown(() => {
+		disposables.clear();
+		rmSync(testDir, { recursive: true, force: true });
+	});
+	ensureNoDisposablesAreLeakedInTestSuite();
+
+	test('returns a functional database reference', async () => {
+		const session = AgentSession.uri('copilot', 'ref-test');
+		const ref = service.openDatabase(session);
+		disposables.add(ref);
+
+		await ref.object.createTurn('turn-1');
+		const edits = await ref.object.getFileEdits([]);
+		assert.deepStrictEqual(edits, []);
+	});
+
+	test('multiple references share the same database', async () => {
+		const session = AgentSession.uri('copilot', 'shared-test');
+		const ref1 = service.openDatabase(session);
+		const ref2 = service.openDatabase(session);
+
+		assert.strictEqual(ref1.object, ref2.object);
+
+		ref1.dispose();
+		ref2.dispose();
+	});
+
+	test('database remains usable until last reference is disposed', async () => {
+		const session = AgentSession.uri('copilot', 'refcount-test');
+		const ref1 = service.openDatabase(session);
+		const ref2 = service.openDatabase(session);
+
+		ref1.dispose();
+
+		// ref2 still works
+		await ref2.object.createTurn('turn-1');
+
+		ref2.dispose();
+	});
+
+	test('new reference after all disposed gets a fresh database', async () => {
+		const session = AgentSession.uri('copilot', 'reopen-test');
+		const ref1 = service.openDatabase(session);
+		const db1 = ref1.object;
+		ref1.dispose();
+
+		const ref2 = service.openDatabase(session);
+		disposables.add(ref2);
+		// New reference — may or may not be the same object, but must be functional
+		await ref2.object.createTurn('turn-1');
+		assert.notStrictEqual(ref2.object, db1);
+	});
+});

src/vs/platform/agentHost/test/node/sessionDatabase.test.ts

@@ -0,0 +1,421 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+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';
+
+suite('SessionDatabase', () => {
+
+	const disposables = new DisposableStore();
+	let testDir: string;
+
+	setup(() => {
+		testDir = join(tmpdir(), `vscode-session-db-test-${randomUUID()}`);
+		mkdirSync(testDir, { recursive: true });
+	});
+
+	teardown(() => {
+		disposables.clear();
+		rmSync(testDir, { recursive: true, force: true });
+	});
+	ensureNoDisposablesAreLeakedInTestSuite();
+
+	function dbPath(name = 'session.db'): string {
+		return join(testDir, name);
+	}
+
+	// ---- Migration system -----------------------------------------------
+
+	suite('migrations', () => {
+
+		test('applies all migrations on a fresh database', async () => {
+			const migrations: ISessionDatabaseMigration[] = [
+				{ version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' },
+				{ version: 2, sql: 'CREATE TABLE t2 (id INTEGER PRIMARY KEY)' },
+			];
+
+			const db = disposables.add(await SessionDatabase.open(dbPath(), migrations));
+
+			const tables = (await db.getAllTables()).sort();
+			assert.deepStrictEqual(tables, ['t1', 't2']);
+		});
+
+		test('reopening with same migrations is a no-op', async () => {
+			const migrations: ISessionDatabaseMigration[] = [
+				{ version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' },
+			];
+
+			const db1 = await SessionDatabase.open(dbPath(), migrations);
+			db1.dispose();
+
+			// Reopen — should not throw (table already exists, migration skipped)
+			const db2 = disposables.add(await SessionDatabase.open(dbPath(), migrations));
+			assert.deepStrictEqual(await db2.getAllTables(), ['t1']);
+		});
+
+		test('only applies new migrations on reopen', async () => {
+			const v1: ISessionDatabaseMigration[] = [
+				{ version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' },
+			];
+			const db1 = await SessionDatabase.open(dbPath(), v1);
+			db1.dispose();
+
+			const v2: ISessionDatabaseMigration[] = [
+				...v1,
+				{ version: 2, sql: 'CREATE TABLE t2 (id INTEGER PRIMARY KEY)' },
+			];
+			const db2 = disposables.add(await SessionDatabase.open(dbPath(), v2));
+
+			const tables = (await db2.getAllTables()).sort();
+			assert.deepStrictEqual(tables, ['t1', 't2']);
+		});
+
+		test('rolls back on migration failure', async () => {
+			const migrations: ISessionDatabaseMigration[] = [
+				{ version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' },
+				{ version: 2, sql: 'THIS IS INVALID SQL' },
+			];
+
+			await assert.rejects(() => SessionDatabase.open(dbPath(), migrations));
+
+			// Reopen with only v1 — t1 should not exist because the whole
+			// transaction was rolled back
+			const db = disposables.add(await SessionDatabase.open(dbPath(), [
+				{ version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' },
+			]));
+			assert.deepStrictEqual(await db.getAllTables(), ['t1']);
+		});
+	});
+
+	// ---- File edits -----------------------------------------------------
+
+	suite('file edits', () => {
+
+		test('store and retrieve a file edit', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/file.ts',
+				beforeContent: new TextEncoder().encode('before'),
+				afterContent: new TextEncoder().encode('after'),
+				addedLines: 5,
+				removedLines: 2,
+			});
+
+			const edits = await db.getFileEdits(['tc-1']);
+			assert.deepStrictEqual(edits, [{
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/file.ts',
+				addedLines: 5,
+				removedLines: 2,
+			}]);
+		});
+
+		test('retrieve multiple edits for a single tool call', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/a.ts',
+				beforeContent: new TextEncoder().encode('a-before'),
+				afterContent: new TextEncoder().encode('a-after'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/b.ts',
+				beforeContent: new TextEncoder().encode('b-before'),
+				afterContent: new TextEncoder().encode('b-after'),
+				addedLines: 1,
+				removedLines: 0,
+			});
+
+			const edits = await db.getFileEdits(['tc-1']);
+			assert.strictEqual(edits.length, 2);
+			assert.strictEqual(edits[0].filePath, '/workspace/a.ts');
+			assert.strictEqual(edits[1].filePath, '/workspace/b.ts');
+		});
+
+		test('retrieve edits across multiple tool calls', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/a.ts',
+				beforeContent: new Uint8Array(0),
+				afterContent: new TextEncoder().encode('hello'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-2',
+				filePath: '/workspace/b.ts',
+				beforeContent: new Uint8Array(0),
+				afterContent: new TextEncoder().encode('world'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+
+			const edits = await db.getFileEdits(['tc-1', 'tc-2']);
+			assert.strictEqual(edits.length, 2);
+
+			// Only tc-2
+			const edits2 = await db.getFileEdits(['tc-2']);
+			assert.strictEqual(edits2.length, 1);
+			assert.strictEqual(edits2[0].toolCallId, 'tc-2');
+		});
+
+		test('returns empty array for unknown tool call IDs', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+			const edits = await db.getFileEdits(['nonexistent']);
+			assert.deepStrictEqual(edits, []);
+		});
+
+		test('returns empty array when given empty array', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+			const edits = await db.getFileEdits([]);
+			assert.deepStrictEqual(edits, []);
+		});
+
+		test('replace on conflict (same toolCallId + filePath)', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/file.ts',
+				beforeContent: new TextEncoder().encode('v1'),
+				afterContent: new TextEncoder().encode('v1-after'),
+				addedLines: 1,
+				removedLines: 0,
+			});
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/file.ts',
+				beforeContent: new TextEncoder().encode('v2'),
+				afterContent: new TextEncoder().encode('v2-after'),
+				addedLines: 3,
+				removedLines: 1,
+			});
+
+			const edits = await db.getFileEdits(['tc-1']);
+			assert.strictEqual(edits.length, 1);
+			assert.strictEqual(edits[0].addedLines, 3);
+
+			const content = await db.readFileEditContent('tc-1', '/workspace/file.ts');
+			assert.ok(content);
+			assert.deepStrictEqual(new TextDecoder().decode(content.beforeContent), 'v2');
+		});
+
+		test('readFileEditContent returns content on demand', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/file.ts',
+				beforeContent: new TextEncoder().encode('before'),
+				afterContent: new TextEncoder().encode('after'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+
+			const content = await db.readFileEditContent('tc-1', '/workspace/file.ts');
+			assert.ok(content);
+			assert.deepStrictEqual(content.beforeContent, new TextEncoder().encode('before'));
+			assert.deepStrictEqual(content.afterContent, new TextEncoder().encode('after'));
+		});
+
+		test('readFileEditContent returns undefined for missing edit', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+			const content = await db.readFileEditContent('tc-missing', '/no/such/file');
+			assert.strictEqual(content, undefined);
+		});
+
+		test('persists binary content correctly', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+			const binary = new Uint8Array([0, 1, 2, 255, 128, 64]);
+
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-bin',
+				filePath: '/workspace/image.png',
+				beforeContent: new Uint8Array(0),
+				afterContent: binary,
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+
+			const content = await db.readFileEditContent('tc-bin', '/workspace/image.png');
+			assert.ok(content);
+			assert.deepStrictEqual(content.afterContent, binary);
+		});
+
+		test('auto-creates turn if it does not exist', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			// storeFileEdit should succeed even without a prior createTurn call
+			await db.storeFileEdit({
+				turnId: 'auto-turn',
+				toolCallId: 'tc-1',
+				filePath: '/x',
+				beforeContent: new Uint8Array(0),
+				afterContent: new Uint8Array(0),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+
+			const edits = await db.getFileEdits(['tc-1']);
+			assert.strictEqual(edits.length, 1);
+			assert.strictEqual(edits[0].turnId, 'auto-turn');
+		});
+	});
+
+	// ---- Turns ----------------------------------------------------------
+
+	suite('turns', () => {
+
+		test('createTurn is idempotent', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+			await db.createTurn('turn-1');
+			await db.createTurn('turn-1'); // should not throw
+		});
+
+		test('deleteTurn cascades to file edits', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			await db.createTurn('turn-1');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/a.ts',
+				beforeContent: new TextEncoder().encode('before'),
+				afterContent: new TextEncoder().encode('after'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+
+			// Edits exist
+			assert.strictEqual((await db.getFileEdits(['tc-1'])).length, 1);
+
+			// Delete the turn — edits should be gone
+			await db.deleteTurn('turn-1');
+			assert.deepStrictEqual(await db.getFileEdits(['tc-1']), []);
+		});
+
+		test('deleteTurn only removes its own edits', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+
+			await db.createTurn('turn-1');
+			await db.createTurn('turn-2');
+			await db.storeFileEdit({
+				turnId: 'turn-1',
+				toolCallId: 'tc-1',
+				filePath: '/workspace/a.ts',
+				beforeContent: new Uint8Array(0),
+				afterContent: new TextEncoder().encode('a'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+			await db.storeFileEdit({
+				turnId: 'turn-2',
+				toolCallId: 'tc-2',
+				filePath: '/workspace/b.ts',
+				beforeContent: new Uint8Array(0),
+				afterContent: new TextEncoder().encode('b'),
+				addedLines: undefined,
+				removedLines: undefined,
+			});
+
+			await db.deleteTurn('turn-1');
+
+			assert.deepStrictEqual(await db.getFileEdits(['tc-1']), []);
+			assert.strictEqual((await db.getFileEdits(['tc-2'])).length, 1);
+		});
+
+		test('deleteTurn is a no-op for unknown turn', async () => {
+			const db = disposables.add(await SessionDatabase.open(dbPath()));
+			await db.deleteTurn('nonexistent'); // should not throw
+		});
+	});
+
+	// ---- Dispose --------------------------------------------------------
+
+	suite('dispose', () => {
+
+		test('methods throw after dispose', async () => {
+			const db = await SessionDatabase.open(dbPath());
+			db.dispose();
+
+			await assert.rejects(
+				() => db.createTurn('turn-1'),
+				/disposed/,
+			);
+		});
+
+		test('double dispose is safe', async () => {
+			const db = await SessionDatabase.open(dbPath());
+			db.dispose();
+			db.dispose(); // should not throw
+		});
+	});
+
+	// ---- Lazy open ------------------------------------------------------
+
+	suite('lazy open', () => {
+
+		test('constructor does not open the database', () => {
+			// Should not throw even if path does not exist yet
+			const db = new SessionDatabase(join(testDir, 'lazy', 'session.db'));
+			disposables.add(db);
+			// No error — the database is not opened until first use
+		});
+
+		test('first async call opens and migrates the database', async () => {
+			const db = disposables.add(new SessionDatabase(dbPath()));
+			// Database file may not exist yet — first call triggers open
+			await db.createTurn('turn-1');
+			const edits = await db.getFileEdits(['nonexistent']);
+			assert.deepStrictEqual(edits, []);
+		});
+
+		test('multiple concurrent calls share the same open promise', async () => {
+			const db = disposables.add(new SessionDatabase(dbPath()));
+			// Fire multiple calls concurrently — all should succeed
+			await Promise.all([
+				db.createTurn('turn-1'),
+				db.createTurn('turn-2'),
+				db.getFileEdits([]),
+			]);
+		});
+
+		test('dispose during open rejects subsequent calls', async () => {
+			const db = new SessionDatabase(dbPath());
+			db.dispose();
+			await assert.rejects(() => db.createTurn('turn-1'), /disposed/);
+		});
+	});
+});

src/vs/platform/agentHost/test/node/sessionStateManager.test.ts

@@ -121,7 +121,7 @@ suite('SessionStateManager', () => {
 		assert.deepStrictEqual(envelopes[0].origin, origin);
 	});
 
-	test('removeSession clears state and emits notification', () => {
+	test('removeSession clears state without notification', () => {
 		manager.createSession(makeSessionSummary());
 
 		const notifications: INotification[] = [];
@@ -129,6 +129,19 @@ suite('SessionStateManager', () => {
 
 		manager.removeSession(sessionUri);
 
+		assert.strictEqual(manager.getSessionState(sessionUri), undefined);
+		assert.strictEqual(manager.getSnapshot(sessionUri), undefined);
+		assert.strictEqual(notifications.length, 0);
+	});
+
+	test('deleteSession clears state and emits notification', () => {
+		manager.createSession(makeSessionSummary());
+
+		const notifications: INotification[] = [];
+		disposables.add(manager.onDidEmitNotification(n => notifications.push(n)));
+
+		manager.deleteSession(sessionUri);
+
 		assert.strictEqual(manager.getSessionState(sessionUri), undefined);
 		assert.strictEqual(manager.getSnapshot(sessionUri), undefined);
 		assert.strictEqual(notifications.length, 1);