diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostCustomizationHarness.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostCustomizationHarness.test.ts index 2c3711e852f4c3..2941152ef47d05 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostCustomizationHarness.test.ts +++ b/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; }, diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentCustomizationItemProvider.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentCustomizationItemProvider.ts index b69a3520c8d980..fda5db3ccf5191 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentCustomizationItemProvider.ts +++ b/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 { + async provideSourceFolders(sessionResource: URI, type: PromptsType, token: CancellationToken): Promise { + // 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[] = []; diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts index c5fce8067b5bf0..a700b46f9be4c5 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts +++ b/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; + /** * 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 { + 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 { + 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; + readiness?: Promise; +} + export class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCustomizationService { - private readonly _sessionStateSubscriptions = this._register(new DisposableResourceMap }>()); + private readonly _sessionStateSubscriptions = this._register(new DisposableResourceMap()); + + /** 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,6 +569,50 @@ 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 { + 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): Promise { + const store = new DisposableStore(); + try { + const firstSnapshot = new Promise(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; @@ -535,7 +620,7 @@ export class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCus return value instanceof Error ? subscription?.verifiedValue : value; } - private _ensureSessionStateSubscription(sessionResource: URI, target: IAgentHostSessionResolution): (IDisposable & { readonly connection: IAgentConnection; readonly backendSession: URI; readonly sub: IAgentSubscription }) | 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, diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostCustomizationService.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostCustomizationService.test.ts index 825c6011c1fdaf..d4d67814e9f22c 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostCustomizationService.test.ts +++ b/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>() { + private readonly _onDidChange = new Emitter(); + /** Number of listeners installed on this subscription, including readiness waits. */ + listenerCount = 0; + override readonly onDidChange: Event = (listener, thisArgs?, disposables?) => { + this.listenerCount++; + return this._onDidChange.event(listener, thisArgs, disposables); + }; + private readonly _onDidError = new Emitter(); + 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() { + 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; + + override getSubscription(_kind: StateComponents): IReference> { + return { + object: subscription as unknown as IAgentSubscription, + 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() { + override readonly ambientConnection = connection; + }(), + new class extends mock() { + override readonly onDidChange = Event.None; + override get(): URI { + return backendSession; + } + override getProvisionalWorkingDirectories(): readonly URI[] { + return []; + } + }(), + instantiationService, + new NullLogService(), + new class extends mock() { + override readonly onDidDisposeSession = Event.None; + }(), + new class extends mock() { }(), + )); + 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, + }); + }); });