Skip to content

Commit d28df09

Browse files
ashishrp-awsashishrp-aws
authored
fix(amazonq): fix for mcp servers operations to edit server config only (aws#2165)
* fix(amazonq): fix for mcp servers operations to edit server specific config * fix(amazonq): additional mcp server config fixes * fix: resolve test failures * fix: update MCP manager configuration
1 parent 558bc1a commit d28df09

File tree

4 files changed

+306
-34
lines changed

4 files changed

+306
-34
lines changed

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

Lines changed: 19 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -239,12 +239,12 @@ describe('callTool()', () => {
239239
describe('addServer()', () => {
240240
let loadStub: sinon.SinonStub
241241
let initOneStub: sinon.SinonStub
242-
let saveAgentConfigStub: sinon.SinonStub
242+
let saveServerSpecificAgentConfigStub: sinon.SinonStub
243243

244244
beforeEach(() => {
245245
loadStub = stubAgentConfig()
246246
initOneStub = stubInitOneServer()
247-
saveAgentConfigStub = sinon.stub(mcpUtils, 'saveAgentConfig').resolves()
247+
saveServerSpecificAgentConfigStub = sinon.stub(mcpUtils, 'saveServerSpecificAgentConfig').resolves()
248248
})
249249

250250
afterEach(async () => {
@@ -268,7 +268,7 @@ describe('addServer()', () => {
268268

269269
await mgr.addServer('newS', newCfg, 'path.json')
270270

271-
expect(saveAgentConfigStub.calledOnce).to.be.true
271+
expect(saveServerSpecificAgentConfigStub.calledOnce).to.be.true
272272
expect(initOneStub.calledOnceWith('newS', sinon.match(newCfg))).to.be.true
273273
})
274274

@@ -301,14 +301,14 @@ describe('addServer()', () => {
301301

302302
await mgr.addServer('httpSrv', httpCfg, 'http.json')
303303

304-
expect(saveAgentConfigStub.calledOnce).to.be.true
304+
expect(saveServerSpecificAgentConfigStub.calledOnce).to.be.true
305305
expect(initOneStub.calledOnceWith('httpSrv', sinon.match(httpCfg))).to.be.true
306306
})
307307
})
308308

309309
describe('removeServer()', () => {
310310
let loadStub: sinon.SinonStub
311-
let saveAgentConfigStub: sinon.SinonStub
311+
let saveServerSpecificAgentConfigStub: sinon.SinonStub
312312
let existsStub: sinon.SinonStub
313313
let readFileStub: sinon.SinonStub
314314
let writeFileStub: sinon.SinonStub
@@ -318,7 +318,7 @@ describe('removeServer()', () => {
318318

319319
beforeEach(() => {
320320
loadStub = stubAgentConfig()
321-
saveAgentConfigStub = sinon.stub(mcpUtils, 'saveAgentConfig').resolves()
321+
saveServerSpecificAgentConfigStub = sinon.stub(mcpUtils, 'saveServerSpecificAgentConfig').resolves()
322322
existsStub = sinon.stub(fakeWorkspace.fs, 'exists').resolves(true)
323323
readFileStub = sinon
324324
.stub(fakeWorkspace.fs, 'readFile')
@@ -364,7 +364,7 @@ describe('removeServer()', () => {
364364
}
365365

366366
await mgr.removeServer('x')
367-
expect(saveAgentConfigStub.calledOnce).to.be.true
367+
expect(saveServerSpecificAgentConfigStub.calledOnce).to.be.true
368368
expect((mgr as any).clients.has('x')).to.be.false
369369
})
370370

@@ -395,8 +395,8 @@ describe('removeServer()', () => {
395395

396396
await mgr.removeServer('x')
397397

398-
// Verify that saveAgentConfig was called
399-
expect(saveAgentConfigStub.calledOnce).to.be.true
398+
// Verify that saveServerSpecificAgentConfig was called
399+
expect(saveServerSpecificAgentConfigStub.calledOnce).to.be.true
400400
expect((mgr as any).clients.has('x')).to.be.false
401401

402402
// Verify server was removed from agent config
@@ -472,11 +472,11 @@ describe('mutateConfigFile()', () => {
472472
describe('updateServer()', () => {
473473
let loadStub: sinon.SinonStub
474474
let initOneStub: sinon.SinonStub
475-
let saveAgentConfigStub: sinon.SinonStub
475+
let saveServerSpecificAgentConfigStub: sinon.SinonStub
476476

477477
beforeEach(() => {
478478
initOneStub = stubInitOneServer()
479-
saveAgentConfigStub = sinon.stub(mcpUtils, 'saveAgentConfig').resolves()
479+
saveServerSpecificAgentConfigStub = sinon.stub(mcpUtils, 'saveServerSpecificAgentConfig').resolves()
480480
})
481481

482482
afterEach(async () => {
@@ -519,11 +519,11 @@ describe('updateServer()', () => {
519519

520520
const closeStub = sinon.stub(fakeClient, 'close').resolves()
521521
initOneStub.resetHistory()
522-
saveAgentConfigStub.resetHistory()
522+
saveServerSpecificAgentConfigStub.resetHistory()
523523

524524
await mgr.updateServer('u1', { timeout: 999 }, 'u.json')
525525

526-
expect(saveAgentConfigStub.calledOnce).to.be.true
526+
expect(saveServerSpecificAgentConfigStub.calledOnce).to.be.true
527527
expect(closeStub.calledOnce).to.be.true
528528
expect(initOneStub.calledOnceWith('u1', sinon.match.has('timeout', 999))).to.be.true
529529
})
@@ -559,11 +559,11 @@ describe('updateServer()', () => {
559559
const mgr = McpManager.instance
560560

561561
initOneStub.resetHistory()
562-
saveAgentConfigStub.resetHistory()
562+
saveServerSpecificAgentConfigStub.resetHistory()
563563

564564
await mgr.updateServer('srv', { command: undefined, url: 'https://new.host/mcp' }, 'z.json')
565565

566-
expect(saveAgentConfigStub.calledOnce).to.be.true
566+
expect(saveServerSpecificAgentConfigStub.calledOnce).to.be.true
567567
expect(initOneStub.calledOnceWith('srv', sinon.match({ url: 'https://new.host/mcp' }))).to.be.true
568568
})
569569
})
@@ -1061,9 +1061,11 @@ describe('listServersAndTools()', () => {
10611061

10621062
describe('updateServerPermission()', () => {
10631063
let saveAgentConfigStub: sinon.SinonStub
1064+
let saveServerSpecificAgentConfigStub: sinon.SinonStub
10641065

10651066
beforeEach(() => {
10661067
saveAgentConfigStub = sinon.stub(mcpUtils, 'saveAgentConfig').resolves()
1068+
saveServerSpecificAgentConfigStub = sinon.stub(mcpUtils, 'saveServerSpecificAgentConfig').resolves()
10671069
})
10681070

10691071
afterEach(async () => {
@@ -1112,8 +1114,8 @@ describe('updateServerPermission()', () => {
11121114
__configPath__: '/p',
11131115
})
11141116

1115-
// Verify saveAgentConfig was called
1116-
expect(saveAgentConfigStub.calledOnce).to.be.true
1117+
// Verify saveServerSpecificAgentConfig was called
1118+
expect(saveServerSpecificAgentConfigStub.calledOnce).to.be.true
11171119

11181120
// Verify the tool permission was updated
11191121
expect(mgr.requiresApproval('srv', 'tool1')).to.be.false

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

Lines changed: 77 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import {
2626
isEmptyEnv,
2727
loadAgentConfig,
2828
saveAgentConfig,
29+
saveServerSpecificAgentConfig,
2930
sanitizeName,
3031
getGlobalAgentConfigPath,
3132
getWorkspaceMcpConfigPaths,
@@ -730,8 +731,23 @@ export class McpManager {
730731
this.agentConfig.tools.push(serverPrefix)
731732
}
732733

733-
// Save agent config once with all changes
734-
await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, agentPath)
734+
// Save server-specific changes to agent config
735+
const serverTools = this.agentConfig.tools.filter(
736+
tool => tool === serverPrefix || tool.startsWith(`${serverPrefix}/`)
737+
)
738+
const serverAllowedTools = this.agentConfig.allowedTools.filter(
739+
tool => tool === serverPrefix || tool.startsWith(`${serverPrefix}/`)
740+
)
741+
742+
await saveServerSpecificAgentConfig(
743+
this.features.workspace,
744+
this.features.logging,
745+
serverName,
746+
serverConfig,
747+
serverTools,
748+
serverAllowedTools,
749+
agentPath
750+
)
735751

736752
// Add server tools to tools list after initialization
737753
await this.initOneServer(sanitizedName, newCfg, AuthIntent.Interactive)
@@ -794,8 +810,16 @@ export class McpManager {
794810
return true
795811
})
796812

797-
// Save agent config
798-
await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, cfg.__configPath__)
813+
// Save server removal to agent config
814+
await saveServerSpecificAgentConfig(
815+
this.features.workspace,
816+
this.features.logging,
817+
unsanitizedName,
818+
null, // null indicates server should be removed
819+
[],
820+
[],
821+
cfg.__configPath__
822+
)
799823
}
800824

801825
this.mcpServers.delete(serverName)
@@ -853,8 +877,24 @@ export class McpManager {
853877
}
854878
this.agentConfig.mcpServers[unsanitizedServerName] = updatedConfig
855879

856-
// Save agent config
857-
await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, agentPath)
880+
// Save server-specific changes to agent config
881+
const serverPrefix = `@${unsanitizedServerName}`
882+
const serverTools = this.agentConfig.tools.filter(
883+
tool => tool === serverPrefix || tool.startsWith(`${serverPrefix}/`)
884+
)
885+
const serverAllowedTools = this.agentConfig.allowedTools.filter(
886+
tool => tool === serverPrefix || tool.startsWith(`${serverPrefix}/`)
887+
)
888+
889+
await saveServerSpecificAgentConfig(
890+
this.features.workspace,
891+
this.features.logging,
892+
unsanitizedServerName,
893+
updatedConfig,
894+
serverTools,
895+
serverAllowedTools,
896+
agentPath
897+
)
858898
}
859899

860900
const newCfg: MCPServerConfig = {
@@ -1057,6 +1097,12 @@ export class McpManager {
10571097
}
10581098
}
10591099

1100+
// Update mcpServerPermissions map immediately to reflect changes
1101+
this.mcpServerPermissions.set(serverName, {
1102+
enabled: perm.enabled,
1103+
toolPerms: perm.toolPerms || {},
1104+
})
1105+
10601106
// Update server enabled/disabled state in agent config
10611107
if (this.agentConfig.mcpServers[unsanitizedServerName]) {
10621108
this.agentConfig.mcpServers[unsanitizedServerName].disabled = !perm.enabled
@@ -1067,17 +1113,28 @@ export class McpManager {
10671113
serverConfig.disabled = !perm.enabled
10681114
}
10691115

1070-
// Save agent config
1116+
// Save only server-specific changes to agent config
10711117
const agentPath = perm.__configPath__
10721118
if (agentPath) {
1073-
await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, agentPath)
1074-
}
1119+
// Collect server-specific tools and allowedTools
1120+
const serverPrefix = `@${unsanitizedServerName}`
1121+
const serverTools = this.agentConfig.tools.filter(
1122+
tool => tool === serverPrefix || tool.startsWith(`${serverPrefix}/`)
1123+
)
1124+
const serverAllowedTools = this.agentConfig.allowedTools.filter(
1125+
tool => tool === serverPrefix || tool.startsWith(`${serverPrefix}/`)
1126+
)
10751127

1076-
// Update mcpServerPermissions map
1077-
this.mcpServerPermissions.set(serverName, {
1078-
enabled: perm.enabled,
1079-
toolPerms: perm.toolPerms || {},
1080-
})
1128+
await saveServerSpecificAgentConfig(
1129+
this.features.workspace,
1130+
this.features.logging,
1131+
unsanitizedServerName,
1132+
this.agentConfig.mcpServers[unsanitizedServerName],
1133+
serverTools,
1134+
serverAllowedTools,
1135+
agentPath
1136+
)
1137+
}
10811138

10821139
// enable/disable server
10831140
if (this.isServerDisabled(serverName)) {
@@ -1184,11 +1241,14 @@ export class McpManager {
11841241
return true
11851242
})
11861243

1187-
// Save agent config
1188-
await saveAgentConfig(
1244+
// Save server removal to agent config
1245+
await saveServerSpecificAgentConfig(
11891246
this.features.workspace,
11901247
this.features.logging,
1191-
this.agentConfig,
1248+
unsanitizedName,
1249+
null, // null indicates server should be removed
1250+
[],
1251+
[],
11921252
cfg.__configPath__
11931253
)
11941254
}

0 commit comments

Comments
 (0)