agentHost: bound catalog reconciliation work

Publish passive metadata without awaiting full catalog synchronization, batch discovery timestamp and dirty-marker updates, and rotate persisted bounded verification samples instead of rescanning every catalog row on each startup.

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

Author: sandy081

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

@@ -21,6 +21,10 @@ const DEFAULT_INTERVAL_MS = 5 * 60 * 1000;
 const DEFAULT_FULL_VERIFICATION_INTERVAL_MS = 60 * 60 * 1000;
 const DEFAULT_BACKGROUND_DELAY_MS = 1000;
 const RECONCILIATION_CURSOR_STORAGE_KEY = 'agentHost.catalogReconciliation.cursor';
+const VERIFICATION_CURSOR_STORAGE_KEY = 'agentHost.catalogReconciliation.verificationCursor';
+const VERIFICATION_VERSION_STORAGE_KEY = 'agentHost.catalogReconciliation.verificationVersion';
+const LAST_VERIFICATION_STORAGE_KEY = 'agentHost.catalogReconciliation.lastVerification';
+const CATALOG_VERIFICATION_VERSION = 1;
 type AgentHostCatalogSyncPendingReason = Extract<AgentHostCatalogSyncResult, { status: 'pending' }>['reason'];
 type ScheduledPassKind = 'background' | 'periodic';
 
@@ -80,8 +84,9 @@ export class AgentHostCatalogReconciliationService extends Disposable {
 	private readonly _scheduledPass = this._register(new MutableDisposable<IDisposable>());
 	private _scheduledPassKind: ScheduledPassKind | undefined;
 	private _payloadDirtyMark: Promise<void> | undefined;
-	private _initialPayloadDirtyMarkPending = true;
-	private _lastFullVerification = 0;
+	private _initialPayloadDirtyMarkPending: boolean;
+	private _lastCompatibilityVerification: number;
+	private _verificationCursor: string | undefined;
 	private _running: Promise<IAgentHostCatalogReconciliationReport> | undefined;
 	private _rerunRequested = false;
 	private _periodic = false;
@@ -104,6 +109,11 @@ export class AgentHostCatalogReconciliationService extends Disposable {
 		this._backgroundDelayMs = this._nonNegativeInteger(options.backgroundDelayMs, DEFAULT_BACKGROUND_DELAY_MS, 'backgroundDelayMs');
 		this._schedule = options.schedule ?? ((callback, delay) => disposableTimeout(callback, delay));
 		this._now = options.now ?? Date.now;
+		this._initialPayloadDirtyMarkPending = this._storageService.get<number>(VERIFICATION_VERSION_STORAGE_KEY) !== CATALOG_VERIFICATION_VERSION;
+		const lastVerification = this._storageService.get<number>(LAST_VERIFICATION_STORAGE_KEY);
+		this._lastCompatibilityVerification = typeof lastVerification === 'number' && Number.isFinite(lastVerification) && lastVerification <= this._now() ? lastVerification : 0;
+		const verificationCursor = this._storageService.get<string>(VERIFICATION_CURSOR_STORAGE_KEY);
+		this._verificationCursor = typeof verificationCursor === 'string' ? verificationCursor : undefined;
 	}
 
 	schedule(): void {
@@ -207,21 +217,20 @@ export class AgentHostCatalogReconciliationService extends Disposable {
 
 	private async _runSinglePass(token: CancellationToken): Promise<IAgentHostCatalogReconciliationReport> {
 		await this._ensureInitialPayloadDirtyMark();
-		if (this._now() - this._lastFullVerification >= this._fullVerificationIntervalMs) {
-			await this._markAllPayloadsDirty();
-			this._lastFullVerification = this._now();
+		if (this._now() - this._lastCompatibilityVerification >= this._fullVerificationIntervalMs) {
+			await this._markVerificationSampleDirty(token);
 		}
 		const { sessions, receiptBySession } = await this._listDirtySessions();
 		if (sessions.length === 0) {
-			this._storageService.delete(this._cursorStorageKey);
+			this._tryDeleteStorage(this._cursorStorageKey);
 			return { outcomes: [], cursor: undefined };
 		}
 
 		const selected = this._selectBatch(sessions, this._readCursor());
 		const outcomes = await this._runBatch(selected, receiptBySession, token);
 		const cursor = selected.at(-1)?.session.toString();
 		if (cursor && !token.isCancellationRequested) {
-			this._storageService.set(this._cursorStorageKey, cursor);
+			this._trySetStorage(this._cursorStorageKey, cursor);
 		}
 		return { outcomes, cursor };
 	}
@@ -235,11 +244,11 @@ export class AgentHostCatalogReconciliationService extends Disposable {
 			outcomes.push(...await this._runBatch(selected, receiptBySession, token));
 			cursor = selected.at(-1)?.session.toString();
 			if (cursor && !token.isCancellationRequested) {
-				this._storageService.set(this._cursorStorageKey, cursor);
+				this._trySetStorage(this._cursorStorageKey, cursor);
 			}
 		}
 		if (sessions.length === 0) {
-			this._storageService.delete(this._cursorStorageKey);
+			this._tryDeleteStorage(this._cursorStorageKey);
 		}
 		return { outcomes, cursor };
 	}
@@ -556,13 +565,65 @@ export class AgentHostCatalogReconciliationService extends Disposable {
 		}
 		await this._markAllPayloadsDirty();
 		this._initialPayloadDirtyMarkPending = false;
-		this._lastFullVerification = this._now();
+		this._trySetStorage(VERIFICATION_VERSION_STORAGE_KEY, CATALOG_VERIFICATION_VERSION);
+		this._recordCompatibilityVerification();
 	}
 
 	private async _prepareFullVerification(): Promise<void> {
 		await this._markAllPayloadsDirty();
 		this._initialPayloadDirtyMarkPending = false;
-		this._lastFullVerification = this._now();
+		this._trySetStorage(VERIFICATION_VERSION_STORAGE_KEY, CATALOG_VERIFICATION_VERSION);
+		this._recordCompatibilityVerification();
+	}
+
+	private async _markVerificationSampleDirty(token: CancellationToken): Promise<void> {
+		if (token.isCancellationRequested) {
+			return;
+		}
+		const receipts = (await this._catalogDatabase.listSessionsV2Receipts())
+			.filter(receipt => receipt.payloadDirty === 0)
+			.sort((first, second) => compareSessionKeys(first.session, second.session));
+		const selected = this._selectVerificationSample(receipts, this._verificationCursor);
+		await this._catalogDatabase.markSessionsV2PayloadsDirty(selected.map(receipt => receipt.session));
+		if (token.isCancellationRequested) {
+			return;
+		}
+		const cursor = selected.at(-1)?.session;
+		if (cursor) {
+			this._verificationCursor = cursor;
+			this._trySetStorage(VERIFICATION_CURSOR_STORAGE_KEY, cursor);
+		}
+		this._recordCompatibilityVerification();
+	}
+
+	private _selectVerificationSample(receipts: readonly IAgentHostDatabaseSessionV2Receipt[], cursor: string | undefined): readonly IAgentHostDatabaseSessionV2Receipt[] {
+		if (receipts.length === 0) {
+			return [];
+		}
+		const start = cursor === undefined ? 0 : Math.max(0, receipts.findIndex(receipt => compareSessionKeys(receipt.session, cursor) > 0));
+		const ordered = start === 0 ? receipts : [...receipts.slice(start), ...receipts.slice(0, start)];
+		return ordered.slice(0, this._batchSize);
+	}
+
+	private _recordCompatibilityVerification(): void {
+		this._lastCompatibilityVerification = this._now();
+		this._trySetStorage(LAST_VERIFICATION_STORAGE_KEY, this._lastCompatibilityVerification);
+	}
+
+	private _trySetStorage<T>(key: string, value: T): void {
+		try {
+			this._storageService.set(key, value);
+		} catch (error) {
+			this._logService.warn(`[AgentHostCatalogReconciliation] Failed to persist '${key}'`, error);
+		}
+	}
+
+	private _tryDeleteStorage(key: string): void {
+		try {
+			this._storageService.delete(key);
+		} catch (error) {
+			this._logService.warn(`[AgentHostCatalogReconciliation] Failed to delete '${key}'`, error);
+		}
 	}
 
 	private _markAllPayloadsDirty(): Promise<void> {

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

@@ -123,7 +123,7 @@ export interface IAgentHostDatabase extends IDisposable {
 	updateSessionExternal(updates: readonly IAgentHostDatabaseExternalUpdate[]): Promise<void>;
 	/** Advances the durable last-observed modification time. */
 	updateSessionModifiedTime(session: string, modifiedTime: number): Promise<boolean>;
-	/** Advances the durable last-observed modification time for many sessions in one transaction. */
+	/** Advances modification times and marks changed catalog payloads dirty in one transaction. */
 	updateSessionModifiedTimes(updates: readonly IAgentHostDatabaseModifiedTimeUpdate[]): Promise<void>;
 	getSession(session: string): Promise<IAgentHostDatabaseSession | undefined>;
 	listSessions(): Promise<readonly IAgentHostDatabaseSession[]>;
@@ -207,6 +207,8 @@ export interface IAgentHostDatabase extends IDisposable {
 	getSessionV2PayloadDirty(session: string): Promise<number | undefined>;
 	/** Marks every cached payload dirty once so mutations made by older builds are rechecked. */
 	markAllSessionsV2PayloadsDirty(): Promise<void>;
+	/** Marks selected cached payloads dirty in one transaction. */
+	markSessionsV2PayloadsDirty(sessions: readonly string[]): Promise<void>;
 	/** Clears a dirty marker only when no newer mutation superseded it. */
 	markSessionV2PayloadClean(session: string, expectedDirty: number): Promise<boolean>;
 	upsertSessionV2(envelope: IAgentHostDatabaseSessionV2Envelope, expectedSessionGeneration: string | undefined): Promise<AgentHostDatabaseSessionV2UpsertResult>;
@@ -389,6 +391,7 @@ function sessionsV2BackfillKey(provider: AgentProvider, payloadVersion: number):
 
 const sessionsV2ExcludedKeyPrefix = 'sessionsV2Excluded:';
 const sessionsV2PayloadDirtyKeyPrefix = 'sessionsV2PayloadDirty:';
+const MODIFIED_TIME_UPDATE_BATCH_SIZE = 400;
 const sessionChatCatalogLegacyMirrorKeyPrefix = 'sessionChatCatalogLegacyMirror:';
 
 function sessionsV2ExcludedProviderPrefix(provider: AgentProvider): string {
@@ -558,10 +561,36 @@ export class AgentHostDatabase implements IAgentHostDatabase {
 			const database = await this._ensureDatabase();
 			await exec(database, 'BEGIN IMMEDIATE');
 			try {
-				for (const { session, modifiedTime } of updates) {
-					await run(database, 'UPDATE sessions_v2 SET modified_time = ? WHERE session_uri = ? AND modified_time < ?', [modifiedTime, session, modifiedTime]);
-					await run(database, 'UPDATE sessions SET modified_time = ? WHERE session_uri = ? AND modified_time < ?', [modifiedTime, session, modifiedTime]);
+				await exec(database, `CREATE TEMP TABLE IF NOT EXISTS session_modified_time_updates (
+					session_uri TEXT PRIMARY KEY NOT NULL,
+					modified_time INTEGER NOT NULL
+				);
+				DELETE FROM session_modified_time_updates`);
+				for (let offset = 0; offset < updates.length; offset += MODIFIED_TIME_UPDATE_BATCH_SIZE) {
+					const batch = updates.slice(offset, offset + MODIFIED_TIME_UPDATE_BATCH_SIZE);
+					await run(database, `INSERT OR REPLACE INTO session_modified_time_updates (session_uri, modified_time) VALUES ${batch.map(() => '(?, ?)').join(', ')}`,
+						batch.flatMap(({ session, modifiedTime }) => [session, modifiedTime]));
 				}
+				await run(database, `INSERT INTO metadata (key, value)
+					SELECT '${sessionsV2PayloadDirtyKeyPrefix}' || sessions_v2.session_uri, '1'
+					FROM sessions_v2
+					INNER JOIN session_modified_time_updates AS updates ON updates.session_uri = sessions_v2.session_uri
+					WHERE sessions_v2.modified_time < updates.modified_time
+					ON CONFLICT(key) DO UPDATE SET value = CAST(value AS INTEGER) + 1`, []);
+				await run(database, `UPDATE sessions_v2 SET modified_time = (
+					SELECT updates.modified_time FROM session_modified_time_updates AS updates
+					WHERE updates.session_uri = sessions_v2.session_uri
+				) WHERE EXISTS (
+					SELECT 1 FROM session_modified_time_updates AS updates
+					WHERE updates.session_uri = sessions_v2.session_uri AND sessions_v2.modified_time < updates.modified_time
+				)`, []);
+				await run(database, `UPDATE sessions SET modified_time = (
+					SELECT updates.modified_time FROM session_modified_time_updates AS updates
+					WHERE updates.session_uri = sessions.session_uri
+				) WHERE EXISTS (
+					SELECT 1 FROM session_modified_time_updates AS updates
+					WHERE updates.session_uri = sessions.session_uri AND sessions.modified_time < updates.modified_time
+				)`, []);
 				await exec(database, 'COMMIT');
 			} catch (error) {
 				await this._rollback(database, error, 'Failed to update session modified times');
@@ -1108,8 +1137,7 @@ export class AgentHostDatabase implements IAgentHostDatabase {
 			try {
 				const exists = await get(database, 'SELECT 1 AS present FROM sessions_v2 WHERE session_uri = ?', [session]);
 				if (exists) {
-					await run(database, `INSERT INTO metadata (key, value) VALUES (?, '1')
-						ON CONFLICT(key) DO UPDATE SET value = CAST(value AS INTEGER) + 1`, [sessionsV2PayloadDirtyKey(session)]);
+					await this._markSessionV2PayloadDirty(database, session);
 				}
 				const row = exists
 					? await get(database, 'SELECT CAST(value AS INTEGER) AS payload_dirty FROM metadata WHERE key = ?', [sessionsV2PayloadDirtyKey(session)])
@@ -1144,6 +1172,32 @@ export class AgentHostDatabase implements IAgentHostDatabase {
 		});
 	}
 
+	async markSessionsV2PayloadsDirty(sessions: readonly string[]): Promise<void> {
+		if (sessions.length === 0) {
+			return;
+		}
+		return this._transactionSequencer.queue(async () => {
+			const database = await this._ensureDatabase();
+			await exec(database, 'BEGIN IMMEDIATE');
+			try {
+				for (const session of sessions) {
+					const exists = await get(database, 'SELECT 1 AS present FROM sessions_v2 WHERE session_uri = ?', [session]);
+					if (exists) {
+						await this._markSessionV2PayloadDirty(database, session);
+					}
+				}
+				await exec(database, 'COMMIT');
+			} catch (error) {
+				await this._rollback(database, error, 'Failed to mark selected sessions_v2 payloads dirty');
+			}
+		});
+	}
+
+	private _markSessionV2PayloadDirty(database: Database, session: string): Promise<void> {
+		return run(database, `INSERT INTO metadata (key, value) VALUES (?, '1')
+			ON CONFLICT(key) DO UPDATE SET value = CAST(value AS INTEGER) + 1`, [sessionsV2PayloadDirtyKey(session)]);
+	}
+
 	async markSessionV2PayloadClean(session: string, expectedDirty: number): Promise<boolean> {
 		this._validatePayloadDirty(expectedDirty);
 		return this._transactionSequencer.queue(async () => {