Await pending session DB writes on agent host shutdown (#311432)

* Await pending session DB writes on agent host shutdown

The agent host server's SIGTERM/SIGINT handler called process.exit(0)
synchronously, abandoning any fire-and-forget SQLite writes that were
in flight (configValues, customTitle, isRead/isDone, diffs). Under CI
load this caused the 'Session Config persistence across restarts'
integration test to  the most recent SessionConfigChanged writeflake
could lose the race against shutdown, leaving the previous value
persisted instead.

Track in-flight writes inside SessionDatabase via a _pendingWrites set
populated by every public mutating method (the outermost wrap is
required so the await this._ensureDb() window is also covered).
SessionDataService aggregates whenIdle() across all live per-session
DBs. The server's shutdown handler now awaits this with a 3s
raceTimeout before disposing.

Removes the await timeout(500) hack the test previously needed.

(Written by Copilot)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR review: close ws server first, simplify whenIdle loop

- Close the WebSocket server before awaiting whenIdle() so no further
  actions can be dispatched during the flush window.
- Simplify SessionDataService.whenIdle(): per-DB whenIdle() already
  drains writes against existing DBs, so the outer loop only needs to
  re-pass when a NEW DB was opened during the await. Comment now
  matches the code.

(Written by Copilot)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix Windows: trigger graceful shutdown via stdin close in test

On Windows, child.kill() (SIGTERM) terminates the process unconditionally
without invoking the SIGTERM  so the in-flight setMetadata writehandler
never reaches SQLite and the second phase sees no persisted config at all.
Closing the child's stdin fires process.stdin.on('end', shutdown) on every
platform, exercising the same graceful flush path.

(Written by Copilot)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: roblourens

src/vs/platform/agentHost/common/sessionDataService.ts

@@ -173,6 +173,13 @@ export interface ISessionDatabase extends IDisposable {
 	 */
 	remapTurnIds(mapping: ReadonlyMap<string, string>): Promise<void>;
 
+	/**
+	 * Resolves once all in-flight write operations on this database have
+	 * settled. Used by graceful shutdown to flush fire-and-forget writes
+	 * before the process exits.
+	 */
+	whenIdle(): Promise<void>;
+
 	/**
 	 * Close the database connection. After calling this method, the object is
 	 * considered disposed and all other methods will reject with an error.
@@ -235,4 +242,12 @@ export interface ISessionDataService {
 	 * Called at startup; safe to call multiple times.
 	 */
 	cleanupOrphanedData(knownSessionIds: Set<string>): Promise<void>;
+
+	/**
+	 * Resolves once all in-flight write operations across every currently
+	 * open per-session database have settled. Intended for graceful
+	 * shutdown — fire-and-forget writes (e.g. metadata persistence) would
+	 * otherwise be lost when the process exits.
+	 */
+	whenIdle(): Promise<void>;
 }

src/vs/platform/agentHost/node/agentHostServerMain.ts

@@ -16,6 +16,7 @@ globalThis._VSCODE_FILE_ROOT = fileURLToPath(new URL('../../../..', import.meta.
 import * as fs from 'fs';
 import * as os from 'os';
 import { DisposableStore } from '../../../base/common/lifecycle.js';
+import { raceTimeout } from '../../../base/common/async.js';
 import { URI } from '../../../base/common/uri.js';
 import { generateUuid } from '../../../base/common/uuid.js';
 import { localize } from '../../../nls.js';
@@ -253,12 +254,31 @@ async function main(): Promise<void> {
 
 	// Keep alive until stdin closes or signal
 	process.stdin.resume();
-	process.stdin.on('end', shutdown);
-	process.on('SIGTERM', shutdown);
-	process.on('SIGINT', shutdown);
-
-	function shutdown(): void {
+	process.stdin.on('end', () => { void shutdown(); });
+	process.on('SIGTERM', () => { void shutdown(); });
+	process.on('SIGINT', () => { void shutdown(); });
+
+	let shuttingDown = false;
+	async function shutdown(): Promise<void> {
+		if (shuttingDown) {
+			return;
+		}
+		shuttingDown = true;
 		logService.info('[AgentHostServer] Shutting down...');
+		// Close the WebSocket server first so no further actions can be
+		// dispatched while we wait for in-flight writes to flush — otherwise
+		// a late-arriving action could keep queuing DB writes and either
+		// undermine the flush or push us past the timeout.
+		wsServer.dispose();
+		// Wait for in-flight persistence writes to flush to the per-session
+		// SQLite databases. Without this, a SIGTERM arriving while a
+		// `setMetadata` write (configValues, customTitle, isRead, isDone,
+		// diffs) is in flight can drop the latest value — see the
+		// "Session Config persistence across restarts" integration test.
+		// Capped so a stuck write cannot hang shutdown indefinitely.
+		await raceTimeout(sessionDataService.whenIdle(), 3000, () => {
+			logService.warn('[AgentHostServer] Timed out waiting for session database writes to flush; exiting anyway.');
+		});
 		disposables.dispose();
 		loggerService?.dispose();
 		process.exit(0);

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

@@ -12,6 +12,14 @@ import { ISessionDatabase, ISessionDataService, SESSION_DB_FILENAME } from '../c
 import { SessionDatabase } from './sessionDatabase.js';
 
 class SessionDatabaseCollection extends ReferenceCollection<ISessionDatabase> {
+
+	/**
+	 * The set of currently-open databases. Mirrors what's held by the
+	 * underlying ref-counted map, but exposed so {@link SessionDataService.whenIdle}
+	 * can iterate without reaching into private state.
+	 */
+	readonly liveDatabases = new Set<ISessionDatabase>();
+
 	constructor(
 		private readonly _getDbPath: (key: string) => string,
 		private readonly _logService: ILogService,
@@ -22,10 +30,13 @@ class SessionDatabaseCollection extends ReferenceCollection<ISessionDatabase> {
 	protected createReferencedObject(key: string): ISessionDatabase {
 		const dbPath = this._getDbPath(key);
 		this._logService.trace(`[SessionDataService] Opening database: ${dbPath}`);
-		return new SessionDatabase(dbPath);
+		const db = new SessionDatabase(dbPath);
+		this.liveDatabases.add(db);
+		return db;
 	}
 
 	protected destroyReferencedObject(_key: string, object: ISessionDatabase): void {
+		this.liveDatabases.delete(object);
 		object.dispose();
 	}
 }
@@ -124,4 +135,22 @@ export class SessionDataService implements ISessionDataService {
 			this._logService.warn('[SessionDataService] Failed to run orphan cleanup', err);
 		}
 	}
+
+	async whenIdle(): Promise<void> {
+		// Each `SessionDatabase.whenIdle()` already loops internally until
+		// that DB is quiescent, so the outer loop only needs to handle the
+		// case where a new DB was opened (and writes queued against it)
+		// while we were awaiting an earlier pass.
+		while (true) {
+			const dbs = [...this._databases.liveDatabases];
+			if (dbs.length === 0) {
+				return;
+			}
+			await Promise.all(dbs.map(db => db.whenIdle()));
+			const newOnes = [...this._databases.liveDatabases].filter(db => !dbs.includes(db));
+			if (newOnes.length === 0) {
+				return;
+			}
+		}
+	}
 }

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

@@ -195,6 +195,17 @@ export class SessionDatabase implements ISessionDatabase {
 	protected _closed: Promise<void> | true | undefined;
 	private readonly _fileEditSequencer = new SequencerByKey<string>();
 
+	/**
+	 * In-flight write operations. Tracked so {@link whenIdle} can await them
+	 * before the process exits — without this, a `SIGTERM` arriving between
+	 * a fire-and-forget mutating call (e.g. `setMetadata`) being invoked and
+	 * its underlying SQLite query completing would silently drop the write.
+	 * Every public mutating method routes its returned promise through
+	 * {@link _track}; reads (`getMetadata`, `getFileEdits`, ...) skip
+	 * tracking since shutdown does not need to wait for them.
+	 */
+	private readonly _pendingWrites = new Set<Promise<unknown>>();
+
 	constructor(
 		private readonly _path: string,
 		private readonly _migrations: readonly ISessionDatabaseMigration[] = sessionDatabaseMigrations,
@@ -251,23 +262,29 @@ export class SessionDatabase implements ISessionDatabase {
 
 	// ---- Turns ----------------------------------------------------------
 
-	async createTurn(turnId: string): Promise<void> {
-		const db = await this._ensureDb();
-		await dbRun(db, 'INSERT OR IGNORE INTO turns (id) VALUES (?)', [turnId]);
+	createTurn(turnId: string): Promise<void> {
+		return this._track(async () => {
+			const db = await this._ensureDb();
+			await dbRun(db, 'INSERT OR IGNORE INTO turns (id) VALUES (?)', [turnId]);
+		});
 	}
 
-	async deleteTurn(turnId: string): Promise<void> {
-		const db = await this._ensureDb();
-		await dbRun(db, 'DELETE FROM turns WHERE id = ?', [turnId]);
+	deleteTurn(turnId: string): Promise<void> {
+		return this._track(async () => {
+			const db = await this._ensureDb();
+			await dbRun(db, 'DELETE FROM turns WHERE id = ?', [turnId]);
+		});
 	}
 
-	async setTurnEventId(turnId: string, eventId: string): Promise<void> {
-		const db = await this._ensureDb();
-		await dbRun(db, 'INSERT OR IGNORE INTO turns (id) VALUES (?)', [turnId]);
-		// Only set the event ID if not already set — steering messages
-		// trigger additional user.message events within the same turn,
-		// and we must preserve the first (boundary) event ID.
-		await dbRun(db, 'UPDATE turns SET event_id = ? WHERE id = ? AND event_id IS NULL', [eventId, turnId]);
+	setTurnEventId(turnId: string, eventId: string): Promise<void> {
+		return this._track(async () => {
+			const db = await this._ensureDb();
+			await dbRun(db, 'INSERT OR IGNORE INTO turns (id) VALUES (?)', [turnId]);
+			// Only set the event ID if not already set — steering messages
+			// trigger additional user.message events within the same turn,
+			// and we must preserve the first (boundary) event ID.
+			await dbRun(db, 'UPDATE turns SET event_id = ? WHERE id = ? AND event_id IS NULL', [eventId, turnId]);
+		});
 	}
 
 	async getTurnEventId(turnId: string): Promise<string | undefined> {
@@ -292,36 +309,42 @@ export class SessionDatabase implements ISessionDatabase {
 		return row?.event_id as string | undefined ?? undefined;
 	}
 
-	async truncateFromTurn(turnId: string): Promise<void> {
-		const db = await this._ensureDb();
-		// Delete the target turn and all turns inserted after it (by rowid order).
-		// File edits cascade-delete via the foreign key constraint.
-		await dbRun(db,
-			`DELETE FROM turns WHERE rowid >= (SELECT rowid FROM turns WHERE id = ?)`,
-			[turnId],
-		);
+	truncateFromTurn(turnId: string): Promise<void> {
+		return this._track(async () => {
+			const db = await this._ensureDb();
+			// Delete the target turn and all turns inserted after it (by rowid order).
+			// File edits cascade-delete via the foreign key constraint.
+			await dbRun(db,
+				`DELETE FROM turns WHERE rowid >= (SELECT rowid FROM turns WHERE id = ?)`,
+				[turnId],
+			);
+		});
 	}
 
-	async deleteTurnsAfter(turnId: string): Promise<void> {
-		const db = await this._ensureDb();
-		// Delete all turns inserted after the given turn (by rowid order),
-		// keeping the given turn itself.
-		// File edits cascade-delete via the foreign key constraint.
-		await dbRun(db,
-			`DELETE FROM turns WHERE rowid > (SELECT rowid FROM turns WHERE id = ?)`,
-			[turnId],
-		);
+	deleteTurnsAfter(turnId: string): Promise<void> {
+		return this._track(async () => {
+			const db = await this._ensureDb();
+			// Delete all turns inserted after the given turn (by rowid order),
+			// keeping the given turn itself.
+			// File edits cascade-delete via the foreign key constraint.
+			await dbRun(db,
+				`DELETE FROM turns WHERE rowid > (SELECT rowid FROM turns WHERE id = ?)`,
+				[turnId],
+			);
+		});
 	}
 
-	async deleteAllTurns(): Promise<void> {
-		const db = await this._ensureDb();
-		await dbExec(db, 'DELETE FROM turns');
+	deleteAllTurns(): Promise<void> {
+		return this._track(async () => {
+			const db = await this._ensureDb();
+			await dbExec(db, 'DELETE FROM turns');
+		});
 	}
 
 	// ---- File edits -----------------------------------------------------
 
-	async storeFileEdit(edit: IFileEditRecord & IFileEditContent): Promise<void> {
-		return this._fileEditSequencer.queue(edit.filePath, async () => {
+	storeFileEdit(edit: IFileEditRecord & IFileEditContent): Promise<void> {
+		return this._track(() => this._fileEditSequencer.queue(edit.filePath, async () => {
 			const db = await this._ensureDb();
 			// Ensure the turn exists — lazily insert since the turn record
 			// may not have been created by an explicit createTurn() call.
@@ -343,7 +366,7 @@ export class SessionDatabase implements ISessionDatabase {
 					edit.removedLines ?? null,
 				],
 			);
-		});
+		}));
 	}
 
 	async getFileEdits(toolCallIds: string[]): Promise<IFileEditRecord[]> {
@@ -459,42 +482,74 @@ export class SessionDatabase implements ISessionDatabase {
 		return result;
 	}
 
-	async setMetadata(key: string, value: string): Promise<void> {
-		const db = await this._ensureDb();
-		await dbRun(db, 'INSERT OR REPLACE INTO session_metadata (key, value) VALUES (?, ?)', [key, value]);
+	setMetadata(key: string, value: string): Promise<void> {
+		return this._track(async () => {
+			const db = await this._ensureDb();
+			await dbRun(db, 'INSERT OR REPLACE INTO session_metadata (key, value) VALUES (?, ?)', [key, value]);
+		});
 	}
 
-	async remapTurnIds(mapping: ReadonlyMap<string, string>): Promise<void> {
-		const db = await this._ensureDb();
-		// Defer FK checks to commit time so we can update turns.id and
-		// file_edits.turn_id in any order without mid-statement violations.
-		// This pragma auto-resets after the transaction ends.
-		await dbExec(db, 'PRAGMA defer_foreign_keys = ON');
-		await dbExec(db, 'BEGIN TRANSACTION');
-		try {
-			// Delete turns not present in the mapping (e.g. turns beyond
-			// the fork point). File edits cascade-delete via FK.
-			const oldIds = [...mapping.keys()];
-			if (oldIds.length > 0) {
-				const placeholders = oldIds.map(() => '?').join(',');
-				await dbRun(db,
-					`DELETE FROM turns WHERE id NOT IN (${placeholders})`,
-					oldIds,
-				);
-			}
+	remapTurnIds(mapping: ReadonlyMap<string, string>): Promise<void> {
+		return this._track(async () => {
+			const db = await this._ensureDb();
+			// Defer FK checks to commit time so we can update turns.id and
+			// file_edits.turn_id in any order without mid-statement violations.
+			// This pragma auto-resets after the transaction ends.
+			await dbExec(db, 'PRAGMA defer_foreign_keys = ON');
+			await dbExec(db, 'BEGIN TRANSACTION');
+			try {
+				// Delete turns not present in the mapping (e.g. turns beyond
+				// the fork point). File edits cascade-delete via FK.
+				const oldIds = [...mapping.keys()];
+				if (oldIds.length > 0) {
+					const placeholders = oldIds.map(() => '?').join(',');
+					await dbRun(db,
+						`DELETE FROM turns WHERE id NOT IN (${placeholders})`,
+						oldIds,
+					);
+				}
 
-			// Remap the remaining turn IDs to their new values
-			for (const [oldId, newId] of mapping) {
-				await dbRun(db, 'UPDATE turns SET id = ? WHERE id = ?', [newId, oldId]);
-				await dbRun(db, 'UPDATE file_edits SET turn_id = ? WHERE turn_id = ?', [newId, oldId]);
+				// Remap the remaining turn IDs to their new values
+				for (const [oldId, newId] of mapping) {
+					await dbRun(db, 'UPDATE turns SET id = ? WHERE id = ?', [newId, oldId]);
+					await dbRun(db, 'UPDATE file_edits SET turn_id = ? WHERE turn_id = ?', [newId, oldId]);
+				}
+				await dbExec(db, 'COMMIT');
+			} catch (err) {
+				await dbExec(db, 'ROLLBACK');
+				throw err;
 			}
-			await dbExec(db, 'COMMIT');
-		} catch (err) {
-			await dbExec(db, 'ROLLBACK');
-			throw err;
+		});
+	}
+
+	/**
+	 * Resolves once all currently in-flight write operations have settled.
+	 * Used by graceful shutdown to flush pending fire-and-forget writes
+	 * before the process exits. Should be called from a path where no
+	 * further writes are expected; loops until idle to also drain any
+	 * writes that get queued while we're awaiting.
+	 */
+	async whenIdle(): Promise<void> {
+		while (this._pendingWrites.size > 0) {
+			await Promise.allSettled([...this._pendingWrites]);
 		}
 	}
 
+	/**
+	 * Wrap a mutating operation's promise so {@link whenIdle} can await it.
+	 * Invoke at the **outermost** layer of every public mutating method so
+	 * that any internal awaits (notably `_ensureDb()`) are covered too —
+	 * tracking only the leaf `dbRun`/`dbExec` would miss the window
+	 * between the method being called and the query actually being queued.
+	 */
+	private _track<T>(fn: () => Promise<T>): Promise<T> {
+		const p = fn();
+		this._pendingWrites.add(p);
+		const untrack = () => { this._pendingWrites.delete(p); };
+		p.then(untrack, untrack);
+		return p;
+	}
+
 	async close() {
 		await (this._closed ??= this._dbPromise?.then(db => db.close()).catch(() => { }) || true);
 	}

src/vs/platform/agentHost/test/common/sessionTestHelpers.ts

@@ -90,6 +90,8 @@ export class TestSessionDatabase implements ISessionDatabase {
 
 	async remapTurnIds(_mapping: ReadonlyMap<string, string>): Promise<void> { }
 
+	async whenIdle(): Promise<void> { }
+
 	private _toEditRecords(edits: (IFileEditRecord & IFileEditContent)[]): IFileEditRecord[] {
 		return edits.map(({ beforeContent: _, afterContent: _2, ...metadata }) => metadata);
 	}
@@ -130,6 +132,7 @@ export function createSessionDataService(database: ISessionDatabase = new TestSe
 		tryOpenDatabase: async () => createReference(database),
 		deleteSessionData: async () => { },
 		cleanupOrphanedData: async () => { },
+		whenIdle: async () => { },
 	};
 }
 
@@ -142,6 +145,7 @@ export function createNullSessionDataService(): ISessionDataService {
 		tryOpenDatabase: async () => undefined,
 		deleteSessionData: async () => { },
 		cleanupOrphanedData: async () => { },
+		whenIdle: async () => { },
 	};
 }