From f8806ccf95596bb7a92ca2f599c0a370920e93e7 Mon Sep 17 00:00:00 2001 From: Shruti Sinha Date: Wed, 7 Jan 2026 15:22:40 -0800 Subject: [PATCH 1/7] fix: fix mcp not resetting on switching --- core/aws-lsp-core/src/util/retryUtils.ts | 6 +++--- .../agenticChat/tools/toolServer.ts | 4 ++-- .../AmazonQTokenServiceManager.ts | 20 ++++++++----------- 3 files changed, 13 insertions(+), 17 deletions(-) diff --git a/core/aws-lsp-core/src/util/retryUtils.ts b/core/aws-lsp-core/src/util/retryUtils.ts index dc135ce23d..9a74c0bc80 100644 --- a/core/aws-lsp-core/src/util/retryUtils.ts +++ b/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 ) } diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts index 4e4896f115..883fc105c4 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts +++ b/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') diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts index 44766ebd2b..79a6c7e537 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts @@ -37,6 +37,7 @@ 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' +import { McpManager } from '../../language-server/agenticChat/tools/mcp/mcpManager' /** * AmazonQTokenServiceManager manages state and provides centralized access to @@ -160,6 +161,10 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< // Reset MCP state cache when auth changes ProfileStatusMonitor.resetMcpState() + if (McpManager.isInitialized()) { + McpManager.instance.setRegistryActive(false) + void McpManager.instance.close(true) + } } public handleOnCredentialsUpdated(type: CredentialsType): void { @@ -264,9 +269,9 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< ) this.state = 'INITIALIZED' this.logging.log(`Initialized Amazon Q service with ${newConnectionType} connection`) - - // Emit auth success event - ProfileStatusMonitor.emitAuthSuccess() + if (McpManager.isInitialized()) { + void McpManager.instance.discoverAllServers() + } return } @@ -290,9 +295,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 +416,6 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< this.activeIdcProfile = newProfile this.state = 'INITIALIZED' - // Emit auth success event - ProfileStatusMonitor.emitAuthSuccess() - return } @@ -440,9 +439,6 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< this.cachedStreamingClient.profileArn = newProfile.arn } - // Emit auth success event - ProfileStatusMonitor.emitAuthSuccess() - return } From ef5ac44b588080edd82ef2eb030dc1cd8fa45b2e Mon Sep 17 00:00:00 2001 From: Shruti Sinha Date: Wed, 7 Jan 2026 16:04:50 -0800 Subject: [PATCH 2/7] chore: reset registry --- .../src/language-server/agenticChat/tools/mcp/mcpManager.ts | 4 ++++ .../amazonQServiceManager/AmazonQTokenServiceManager.ts | 1 + 2 files changed, 5 insertions(+) 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 index 24586d3331..4b18320043 100644 --- 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 @@ -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 diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts index 79a6c7e537..2fa8eb3eac 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts @@ -163,6 +163,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< ProfileStatusMonitor.resetMcpState() if (McpManager.isInitialized()) { McpManager.instance.setRegistryActive(false) + McpManager.instance.resetRegistryService() void McpManager.instance.close(true) } } From 40046758dab3cd809371308a281a9a5a9ea42c1d Mon Sep 17 00:00:00 2001 From: Shruti Sinha Date: Thu, 8 Jan 2026 14:40:08 -0800 Subject: [PATCH 3/7] chore: refactor mcpmanager calls to profilestatusmonitor --- .../tools/mcp/profileStatusMonitor.ts | 19 +++++++++++++++++++ .../AmazonQTokenServiceManager.ts | 13 +++---------- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts index e69829b8e9..ddf1238e57 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts @@ -14,6 +14,7 @@ import { EventEmitter } from 'events' import { McpRegistryService } from './mcpRegistryService' import { McpRegistryData } from './mcpTypes' import { GetProfileResponse } from '@amzn/codewhisperer-runtime' +import { McpManager } from './mcpManager' export const AUTH_SUCCESS_EVENT = 'authSuccess' @@ -182,6 +183,24 @@ export class ProfileStatusMonitor { ProfileStatusMonitor.setMcpState(true) } + static resetMcpManager(): void { + if (McpManager.isInitialized()) { + McpManager.instance.setRegistryActive(false) + McpManager.instance.resetRegistryService() + void McpManager.instance.close(true) + } + } + + static discoverMcpServers(): void { + if (McpManager.isInitialized()) { + void McpManager.instance.discoverAllServers() + } + } + + static isMcpManagerInitialized(): boolean { + return McpManager.isInitialized() + } + static emitAuthSuccess(): void { ProfileStatusMonitor.eventEmitter.emit(AUTH_SUCCESS_EVENT) } diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts index 2fa8eb3eac..61ea49b7a1 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts @@ -37,8 +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' -import { McpManager } from '../../language-server/agenticChat/tools/mcp/mcpManager' - /** * AmazonQTokenServiceManager manages state and provides centralized access to * instance of CodeWhispererServiceToken SDK client to any consuming code. @@ -161,11 +159,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< // Reset MCP state cache when auth changes ProfileStatusMonitor.resetMcpState() - if (McpManager.isInitialized()) { - McpManager.instance.setRegistryActive(false) - McpManager.instance.resetRegistryService() - void McpManager.instance.close(true) - } + ProfileStatusMonitor.resetMcpManager() } public handleOnCredentialsUpdated(type: CredentialsType): void { @@ -270,9 +264,8 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< ) this.state = 'INITIALIZED' this.logging.log(`Initialized Amazon Q service with ${newConnectionType} connection`) - if (McpManager.isInitialized()) { - void McpManager.instance.discoverAllServers() - } + + ProfileStatusMonitor.discoverServersWhenNoProfiles() return } From db5f4744638a3f86996eefcd3728e5e4ebc63091 Mon Sep 17 00:00:00 2001 From: Shruti Sinha Date: Thu, 8 Jan 2026 14:44:48 -0800 Subject: [PATCH 4/7] fix: rename function --- .../agenticChat/tools/mcp/profileStatusMonitor.ts | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts index ddf1238e57..381fe4b968 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts @@ -191,16 +191,16 @@ export class ProfileStatusMonitor { } } - static discoverMcpServers(): void { + /** + * When switching between connections especially builder id + * re-discover mcp servers for non-profile connections + */ + static discoverServersWhenNoProfiles(): void { if (McpManager.isInitialized()) { void McpManager.instance.discoverAllServers() } } - static isMcpManagerInitialized(): boolean { - return McpManager.isInitialized() - } - static emitAuthSuccess(): void { ProfileStatusMonitor.eventEmitter.emit(AUTH_SUCCESS_EVENT) } From 250b7eb4d5dd40aa36fbc33b05261edf0dfda9e8 Mon Sep 17 00:00:00 2001 From: Shruti Sinha Date: Thu, 8 Jan 2026 15:21:22 -0800 Subject: [PATCH 5/7] fix: use dynamic require to satisfy webpack --- .../tools/mcp/profileStatusMonitor.ts | 25 +++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts index 381fe4b968..141500074c 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts @@ -14,7 +14,6 @@ import { EventEmitter } from 'events' import { McpRegistryService } from './mcpRegistryService' import { McpRegistryData } from './mcpTypes' import { GetProfileResponse } from '@amzn/codewhisperer-runtime' -import { McpManager } from './mcpManager' export const AUTH_SUCCESS_EVENT = 'authSuccess' @@ -184,10 +183,16 @@ export class ProfileStatusMonitor { } static resetMcpManager(): void { - if (McpManager.isInitialized()) { - McpManager.instance.setRegistryActive(false) - McpManager.instance.resetRegistryService() - void McpManager.instance.close(true) + try { + // note: use dynamic require to satisfy webpack requirements for webworker + const { McpManager } = require('./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}`) } } @@ -196,8 +201,14 @@ export class ProfileStatusMonitor { * re-discover mcp servers for non-profile connections */ static discoverServersWhenNoProfiles(): void { - if (McpManager.isInitialized()) { - void McpManager.instance.discoverAllServers() + try { + // note: use dynamic require to satisfy webpack requirements for webworker + const { McpManager } = require('./mcpManager') + if (McpManager.isInitialized()) { + void McpManager.instance.discoverAllServers() + } + } catch (error) { + ProfileStatusMonitor.logging?.error(`Failed to discover mcp servers when no profiles available: ${error}`) } } From 31db52296f579699751a3015aed9a1a072efd345 Mon Sep 17 00:00:00 2001 From: Shruti Sinha Date: Thu, 8 Jan 2026 15:36:04 -0800 Subject: [PATCH 6/7] fix: try eval require --- .../agenticChat/tools/mcp/profileStatusMonitor.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts index 141500074c..d7f1ae2466 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts @@ -185,7 +185,7 @@ export class ProfileStatusMonitor { static resetMcpManager(): void { try { // note: use dynamic require to satisfy webpack requirements for webworker - const { McpManager } = require('./mcpManager') + const McpManager = eval('require')('./mcpManager').McpManager if (McpManager.isInitialized()) { McpManager.instance.setRegistryActive(false) McpManager.instance.resetRegistryService() @@ -203,7 +203,7 @@ export class ProfileStatusMonitor { static discoverServersWhenNoProfiles(): void { try { // note: use dynamic require to satisfy webpack requirements for webworker - const { McpManager } = require('./mcpManager') + const McpManager = eval('require')('./mcpManager').McpManager if (McpManager.isInitialized()) { void McpManager.instance.discoverAllServers() } From 9f406c8620c17951a3160cdd60b12e95721864fc Mon Sep 17 00:00:00 2001 From: Shruti Sinha Date: Tue, 13 Jan 2026 10:03:50 -0800 Subject: [PATCH 7/7] fix: address issue where mcp registry servers installed did not load on connection switch --- .../agenticChat/tools/mcp/mcpManager.ts | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) 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 index 4b18320043..5a01b86b89 100644 --- 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 @@ -1853,8 +1853,6 @@ export class McpManager { this.registryService = new McpRegistryService(this.features.logging) } - const wasActive = this.registryUrlProvided - try { const registry = await this.registryService.fetchRegistry(registryUrl) if (registry) { @@ -1862,19 +1860,13 @@ export class McpManager { // 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