feat: update MCP manager and utilities (aws#2158)

ashishrp-aws · bywang56 · web-flow · commit b99df82826d0 · 2025-08-25T14:31:41.000-07:00
Co-authored-by: Boyu &lt;bywang@amazon.com&gt;
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.test.ts
@@ -368,7 +368,7 @@ describe('removeServer()', () => {
         expect((mgr as any).clients.has('x')).to.be.false
     })
 
-    it('removes server from all config files', async () => {
+    it('removes server from agent config', async () => {
         const mgr = await McpManager.init([], features)
         const dummy = new Client({ name: 'c', version: 'v' })
         ;(mgr as any).clients.set('x', dummy)
@@ -395,14 +395,13 @@ describe('removeServer()', () => {
 
         await mgr.removeServer('x')
 
-        // Verify that writeFile was called for each config path (2 workspace + 1 global)
-        expect(writeFileStub.callCount).to.equal(3)
+        // Verify that saveAgentConfig was called
+        expect(saveAgentConfigStub.calledOnce).to.be.true
+        expect((mgr as any).clients.has('x')).to.be.false
 
-        // Verify the content of the writes (should have removed the server)
-        writeFileStub.getCalls().forEach(call => {
-            const content = JSON.parse(call.args[1])
-            expect(content.mcpServers).to.not.have.property('x')
-        })
+        // Verify server was removed from agent config
+        expect((mgr as any).agentConfig.mcpServers).to.not.have.property('x')
+        expect((mgr as any).agentConfig.tools).to.not.include('@x')
     })
 })
 
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.ts
@@ -796,30 +796,6 @@ export class McpManager {
 
             // Save agent config
             await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, cfg.__configPath__)
-
-            // Get all config paths and delete the server from each one
-            const wsUris = this.features.workspace.getAllWorkspaceFolders()?.map(f => f.uri) ?? []
-            const wsConfigPaths = getWorkspaceMcpConfigPaths(wsUris)
-            const globalConfigPath = getGlobalMcpConfigPath(this.features.workspace.fs.getUserHomeDir())
-            const allConfigPaths = [...wsConfigPaths, globalConfigPath]
-
-            // Delete the server from all config files
-            for (const configPath of allConfigPaths) {
-                try {
-                    await this.mutateConfigFile(configPath, json => {
-                        if (json.mcpServers && json.mcpServers[unsanitizedName]) {
-                            delete json.mcpServers[unsanitizedName]
-                            this.features.logging.info(
-                                `Deleted server '${unsanitizedName}' from config file: ${configPath}`
-                            )
-                        }
-                    })
-                } catch (err) {
-                    this.features.logging.warn(
-                        `Failed to delete server '${unsanitizedName}' from config file ${configPath}: ${err}`
-                    )
-                }
-            }
         }
 
         this.mcpServers.delete(serverName)
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.test.ts
@@ -753,6 +753,12 @@ describe('migrateToAgentConfig', () => {
     })
 
     it('migrates when no existing configs exist', async () => {
+        // Create empty MCP config to trigger migration
+        const mcpDir = path.join(tmpDir, '.aws', 'amazonq')
+        fs.mkdirSync(mcpDir, { recursive: true })
+        const mcpPath = path.join(mcpDir, 'mcp.json')
+        fs.writeFileSync(mcpPath, JSON.stringify({ mcpServers: {} }))
+
         await migrateToAgentConfig(workspace, logger, mockAgent)
 
         // Should create default agent config
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts
@@ -175,7 +175,9 @@ const DEFAULT_AGENT_RAW = `{
     "README.md",
     ".amazonq/rules/**/*.md"
   ],
-  "resources": []
+  "resources": [],
+  "createHooks": [],
+  "promptHooks": []
 }`
 
 const DEFAULT_PERSONA_RAW = `{
@@ -840,48 +842,40 @@ async function migrateConfigToAgent(
     const normalizedPersonaPath = normalizePathFromUri(personaPath, logging)
     agentPath = normalizePathFromUri(agentPath)
 
-    // Check if agent config exists
+    // Check if config and agent files exist
+    const configExists = await workspace.fs.exists(normalizedConfigPath).catch(() => false)
     const agentExists = await workspace.fs.exists(agentPath).catch(() => false)
 
-    // Load existing agent config if it exists
-    let existingAgentConfig: AgentConfig | undefined
+    // Only migrate if agent file does not exist
+    // If config exists, migrate from it; if not, create default agent config
     if (agentExists) {
-        try {
-            const raw = (await workspace.fs.readFile(agentPath)).toString().trim()
-            existingAgentConfig = raw ? JSON.parse(raw) : undefined
-        } catch (err) {
-            logging.warn(`Failed to read existing agent config at ${agentPath}: ${err}`)
-        }
+        return
     }
 
     // Read MCP server configs directly from file
     const serverConfigs: Record<string, MCPServerConfig> = {}
     try {
-        const configExists = await workspace.fs.exists(normalizedConfigPath)
-
-        if (configExists) {
-            const raw = (await workspace.fs.readFile(normalizedConfigPath)).toString().trim()
-            if (raw) {
-                const config = JSON.parse(raw)
-
-                if (config.mcpServers && typeof config.mcpServers === 'object') {
-                    // Add each server to the serverConfigs
-                    for (const [name, serverConfig] of Object.entries(config.mcpServers)) {
-                        serverConfigs[name] = {
-                            command: (serverConfig as any).command,
-                            args: Array.isArray((serverConfig as any).args) ? (serverConfig as any).args : undefined,
-                            env: typeof (serverConfig as any).env === 'object' ? (serverConfig as any).env : undefined,
-                            initializationTimeout:
-                                typeof (serverConfig as any).initializationTimeout === 'number'
-                                    ? (serverConfig as any).initializationTimeout
-                                    : undefined,
-                            timeout:
-                                typeof (serverConfig as any).timeout === 'number'
-                                    ? (serverConfig as any).timeout
-                                    : undefined,
-                        }
-                        logging.info(`Added server ${name} to serverConfigs`)
+        const raw = (await workspace.fs.readFile(normalizedConfigPath)).toString().trim()
+        if (raw) {
+            const config = JSON.parse(raw)
+
+            if (config.mcpServers && typeof config.mcpServers === 'object') {
+                // Add each server to the serverConfigs
+                for (const [name, serverConfig] of Object.entries(config.mcpServers)) {
+                    serverConfigs[name] = {
+                        command: (serverConfig as any).command,
+                        args: Array.isArray((serverConfig as any).args) ? (serverConfig as any).args : undefined,
+                        env: typeof (serverConfig as any).env === 'object' ? (serverConfig as any).env : undefined,
+                        initializationTimeout:
+                            typeof (serverConfig as any).initializationTimeout === 'number'
+                                ? (serverConfig as any).initializationTimeout
+                                : undefined,
+                        timeout:
+                            typeof (serverConfig as any).timeout === 'number'
+                                ? (serverConfig as any).timeout
+                                : undefined,
                     }
+                    logging.info(`Added server ${name} to serverConfigs`)
                 }
             }
         }
@@ -908,46 +902,24 @@ async function migrateConfigToAgent(
     }
 
     // Convert to agent config
-    const newAgentConfig = convertPersonaToAgent(personaConfig, serverConfigs, agent)
-    newAgentConfig.includedFiles = ['AmazonQ.md', 'README.md', '.amazonq/rules/**/*.md']
-    newAgentConfig.resources = [] // Initialize with empty array
-
-    // Merge with existing config if available
-    let finalAgentConfig: AgentConfig
-    if (existingAgentConfig) {
-        // Keep existing metadata
-        finalAgentConfig = {
-            ...existingAgentConfig,
-            // Merge MCP servers, keeping existing ones if they exist
-            mcpServers: {
-                ...newAgentConfig.mcpServers,
-                ...existingAgentConfig.mcpServers,
-            },
-            // Merge tools lists without duplicates
-            tools: [...new Set([...existingAgentConfig.tools, ...newAgentConfig.tools])],
-            allowedTools: [...new Set([...existingAgentConfig.allowedTools, ...newAgentConfig.allowedTools])],
-            // Merge tool settings, preferring existing ones
-            toolsSettings: {
-                ...newAgentConfig.toolsSettings,
-                ...existingAgentConfig.toolsSettings,
-            },
-            // Keep other properties from existing config
-            includedFiles: existingAgentConfig.includedFiles || newAgentConfig.includedFiles,
-            createHooks: existingAgentConfig.createHooks || newAgentConfig.createHooks,
-            promptHooks: [
-                ...new Set([...(existingAgentConfig.promptHooks || []), ...(newAgentConfig.promptHooks || [])]),
-            ],
-            resources: [...new Set([...(existingAgentConfig.resources || []), ...(newAgentConfig.resources || [])])],
-        }
-    } else {
-        finalAgentConfig = newAgentConfig
-        logging.info(`Using new config (no existing config to merge)`)
-    }
+    const agentConfig = convertPersonaToAgent(personaConfig, serverConfigs, agent)
+
+    // Parse default values from DEFAULT_AGENT_RAW
+    const defaultAgent = JSON.parse(DEFAULT_AGENT_RAW)
+
+    // Add complete agent format sections using default values
+    agentConfig.name = defaultAgent.name
+    agentConfig.description = defaultAgent.description
+    agentConfig.version = defaultAgent.version
+    agentConfig.includedFiles = defaultAgent.includedFiles
+    agentConfig.resources = defaultAgent.resources
+    agentConfig.createHooks = defaultAgent.createHooks
+    agentConfig.promptHooks = defaultAgent.promptHooks
 
     // Save agent config
     try {
-        await saveAgentConfig(workspace, logging, finalAgentConfig, agentPath)
-        logging.info(`Successfully ${existingAgentConfig ? 'updated' : 'created'} agent config at ${agentPath}`)
+        await saveAgentConfig(workspace, logging, agentConfig, agentPath)
+        logging.info(`Successfully created agent config at ${agentPath}`)
     } catch (err) {
         logging.error(`Failed to save agent config to ${agentPath}: ${err}`)
         throw err
