Await pending session DB writes on agent host shutdown

[roblourens]

Fixes a flake in the Protocol WebSocket - Session Config persistence across restarts > persisted config values are restored on subscribe after server restart integration test, where the assertion saw the previous branch: 'main' value instead of the updated branch: 'release'.

Root cause

The agent host server's SIGTERM/SIGINT handler called process.exit(0) synchronously, abandoning any fire-and-forget SQLite writes that were still in flight — configValues, customTitle, isRead/isDone, and per-session diffs. Under CI load the most recent SessionConfigChanged write could lose the race against shutdown, leaving the prior value persisted.

Fix

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 — tracking only the leaf dbRun/dbExec would miss the gap between the method being called and the query actually being queued.

SessionDataService aggregates whenIdle() across all live per-session DBs (tracked via the existing ref-counted collection). The server's shutdown handler now awaits this with a 3s raceTimeout before disposing, so a stuck write cannot hang shutdown indefinitely.

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

Verification

• 632/632 platform/agentHost tests pass locally.
• 10/10 stability runs of the previously-flaky test pass with the fix.
• 5/5 runs fail when the await raceTimeout(...) line is commented out, confirming the test now deterministically catches the race rather than masking it with a sleep.
• npm run compile-check-ts-native clean.

(Written by Copilot)

skipped in https://github.com/microsoft/vscode/pull/311374/changes

[github-actions]

Screenshot Changes

Base: 3aeae51b Current: 063b0e92

Changed (1)

chat/aiCustomizations/aiCustomizationManagementEditor/McpBrowseMode/Light

Before	After

[Copilot]

Pull request overview

This PR fixes a shutdown race in the Agent Host where fire-and-forget SQLite writes could be dropped on SIGTERM/SIGINT, causing flaky restart/persistence integration tests.

Changes:

• Track in-flight write promises in SessionDatabase and expose whenIdle() to await pending writes.
• Add ISessionDataService.whenIdle() to aggregate whenIdle() across all live per-session databases.
• Await whenIdle() (with a 3s raceTimeout) during agent host shutdown; remove the test’s timeout(500) workaround.

Show a summary per file

File	Description
src/vs/platform/agentHost/node/sessionDatabase.ts	Adds pending-write tracking and whenIdle(); wraps mutating DB methods with _track().
src/vs/platform/agentHost/node/sessionDataService.ts	Tracks live DB instances and implements whenIdle() across open databases.
src/vs/platform/agentHost/node/agentHostServerMain.ts	Makes shutdown async and awaits DB idleness with a bounded timeout before exiting.
src/vs/platform/agentHost/common/sessionDataService.ts	Extends ISessionDatabase/ISessionDataService interfaces with whenIdle().
src/vs/platform/agentHost/test/node/protocol/sessionConfig.integrationTest.ts	Removes sleep-based flake workaround; relies on graceful shutdown flush.
src/vs/platform/agentHost/test/node/copilotAgent.test.ts	Updates test ISessionDataService stub to include whenIdle().
src/vs/platform/agentHost/test/node/agentService.test.ts	Updates session data service mocks to include whenIdle().
src/vs/platform/agentHost/test/common/sessionTestHelpers.ts	Updates ISessionDatabase/ISessionDataService test helpers with whenIdle().

Copilot's findings

• Files reviewed: 8/8 changed files
• Comments generated: 2