fix(amazonq): Fixes mcp server registry not clearing on switching connections and reduces GetProfile calls

core/aws-lsp-core/src/util/retryUtils.ts

@@ -34,18 +34,18 @@ export interface RetryOptions {
  */
 function defaultIsRetryable(error: any): boolean {
     const errorCode = error.code || error.name
-    const statusCode = error.statusCode
+    const statusCode = error.statusCode || error?.$metadata?.httpStatusCode
 
     // Fast fail on non-retryable client errors (except throttling)
     if (statusCode >= CLIENT_ERROR_MIN && statusCode < CLIENT_ERROR_MAX && errorCode !== THROTTLING_EXCEPTION) {
         return false
     }
 
-    // Retry on throttling, server errors, and specific status codes
+    // Retry on throttling, server errors (5xx), and specific status codes
     return (
         errorCode === THROTTLING_EXCEPTION ||
         errorCode === INTERNAL_SERVER_EXCEPTION ||
-        statusCode === INTERNAL_SERVER_ERROR ||
+        statusCode >= 500 ||
         statusCode === SERVICE_UNAVAILABLE
     )
 }

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

@@ -1840,6 +1840,10 @@ export class McpManager {
         }
     }
 
+    public resetRegistryService(): void {
+        this.registryService = undefined
+    }
+
     /**
      * Update registry URL and refetch registry
      * @throws Error if registry fetch or validation fails
@@ -1849,28 +1853,20 @@ export class McpManager {
             this.registryService = new McpRegistryService(this.features.logging)
         }
 
-        const wasActive = this.registryUrlProvided
-
         try {
             const registry = await this.registryService.fetchRegistry(registryUrl)
             if (registry) {
                 this.currentRegistry = registry
                 // Clear any previous registry errors on success
                 this.configLoadErrors.delete('registry')
 
-                if (!wasActive) {
+                this.features.logging.info(`MCP Registry: Registry mode ACTIVATED - ${registry.servers.length} servers`)
+                // Only discover servers when registry is newly activated and not during periodic sync
+                if (!isPeriodicSync) {
+                    await this.discoverAllServers()
                     this.features.logging.info(
-                        `MCP Registry: Registry mode ACTIVATED - ${registry.servers.length} servers`
+                        `MCP: discovered ${this.getAllTools().length} tools after registry activation`
                     )
-                    // Only discover servers when registry is newly activated and not during periodic sync
-                    if (!isPeriodicSync) {
-                        await this.discoverAllServers()
-                        this.features.logging.info(
-                            `MCP: discovered ${this.getAllTools().length} tools after registry activation`
-                        )
-                    }
-                } else {
-                    this.features.logging.info(`MCP Registry: Updated registry with ${registry.servers.length} servers`)
                 }
 
                 // Only sync during periodic updates, not at startup

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

@@ -182,6 +182,36 @@ export class ProfileStatusMonitor {
         ProfileStatusMonitor.setMcpState(true)
     }
 
+    static resetMcpManager(): void {
+        try {
+            // note: use dynamic require to satisfy webpack requirements for webworker
+            const McpManager = eval('require')('./mcpManager').McpManager
+            if (McpManager.isInitialized()) {
+                McpManager.instance.setRegistryActive(false)
+                McpManager.instance.resetRegistryService()
+                void McpManager.instance.close(true)
+            }
+        } catch (error) {
+            ProfileStatusMonitor.logging?.error(`Failed to reset MCP manager: ${error}`)
+        }
+    }
+
+    /**
+     * When switching between connections especially builder id
+     * re-discover mcp servers for non-profile connections
+     */
+    static discoverServersWhenNoProfiles(): void {
+        try {
+            // note: use dynamic require to satisfy webpack requirements for webworker
+            const McpManager = eval('require')('./mcpManager').McpManager
+            if (McpManager.isInitialized()) {
+                void McpManager.instance.discoverAllServers()
+            }
+        } catch (error) {
+            ProfileStatusMonitor.logging?.error(`Failed to discover mcp servers when no profiles available: ${error}`)
+        }
+    }
+
     static emitAuthSuccess(): void {
         ProfileStatusMonitor.eventEmitter.emit(AUTH_SUCCESS_EVENT)
     }

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts

@@ -461,8 +461,8 @@ export const McpToolsServer: Server = ({
             // Wait for profile ARN to be available before checking MCP state
             const checkAndInitialize = async () => {
                 try {
-                    // Check if MCP is enabled via isMcpEnabled check (only call once)
-                    const mcpEnabled = await profileStatusMonitor!.checkInitialState()
+                    // Check if MCP is enabled via get Mcp state
+                    const mcpEnabled = ProfileStatusMonitor.getMcpState()
 
                     if (mcpEnabled) {
                         logging.info('MCP is enabled, discovering servers')

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts

@@ -37,7 +37,6 @@ import { parse } from '@aws-sdk/util-arn-parser'
 import { ChatDatabase } from '../../language-server/agenticChat/tools/chatDb/chatDb'
 import { ProfileStatusMonitor } from '../../language-server/agenticChat/tools/mcp/profileStatusMonitor'
 import { UserContext } from '@amzn/codewhisperer-runtime'
-
 /**
  * AmazonQTokenServiceManager manages state and provides centralized access to
  * instance of CodeWhispererServiceToken SDK client to any consuming code.
@@ -160,6 +159,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
 
         // Reset MCP state cache when auth changes
         ProfileStatusMonitor.resetMcpState()
+        ProfileStatusMonitor.resetMcpManager()
     }
 
     public handleOnCredentialsUpdated(type: CredentialsType): void {
@@ -265,8 +265,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
             this.state = 'INITIALIZED'
             this.logging.log(`Initialized Amazon Q service with ${newConnectionType} connection`)
 
-            // Emit auth success event
-            ProfileStatusMonitor.emitAuthSuccess()
+            ProfileStatusMonitor.discoverServersWhenNoProfiles()
 
             return
         }
@@ -290,9 +289,6 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
             this.state = 'INITIALIZED'
             this.logging.log('Initialized Amazon Q service with identityCenter connection')
 
-            // Emit auth success event
-            ProfileStatusMonitor.emitAuthSuccess()
-
             return
         }
 
@@ -414,9 +410,6 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
             this.activeIdcProfile = newProfile
             this.state = 'INITIALIZED'
 
-            // Emit auth success event
-            ProfileStatusMonitor.emitAuthSuccess()
-
             return
         }
 
@@ -440,9 +433,6 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
                 this.cachedStreamingClient.profileArn = newProfile.arn
             }
 
-            // Emit auth success event
-            ProfileStatusMonitor.emitAuthSuccess()
-
             return
         }