Await session state before computing the customization migration hint

src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostCustomizationHarness.test.ts

@@ -163,6 +163,7 @@ function createTestCustomAgentsService(connection: MockAgentConnection, rootCust
 			return [...rootCustomizations, ...(sessionState.customizations ?? [])];
 		},
 		getFolderPickerDecision: () => undefined,
+		whenCustomizationsReady: () => Promise.resolve(),
 		getWorkingDirectory(sessionResource: URI): string | undefined {
 			return undefined;
 		},

src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentCustomizationItemProvider.ts

@@ -192,7 +192,11 @@ export class AgentCustomizationItemProvider extends Disposable implements ICusto
 		};
 	}
 
-	async provideSourceFolders(sessionResource: URI, type: PromptsType, _token: CancellationToken): Promise<readonly ICustomizationSourceFolder[]> {
+	async provideSourceFolders(sessionResource: URI, type: PromptsType, token: CancellationToken): Promise<readonly ICustomizationSourceFolder[]> {
+		// One-shot callers (the migration hint) must not read the empty
+		// placeholder a still-loading session reports, or they conclude there is
+		// nothing to migrate.
+		await this._customAgentsService.whenCustomizationsReady(sessionResource, token);
 		const workingDirectories = this._customAgentsService.getWorkingDirectories(sessionResource);
 
 		const folders: ICustomizationSourceFolder[] = [];

src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts

@@ -4,9 +4,11 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { URI } from '../../../../../../base/common/uri.js';
+import { raceCancellation, raceTimeout } from '../../../../../../base/common/async.js';
+import { CancellationToken } from '../../../../../../base/common/cancellation.js';
 import { Emitter, Event } from '../../../../../../base/common/event.js';
 import { StringSHA1 } from '../../../../../../base/common/hash.js';
-import { Disposable, DisposableResourceMap, IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js';
+import { Disposable, DisposableResourceMap, DisposableStore, IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js';
 import { ResourceSet } from '../../../../../../base/common/map.js';
 import { AgentHostMcpServers, AgentHostMcpServersConfigKey } from '../../../../../../platform/agentHost/common/agentHostSchema.js';
 import { IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js';
@@ -42,6 +44,13 @@ export interface IAgentHostCustomizationService {
 
 	getCustomizations(sessionResource: URI): readonly Customization[];
 
+	/**
+	 * Waits up to two seconds for {@link getCustomizations} to reflect the session's first state snapshot; it may resolve earlier on cancellation, failure, or when no agent-host session exists.
+	 * The wait is shared per session, so repeated calls observe one deadline rather than restarting it, and resolve immediately once it has elapsed.
+	 * Intended for one-shot reads; reactive callers should continue listening to {@link onDidChangeCustomizations}.
+	 */
+	whenCustomizationsReady(sessionResource: URI, token?: CancellationToken): Promise<void>;
+
 	/**
 	 * The harness-owned decision about the multi-root Folder picker for a
 	 * session (or `undefined` when the provider expressed no opinion). Read from
@@ -108,6 +117,9 @@ export class NullAgentHostCustomizationService implements IAgentHostCustomizatio
 	getCustomizations(_sessionResource: URI): readonly Customization[] {
 		return [];
 	}
+	whenCustomizationsReady(_sessionResource: URI, _token?: CancellationToken): Promise<void> {
+		return Promise.resolve();
+	}
 	getFolderPickerDecision(_sessionResource: URI): ISessionFolderPickerDecision | undefined {
 		return undefined;
 	}
@@ -185,6 +197,14 @@ export abstract class AbstractAgentHostCustomizationService extends Disposable i
 		return this._resolveTarget(sessionResource)?.customizations ?? [];
 	}
 
+	/**
+	 * Targets resolved by this base are backed by already-materialized provider
+	 * state, so a snapshot is available as soon as the target resolves.
+	 */
+	whenCustomizationsReady(_sessionResource: URI, _token?: CancellationToken): Promise<void> {
+		return Promise.resolve();
+	}
+
 	getFolderPickerDecision(sessionResource: URI): ISessionFolderPickerDecision | undefined {
 		return this._resolveTarget(sessionResource)?.folderPickerDecision;
 	}
@@ -434,9 +454,30 @@ export function getPresentableMcpServerCustomizations(customizations: readonly C
 	return entries.filter(entry => entry.isTopLevel || !topLevelNames.has(entry.server.name));
 }
 
+/**
+ * Upper bound on how long {@link WorkbenchAgentHostCustomizationService.whenCustomizationsReady}
+ * waits for a session's first state snapshot.
+ */
+const SESSION_STATE_SNAPSHOT_TIMEOUT_MS = 2000;
+
+/**
+ * A live session-state subscription plus the memoized readiness wait shared by
+ * every {@link WorkbenchAgentHostCustomizationService.whenCustomizationsReady}
+ * caller for that subscription.
+ */
+interface ISessionStateSubscriptionEntry extends IDisposable {
+	readonly connection: IAgentConnection;
+	readonly backendSession: URI;
+	readonly sub: IAgentSubscription<SessionState>;
+	readiness?: Promise<void>;
+}
+
 export class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCustomizationService {
 
-	private readonly _sessionStateSubscriptions = this._register(new DisposableResourceMap<IDisposable & { readonly connection: IAgentConnection; readonly backendSession: URI; readonly sub: IAgentSubscription<SessionState> }>());
+	private readonly _sessionStateSubscriptions = this._register(new DisposableResourceMap<ISessionStateSubscriptionEntry>());
+
+	/** Overridable so tests can exercise the timeout without real-time waits. */
+	protected readonly _snapshotTimeoutMs: number = SESSION_STATE_SNAPSHOT_TIMEOUT_MS;
 
 	constructor(
 		@IAgentHostConnectionsService private readonly _connectionsService: IAgentHostConnectionsService,
@@ -528,14 +569,58 @@ export class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCus
 		};
 	}
 
+	/**
+	 * Session state arrives asynchronously over the protocol, so a freshly
+	 * created subscription reports `undefined` until its first snapshot lands.
+	 *
+	 * The wait is memoized per subscription so that the many source-folder
+	 * queries behind a single migration hint observe one shared deadline rather
+	 * than restarting it per prompt type. It is bounded because the chat request
+	 * path blocks on this before sending the user's message: once it elapses,
+	 * callers fall back to the current (possibly empty) snapshot rather than
+	 * stalling the send again on every subsequent query.
+	 */
+	override async whenCustomizationsReady(sessionResource: URI, token: CancellationToken = CancellationToken.None): Promise<void> {
+		const target = this._resolveSessionTarget(sessionResource);
+		if (!target) {
+			return;
+		}
+		const entry = this._ensureSessionStateSubscription(sessionResource, target);
+		// An `Error` value counts as resolved: the subscription settled, just not with a snapshot.
+		if (!entry || entry.sub.value !== undefined) {
+			return;
+		}
+
+		// Each caller races the shared wait against its own token, so one
+		// cancellation cannot settle the wait for the others.
+		entry.readiness ??= this._awaitFirstSnapshot(entry.sub);
+		await raceCancellation(entry.readiness, token);
+	}
+
+	private async _awaitFirstSnapshot(subscription: IAgentSubscription<SessionState>): Promise<void> {
+		const store = new DisposableStore();
+		try {
+			const firstSnapshot = new Promise<void>(resolve => {
+				store.add(subscription.onDidChange(() => resolve()));
+				const onDidError = subscription.onDidError;
+				if (onDidError) {
+					store.add(onDidError(() => resolve()));
+				}
+			});
+			await raceTimeout(firstSnapshot, this._snapshotTimeoutMs);
+		} finally {
+			store.dispose();
+		}
+	}
+
 	private _readSessionState(sessionResource: URI): SessionState | undefined {
 		const target = this._resolveSessionTarget(sessionResource);
 		const subscription = target ? this._ensureSessionStateSubscription(sessionResource, target)?.sub : undefined;
 		const value = subscription?.value;
 		return value instanceof Error ? subscription?.verifiedValue : value;
 	}
 
-	private _ensureSessionStateSubscription(sessionResource: URI, target: IAgentHostSessionResolution): (IDisposable & { readonly connection: IAgentConnection; readonly backendSession: URI; readonly sub: IAgentSubscription<SessionState> }) | undefined {
+	private _ensureSessionStateSubscription(sessionResource: URI, target: IAgentHostSessionResolution): ISessionStateSubscriptionEntry | undefined {
 		const existing = this._sessionStateSubscriptions.get(sessionResource);
 		if (existing?.backendSession.toString() === target.backendSession.toString() && existing.connection === target.connection) {
 			return existing;
@@ -547,7 +632,9 @@ export class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCus
 			this._fireCustomizationsChanged();
 			this._fireCustomAgentsChanged();
 		});
-		const entry = {
+		// A new generation starts with no memoized readiness, so the untitled →
+		// real rebind that backs a first send always gets a full wait.
+		const entry: ISessionStateSubscriptionEntry = {
 			connection: target.connection,
 			backendSession: target.backendSession,
 			sub,

src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostCustomizationService.test.ts

@@ -4,7 +4,8 @@
  *--------------------------------------------------------------------------------------------*/
 
 import assert from 'assert';
-import { Event } from '../../../../../../base/common/event.js';
+import { timeout } from '../../../../../../base/common/async.js';
+import { Emitter, Event } from '../../../../../../base/common/event.js';
 import { IReference } from '../../../../../../base/common/lifecycle.js';
 import { ResourceMap } from '../../../../../../base/common/map.js';
 import { URI } from '../../../../../../base/common/uri.js';
@@ -390,4 +391,181 @@ suite('WorkbenchAgentHostCustomizationService', () => {
 			afterError: [retainedRoot.toString()],
 		});
 	});
+
+	/**
+	 * A subscription whose snapshot arrives after the fact, so tests can observe
+	 * the window in which `value` is still `undefined`.
+	 */
+	class LiveSessionSubscription extends mock<IAgentSubscription<SessionState>>() {
+		private readonly _onDidChange = new Emitter<SessionState>();
+		/** Number of listeners installed on this subscription, including readiness waits. */
+		listenerCount = 0;
+		override readonly onDidChange: Event<SessionState> = (listener, thisArgs?, disposables?) => {
+			this.listenerCount++;
+			return this._onDidChange.event(listener, thisArgs, disposables);
+		};
+		private readonly _onDidError = new Emitter<Error>();
+		override readonly onDidError = this._onDidError.event;
+		private current: SessionState | Error | undefined;
+		private confirmed: SessionState | undefined;
+
+		override get value(): SessionState | Error | undefined {
+			return this.current;
+		}
+
+		override get verifiedValue(): SessionState | undefined {
+			return this.confirmed;
+		}
+
+		setSnapshot(state: SessionState): void {
+			this.current = state;
+			this.confirmed = state;
+			this._onDidChange.fire(state);
+		}
+
+		setError(error: Error): void {
+			this.current = error;
+			this._onDidError.fire(error);
+		}
+
+		dispose(): void {
+			this._onDidChange.dispose();
+			this._onDidError.dispose();
+		}
+	}
+
+	function createReadinessSut() {
+		/** Keeps the bounded wait short so timeout coverage costs no real time. */
+		class TestTimeoutCustomizationService extends WorkbenchAgentHostCustomizationService {
+			protected override readonly _snapshotTimeoutMs = 20;
+		}
+		const sessionResource = URI.parse('untitled:chat');
+		const backendSession = URI.parse('copilot:/session');
+		const subscription = store.add(new LiveSessionSubscription());
+		const connection = new class extends mock<IAgentConnection>() {
+			override readonly resourceUris = identityAgentHostResourceUriMapper;
+			override readonly onDidAction = Event.None;
+			override readonly rootState = {
+				value: undefined,
+				verifiedValue: undefined,
+				onDidChange: Event.None,
+				onWillApplyAction: Event.None,
+				onDidApplyAction: Event.None,
+			} satisfies IAgentSubscription<RootState>;
+
+			override getSubscription<T>(_kind: StateComponents): IReference<IAgentSubscription<T>> {
+				return {
+					object: subscription as unknown as IAgentSubscription<T>,
+					dispose: () => { },
+				};
+			}
+		}();
+		const instantiationService = store.add(new TestInstantiationService());
+		instantiationService.stub(ILoggerService, store.add(new NullLoggerService()));
+		instantiationService.stub(IOutputService, {
+			getChannel: () => undefined,
+			getChannelDescriptor: () => undefined,
+			showChannel: async () => { },
+		});
+		const service = store.add(new TestTimeoutCustomizationService(
+			new class extends mock<IAgentHostConnectionsService>() {
+				override readonly ambientConnection = connection;
+			}(),
+			new class extends mock<IAgentHostUntitledProvisionalSessionService>() {
+				override readonly onDidChange = Event.None;
+				override get(): URI {
+					return backendSession;
+				}
+				override getProvisionalWorkingDirectories(): readonly URI[] {
+					return [];
+				}
+			}(),
+			instantiationService,
+			new NullLogService(),
+			new class extends mock<IChatService>() {
+				override readonly onDidDisposeSession = Event.None;
+			}(),
+			new class extends mock<IAgentHostActiveClientService>() { }(),
+		));
+		const directory: Customization = {
+			type: CustomizationType.Directory,
+			id: 'dir-1',
+			uri: 'file:///workspace/.github/skills',
+			name: 'skills',
+			contents: CustomizationType.Skill,
+			writable: true,
+			children: [],
+		} as unknown as Customization;
+		const stateWithDirectory: SessionState = {
+			...createSessionState({
+				resource: backendSession.toString(),
+				provider: 'copilot',
+				title: 'Session',
+				status: SessionStatus.Idle,
+				createdAt: new Date(0).toISOString(),
+				modifiedAt: new Date(0).toISOString(),
+			}),
+			customizations: [directory],
+		};
+		return { service, subscription, sessionResource, stateWithDirectory };
+	}
+
+	test('whenCustomizationsReady defers until the first snapshot rather than reporting no customizations', async () => {
+		const { service, subscription, sessionResource, stateWithDirectory } = createReadinessSut();
+
+		let resolved = false;
+		const ready = service.whenCustomizationsReady(sessionResource).then(() => { resolved = true; });
+		await timeout(0);
+		const whileLoading = { resolved, customizations: service.getCustomizations(sessionResource).map(c => c.id) };
+
+		subscription.setSnapshot(stateWithDirectory);
+		await ready;
+		const afterSnapshot = { resolved, customizations: service.getCustomizations(sessionResource).map(c => c.id) };
+
+		let resolvedAgain = false;
+		service.whenCustomizationsReady(sessionResource).then(() => { resolvedAgain = true; });
+		await timeout(0);
+
+		assert.deepStrictEqual({ whileLoading, afterSnapshot, resolvedAgain }, {
+			whileLoading: { resolved: false, customizations: [] },
+			afterSnapshot: { resolved: true, customizations: ['dir-1'] },
+			resolvedAgain: true,
+		});
+	});
+
+	test('whenCustomizationsReady stops waiting when the subscription fails', async () => {
+		const { service, subscription, sessionResource } = createReadinessSut();
+
+		let resolved = false;
+		const ready = service.whenCustomizationsReady(sessionResource).then(() => { resolved = true; });
+		subscription.setError(new Error('subscription failed'));
+		await ready;
+
+		assert.strictEqual(resolved, true);
+	});
+
+	test('whenCustomizationsReady shares one bounded wait across every prompt-type query', async () => {
+		const { service, subscription, sessionResource } = createReadinessSut();
+
+		// `createFileMigration` queries source folders once per target prompt
+		// type, sequentially, so a never-hydrating subscription must cost one
+		// deadline for the whole hint rather than one per type. Counting
+		// listeners keeps this deterministic; a wall-clock bound would be flaky.
+		// The expected two are the subscription entry's own listener plus the
+		// single shared readiness wait; the point is that it stops growing.
+		await service.whenCustomizationsReady(sessionResource);
+		const afterFirstQuery = subscription.listenerCount;
+		await service.whenCustomizationsReady(sessionResource);
+		await service.whenCustomizationsReady(sessionResource);
+
+		assert.deepStrictEqual({
+			afterFirstQuery,
+			afterThreeQueries: subscription.listenerCount,
+			stillUnresolved: subscription.value === undefined,
+		}, {
+			afterFirstQuery: 2,
+			afterThreeQueries: 2,
+			stillUnresolved: true,
+		});
+	});
 });