Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,7 @@ export class AgenticChatController implements ChatHandlers {
this.#features.workspace,
this.#features.lsp
)
this.#mcpEventHandler = new McpEventHandler(features)
this.#mcpEventHandler = new McpEventHandler(features, telemetryService)
}

async onButtonClick(params: ButtonClickParams): Promise<ButtonClickResult> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,12 @@ import * as sinon from 'sinon'
import { McpEventHandler } from './mcpEventHandler'
import { McpManager } from './mcpManager'
import * as mcpUtils from './mcpUtils'
import { TelemetryService } from '../../../../shared/telemetry/telemetryService'

describe('McpEventHandler error handling', () => {
let eventHandler: McpEventHandler
let features: any
let telemetryService: TelemetryService
let loadStub: sinon.SinonStub

beforeEach(() => {
Expand Down Expand Up @@ -44,8 +46,16 @@ describe('McpEventHandler error handling', () => {
lsp: {},
}

// Create mock telemetry service
telemetryService = {
emitUserTriggerDecision: sinon.stub(),
emitChatInteractWithMessage: sinon.stub(),
emitUserModificationEvent: sinon.stub(),
emitCodeCoverageEvent: sinon.stub(),
} as unknown as TelemetryService

// Create the event handler
eventHandler = new McpEventHandler(features)
eventHandler = new McpEventHandler(features, telemetryService)

// Stub loadPersonaPermissions
sinon
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { Features } from '../../../types'
import { MCP_SERVER_STATUS_CHANGED, McpManager } from './mcpManager'
import { ChatTelemetryController } from '../../../chat/telemetry/chatTelemetryController'
import {
DetailedListGroup,
DetailedListItem,
Expand All @@ -22,6 +23,7 @@ import {
McpServerRuntimeState,
McpServerStatus,
} from './mcpTypes'
import { TelemetryService } from '../../../../shared/telemetry/telemetryService'

interface PermissionOption {
label: string
Expand All @@ -33,12 +35,14 @@ export class McpEventHandler {
#eventListenerRegistered: boolean
#currentEditingServerName: string | undefined
#shouldDisplayListMCPServers: boolean
#telemetryController: ChatTelemetryController

constructor(features: Features) {
constructor(features: Features, telemetryService: TelemetryService) {
this.#features = features
this.#eventListenerRegistered = false
this.#currentEditingServerName = undefined
this.#shouldDisplayListMCPServers = true
this.#telemetryController = new ChatTelemetryController(features, telemetryService)
}

/**
Expand Down Expand Up @@ -635,9 +639,28 @@ export class McpEventHandler {
if (isEditMode && originalServerName) {
await McpManager.instance.removeServer(originalServerName)
await McpManager.instance.addServer(serverName, config, configPath, personaPath)
// Emit server initialize event after updating server
this.#telemetryController?.emitMCPServerInitializeEvent({
source: 'updateServer',
command: config.command,
enabled: true,
numTools: McpManager.instance.getAllToolsWithPermissions(serverName).length,
scope: params.optionsValues['scope'] === 'global' ? 'global' : 'workspace',
transportType: 'stdio',
languageServerVersion: this.#features.runtime.serverInfo.version,
})
} else {
// Create new server
await McpManager.instance.addServer(serverName, config, configPath, personaPath)
this.#telemetryController?.emitMCPServerInitializeEvent({
source: 'addServer',
command: config.command,
enabled: true,
numTools: McpManager.instance.getAllToolsWithPermissions(serverName).length,
scope: params.optionsValues['scope'] === 'global' ? 'global' : 'workspace',
transportType: 'stdio',
languageServerVersion: this.#features.runtime.serverInfo.version,
})
}

this.#currentEditingServerName = undefined
Expand Down Expand Up @@ -755,6 +778,7 @@ export class McpEventHandler {

try {
await McpManager.instance.updateServerPermission(serverName, perm)
this.#emitMCPConfigEvent()
} catch (error) {
this.#features.logging.error(`Failed to enable MCP server: ${error}`)
}
Expand All @@ -781,6 +805,7 @@ export class McpEventHandler {

try {
await McpManager.instance.updateServerPermission(serverName, perm)
this.#emitMCPConfigEvent()
} catch (error) {
this.#features.logging.error(`Failed to disable MCP server: ${error}`)
}
Expand Down Expand Up @@ -973,20 +998,100 @@ export class McpEventHandler {
const mcpServerPermission = await this.#processPermissionUpdates(updatedPermissionConfig)

await McpManager.instance.updateServerPermission(serverName, mcpServerPermission)
this.#emitMCPConfigEvent()

// Get server config to emit telemetry
const serverConfig = McpManager.instance.getAllServerConfigs().get(serverName)
if (serverConfig) {
// Emit server initialize event after permission change
this.#telemetryController?.emitMCPServerInitializeEvent({
source: 'updatePermission',
command: serverConfig.command,
enabled: true,
numTools: McpManager.instance.getAllToolsWithPermissions(serverName).length,
scope:
serverConfig?.__configPath__ ===
getGlobalMcpConfigPath(this.#features.workspace.fs.getUserHomeDir())
? 'global'
: 'workspace',
transportType: 'stdio',
languageServerVersion: this.#features.runtime.serverInfo.version,
})
}
return { id: params.id }
} catch (error) {
this.#features.logging.error(`Failed to update MCP permissions: ${error}`)
return { id: params.id }
}
}

#emitMCPConfigEvent() {
// Emit MCP config event after reinitialization
const mcpManager = McpManager.instance
const serverConfigs = mcpManager.getAllServerConfigs()
const activeServers = Array.from(serverConfigs.entries()).filter(
([name, _]) => !mcpManager.isServerDisabled(name)
)

// Count global vs project servers
const globalServers = Array.from(serverConfigs.entries()).filter(
([_, config]) =>
config?.__configPath__ === getGlobalMcpConfigPath(this.#features.workspace.fs.getUserHomeDir())
).length
const projectServers = serverConfigs.size - globalServers

// Count tools by permission
let toolsAlwaysAllowed = 0
let toolsDenied = 0

for (const [serverName, _] of activeServers) {
const toolsWithPermissions = mcpManager.getAllToolsWithPermissions(serverName)
toolsWithPermissions.forEach(item => {
if (item.permission === McpPermissionType.alwaysAllow) {
toolsAlwaysAllowed++
} else if (item.permission === McpPermissionType.deny) {
toolsDenied++
}
})
}

this.#telemetryController?.emitMCPConfigEvent({
numActiveServers: activeServers.length,
numGlobalServers: globalServers,
numProjectServers: projectServers,
numToolsAlwaysAllowed: toolsAlwaysAllowed,
numToolsDenied: toolsDenied,
languageServerVersion: this.#features.runtime.serverInfo.version,
})

// Emit server initialize events for all active servers
for (const [serverName, config] of serverConfigs.entries()) {
const enabled = !mcpManager.isServerDisabled(serverName)
if (enabled) {
this.#telemetryController?.emitMCPServerInitializeEvent({
source: 'reload',
command: config.command,
enabled,
numTools: mcpManager.getAllToolsWithPermissions(serverName).length,
scope:
config?.__configPath__ === getGlobalMcpConfigPath(this.#features.workspace.fs.getUserHomeDir())
? 'global'
: 'workspace',
transportType: 'stdio',
languageServerVersion: this.#features.runtime.serverInfo.version,
})
}
}
}

/**
* Handled refresh MCP list events
*/
async #handleRefreshMCPList(params: McpServerClickParams) {
this.#shouldDisplayListMCPServers = true
try {
await McpManager.instance.reinitializeMcpServers()
this.#emitMCPConfigEvent()
return {
id: params.id,
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
*/

import type { Features } from '@aws/language-server-runtimes/server-interface/server'
import { ChatTelemetryEventName } from '../../../../shared/telemetry/types'
import { getGlobalMcpConfigPath } from './mcpUtils'
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
import {
Expand Down Expand Up @@ -45,7 +47,10 @@ export class McpManager {
private constructor(
private configPaths: string[],
private personaPaths: string[],
private features: Pick<Features, 'logging' | 'workspace' | 'lsp'>
private features: Pick<
Features,
'logging' | 'workspace' | 'lsp' | 'telemetry' | 'credentialsProvider' | 'runtime'
>
) {
this.mcpTools = []
this.clients = new Map<string, Client>()
Expand All @@ -58,19 +63,58 @@ export class McpManager {
this.toolNameMapping = new Map<string, { serverName: string; toolName: string }>()
}

/**
* Initialize or return existing manager, then discover all servers.
*/
public static async init(
configPaths: string[],
personaPaths: string[],
features: Pick<Features, 'logging' | 'workspace' | 'lsp'>
features: Pick<Features, 'logging' | 'workspace' | 'lsp' | 'telemetry' | 'credentialsProvider' | 'runtime'>
): Promise<McpManager> {
if (!McpManager.#instance) {
const mgr = new McpManager(configPaths, personaPaths, features)
McpManager.#instance = mgr
await mgr.discoverAllServers()
features.logging.info(`MCP: discovered ${mgr.mcpTools.length} tools across all servers`)

// Emit MCP configuration metrics
const serverConfigs = mgr.getAllServerConfigs()
const activeServers = Array.from(serverConfigs.entries()).filter(([name, _]) => !mgr.isServerDisabled(name))

// Count global vs project servers
const globalServers = Array.from(serverConfigs.entries()).filter(
([_, config]) =>
config?.__configPath__ === getGlobalMcpConfigPath(features.workspace.fs.getUserHomeDir())
).length
const projectServers = serverConfigs.size - globalServers

// Count tools by permission
let toolsAlwaysAllowed = 0
let toolsDenied = 0

for (const [serverName, _] of activeServers) {
const toolsWithPermissions = mgr.getAllToolsWithPermissions(serverName)
toolsWithPermissions.forEach(item => {
if (item.permission === McpPermissionType.alwaysAllow) {
toolsAlwaysAllowed++
} else if (item.permission === McpPermissionType.deny) {
toolsDenied++
}
})
}

// Emit MCP configuration metrics
if (features.telemetry) {
features.telemetry.emitMetric({
name: ChatTelemetryEventName.MCPConfig,
data: {
credentialStartUrl: features.credentialsProvider?.getConnectionMetadata()?.sso?.startUrl,
languageServerVersion: features.runtime?.serverInfo.version,
numActiveServers: activeServers.length,
numGlobalServers: globalServers,
numProjectServers: projectServers,
numToolsAlwaysAllowed: toolsAlwaysAllowed,
numToolsDenied: toolsDenied,
},
})
}
}
return McpManager.#instance
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,11 @@ describe('McpTool', () => {
},
},
lsp: {},
credentialsProvider: {},
telemetry: { record: () => {} },
} as unknown as Pick<
import('@aws/language-server-runtimes/server-interface/server').Features,
'logging' | 'workspace' | 'lsp'
'logging' | 'workspace' | 'lsp' | 'credentialsProvider' | 'telemetry' | 'runtime'
>

const definition: McpToolDefinition = {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ export const LspToolsServer: Server = ({ workspace, logging, lsp, agent }) => {
return () => {}
}

export const McpToolsServer: Server = ({ credentialsProvider, workspace, logging, lsp, agent }) => {
export const McpToolsServer: Server = ({ credentialsProvider, workspace, logging, lsp, agent, telemetry, runtime }) => {
const registered: Record<string, string[]> = {}

const allNamespacedTools = new Set<string>()
Expand Down Expand Up @@ -141,7 +141,14 @@ export const McpToolsServer: Server = ({ credentialsProvider, workspace, logging
const globalPersonaPath = getGlobalPersonaConfigPath(workspace.fs.getUserHomeDir())
const allPersonaPaths = [...wsPersonaPaths, globalPersonaPath]

const mgr = await McpManager.init(allConfigPaths, allPersonaPaths, { logging, workspace, lsp })
const mgr = await McpManager.init(allConfigPaths, allPersonaPaths, {
logging,
workspace,
lsp,
telemetry,
credentialsProvider,
runtime,
})

// Clear tool name mapping before registering all tools to avoid conflicts from previous registrations
McpManager.instance.clearToolNameMapping()
Expand Down
Loading