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/mcp/mcpManager.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.ts index 24586d3331..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 @@ -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,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) { @@ -1858,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 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..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 @@ -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) } 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..61ea49b7a1 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts +++ b/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 }