Merge pull request #311129 from microsoft/connor4312/eager-activeclient

connor4312 · web-flow · commit abdc25ae7b2c · 2026-04-20T08:58:01.000-07:00
agentHost: adopt eager activeClient announcement
diff --git a/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts b/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts
@@ -188,6 +188,7 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC
 			model: config?.model,
 			workingDirectory: config?.workingDirectory ? fromAgentHostUri(config.workingDirectory).toString() : undefined,
 			config: config?.config,
+			activeClient: config?.activeClient,
 		});
 		return session;
 	}
diff --git a/src/vs/platform/agentHost/common/agentService.ts b/src/vs/platform/agentHost/common/agentService.ts
@@ -12,7 +12,7 @@ import { createDecorator } from '../../instantiation/common/instantiation.js';
 import type { ISyncedCustomization } from './agentPluginManager.js';
 import type { IAgentSubscription } from './state/agentSubscription.js';
 import type { ICreateTerminalParams, IResolveSessionConfigResult, ISessionConfigCompletionsResult } from './state/protocol/commands.js';
-import { IProtectedResourceMetadata, type IConfigSchema, type IFileEdit, type IModelSelection, type IToolDefinition } from './state/protocol/state.js';
+import { IProtectedResourceMetadata, type IConfigSchema, type IFileEdit, type IModelSelection, type ISessionActiveClient, type IToolDefinition } from './state/protocol/state.js';
 import type { IActionEnvelope, INotification, ISessionAction, ITerminalAction } from './state/sessionActions.js';
 import type { IResourceCopyParams, IResourceCopyResult, IResourceDeleteParams, IResourceDeleteResult, IResourceListResult, IResourceMoveParams, IResourceMoveResult, IResourceReadResult, IResourceWriteParams, IResourceWriteResult, IStateSnapshot } from './state/sessionProtocol.js';
 import { AttachmentType, ComponentToState, SessionInputResponseKind, SessionStatus, StateComponents, type ICustomizationRef, type IPendingMessage, type IRootState, type ISessionInputAnswer, type ISessionInputRequest, type IToolCallResult, type IToolResultContent, type PolicyState, type StringOrMarkdown } from './state/sessionState.js';
@@ -125,6 +125,14 @@ export interface IAgentCreateSessionConfig {
 	readonly session?: URI;
 	readonly workingDirectory?: URI;
 	readonly config?: Record<string, string>;
+	/**
+	 * Eagerly claim the active client role for the new session. When provided,
+	 * the server initializes the session with this client as the active
+	 * client, equivalent to dispatching a `session/activeClientChanged`
+	 * action immediately after creation. The `clientId` MUST match the
+	 * connection's own `clientId`.
+	 */
+	readonly activeClient?: ISessionActiveClient;
 	/** Fork from an existing session at a specific turn. */
 	readonly fork?: {
 		readonly session: URI;
diff --git a/src/vs/platform/agentHost/common/state/protocol/.ahp-version b/src/vs/platform/agentHost/common/state/protocol/.ahp-version
@@ -1 +1 @@
-ab467b2
+0947b17
diff --git a/src/vs/platform/agentHost/common/state/protocol/commands.ts b/src/vs/platform/agentHost/common/state/protocol/commands.ts
@@ -6,7 +6,7 @@
 // allow-any-unicode-comment-file
 // DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts
 
-import type { URI, ISnapshot, ISessionConfigSchema, ISessionSummary, IModelSelection, ITurn, ITerminalClaim } from './state.js';
+import type { URI, ISnapshot, ISessionConfigSchema, ISessionSummary, IModelSelection, ITurn, ITerminalClaim, ISessionActiveClient } from './state.js';
 import type { IActionEnvelope, IStateAction } from './actions.js';
 
 export type { IConfigPropertySchema, IConfigSchema, ISessionConfigPropertySchema, ISessionConfigSchema } from './state.js';
@@ -198,6 +198,15 @@ export interface ICreateSessionParams {
 	 * Keys and values correspond to the schema returned by the server.
 	 */
 	config?: Record<string, string>;
+	/**
+	 * Eagerly claim the active client role for the new session.
+	 *
+	 * When provided, the server initializes the session with this client as the
+	 * active client, equivalent to dispatching a `session/activeClientChanged`
+	 * action immediately after creation. The `clientId` MUST match the
+	 * `clientId` the creating client supplied in `initialize`.
+	 */
+	activeClient?: ISessionActiveClient;
 }
 
 // ─── disposeSession ──────────────────────────────────────────────────────────
diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts
@@ -263,6 +263,7 @@ export class AgentService extends Disposable implements IAgentService {
 			const state = this._stateManager.createSession(summary);
 			state.config = sessionConfig;
 			state.turns = sourceTurns;
+			state.activeClient = config.activeClient;
 		} else {
 			// Create empty state for new sessions
 			const summary: ISessionSummary = {
@@ -278,6 +279,7 @@ export class AgentService extends Disposable implements IAgentService {
 			};
 			const state = this._stateManager.createSession(summary);
 			state.config = sessionConfig;
+			state.activeClient = config?.activeClient;
 		}
 		// Persist initial config values so a subsequent `restoreSession` can
 		// re-hydrate them. We persist the full resolved values (not just the
diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts
@@ -495,6 +495,15 @@ export class CopilotAgent extends Disposable implements IAgent {
 
 		const sessionId = config?.session ? AgentSession.id(config.session) : generateUuid();
 		const sessionUri = AgentSession.uri(this.id, sessionId);
+		let seededActiveClient = false;
+		if (config?.activeClient) {
+			const ac = this._getOrCreateActiveClient(sessionUri);
+			seededActiveClient = true;
+			ac.updateTools(config.activeClient.clientId, config.activeClient.tools);
+			if (config.activeClient.customizations !== undefined) {
+				await this._plugins.sync(config.activeClient.clientId, config.activeClient.customizations);
+			}
+		}
 		const activeClient = this._activeClients.get(sessionUri);
 		const snapshot = activeClient ? await activeClient.snapshot() : undefined;
 		const workingDirectory = await this._resolveSessionWorkingDirectory(config, sessionId);
@@ -518,6 +527,9 @@ export class CopilotAgent extends Disposable implements IAgent {
 			agentSession = this._createAgentSession(factory, sessionId, shellManager, snapshot);
 			await agentSession.initializeSession();
 		} catch (error) {
+			if (seededActiveClient) {
+				this._activeClients.delete(sessionUri);
+			}
 			await this._removeCreatedWorktree(sessionId);
 			throw error;
 		}
diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts
@@ -23,6 +23,7 @@ import {
 	isJsonRpcNotification,
 	isJsonRpcRequest,
 	JSON_RPC_INTERNAL_ERROR,
+	JsonRpcErrorCodes,
 	ProtocolError,
 	type IAhpServerNotification,
 	type IInitializeParams,
@@ -349,6 +350,11 @@ export class ProtocolServerHandler extends Disposable {
 				}
 				fork = { session: URI.parse(params.fork.session), turnIndex, turnId: params.fork.turnId };
 			}
+			// If the client eagerly claimed the active client role, validate
+			// the clientId matches the connection before forwarding.
+			if (params.activeClient && params.activeClient.clientId !== _client.clientId) {
+				throw new ProtocolError(JsonRpcErrorCodes.InvalidParams, `createSession.activeClient.clientId must match the connection's clientId`);
+			}
 			try {
 				createdSession = await this._agentService.createSession({
 					provider: params.provider,
@@ -357,6 +363,7 @@ export class ProtocolServerHandler extends Disposable {
 					session: URI.parse(params.session),
 					fork,
 					config: params.config,
+					activeClient: params.activeClient,
 				});
 			} catch (err) {
 				if (err instanceof ProtocolError) {
diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts
@@ -19,7 +19,7 @@ import { AgentSession } from '../../common/agentService.js';
 import { ISessionDatabase, ISessionDataService } from '../../common/sessionDataService.js';
 import { SessionDatabase } from '../../node/sessionDatabase.js';
 import { ActionType, IActionEnvelope } from '../../common/state/sessionActions.js';
-import { ResponsePartKind, SessionLifecycle, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, TurnState, buildSubagentSessionUri, type IMarkdownResponsePart, type IToolCallCompletedState, type IToolCallResponsePart } from '../../common/state/sessionState.js';
+import { ISessionActiveClient, ResponsePartKind, SessionLifecycle, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, TurnState, buildSubagentSessionUri, type IMarkdownResponsePart, type IToolCallCompletedState, type IToolCallResponsePart } from '../../common/state/sessionState.js';
 import { IProductService } from '../../../product/common/productService.js';
 import { AgentService } from '../../node/agentService.js';
 import { MockAgent } from './mockAgent.js';
@@ -248,6 +248,36 @@ suite('AgentService (node dispatcher)', () => {
 
 			assert.deepStrictEqual(service.stateManager.getSessionState(session.toString())?.config?.values, config);
 		});
+
+		test('seeds activeClient into the initial session state when provided', async () => {
+			service.registerProvider(copilotAgent);
+
+			const envelopes: IActionEnvelope[] = [];
+			disposables.add(service.onDidAction(env => envelopes.push(env)));
+
+			const activeClient: ISessionActiveClient = {
+				clientId: 'client-eager',
+				tools: [{ name: 't1', description: 'd', inputSchema: { type: 'object' } }],
+				customizations: [{ uri: 'file:///plugin-a', displayName: 'A' }],
+			};
+			const session = await service.createSession({ provider: 'copilot', activeClient });
+
+			assert.deepStrictEqual({
+				activeClient: service.stateManager.getSessionState(session.toString())?.activeClient,
+				dispatchedActiveClientChanged: envelopes.some(e => e.action.type === ActionType.SessionActiveClientChanged),
+			}, {
+				activeClient,
+				dispatchedActiveClientChanged: false,
+			});
+		});
+
+		test('omits activeClient from the initial session state when not provided', async () => {
+			service.registerProvider(copilotAgent);
+
+			const session = await service.createSession({ provider: 'copilot' });
+
+			assert.strictEqual(service.stateManager.getSessionState(session.toString())?.activeClient, undefined);
+		});
 	});
 
 	// ---- authenticate ---------------------------------------------------
diff --git a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts
@@ -188,14 +188,14 @@ class TestableCopilotAgent extends CopilotAgent {
 	}
 }
 
-function createTestAgentContext(disposables: Pick<DisposableStore, 'add'>, options?: { sessionDataService?: ISessionDataService; copilotClient?: ICopilotClient; gitService?: TestAgentHostGitService; environmentServiceRegistration?: 'native' | 'none' }): { agent: CopilotAgent; instantiationService: IInstantiationService } {
+function createTestAgentContext(disposables: Pick<DisposableStore, 'add'>, options?: { sessionDataService?: ISessionDataService; copilotClient?: ICopilotClient; gitService?: TestAgentHostGitService; environmentServiceRegistration?: 'native' | 'none'; pluginManager?: IAgentPluginManager }): { agent: CopilotAgent; instantiationService: IInstantiationService } {
 	const services = new ServiceCollection();
 	const logService = new NullLogService();
 	const fileService = disposables.add(new FileService(logService));
 	services.set(ILogService, logService);
 	services.set(IFileService, fileService);
 	services.set(ISessionDataService, options?.sessionDataService ?? createNullSessionDataService());
-	services.set(IAgentPluginManager, new TestAgentPluginManager());
+	services.set(IAgentPluginManager, options?.pluginManager ?? new TestAgentPluginManager());
 	services.set(IAgentHostGitService, options?.gitService ?? new TestAgentHostGitService());
 	services.set(IAgentHostTerminalManager, new TestAgentHostTerminalManager());
 	if (options?.environmentServiceRegistration !== 'none') {
@@ -213,7 +213,7 @@ function createTestAgentContext(disposables: Pick<DisposableStore, 'add'>, optio
 	return { agent, instantiationService };
 }
 
-function createTestAgent(disposables: Pick<DisposableStore, 'add'>, options?: { sessionDataService?: ISessionDataService; copilotClient?: ICopilotClient; gitService?: TestAgentHostGitService; environmentServiceRegistration?: 'native' | 'none' }): CopilotAgent {
+function createTestAgent(disposables: Pick<DisposableStore, 'add'>, options?: { sessionDataService?: ISessionDataService; copilotClient?: ICopilotClient; gitService?: TestAgentHostGitService; environmentServiceRegistration?: 'native' | 'none'; pluginManager?: IAgentPluginManager }): CopilotAgent {
 	return createTestAgentContext(disposables, options).agent;
 }
 
@@ -386,6 +386,75 @@ suite('CopilotAgent', () => {
 		}
 	});
 
+	suite('createSession activeClient eager-claim', () => {
+
+		class SpyingPluginManager extends TestAgentPluginManager {
+			public readonly calls: { clientId: string; customizations: ICustomizationRef[] }[] = [];
+
+			override async syncCustomizations(clientId: string, customizations: ICustomizationRef[], _progress?: (status: ISessionCustomization[]) => void): Promise<ISyncedCustomization[]> {
+				this.calls.push({ clientId, customizations: [...customizations] });
+				return [];
+			}
+		}
+
+		test('createSession seeds activeClient tools and syncs customizations', async () => {
+			const sessionDataService = disposables.add(new TestSessionDataService());
+			const client = new TestCopilotClient([]);
+			const pluginManager = new SpyingPluginManager();
+			// Fail fast inside the SDK factory so we don't need to wire up a
+			// real raw session. The seeding of activeClient and the plugin
+			// sync both happen before `client.createSession` is invoked.
+			client.createSession = async () => { throw new Error('sentinel'); };
+
+			const agent = createTestAgent(disposables, { sessionDataService, copilotClient: client, pluginManager });
+			try {
+				await agent.authenticate('https://api.github.com', 'token');
+
+				const customizations: ICustomizationRef[] = [{ uri: 'file:///plugin-a', displayName: 'Plugin A' }];
+				await assert.rejects(
+					agent.createSession({
+						session: AgentSession.uri('copilot', 'test-session'),
+						workingDirectory: URI.file('/workspace'),
+						activeClient: {
+							clientId: 'client-1',
+							tools: [{ name: 't1', description: 'd', inputSchema: { type: 'object' } }],
+							customizations,
+						},
+					}),
+					(err: Error) => /sentinel/.test(err.message),
+				);
+
+				assert.deepStrictEqual(pluginManager.calls, [{ clientId: 'client-1', customizations }]);
+			} finally {
+				await disposeAgent(agent);
+			}
+		});
+
+		test('createSession without activeClient does not sync customizations', async () => {
+			const sessionDataService = disposables.add(new TestSessionDataService());
+			const client = new TestCopilotClient([]);
+			const pluginManager = new SpyingPluginManager();
+			client.createSession = async () => { throw new Error('sentinel'); };
+
+			const agent = createTestAgent(disposables, { sessionDataService, copilotClient: client, pluginManager });
+			try {
+				await agent.authenticate('https://api.github.com', 'token');
+
+				await assert.rejects(
+					agent.createSession({
+						session: AgentSession.uri('copilot', 'test-session-2'),
+						workingDirectory: URI.file('/workspace'),
+					}),
+					(err: Error) => /sentinel/.test(err.message),
+				);
+
+				assert.deepStrictEqual(pluginManager.calls, []);
+			} finally {
+				await disposeAgent(agent);
+			}
+		});
+	});
+
 	suite('worktree announcement', () => {
 		// Drives a real session through worktree creation (calling the
 		// agent's _resolveSessionWorkingDirectory via a test seam so we don't
diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts
@@ -72,6 +72,7 @@ class MockAgentService implements IAgentService {
 	readonly browsedUris: URI[] = [];
 	readonly browseErrors = new Map<string, Error>();
 	readonly listedSessions: IAgentSessionMetadata[] = [];
+	readonly createSessionConfigs: (IAgentCreateSessionConfig | undefined)[] = [];
 
 	private readonly _onDidAction = new Emitter<import('../../common/state/sessionActions.js').IActionEnvelope>();
 	readonly onDidAction = this._onDidAction.event;
@@ -91,6 +92,7 @@ class MockAgentService implements IAgentService {
 		this._stateManager.dispatchClientAction(action, origin);
 	}
 	async createSession(config?: IAgentCreateSessionConfig): Promise<URI> {
+		this.createSessionConfigs.push(config);
 		const session = config?.session ?? URI.parse('copilot:///new-session');
 		this._stateManager.createSession({
 			resource: session.toString(),
@@ -577,4 +579,62 @@ suite('ProtocolServerHandler', () => {
 		transport2.simulateClose();
 		assert.deepStrictEqual(counts, [1, 1, 0]);
 	});
+
+	// ---- createSession activeClient -------------------------------------
+
+	suite('createSession activeClient', () => {
+
+		test('forwards activeClient to the agent service', async () => {
+			const newSession = URI.parse('copilot:///eager-session').toString();
+
+			const transport = connectClient('client-1');
+			transport.sent.length = 0;
+
+			const responsePromise = waitForResponse(transport, 2);
+			transport.simulateMessage(request(2, 'createSession', {
+				session: newSession,
+				provider: 'copilot',
+				activeClient: {
+					clientId: 'client-1',
+					tools: [{ name: 't1', description: 'd', inputSchema: { type: 'object' } }],
+					customizations: [{ uri: 'file:///plugin-a', displayName: 'A' }],
+				},
+			}));
+			const resp = await responsePromise as { result?: unknown; error?: unknown };
+
+			assert.strictEqual(resp.error, undefined, 'createSession should succeed');
+			const config = agentService.createSessionConfigs.at(-1);
+			assert.deepStrictEqual({
+				clientId: config?.activeClient?.clientId,
+				toolName: config?.activeClient?.tools[0]?.name,
+				customizationUri: config?.activeClient?.customizations?.[0].uri,
+			}, {
+				clientId: 'client-1',
+				toolName: 't1',
+				customizationUri: 'file:///plugin-a',
+			});
+		});
+
+		test('rejects createSession when activeClient.clientId mismatches', async () => {
+			const newSession = URI.parse('copilot:///mismatch-session').toString();
+
+			const transport = connectClient('client-1');
+			transport.sent.length = 0;
+
+			const responsePromise = waitForResponse(transport, 2);
+			transport.simulateMessage(request(2, 'createSession', {
+				session: newSession,
+				provider: 'copilot',
+				activeClient: {
+					clientId: 'other-client',
+					tools: [],
+				},
+			}));
+			const resp = await responsePromise as { result?: unknown; error?: { code: number; message: string } };
+
+			assert.ok(resp.error, 'response should be an error');
+			assert.strictEqual(resp.result, undefined);
+			assert.strictEqual(agentService.createSessionConfigs.length, 0, 'agent service should not have been called');
+		});
+	});
 });
diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts
@@ -2249,6 +2249,12 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 			}
 		}
 
+		const activeClient = {
+			clientId: this._config.connection.clientId,
+			tools: this._clientToolsObs.get().map(toolDataToDefinition),
+			customizations: this._config.customizations?.get() ?? [],
+		};
+
 		let session: URI;
 		try {
 			session = await this._config.connection.createSession({
@@ -2257,6 +2263,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 				workingDirectory,
 				fork,
 				config,
+				activeClient,
 			});
 		} catch (err) {
 			// If authentication is required (e.g. token expired), try interactive auth and retry once
@@ -2270,6 +2277,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 						workingDirectory,
 						fork,
 						config,
+						activeClient,
 					});
 				} else {
 					throw new Error(localize('agentHost.authRequired', "Authentication is required to start a session. Please sign in and try again."));
@@ -2290,10 +2298,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 			});
 		}
 
-		// Claim the active client role with current customizations
-		const customizations = this._config.customizations?.get() ?? [];
-		this._dispatchActiveClient(session, customizations);
-
 		// Start syncing the chat model's pending requests to the protocol
 		this._ensurePendingMessageSubscription(sessionResource, session);
 
diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/loggingAgentConnection.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/loggingAgentConnection.ts
diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts

Original file line number	Diff line number	Diff line change
`@@ -188,6 +188,7 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC`
`188`	`188`	`model: config?.model,`
`189`	`189`	`workingDirectory: config?.workingDirectory ? fromAgentHostUri(config.workingDirectory).toString() : undefined,`
`190`	`190`	`config: config?.config,`
	`191`	`+ activeClient: config?.activeClient,`
`191`	`192`	`});`
`192`	`193`	`return session;`
`193`	`194`	`}`