feat(amazonq): adding feature flag and API for mcp admin configuration

Author: ashishrp-aws

packages/amazonq/src/extension.ts

@@ -7,6 +7,7 @@ import { AuthUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-c
 import { activate as activateCodeWhisperer, shutdown as shutdownCodeWhisperer } from 'aws-core-vscode/codewhisperer'
 import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
 import { CommonAuthWebview } from 'aws-core-vscode/login'
+import { checkMcpConfiguration } from 'aws-core-vscode/codewhisperer'
 import {
     amazonQDiffScheme,
     DefaultAWSClientBuilder,
@@ -48,6 +49,16 @@ import { hasGlibcPatch } from './lsp/client'
 
 export const amazonQContextPrefix = 'amazonq'
 
+// Global variable to store mcpAdmin feature flag
+export let mcpAdmin = true
+
+/**
+ * Safely checks MCP configuration from user profile and sets mcpAdmin feature flag
+ */
+async function setMcpAdminFlag(): Promise<void> {
+    mcpAdmin = await checkMcpConfiguration()
+}
+
 /**
  * Activation code for Amazon Q that will we want in all environments (eg Node.js, web mode)
  */
@@ -126,6 +137,9 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     // This contains every lsp agnostic things (auth, security scan, code scan)
     await activateCodeWhisperer(extContext as ExtContext)
 
+    // Check MCP configuration from profile before activation
+    await setMcpAdminFlag()
+
     if (!isAmazonLinux2() || hasGlibcPatch()) {
         // Activate Amazon Q LSP for everyone unless they're using AL2 without the glibc patch
         await activateAmazonqLsp(context)

packages/amazonq/src/lsp/client.ts

@@ -50,6 +50,7 @@ import { SessionManager } from '../app/inline/sessionManager'
 import { LineTracker } from '../app/inline/stateTracker/lineTracker'
 import { InlineTutorialAnnotation } from '../app/inline/tutorials/inlineTutorialAnnotation'
 import { InlineChatTutorialAnnotation } from '../app/inline/tutorials/inlineChatTutorialAnnotation'
+import { mcpAdmin } from '../extension'
 
 const localize = nls.loadMessageBundle()
 const logger = getLogger('amazonqLsp.lspClient')
@@ -165,6 +166,7 @@ export async function startLanguageServer(
                         pinnedContextEnabled: true,
                         imageContextEnabled: true,
                         mcp: true,
+                        mcpAdmin: mcpAdmin,
                         shortcut: true,
                         reroute: true,
                         modelSelection: true,

packages/core/src/codewhisperer/client/user-service-2.json

@@ -202,6 +202,37 @@
                 { "shape": "AccessDeniedException" }
             ]
         },
+        "GetProfile": {
+            "name": "GetProfile",
+            "http": {
+                "method": "POST",
+                "requestUri": "/"
+            },
+            "input": {
+                "shape": "GetProfileRequest"
+            },
+            "output": {
+                "shape": "GetProfileResponse"
+            },
+            "errors": [
+                {
+                    "shape": "ThrottlingException"
+                },
+                {
+                    "shape": "ResourceNotFoundException"
+                },
+                {
+                    "shape": "InternalServerException"
+                },
+                {
+                    "shape": "ValidationException"
+                },
+                {
+                    "shape": "AccessDeniedException"
+                }
+            ],
+            "documentation": "<p>Get the requested CodeWhisperer profile.</p>"
+        },
         "GetTaskAssistCodeGeneration": {
             "name": "GetTaskAssistCodeGeneration",
             "http": {
@@ -1997,6 +2028,24 @@
                 "suggestedFix": { "shape": "SuggestedFix" }
             }
         },
+        "GetProfileRequest": {
+            "type": "structure",
+            "required": ["profileArn"],
+            "members": {
+                "profileArn": {
+                    "shape": "ProfileArn"
+                }
+            }
+        },
+        "GetProfileResponse": {
+            "type": "structure",
+            "required": ["profile"],
+            "members": {
+                "profile": {
+                    "shape": "ProfileInfo"
+                }
+            }
+        },
         "GetTaskAssistCodeGenerationRequest": {
             "type": "structure",
             "required": ["conversationId", "codeGenerationId"],
@@ -2426,6 +2475,15 @@
             "type": "long",
             "box": true
         },
+        "MCPConfiguration": {
+            "type": "structure",
+            "required": ["toggle"],
+            "members": {
+                "toggle": {
+                    "shape": "OptInFeatureToggle"
+                }
+            }
+        },
         "MemoryEntry": {
             "type": "structure",
             "required": ["id", "memoryEntryString", "metadata"],
@@ -2532,7 +2590,8 @@
                 "byUserAnalytics": { "shape": "ByUserAnalytics" },
                 "dashboardAnalytics": { "shape": "DashboardAnalytics" },
                 "notifications": { "shape": "Notifications" },
-                "workspaceContext": { "shape": "WorkspaceContext" }
+                "workspaceContext": { "shape": "WorkspaceContext" },
+                "mcpConfiguration": { "shape": "MCPConfiguration" }
             }
         },
         "OptOutPreference": {
@@ -2682,6 +2741,30 @@
             "min": 1,
             "pattern": "[\\sa-zA-Z0-9_-]*"
         },
+        "ProfileInfo": {
+            "type": "structure",
+            "required": ["arn"],
+            "members": {
+                "arn": {
+                    "shape": "ProfileArn"
+                },
+                "profileName": {
+                    "shape": "ProfileName"
+                },
+                "description": {
+                    "shape": "ProfileDescription"
+                },
+                "status": {
+                    "shape": "ProfileStatus"
+                },
+                "profileType": {
+                    "shape": "ProfileType"
+                },
+                "optInFeatures": {
+                    "shape": "OptInFeatures"
+                }
+            }
+        },
         "ProfileList": {
             "type": "list",
             "member": { "shape": "Profile" }

packages/core/src/codewhisperer/index.ts

@@ -86,3 +86,4 @@ export * from './util/gitUtil'
 export * from './ui/prompters'
 export { UserWrittenCodeTracker } from './tracker/userWrittenCodeTracker'
 export { RegionProfileManager, defaultServiceConfig } from './region/regionProfileManager'
+export { checkMcpConfiguration } from './util/profileUtils'

packages/core/src/codewhisperer/util/profileUtils.ts

@@ -0,0 +1,71 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { getLogger } from '../../shared/logger/logger'
+import { codeWhispererClient } from '../client/codewhisperer'
+import { AuthUtil } from './authUtil'
+
+export async function checkMcpConfiguration(): Promise<boolean> {
+    try {
+        if (!AuthUtil.instance.isConnected()) {
+            return true
+        }
+
+        const userClient = await codeWhispererClient.createUserSdkClient()
+        const profileArn = AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn
+        if (!profileArn) {
+            return true
+        }
+
+        const response = await retryWithBackoff(() => userClient.getProfile({ profileArn }).promise())
+        const mcpConfig = response.profile?.optInFeatures?.mcpConfiguration?.toggle
+        const isMcpEnabled = mcpConfig === 'ON'
+
+        getLogger().debug(`MCP configuration toggle: ${mcpConfig}, mcpAdmin flag set to: ${isMcpEnabled}`)
+        return isMcpEnabled
+    } catch (error) {
+        getLogger().debug(`Failed to check MCP configuration from profile: ${error}. Setting mcpAdmin to false.`)
+        return true
+    }
+}
+
+async function retryWithBackoff<T>(fn: () => Promise<T>, maxRetries = 3): Promise<T> {
+    let lastError: any
+
+    for (let attempt = 0; attempt < maxRetries; attempt++) {
+        try {
+            return await fn()
+        } catch (error) {
+            lastError = error
+
+            // Only retry on specific retryable exceptions
+            const errorCode = (error as any).code || (error as any).name
+            const statusCode = (error as any).statusCode
+
+            // Don't retry on client errors (4xx) except ThrottlingException
+            if (statusCode >= 400 && statusCode < 500 && errorCode !== 'ThrottlingException') {
+                throw error
+            }
+
+            // Only retry on retryable exceptions
+            const retryableExceptions = [
+                'ThrottlingException',
+                'InternalServerException',
+                'ServiceUnavailableException',
+            ]
+            if (!retryableExceptions.includes(errorCode) && statusCode !== 500 && statusCode !== 503) {
+                throw error
+            }
+
+            if (attempt < maxRetries - 1) {
+                const delay = Math.min(1000 * Math.pow(2, attempt), 3000) // Cap at 3s
+                getLogger().debug(`GetProfile attempt ${attempt + 1} failed, retrying in ${delay}ms: ${error}`)
+                await new Promise((resolve) => setTimeout(resolve, delay))
+            }
+        }
+    }
+
+    throw lastError
+}