pr comments

Author: connor4312

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

@@ -265,6 +265,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 = {
@@ -280,6 +281,7 @@ export class AgentService extends Disposable implements IAgentService {
 			};
 			const state = this._stateManager.createSession(summary);
 			state.config = sessionConfig;
+			state.activeClient = config?.activeClient;
 		}
 		this._stateManager.dispatchServerAction({ type: ActionType.SessionReady, session: session.toString() });
 

src/vs/platform/agentHost/node/copilot/copilotAgent.ts

@@ -449,10 +449,12 @@ 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?.length) {
+			if (config.activeClient.customizations !== undefined) {
 				await this._plugins.sync(config.activeClient.clientId, config.activeClient.customizations);
 			}
 		}
@@ -479,6 +481,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;
 		}

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

@@ -12,7 +12,6 @@ import { ILogService } from '../../log/common/log.js';
 import { AHPFileSystemProvider } from '../common/agentHostFileSystemProvider.js';
 import { AgentSession, type IAgentService } from '../common/agentService.js';
 import type { ICommandMap } from '../common/state/protocol/messages.js';
-import { ActionType } from '../common/state/protocol/actions.js';
 import { IActionEnvelope, INotification, isSessionAction, isTerminalAction, type ISessionAction } from '../common/state/sessionActions.js';
 import { MIN_PROTOCOL_VERSION, PROTOCOL_VERSION } from '../common/state/sessionCapabilities.js';
 import {
@@ -24,6 +23,7 @@ import {
 	isJsonRpcNotification,
 	isJsonRpcRequest,
 	JSON_RPC_INTERNAL_ERROR,
+	JsonRpcErrorCodes,
 	ProtocolError,
 	type IAhpServerNotification,
 	type IInitializeParams,
@@ -350,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,
@@ -358,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) {
@@ -369,19 +375,6 @@ export class ProtocolServerHandler extends Disposable {
 			if (createdSession.toString() !== URI.parse(params.session).toString()) {
 				this._logService.warn(`[ProtocolServer] createSession: provider returned URI ${createdSession.toString()} but client requested ${params.session}`);
 			}
-			// If the client eagerly claimed the active client role, dispatch
-			// `session/activeClientChanged` now so the claim is atomic with
-			// creation.
-			if (params.activeClient) {
-				if (params.activeClient.clientId !== _client.clientId) {
-					throw new ProtocolError(JSON_RPC_INTERNAL_ERROR, `createSession.activeClient.clientId must match the connection's clientId`);
-				}
-				this._agentService.dispatchAction({
-					type: ActionType.SessionActiveClientChanged,
-					session: createdSession.toString(),
-					activeClient: params.activeClient,
-				}, _client.clientId, 0);
-			}
 			return null;
 		},
 		disposeSession: async (_client, params) => {

src/vs/platform/agentHost/test/node/agentService.test.ts

@@ -247,6 +247,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 = {
+				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 ---------------------------------------------------

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(),
@@ -582,14 +584,10 @@ suite('ProtocolServerHandler', () => {
 
 	suite('createSession activeClient', () => {
 
-		test('forwards activeClient as SessionActiveClientChanged', async () => {
+		test('forwards activeClient to the agent service', async () => {
 			const newSession = URI.parse('copilot:///eager-session').toString();
-			// Pre-create the session so the client can subscribe at handshake
-			// time. MockAgentService.createSession below is idempotent against
-			// state-manager duplicates.
-			stateManager.createSession(makeSessionSummary(newSession));
 
-			const transport = connectClient('client-1', [newSession]);
+			const transport = connectClient('client-1');
 			transport.sent.length = 0;
 
 			const responsePromise = waitForResponse(transport, 2);
@@ -602,21 +600,15 @@ suite('ProtocolServerHandler', () => {
 					customizations: [{ uri: 'file:///plugin-a', displayName: 'A' }],
 				},
 			}));
-			await responsePromise;
+			const resp = await responsePromise as { result?: unknown; error?: unknown };
 
-			const activeClientMsgs = findNotifications(transport.sent, 'action').filter(m => {
-				const envelope = m.params as unknown as { action: { type: string } };
-				return envelope.action.type === ActionType.SessionActiveClientChanged;
-			});
-			assert.strictEqual(activeClientMsgs.length, 1, 'should emit exactly one SessionActiveClientChanged');
-			const envelope = activeClientMsgs[0].params as unknown as { action: { session: string; activeClient: { clientId: string; tools: { name: string }[]; customizations?: { uri: string }[] } } };
+			assert.strictEqual(resp.error, undefined, 'createSession should succeed');
+			const config = agentService.createSessionConfigs.at(-1);
 			assert.deepStrictEqual({
-				session: envelope.action.session,
-				clientId: envelope.action.activeClient.clientId,
-				toolName: envelope.action.activeClient.tools[0].name,
-				customizationUri: envelope.action.activeClient.customizations?.[0].uri,
+				clientId: config?.activeClient?.clientId,
+				toolName: config?.activeClient?.tools[0]?.name,
+				customizationUri: config?.activeClient?.customizations?.[0].uri,
 			}, {
-				session: newSession,
 				clientId: 'client-1',
 				toolName: 't1',
 				customizationUri: 'file:///plugin-a',
@@ -625,9 +617,8 @@ suite('ProtocolServerHandler', () => {
 
 		test('rejects createSession when activeClient.clientId mismatches', async () => {
 			const newSession = URI.parse('copilot:///mismatch-session').toString();
-			stateManager.createSession(makeSessionSummary(newSession));
 
-			const transport = connectClient('client-1', [newSession]);
+			const transport = connectClient('client-1');
 			transport.sent.length = 0;
 
 			const responsePromise = waitForResponse(transport, 2);
@@ -643,11 +634,7 @@ suite('ProtocolServerHandler', () => {
 
 			assert.ok(resp.error, 'response should be an error');
 			assert.strictEqual(resp.result, undefined);
-			const activeClientMsgs = findNotifications(transport.sent, 'action').filter(m => {
-				const envelope = m.params as unknown as { action: { type: string } };
-				return envelope.action.type === ActionType.SessionActiveClientChanged;
-			});
-			assert.strictEqual(activeClientMsgs.length, 0, 'no activeClient notification should have been emitted');
+			assert.strictEqual(agentService.createSessionConfigs.length, 0, 'agent service should not have been called');
 		});
 	});
 });

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

@@ -93,14 +93,12 @@ class MockAgentHostService extends mock<IAgentHostService>() {
 				createdAt: Date.now(),
 				modifiedAt: Date.now(),
 			};
-			const initial: ISessionState = { ...createSessionState(summary), lifecycle: SessionLifecycle.Ready };
-			const action: ISessionAction = {
-				type: 'session/activeClientChanged' as const,
-				session: session.toString(),
+			const state: ISessionState = {
+				...createSessionState(summary),
+				lifecycle: SessionLifecycle.Ready,
 				activeClient: config.activeClient,
-			} as ISessionAction;
-			const withClient = sessionReducer(initial, action as Parameters<typeof sessionReducer>[1], () => { });
-			this.sessionStates.set(session.toString(), withClient);
+			};
+			this.sessionStates.set(session.toString(), state);
 		}
 		return session;
 	}