From bd590f046f326ca54b9741437034463daf5dae85 Mon Sep 17 00:00:00 2001 From: Ashish Reddy Podduturi Date: Fri, 25 Jul 2025 12:46:34 -0700 Subject: [PATCH] feat(amazonq): adding feature flag and API for mcp admin configuration --- packages/amazonq/src/extension.ts | 14 +++ packages/amazonq/src/lsp/client.ts | 2 + .../codewhisperer/client/user-service-2.json | 85 ++++++++++++++++++- packages/core/src/codewhisperer/index.ts | 1 + .../src/codewhisperer/util/profileUtils.ts | 71 ++++++++++++++++ 5 files changed, 172 insertions(+), 1 deletion(-) create mode 100644 packages/core/src/codewhisperer/util/profileUtils.ts diff --git a/packages/amazonq/src/extension.ts b/packages/amazonq/src/extension.ts index 53d7cd88037..5d8e7f6250c 100644 --- a/packages/amazonq/src/extension.ts +++ b/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 { + 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) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 4d052912c8e..188a9bef447 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/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, diff --git a/packages/core/src/codewhisperer/client/user-service-2.json b/packages/core/src/codewhisperer/client/user-service-2.json index 714937ed402..ee3831401ff 100644 --- a/packages/core/src/codewhisperer/client/user-service-2.json +++ b/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": "

Get the requested CodeWhisperer profile.

" + }, "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" } diff --git a/packages/core/src/codewhisperer/index.ts b/packages/core/src/codewhisperer/index.ts index d782b2abefe..b6adb84eebd 100644 --- a/packages/core/src/codewhisperer/index.ts +++ b/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' diff --git a/packages/core/src/codewhisperer/util/profileUtils.ts b/packages/core/src/codewhisperer/util/profileUtils.ts new file mode 100644 index 00000000000..a30864be475 --- /dev/null +++ b/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 { + 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(fn: () => Promise, maxRetries = 3): Promise { + 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 +}