agentHost: publish client tools before starting turns (#334385)

* agentHost: publish client tools before starting turns Flush and await the latest active-client snapshot before dispatching new or resumed turns so tools registered during MCP startup are available on the first provider request. Keep the barrier cancellation-aware and cover ordering and cancellation regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * agentHost: harden active client turn barrier Avoid retaining cancelled unresolved claims and reconcile before reading resumed-turn state so concurrent resumes are joined instead of duplicated. Strengthen both race regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Paul Wang <paulwang@Pauls-MacBook-Pro.local> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-09-30 18:46:00 +01:00 · 2026-09-05 02:10:11 +00:00
parent 291338d45a
commit 3a8bef768c
2 changed files with 246 additions and 15 deletions
@@ -7,7 +7,7 @@ import { status } from '../../../../../../base/browser/ui/aria/aria.js';
 import { Delayer, disposableTimeout, raceCancellation } from '../../../../../../base/common/async.js';
 import { decodeBase64, encodeBase64, VSBuffer } from '../../../../../../base/common/buffer.js';
 import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js';
-import { getErrorCode, isCancellationError } from '../../../../../../base/common/errors.js';
+import { CancellationError, getErrorCode, isCancellationError } from '../../../../../../base/common/errors.js';
 import { Emitter, Event } from '../../../../../../base/common/event.js';
 import { MarkdownString } from '../../../../../../base/common/htmlContent.js';
 import { getChatErrorDetailsFromMeta, getCopilotPlanFromEntitlement, IChatErrorContext } from '../../../common/chatErrorMessages.js';
@@ -18,6 +18,7 @@ import { equals } from '../../../../../../base/common/objects.js';
 import { autorun, autorunPerKeyedItem, constObservable, derived, derivedOpts, IObservable, ISettableObservable, observableValue, transaction, waitForState } from '../../../../../../base/common/observable.js';
 import { extUriBiasedIgnorePathCase, isEqual } from '../../../../../../base/common/resources.js';
 import { StopWatch } from '../../../../../../base/common/stopwatch.js';
+import { MicrotaskDelay } from '../../../../../../base/common/symbols.js';
 import { Mutable } from '../../../../../../base/common/types.js';
 import { URI } from '../../../../../../base/common/uri.js';
 import { generateUuid } from '../../../../../../base/common/uuid.js';
@@ -933,11 +934,19 @@ class ActiveClientEntry extends Disposable {
 		await waitForState(this._state, state => !state, undefined, this._cancellation.token);
 	}

-	/** Binds the backend session and requests this client join it. */
-	claim(backendSession: URI): void {
+	/** Binds the backend session and reconciles this client's latest snapshot before returning. */
+	async claim(backendSession: URI, cancellationToken: CancellationToken): Promise<void> {
+		await raceCancellation(this._scope.whenResolved(), cancellationToken);
+		if (cancellationToken.isCancellationRequested) {
+			return;
+		}
+		if (!this._scope.isResolved.get()) {
+			throw new CancellationError();
+		}
 		this._backendSession = backendSession;
 		this._claimRequested = true;
-		this._requestReconciliation();
+		this._requestReconciliation(MicrotaskDelay);
+		await raceCancellation(this.whenSettled(), cancellationToken);
 	}

 	/** Binds the backend session and reconciles without claiming it. */
@@ -949,7 +958,7 @@ class ActiveClientEntry extends Disposable {
 		this._requestReconciliation();
 	}

-	private _requestReconciliation(): void {
+	private _requestReconciliation(delay: number | typeof MicrotaskDelay = this._publishDelayer.defaultDelay): void {
 		if (this._cancellation.token.isCancellationRequested) {
 			return;
 		}
@@ -993,7 +1002,7 @@ class ActiveClientEntry extends Disposable {
 					this._state.set(false, undefined);
 				}
 			}
-		}).catch(() => { /* delayer disposed */ });
+		}, delay).catch(() => { /* delayer disposed */ });
 	}
 }

@@ -2276,10 +2285,9 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 		return { ...activeClient, customizations: [] };
 	}

-	private _ensureActiveClient(sessionResource: URI, backendSession: URI): ActiveClientEntry {
+	private async _ensureActiveClient(sessionResource: URI, backendSession: URI, cancellationToken: CancellationToken): Promise<void> {
 		const entry = this._ensureActiveClientEntry(sessionResource);
-		entry.claim(backendSession);
-		return entry;
+		await entry.claim(backendSession, cancellationToken);
 	}

 	private _ensureActiveClientEntry(sessionResource: URI): ActiveClientEntry {
@@ -2991,7 +2999,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 			return;
 		}
 		const turnId = request.requestId;
-		this._clientDispatchedTurnIds.add(turnId);
 		const chatURI = this._getChatURI(request.sessionResource);
 		const turnChannel = chatURI;
 		const messageAttachments = await this._convertVariablesToAttachments(request);
@@ -3003,7 +3010,11 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 		// turn goes out. We only do this on turn start (not on session open)
 		// so that opening a session doesn't eagerly register this client while
 		// another client is in the middle of a turn.
-		this._ensureActiveClient(request.sessionResource, session);
+		await this._ensureActiveClient(request.sessionResource, session, cancellationToken);
+		if (cancellationToken.isCancellationRequested) {
+			return;
+		}
+		this._clientDispatchedTurnIds.add(turnId);

 		// Model and agent selection now travel on the turn message itself rather
 		// than via the removed `session/modelChanged` / `session/agentChanged`
@@ -3102,17 +3113,21 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 		});
 	}

-	private _handleResumedTurn(
+	private async _handleResumedTurn(
 		session: URI,
 		request: IChatAgentRequest,
 		progress: (parts: IChatProgress[]) => void,
 		cancellationToken: CancellationToken,
 	): Promise<Turn | undefined> {
 		if (cancellationToken.isCancellationRequested) {
-			return Promise.resolve(undefined);
+			return;
 		}
 		const turnId = request.requestId;
 		const chatURI = this._getChatURI(request.sessionResource);
+		await this._ensureActiveClient(request.sessionResource, session, cancellationToken);
+		if (cancellationToken.isCancellationRequested) {
+			return;
+		}
 		const state = this._getSessionState(session.toString(), chatURI);
 		const latestTurn = state?.turns.at(-1);
 		const activeTurn = state?.activeTurn?.id === turnId ? state.activeTurn : undefined;
@@ -3129,7 +3144,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 		const shouldDispatchResume = resumableTurn !== undefined;

 		this._clientDispatchedTurnIds.add(turnId);
-		this._ensureActiveClient(request.sessionResource, session);

 		return new Promise<Turn | undefined>((resolve, reject) => {
 			const store = new DisposableStore();
@@ -20,7 +20,7 @@ import { autorun, constObservable, derived, ISettableObservable, observableValue
 import { mock, upcastPartial } from '../../../../../../base/test/common/mock.js';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js';
 import { runWithFakedTimers } from '../../../../../../base/test/common/timeTravelScheduler.js';
-import { DeferredPromise, timeout } from '../../../../../../base/common/async.js';
+import { DeferredPromise, raceTimeout, timeout } from '../../../../../../base/common/async.js';
 import { Range } from '../../../../../../editor/common/core/range.js';
 import { ITextModel } from '../../../../../../editor/common/model.js';
 import { IModelService } from '../../../../../../editor/common/services/model.js';
@@ -1149,6 +1149,7 @@ async function startTurn(
 		agentHostSessionConfig: Record<string, string>;
 		cancellationToken: CancellationToken;
 		agentId: string;
+		beforeInvoke: () => void;
 	}>,
 ) {
 	const agentId = overrides?.agentId ?? 'agent-host-copilot';
@@ -1166,6 +1167,7 @@ async function startTurn(
 	const registered = chatAgentService.registeredAgents.get(agentId);
 	assert.ok(registered, `${agentId} agent should be registered`);

+	overrides?.beforeInvoke?.();
 	const turnPromise = registered.impl.invoke(
 		makeRequest({
 			message: overrides?.message ?? 'Hello',
@@ -6767,6 +6769,115 @@ suite('AgentHostChatContribution', () => {
 			});
 		}));

+		test('a local retry joins a turn resumed while active client scope resolution is pending', () => runWithFakedTimers({ useFakeTimers: true }, async () => {
+			const { instantiationService, agentHostService, chatAgentService, seedActiveClient } = createTestServices(disposables);
+			agentHostService.setRootState({
+				agents: [{
+					provider: 'copilot',
+					displayName: 'Agent Host - Copilot',
+					description: 'test',
+					models: [],
+				}],
+				activeSessions: 1,
+			});
+			const initialResolution = new DeferredPromise<void>();
+			const isResolved = observableValue('pendingRetryActiveClientResolved', false);
+			disposables.add(seedActiveClient('agent-host-copilot', {
+				customizations: constObservable<readonly ClientPluginCustomization[]>([]),
+				isResolved,
+				whenResolved: initialResolution.p,
+			}));
+			const listController = createSessionListController(disposables, instantiationService, agentHostService);
+			const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, {
+				provider: 'copilot',
+				agentId: 'agent-host-copilot',
+				sessionType: 'agent-host-copilot',
+				fullName: 'Agent Host - Copilot',
+				description: 'Copilot SDK agent running in the local agent host process',
+				connection: agentHostService,
+				connectionAuthority: 'local',
+				isNewSession: resource => listController.isNewSession(resource),
+				onSessionMaterialized: resource => listController.notifySessionMaterialized(resource),
+			}));
+			const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/pending-scope-racing-retry' });
+			const backendSession = AgentSession.uri('copilot', 'pending-scope-racing-retry');
+			const summary: SessionSummary = {
+				resource: backendSession.toString(),
+				provider: 'copilot',
+				title: 'Test',
+				status: SessionStatus.Idle,
+				createdAt: new Date().toISOString(),
+				modifiedAt: new Date().toISOString(),
+			};
+			agentHostService.sessionStates.set(backendSession.toString(), {
+				...createSessionState(summary),
+				lifecycle: SessionLifecycle.Ready,
+				activeClients: [],
+			});
+			const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None);
+			disposables.add(toDisposable(() => chatSession.dispose()));
+			const chatUri = buildDefaultChatUri(backendSession);
+			const turnId = 'pending-scope-racing-turn';
+			const fire = (action: ChatAction, serverSeq: number, origin?: { clientId: string; clientSeq: number }) => {
+				agentHostService.fireAction({ channel: chatUri, action, serverSeq, origin });
+			};
+			fire({
+				type: ActionType.ChatTurnStarted,
+				turnId,
+				startedAt: '2025-01-01T00:00:00.000Z',
+				message: { text: 'original request', origin: { kind: MessageKind.User } },
+			}, 1);
+			fire({
+				type: ActionType.ChatResponsePart,
+				turnId,
+				part: { kind: ResponsePartKind.Markdown, id: 'old-part', content: 'partial response' },
+			}, 2);
+			fire({
+				type: ActionType.ChatError,
+				turnId,
+				duration: 100,
+				part: { kind: ResponsePartKind.Error, error: { errorType: 'requestFailed', message: 'failed' }, resumable: true },
+			}, 3);
+
+			agentHostService.dispatchedActions.length = 0;
+			const registered = chatAgentService.registeredAgents.get('agent-host-copilot');
+			assert.ok(registered);
+			const retryProgress: IChatProgress[][] = [];
+			const retryPromise = registered.impl.invoke(
+				makeRequest({
+					sessionResource,
+					requestId: turnId,
+					acceptedConfirmationData: [{ agentHostResumeTurn: true }],
+				}),
+				parts => retryProgress.push(parts),
+				[],
+				CancellationToken.None,
+			);
+			await timeout(10);
+
+			fire({ type: ActionType.ChatTurnResume, turnId }, 100, { clientId: 'other-client', clientSeq: 1 });
+			isResolved.set(true, undefined);
+			initialResolution.complete();
+			await timeout(10);
+			fire({
+				type: ActionType.ChatResponsePart,
+				turnId,
+				part: { kind: ResponsePartKind.Markdown, id: 'new-part', content: 'continued response' },
+			}, 101);
+			fire({ type: ActionType.ChatTurnComplete, turnId, duration: 200 }, 102);
+
+			const retryResult = await retryPromise;
+			assert.deepStrictEqual({
+				errorDetails: retryResult.errorDetails,
+				resumeDispatches: agentHostService.dispatchedActions.filter(entry => entry.action.type === ActionType.ChatTurnResume).length,
+				progress: retryProgress.flat().filter(part => part.kind === 'markdownContent').map(part => (part as IChatMarkdownContent).content.value),
+			}, {
+				errorDetails: undefined,
+				resumeDispatches: 0,
+				progress: ['partial response', 'continued response'],
+			});
+		}));
+
 		test('a rejected local retry keeps observing a concurrently accepted resume', () => runWithFakedTimers({ useFakeTimers: true }, async () => {
 			const { sessionHandler, agentHostService, chatAgentService } = createContribution(disposables);
 			agentHostService.setRootState({
@@ -12847,6 +12958,62 @@ suite('AgentHostChatContribution', () => {
 			);
 		});

+		test('cancels a turn while active client scope resolution is pending', async () => {
+			const { instantiationService, agentHostService, chatAgentService, seedActiveClient } = createTestServices(disposables);
+			const initialResolution = new DeferredPromise<void>();
+			const isResolved = observableValue('pendingActiveClientResolved', false);
+			disposables.add(seedActiveClient('agent-host-copilot', {
+				customizations: constObservable<readonly ClientPluginCustomization[]>([]),
+				isResolved,
+				whenResolved: initialResolution.p,
+			}));
+			const sessionResource = AgentSession.uri('copilot', 'pending-active-client');
+			const summary: SessionSummary = {
+				resource: sessionResource.toString(),
+				provider: 'copilot',
+				title: 'Test',
+				status: SessionStatus.Idle,
+				createdAt: new Date().toISOString(),
+				modifiedAt: new Date().toISOString(),
+			};
+			agentHostService.sessionStates.set(sessionResource.toString(), {
+				...createSessionState(summary),
+				lifecycle: SessionLifecycle.Ready,
+				activeClients: [],
+			});
+			const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, {
+				provider: 'copilot' as const,
+				agentId: 'agent-host-copilot',
+				sessionType: 'agent-host-copilot',
+				fullName: 'Agent Host - Copilot',
+				description: 'test',
+				connection: agentHostService,
+				connectionAuthority: 'local',
+			}));
+			const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None);
+			disposables.add(toDisposable(() => chatSession.dispose()));
+			const cancellation = disposables.add(new CancellationTokenSource());
+			const registered = chatAgentService.registeredAgents.get('agent-host-copilot')!;
+			const turnPromise = registered.impl.invoke(makeRequest({ sessionResource }), () => { }, [], cancellation.token);
+
+			await timeout(10);
+			cancellation.cancel();
+			isResolved.set(true, undefined);
+			initialResolution.complete();
+			const settled = await raceTimeout(turnPromise.then(() => true), 1_000) ?? false;
+			await timeout(10);
+
+			assert.deepStrictEqual({
+				settled,
+				turns: agentHostService.dispatchedActions.filter(({ action }) => action.type === ActionType.ChatTurnStarted),
+				activeClientUpdates: agentHostService.dispatchedActions.filter(({ action }) => action.type === ActionType.SessionActiveClientSet),
+			}, {
+				settled: true,
+				turns: [],
+				activeClientUpdates: [],
+			});
+		});
+
 		test('re-dispatches activeClientSet when customizations observable changes', async () => {
 			const { instantiationService, agentHostService, chatAgentService, seedActiveClient } = createTestServices(disposables);

@@ -12983,6 +13150,56 @@ suite('AgentHostChatContribution', () => {
 			await turnPromise;
 		}));

+		test('publishes newly registered tools before starting the next turn', async () => {
+			const { instantiationService, agentHostService, chatAgentService, seedActiveClient } = createTestServices(disposables);
+			const tools = observableValue<readonly ToolDefinition[]>('lateTools', []);
+			disposables.add(seedActiveClient('agent-host-copilot', {
+				customizations: constObservable<readonly ClientPluginCustomization[]>([]),
+				tools,
+			}));
+			const sessionResource = AgentSession.uri('copilot', 'late-tools');
+			const summary: SessionSummary = {
+				resource: sessionResource.toString(),
+				provider: 'copilot',
+				title: 'Test',
+				status: SessionStatus.Idle,
+				createdAt: new Date().toISOString(),
+				modifiedAt: new Date().toISOString(),
+			};
+			agentHostService.sessionStates.set(sessionResource.toString(), {
+				...createSessionState(summary),
+				lifecycle: SessionLifecycle.Ready,
+				activeClients: [{ clientId: agentHostService.clientId, tools: [], customizations: [] }],
+			});
+			const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, {
+				provider: 'copilot' as const,
+				agentId: 'agent-host-copilot',
+				sessionType: 'agent-host-copilot',
+				fullName: 'Agent Host - Copilot',
+				description: 'test',
+				connection: agentHostService,
+				connectionAuthority: 'local',
+			}));
+
+			const turn = await startTurn(sessionHandler, agentHostService, chatAgentService, disposables, {
+				sessionResource,
+				beforeInvoke: () => tools.set([{ name: 'late_tool' }], undefined),
+			});
+			const orderedActions = agentHostService.dispatchedActions
+				.filter(({ action }) => action.type === ActionType.SessionActiveClientSet || action.type === ActionType.ChatTurnStarted)
+				.map(({ action }) => action.type === ActionType.SessionActiveClientSet
+					? { type: action.type, tools: action.activeClient.tools.map(tool => tool.name) }
+					: { type: action.type });
+
+			turn.fire({ type: ActionType.ChatTurnComplete, endedAt: '2025-01-01T00:00:00.000Z', session: turn.session, turnId: turn.turnId } as ChatAction);
+			await turn.turnPromise;
+
+			assert.deepStrictEqual(orderedActions, [
+				{ type: ActionType.SessionActiveClientSet, tools: ['late_tool'] },
+				{ type: ActionType.ChatTurnStarted },
+			]);
+		});
+
 		test('coalesces customization, custom-agent, and tool changes into one active-client dispatch', async () => {
 			const { instantiationService, agentHostService, seedActiveClient } = createTestServices(disposables);
 			const customizations = observableValue<readonly ClientPluginCustomization[]>('burstCustomizations', []);