feat: Fix MCP server mapping integration in Copilot CLI sessions (#4709)

* feat: Fix MCP server mapping integration in Copilot CLI sessions

* Updates

* Updates

* Updates

---------

Co-authored-by: Raymond Zhao <7199958+rzhao271@users.noreply.github.com>

Author: DonJayamanne

package.json

@@ -6368,5 +6368,5 @@
 		"node-gyp": "npm:node-gyp@10.3.1",
 		"zod": "3.25.76"
 	},
-	"vscodeCommit": "cdf4f2f548f0cd60f7804b6837a0512a15a02889"
+	"vscodeCommit": "e1ab55e1215fa458d1aeabcf8afd1a9478d2c47f"
 }

src/extension/chatSessions/copilotcli/node/copilotCli.ts

@@ -25,6 +25,7 @@ import { IInstantiationService } from '../../../../util/vs/platform/instantiatio
 import { IChatPromptFileService } from '../../common/chatPromptFileService';
 import { getWorkingDirectory, IWorkspaceInfo } from '../../common/workspaceInfo';
 import { getCopilotLogger } from './logger';
+import { remapCustomAgentTools, type McpServerMappings } from './mcpHandler';
 import { ensureNodePtyShim } from './nodePtyShim';
 import { ensureRipgrepShim } from './ripgrepShim';
 
@@ -60,7 +61,7 @@ export class CopilotCLISessionOptions {
 		return this.agent?.name;
 	}
 
-	public toSessionOptions(): Readonly<SessionOptions> {
+	public toSessionOptions(mcpServerMappings?: McpServerMappings): Readonly<SessionOptions> {
 		const allOptions: SessionOptions = {
 			clientName: 'vscode',
 		};
@@ -84,6 +85,9 @@ export class CopilotCLISessionOptions {
 		if (this.skillLocations) {
 			allOptions.skillDirectories = this.skillLocations.map(uri => uri.fsPath);
 		}
+		if (mcpServerMappings?.size && this.customAgents && this.mcpServers) {
+			remapCustomAgentTools(this.customAgents, mcpServerMappings, this.mcpServers, this.agent);
+		}
 		if (this.agent) {
 			allOptions.selectedCustomAgent = this.agent;
 		}
@@ -334,7 +338,7 @@ export class CopilotCLIAgents extends Disposable implements ICopilotCLIAgents {
 			});
 		}
 
-		return this._agentsPromise;
+		return this._agentsPromise.then(agents => agents.map(agent => this.cloneAgent(agent)));
 	}
 
 	async getAgentsImpl(): Promise<Readonly<SweCustomAgent>[]> {

src/extension/chatSessions/copilotcli/node/copilotcliSessionService.ts

@@ -46,7 +46,7 @@ import { CopilotCLISessionOptions, ICopilotCLIAgents, ICopilotCLISDK } from './c
 import { CopilotCliBridgeSpanProcessor } from './copilotCliBridgeSpanProcessor';
 import { CopilotCLISession, ICopilotCLISession } from './copilotcliSession';
 import { ICopilotCLISkills } from './copilotCLISkills';
-import { ICopilotCLIMCPHandler } from './mcpHandler';
+import { ICopilotCLIMCPHandler, McpServerMappings } from './mcpHandler';
 import { IChatDebugFileLoggerService } from '../../../../platform/chat/common/chatDebugFileLoggerService';
 
 const COPILOT_CLI_WORKSPACE_JSON_FILE_KEY = 'github.copilot.cli.workspaceSessionFile';
@@ -84,8 +84,8 @@ export interface ICopilotCLISessionService {
 	renameSession(sessionId: string, title: string): Promise<void>;
 
 	// Session wrapper tracking
-	getSession(options: { sessionId: string; model?: string; workspaceInfo: IWorkspaceInfo; readonly: boolean; agent?: SweCustomAgent }, token: CancellationToken): Promise<IReference<ICopilotCLISession> | undefined>;
-	createSession(options: { model?: string; workspaceInfo: IWorkspaceInfo; agent?: SweCustomAgent; sessionId?: string }, token: CancellationToken): Promise<IReference<ICopilotCLISession>>;
+	getSession(options: { sessionId: string; model?: string; workspaceInfo: IWorkspaceInfo; readonly: boolean; agent?: SweCustomAgent; mcpServerMappings?: McpServerMappings }, token: CancellationToken): Promise<IReference<ICopilotCLISession> | undefined>;
+	createSession(options: { model?: string; workspaceInfo: IWorkspaceInfo; agent?: SweCustomAgent; sessionId?: string; mcpServerMappings?: McpServerMappings }, token: CancellationToken): Promise<IReference<ICopilotCLISession>>;
 	forkSession(sessionId: string, requestId: string | undefined, options: { workspaceInfo: IWorkspaceInfo }, token: CancellationToken): Promise<string>;
 	tryGetPartialSesionHistory(sessionId: string): Promise<readonly (ChatRequestTurn2 | ChatResponseTurn2)[] | undefined>;
 }
@@ -511,13 +511,13 @@ export class CopilotCLISessionService extends Disposable implements ICopilotCLIS
 		}
 	}
 
-	public async createSession({ model, workspaceInfo, agent, sessionId }: { model?: string; workspaceInfo: IWorkspaceInfo; agent?: SweCustomAgent; sessionId?: string }, token: CancellationToken): Promise<RefCountedSession> {
+	public async createSession({ model, workspaceInfo, agent, sessionId, mcpServerMappings }: { model?: string; workspaceInfo: IWorkspaceInfo; agent?: SweCustomAgent; sessionId?: string; mcpServerMappings?: McpServerMappings }, token: CancellationToken): Promise<RefCountedSession> {
 		const { mcpConfig: mcpServers, disposable: mcpGateway } = await this.mcpHandler.loadMcpConfig();
 		try {
 			const copilotUrl = this.configurationService.getConfig(ConfigKey.Shared.DebugOverrideProxyUrl) || undefined;
 			const options = await this.createSessionsOptions({ model, workspaceInfo, mcpServers, agent, copilotUrl });
 			const sessionManager = await raceCancellationError(this.getSessionManager(), token);
-			const sdkSession = await sessionManager.createSession({ ...options.toSessionOptions(), sessionId });
+			const sdkSession = await sessionManager.createSession({ ...options.toSessionOptions(mcpServerMappings), sessionId });
 			this._newSessionIds.delete(sdkSession.sessionId);
 			// After the first session creation, the SDK's OTel TracerProvider is
 			// initialized. Install the bridge processor so SDK-native spans flow
@@ -637,7 +637,7 @@ export class CopilotCLISessionService extends Disposable implements ICopilotCLIS
 		return new CopilotCLISessionOptions({ ...options, customAgents, skillLocations }, this.logService);
 	}
 
-	public async getSession({ sessionId, model, workspaceInfo, readonly, agent }: { sessionId: string; model?: string; workspaceInfo: IWorkspaceInfo; readonly: boolean; agent?: SweCustomAgent }, token: CancellationToken): Promise<RefCountedSession | undefined> {
+	public async getSession({ sessionId, model, workspaceInfo, readonly, agent, mcpServerMappings }: { sessionId: string; model?: string; workspaceInfo: IWorkspaceInfo; readonly: boolean; agent?: SweCustomAgent; mcpServerMappings?: McpServerMappings }, token: CancellationToken): Promise<RefCountedSession | undefined> {
 		// https://github.com/microsoft/vscode/issues/276573
 		const lock = this.sessionMutexForGetSession.get(sessionId) ?? new Mutex();
 		this.sessionMutexForGetSession.set(sessionId, lock);
@@ -668,7 +668,7 @@ export class CopilotCLISessionService extends Disposable implements ICopilotCLIS
 				const copilotUrl = this.configurationService.getConfig(ConfigKey.Shared.DebugOverrideProxyUrl) || undefined;
 				const options = await this.createSessionsOptions({ model, agent, workspaceInfo, mcpServers, copilotUrl }, readonly);
 
-				const sdkSession = await sessionManager.getSession({ ...options.toSessionOptions(), sessionId }, !readonly);
+				const sdkSession = await sessionManager.getSession({ ...options.toSessionOptions(mcpServerMappings), sessionId }, !readonly);
 				if (!sdkSession) {
 					this.logService.error(`[CopilotCLISession] CopilotCLI failed to get session ${sessionId}.`);
 					return undefined;

src/extension/chatSessions/copilotcli/node/mcpHandler.ts

@@ -3,20 +3,33 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import type { Session } from '@github/copilot/sdk';
+import type { Session, SessionOptions, SweCustomAgent } from '@github/copilot/sdk';
 import type { CancellationToken } from 'vscode';
 import { IAuthenticationService } from '../../../../platform/authentication/common/authentication';
 import { ConfigKey, IConfigurationService } from '../../../../platform/configuration/common/configurationService';
 import { ILogService } from '../../../../platform/log/common/logService';
 import { IMcpService } from '../../../../platform/mcp/common/mcpService';
 import { createServiceIdentifier } from '../../../../util/common/services';
 import { Disposable, DisposableStore, IDisposable } from '../../../../util/vs/base/common/lifecycle';
+import { hasKey } from '../../../../util/vs/base/common/types';
 import { URI } from '../../../../util/vs/base/common/uri';
 import { generateUuid } from '../../../../util/vs/base/common/uuid';
+import type { LanguageModelToolInformation } from '../../../../vscodeTypes';
 import { GitHubMcpDefinitionProvider } from '../../../githubMcp/common/githubMcpDefinitionProvider';
 
 const toolInvalidCharRe = /[^a-z0-9_-]/gi;
 
+/** The user-facing display label of an MCP server (from VS Code settings). */
+export type MCPDisplayName = string;
+/** The short server name as used in agent definition files (the prefix of `fullReferenceName`). */
+export type MCPServerName = string;
+
+/**
+ * A mapping from friendly MCP server names (as defined in custom agent files)
+ * to VS Code MCP server display labels.
+ */
+export type McpServerMappings = Map<MCPServerName, MCPDisplayName>;
+
 export type MCPServerConfig = NonNullable<Session['mcpServers']>[string];
 
 export interface ICopilotCLIMCPHandler {
@@ -173,3 +186,85 @@ export class CopilotCLIMCPHandler implements ICopilotCLIMCPHandler {
 		}
 	}
 }
+
+/**
+ * Builds a mapping from friendly MCP server names (as defined in custom agent files)
+ * to VS Code MCP server labels.
+ *
+ * Iterates through tools that have an MCP source (detected via structural typing using
+ * {@link hasKey}) and a `fullReferenceName` in the format `<server name>/<tool name>`,
+ * extracting the server name portion as the key and the source's `label` as the value.
+ */
+export function buildMcpServerMappings(tools: ReadonlyMap<LanguageModelToolInformation, boolean>): McpServerMappings {
+	const mappings = new Map<string, string>();
+	for (const [tool] of tools) {
+		if (!tool.source || !hasKey(tool.source, { name: true }) || !tool.fullReferenceName) {
+			continue;
+		}
+		const slashIndex = tool.fullReferenceName.lastIndexOf('/');
+		if (slashIndex > 0) {
+			const serverName = tool.fullReferenceName.substring(0, slashIndex);
+			if (serverName && !mappings.has(serverName) && tool.source.label) {
+				mappings.set(serverName, tool.source.label);
+			}
+		}
+	}
+	return mappings;
+}
+
+/**
+ * Remaps tool references in custom agents from friendly MCP server names to gateway names.
+ *
+ * Agent definition files reference tools as `<friendly server name>/<tool name>`, but the SDK
+ * expects `<gateway name>/<tool name>` where gateway names are the Record keys in the MCP
+ * server config.
+ *
+ * @param customAgents The list of custom agents whose tools will be remapped in place.
+ * @param mcpServerMappings Maps friendly server names (from agent files) → VS Code MCP display labels.
+ * @param mcpServers The MCP server config, keyed by gateway name.
+ * @param selectedAgent Optional selected agent to also remap.
+ */
+export function remapCustomAgentTools(
+	customAgents: SweCustomAgent[],
+	mcpServerMappings: McpServerMappings,
+	mcpServers: SessionOptions['mcpServers'],
+	selectedAgent: SweCustomAgent | undefined,
+): void {
+	if (!mcpServerMappings.size || !mcpServers) {
+		return;
+	}
+	// Build a map from display name → gateway name (the Record key in mcpServers).
+	const displayNameToGatewayName = new Map<string, string>();
+	for (const [gatewayName, config] of Object.entries(mcpServers)) {
+		if (config.displayName) {
+			displayNameToGatewayName.set(config.displayName, gatewayName);
+		}
+	}
+
+	const agentsToRemap = selectedAgent ? [...customAgents, selectedAgent] : customAgents;
+	for (const agent of agentsToRemap) {
+		if (!agent.tools?.length) {
+			continue;
+		}
+		for (let i = 0; i < agent.tools.length; i++) {
+			const tool = agent.tools[i];
+			const slashIndex = tool.lastIndexOf('/'); // Tool names cannot contain '/', so the last slash separates server from tool
+			if (slashIndex < 1) {
+				continue;
+			}
+			const serverName = tool.substring(0, slashIndex);
+			const toolName = tool.substring(slashIndex + 1);
+			if (!serverName || !toolName) {
+				continue;
+			}
+			// First try: map through mcpServerMappings (friendly name → display name) then to gateway name.
+			const displayName = mcpServerMappings.get(serverName);
+			// Also try to look up the server name directly as a display name in the gateway map.
+			const gatewayName = displayName ? displayNameToGatewayName.get(displayName) : displayNameToGatewayName.get(serverName);
+
+			if (gatewayName) {
+				agent.tools[i] = `${gatewayName}/${toolName}`;
+			}
+		}
+	}
+}