Skip to content

Commit e4e8bbb

Browse files
dungdong-awsbywang56
dungdong-aws
and
bywang56
authored
fix(amazonq): fix regression of mcp config in agent config (aws#2101)
* fix(amazonq): update process-permission-updates to be the same as previous behavior * fix(chat-client): update package.json chat-client to prod * fix(amazonq): retain mcp permissions after disable/enable server * fix: dont call initOneServer on Built-in * fix: deny permission does not persist after restart IDE --------- Co-authored-by: Boyu <[email protected]>
1 parent 971eaa5 commit e4e8bbb

File tree

3 files changed

+98
-166
lines changed

3 files changed

+98
-166
lines changed

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts‎

Lines changed: 28 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import { ProfileStatusMonitor } from './profileStatusMonitor'
3535
interface PermissionOption {
3636
label: string
3737
value: string
38+
description?: string
3839
}
3940

4041
export class McpEventHandler {
@@ -913,20 +914,15 @@ export class McpEventHandler {
913914
}
914915

915916
const mcpManager = McpManager.instance
916-
917917
// Get the appropriate agent path
918918
const agentPath = mcpManager.getAllServerConfigs().get(serverName)?.__configPath__
919-
920-
const perm: MCPServerPermission = {
921-
enabled: true,
922-
toolPerms: {},
923-
__configPath__: agentPath,
924-
}
925-
926919
// Set flag to ignore file changes during permission update
927920
this.#isProgrammaticChange = true
928921

929922
try {
923+
const perm = mcpManager.getMcpServerPermissions(serverName)!
924+
perm.enabled = true
925+
perm.__configPath__ = agentPath
930926
await mcpManager.updateServerPermission(serverName, perm)
931927
this.#emitMCPConfigEvent()
932928
} catch (error) {
@@ -944,21 +940,15 @@ export class McpEventHandler {
944940
if (!serverName) {
945941
return { id: params.id }
946942
}
947-
948943
const mcpManager = McpManager.instance
949-
// Get the appropriate agent path
944+
// Set flag to ignore file changes during permission update
950945
const agentPath = mcpManager.getAllServerConfigs().get(serverName)?.__configPath__
951-
952-
const perm: MCPServerPermission = {
953-
enabled: false,
954-
toolPerms: {},
955-
__configPath__: agentPath,
956-
}
957-
958946
// Set flag to ignore file changes during permission update
959947
this.#isProgrammaticChange = true
960-
961948
try {
949+
const perm = mcpManager.getMcpServerPermissions(serverName)!
950+
perm.enabled = false
951+
perm.__configPath__ = agentPath
962952
await mcpManager.updateServerPermission(serverName, perm)
963953
this.#emitMCPConfigEvent()
964954
} catch (error) {
@@ -1078,17 +1068,16 @@ export class McpEventHandler {
10781068
// Add tool select options
10791069
toolsWithPermissions.forEach(item => {
10801070
const toolName = item.tool.toolName
1081-
const currentPermission = this.#getCurrentPermission(item.permission)
10821071
// For Built-in server, use a special function that doesn't include the 'Deny' option
1083-
const permissionOptions = this.#buildPermissionOptions(item.permission)
1072+
let permissionOptions = this.#buildPermissionOptions()
10841073

10851074
filterOptions.push({
10861075
type: 'select',
10871076
id: `${toolName}`,
10881077
title: toolName,
10891078
description: item.tool.description,
1090-
placeholder: currentPermission,
10911079
options: permissionOptions,
1080+
...{ value: item.permission, boldTitle: true, mandatory: true, hideMandatoryIcon: true },
10921081
})
10931082
})
10941083

@@ -1141,20 +1130,22 @@ export class McpEventHandler {
11411130
/**
11421131
* Builds permission options excluding the current one
11431132
*/
1144-
#buildPermissionOptions(currentPermission: string) {
1133+
#buildPermissionOptions() {
11451134
const permissionOptions: PermissionOption[] = []
11461135

1147-
if (currentPermission !== McpPermissionType.alwaysAllow) {
1148-
permissionOptions.push({ label: 'Always allow', value: McpPermissionType.alwaysAllow })
1149-
}
1136+
permissionOptions.push({
1137+
label: 'Ask',
1138+
value: McpPermissionType.ask,
1139+
description: 'Ask for your approval each time this tool is run',
1140+
})
11501141

1151-
if (currentPermission !== McpPermissionType.ask) {
1152-
permissionOptions.push({ label: 'Ask', value: McpPermissionType.ask })
1153-
}
1142+
permissionOptions.push({
1143+
label: 'Always allow',
1144+
value: McpPermissionType.alwaysAllow,
1145+
description: 'Always allow this tool to run without asking for approval',
1146+
})
11541147

1155-
if (currentPermission !== McpPermissionType.deny) {
1156-
permissionOptions.push({ label: 'Deny', value: McpPermissionType.deny })
1157-
}
1148+
permissionOptions.push({ label: 'Deny', value: McpPermissionType.deny, description: 'Never run this tool' })
11581149

11591150
return permissionOptions
11601151
}
@@ -1203,6 +1194,7 @@ export class McpEventHandler {
12031194
}
12041195

12051196
const mcpServerPermission = await this.#processPermissionUpdates(
1197+
serverName,
12061198
updatedPermissionConfig,
12071199
serverConfig?.__configPath__
12081200
)
@@ -1400,27 +1392,20 @@ export class McpEventHandler {
14001392
/**
14011393
* Processes permission updates from the UI
14021394
*/
1403-
async #processPermissionUpdates(updatedPermissionConfig: any, agentPath: string | undefined) {
1395+
async #processPermissionUpdates(serverName: string, updatedPermissionConfig: any, agentPath: string | undefined) {
1396+
const builtInToolAgentPath = await this.#getAgentPath()
14041397
const perm: MCPServerPermission = {
14051398
enabled: true,
14061399
toolPerms: {},
1407-
__configPath__: agentPath,
1400+
__configPath__: serverName === 'Built-in' ? builtInToolAgentPath : agentPath,
14081401
}
14091402

14101403
// Process each tool permission setting
14111404
for (const [key, val] of Object.entries(updatedPermissionConfig)) {
14121405
if (key === 'scope') continue
14131406

1414-
// // Get the default permission for this tool from McpManager
1415-
// let defaultPermission = McpManager.instance.getToolPerm(serverName, key)
1416-
1417-
// // If no default permission is found, use 'alwaysAllow' for Built-in and 'ask' for MCP servers
1418-
// if (!defaultPermission) {
1419-
// defaultPermission = serverName === 'Built-in' ? 'alwaysAllow' : 'ask'
1420-
// }
1421-
1422-
// If the value is an empty string (''), skip this tool to preserve its existing permission in the persona file
1423-
if (val === '') continue
1407+
const currentPerm = McpManager.instance.getToolPerm(serverName, key)
1408+
if (val === currentPerm) continue
14241409
switch (val) {
14251410
case McpPermissionType.alwaysAllow:
14261411
perm.toolPerms[key] = McpPermissionType.alwaysAllow

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.ts‎

Lines changed: 66 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -178,14 +178,49 @@ export class McpManager {
178178

179179
// Reset permissions map
180180
this.mcpServerPermissions.clear()
181-
182-
// Initialize permissions for servers from agent config
181+
// Create init state
183182
for (const [sanitizedName, _] of this.mcpServers.entries()) {
184-
const name = this.serverNameMapping.get(sanitizedName) || sanitizedName
185-
186183
// Set server status to UNINITIALIZED initially
187184
this.setState(sanitizedName, McpServerStatus.UNINITIALIZED, 0)
185+
}
186+
// Get all servers that need to be initialized
187+
const serversToInit: Array<[string, MCPServerConfig]> = []
188+
189+
for (const [name, cfg] of this.mcpServers.entries()) {
190+
if (this.isServerDisabled(name)) {
191+
this.features.logging.info(`MCP: server '${name}' is disabled by persona settings, skipping`)
192+
this.setState(name, McpServerStatus.DISABLED, 0)
193+
this.emitToolsChanged(name)
194+
continue
195+
}
196+
serversToInit.push([name, cfg])
197+
}
198+
199+
// Process servers in batches of 5 at a time
200+
const MAX_CONCURRENT_SERVERS = 5
201+
const totalServers = serversToInit.length
202+
203+
if (totalServers > 0) {
204+
this.features.logging.info(
205+
`MCP: initializing ${totalServers} servers with max concurrency of ${MAX_CONCURRENT_SERVERS}`
206+
)
207+
208+
// Process servers in batches
209+
for (let i = 0; i < totalServers; i += MAX_CONCURRENT_SERVERS) {
210+
const batch = serversToInit.slice(i, i + MAX_CONCURRENT_SERVERS)
211+
const batchPromises = batch.map(([name, cfg]) => this.initOneServer(name, cfg))
212+
213+
this.features.logging.debug(
214+
`MCP: initializing batch of ${batch.length} servers (${i + 1}-${Math.min(i + MAX_CONCURRENT_SERVERS, totalServers)} of ${totalServers})`
215+
)
216+
await Promise.all(batchPromises)
217+
}
188218

219+
this.features.logging.info(`MCP: completed initialization of ${totalServers} servers`)
220+
}
221+
222+
for (const [sanitizedName, _] of this.mcpServers.entries()) {
223+
const name = this.serverNameMapping.get(sanitizedName) || sanitizedName
189224
// Initialize permissions for this server
190225
const serverPrefix = `@${name}`
191226

@@ -208,9 +243,18 @@ export class McpManager {
208243
})
209244
} else {
210245
// Only specific tools are enabled
246+
// get allTools of this server, if it's not in tools --> it's denied
247+
// have to move the logic after all servers finish init, because that's when we have list of tools
248+
const deniedTools = new Set(
249+
this.getAllTools()
250+
.filter(tool => tool.serverName === name)
251+
.map(tool => tool.toolName)
252+
)
211253
this.agentConfig.tools.forEach(tool => {
212254
if (tool.startsWith(serverPrefix + '/')) {
255+
// remove this from deniedTools
213256
const toolName = tool.substring(serverPrefix.length + 1)
257+
deniedTools.delete(toolName)
214258
if (toolName) {
215259
// Check if tool is in allowedTools
216260
if (this.agentConfig.allowedTools.includes(tool)) {
@@ -221,49 +265,18 @@ export class McpManager {
221265
}
222266
}
223267
})
268+
269+
// update permission to deny for rest of the tools
270+
deniedTools.forEach(tool => {
271+
toolPerms[tool] = McpPermissionType.deny
272+
})
224273
}
225274

226275
this.mcpServerPermissions.set(sanitizedName, {
227276
enabled: true,
228277
toolPerms,
229278
})
230279
}
231-
232-
// Get all servers that need to be initialized
233-
const serversToInit: Array<[string, MCPServerConfig]> = []
234-
235-
for (const [name, cfg] of this.mcpServers.entries()) {
236-
if (this.isServerDisabled(name)) {
237-
this.features.logging.info(`MCP: server '${name}' is disabled by persona settings, skipping`)
238-
this.setState(name, McpServerStatus.DISABLED, 0)
239-
this.emitToolsChanged(name)
240-
continue
241-
}
242-
serversToInit.push([name, cfg])
243-
}
244-
245-
// Process servers in batches of 5 at a time
246-
const MAX_CONCURRENT_SERVERS = 5
247-
const totalServers = serversToInit.length
248-
249-
if (totalServers > 0) {
250-
this.features.logging.info(
251-
`MCP: initializing ${totalServers} servers with max concurrency of ${MAX_CONCURRENT_SERVERS}`
252-
)
253-
254-
// Process servers in batches
255-
for (let i = 0; i < totalServers; i += MAX_CONCURRENT_SERVERS) {
256-
const batch = serversToInit.slice(i, i + MAX_CONCURRENT_SERVERS)
257-
const batchPromises = batch.map(([name, cfg]) => this.initOneServer(name, cfg))
258-
259-
this.features.logging.debug(
260-
`MCP: initializing batch of ${batch.length} servers (${i + 1}-${Math.min(i + MAX_CONCURRENT_SERVERS, totalServers)} of ${totalServers})`
261-
)
262-
await Promise.all(batchPromises)
263-
}
264-
265-
this.features.logging.info(`MCP: completed initialization of ${totalServers} servers`)
266-
}
267280
}
268281

269282
/**
@@ -658,13 +671,7 @@ export class McpManager {
658671
}
659672

660673
// Save agent config once with all changes
661-
await saveAgentConfig(
662-
this.features.workspace,
663-
this.features.logging,
664-
this.agentConfig,
665-
agentPath,
666-
serverName
667-
)
674+
await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, agentPath)
668675

669676
// Add server tools to tools list after initialization
670677
await this.initOneServer(sanitizedName, newCfg)
@@ -728,13 +735,7 @@ export class McpManager {
728735
})
729736

730737
// Save agent config
731-
await saveAgentConfig(
732-
this.features.workspace,
733-
this.features.logging,
734-
this.agentConfig,
735-
cfg.__configPath__,
736-
unsanitizedName
737-
)
738+
await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, cfg.__configPath__)
738739

739740
// Get all config paths and delete the server from each one
740741
const wsUris = this.features.workspace.getAllWorkspaceFolders()?.map(f => f.uri) ?? []
@@ -817,13 +818,7 @@ export class McpManager {
817818
this.agentConfig.mcpServers[unsanitizedServerName] = updatedConfig
818819

819820
// Save agent config
820-
await saveAgentConfig(
821-
this.features.workspace,
822-
this.features.logging,
823-
this.agentConfig,
824-
agentPath,
825-
unsanitizedServerName
826-
)
821+
await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, agentPath)
827822
}
828823

829824
const newCfg: MCPServerConfig = {
@@ -1035,13 +1030,7 @@ export class McpManager {
10351030
// Save agent config
10361031
const agentPath = perm.__configPath__
10371032
if (agentPath) {
1038-
await saveAgentConfig(
1039-
this.features.workspace,
1040-
this.features.logging,
1041-
this.agentConfig,
1042-
agentPath,
1043-
unsanitizedServerName
1044-
)
1033+
await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, agentPath)
10451034
}
10461035

10471036
// Update mcpServerPermissions map
@@ -1059,7 +1048,7 @@ export class McpManager {
10591048
}
10601049
this.setState(serverName, McpServerStatus.DISABLED, 0)
10611050
} else {
1062-
if (!this.clients.has(serverName)) {
1051+
if (!this.clients.has(serverName) && serverName !== 'Built-in') {
10631052
await this.initOneServer(serverName, this.mcpServers.get(serverName)!)
10641053
}
10651054
}
@@ -1087,6 +1076,13 @@ export class McpManager {
10871076
return !this.agentConfig.allowedTools.includes(toolId)
10881077
}
10891078

1079+
/**
1080+
* get server's tool permission
1081+
*/
1082+
public getMcpServerPermissions(serverName: string): MCPServerPermission | undefined {
1083+
return this.mcpServerPermissions.get(serverName)
1084+
}
1085+
10901086
/**
10911087
* Returns any errors that occurred during loading of MCP configuration files
10921088
*/
@@ -1152,8 +1148,7 @@ export class McpManager {
11521148
this.features.workspace,
11531149
this.features.logging,
11541150
this.agentConfig,
1155-
cfg.__configPath__,
1156-
unsanitizedName
1151+
cfg.__configPath__
11571152
)
11581153
}
11591154
} catch (err) {

0 commit comments

Comments
 (0)