diff --git a/chat-client/src/client/mcpMynahUi.test.ts b/chat-client/src/client/mcpMynahUi.test.ts index 9d1dec5407..947e5bc604 100644 --- a/chat-client/src/client/mcpMynahUi.test.ts +++ b/chat-client/src/client/mcpMynahUi.test.ts @@ -107,10 +107,8 @@ describe('McpMynahUi', () => { assert.strictEqual(callArgs.detailedList.header.description, 'Test Description') assert.deepStrictEqual(callArgs.detailedList.header.status, { status: 'success' }) - // Verify the actions in the header - assert.strictEqual(callArgs.detailedList.header.actions.length, 2) - assert.strictEqual(callArgs.detailedList.header.actions[0].id, 'add-new-mcp') - assert.strictEqual(callArgs.detailedList.header.actions[1].id, 'refresh-mcp-list') + // Verify the actions in the header (no default actions are added when header is provided) + assert.strictEqual(callArgs.detailedList.header.actions.length, 0) // Verify the list structure assert.strictEqual(callArgs.detailedList.list.length, 1) diff --git a/chat-client/src/client/mcpMynahUi.ts b/chat-client/src/client/mcpMynahUi.ts index f9216b8168..5ece955dfa 100644 --- a/chat-client/src/client/mcpMynahUi.ts +++ b/chat-client/src/client/mcpMynahUi.ts @@ -272,20 +272,11 @@ export class McpMynahUi { title: params.header.title, description: params.header.description, status: params.header.status, - actions: [ - { - id: MCP_IDS.ADD_NEW, - icon: toMynahIcon('plus'), - status: 'clear', - description: 'Add new MCP', - }, - { - id: MCP_IDS.REFRESH_LIST, - icon: toMynahIcon('refresh'), - status: 'clear', - description: 'Refresh MCP servers', - }, - ], + actions: + params.header.actions?.map(action => ({ + ...action, + icon: action.icon ? toMynahIcon(action.icon) : undefined, + })) || [], } : undefined, filterOptions: params.filterOptions?.map(filter => ({ diff --git a/core/aws-lsp-core/src/index.ts b/core/aws-lsp-core/src/index.ts index 1ad391dc8f..030a4fa2e3 100644 --- a/core/aws-lsp-core/src/index.ts +++ b/core/aws-lsp-core/src/index.ts @@ -19,3 +19,4 @@ export * as workspaceUtils from './util/workspaceUtils' export * as processUtils from './util/processUtils' export * as collectionUtils from './util/collectionUtils' export * as loggingUtils from './util/loggingUtils' +export * as retryUtils from './util/retryUtils' diff --git a/core/aws-lsp-core/src/util/retryUtils.test.ts b/core/aws-lsp-core/src/util/retryUtils.test.ts new file mode 100644 index 0000000000..4223e1f3bf --- /dev/null +++ b/core/aws-lsp-core/src/util/retryUtils.test.ts @@ -0,0 +1,120 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. + * All Rights Reserved. SPDX-License-Identifier: Apache-2.0 + */ + +import { expect } from 'chai' +import * as sinon from 'sinon' +import { retryWithBackoff, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY } from './retryUtils' + +describe('retryUtils', () => { + let clock: sinon.SinonFakeTimers + + beforeEach(() => { + clock = sinon.useFakeTimers() + }) + + afterEach(() => { + clock.restore() + }) + + describe('retryWithBackoff', () => { + it('should return result on first success', async () => { + const fn = sinon.stub().resolves('success') + + const result = await retryWithBackoff(fn) + + expect(result).to.equal('success') + expect(fn.callCount).to.equal(1) + }) + + it('should retry on retryable errors', async () => { + const fn = sinon.stub() + fn.onFirstCall().rejects({ code: 'ThrottlingException' }) + fn.onSecondCall().resolves('success') + + const promise = retryWithBackoff(fn) + await clock.tickAsync(DEFAULT_BASE_DELAY) + const result = await promise + + expect(result).to.equal('success') + expect(fn.callCount).to.equal(2) + }) + + it('should not retry on non-retryable client errors', async () => { + const error = { statusCode: 404 } + const fn = sinon.stub().rejects(error) + + try { + await retryWithBackoff(fn) + expect.fail('Expected function to throw') + } catch (e) { + expect(e).to.equal(error) + } + expect(fn.callCount).to.equal(1) + }) + + it('should retry on server errors', async () => { + const fn = sinon.stub() + fn.onFirstCall().rejects({ statusCode: 500 }) + fn.onSecondCall().resolves('success') + + const promise = retryWithBackoff(fn) + await clock.tickAsync(DEFAULT_BASE_DELAY) + const result = await promise + + expect(result).to.equal('success') + expect(fn.callCount).to.equal(2) + }) + + it('should use exponential backoff by default', async () => { + const fn = sinon.stub() + const error = { code: 'ThrottlingException' } + fn.onFirstCall().rejects(error) + fn.onSecondCall().rejects(error) + + const promise = retryWithBackoff(fn) + + // First retry after baseDelay * 1 + await clock.tickAsync(DEFAULT_BASE_DELAY) + // Second retry after baseDelay * 2 + await clock.tickAsync(DEFAULT_BASE_DELAY * 2) + + try { + await promise + expect.fail('Expected function to throw') + } catch (e) { + expect(e).to.equal(error) + } + expect(fn.callCount).to.equal(DEFAULT_MAX_RETRIES) + }) + + it('should respect custom maxRetries', async () => { + const error = { code: 'ThrottlingException' } + const fn = sinon.stub().rejects(error) + + try { + await retryWithBackoff(fn, { maxRetries: 1 }) + expect.fail('Expected function to throw') + } catch (e) { + expect(e).to.equal(error) + } + expect(fn.callCount).to.equal(1) + }) + + it('should use custom isRetryable function', async () => { + const error = { custom: 'error' } + const fn = sinon.stub().rejects(error) + const isRetryable = sinon.stub().returns(false) + + try { + await retryWithBackoff(fn, { isRetryable }) + expect.fail('Expected function to throw') + } catch (e) { + expect(e).to.equal(error) + } + expect(fn.callCount).to.equal(1) + expect(isRetryable.calledWith(error)).to.equal(true) + }) + }) +}) diff --git a/core/aws-lsp-core/src/util/retryUtils.ts b/core/aws-lsp-core/src/util/retryUtils.ts new file mode 100644 index 0000000000..dc135ce23d --- /dev/null +++ b/core/aws-lsp-core/src/util/retryUtils.ts @@ -0,0 +1,77 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. + * All Rights Reserved. SPDX-License-Identifier: Apache-2.0 + */ + +// Default retry configuration constants +export const DEFAULT_MAX_RETRIES = 2 +export const DEFAULT_BASE_DELAY = 500 +export const DEFAULT_EXPONENTIAL_BACKOFF = true + +// HTTP status code constants +const CLIENT_ERROR_MIN = 400 +const CLIENT_ERROR_MAX = 500 +const INTERNAL_SERVER_ERROR = 500 +const SERVICE_UNAVAILABLE = 503 + +// AWS error code constants +const THROTTLING_EXCEPTION = 'ThrottlingException' +const INTERNAL_SERVER_EXCEPTION = 'InternalServerException' + +export interface RetryOptions { + /** Maximum number of retry attempts (default: DEFAULT_MAX_RETRIES) */ + maxRetries?: number + /** Base delay in milliseconds (default: DEFAULT_BASE_DELAY) */ + baseDelay?: number + /** Whether to use exponential backoff (default: DEFAULT_EXPONENTIAL_BACKOFF) */ + exponentialBackoff?: boolean + /** Custom function to determine if an error is retryable */ + isRetryable?: (error: any) => boolean +} + +/** + * Default AWS error retry logic + */ +function defaultIsRetryable(error: any): boolean { + const errorCode = error.code || error.name + const statusCode = error.statusCode + + // 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 + return ( + errorCode === THROTTLING_EXCEPTION || + errorCode === INTERNAL_SERVER_EXCEPTION || + statusCode === INTERNAL_SERVER_ERROR || + statusCode === SERVICE_UNAVAILABLE + ) +} + +/** + * Executes a function with retry logic and exponential backoff + */ +export async function retryWithBackoff(fn: () => Promise, options: RetryOptions = {}): Promise { + const { + maxRetries = DEFAULT_MAX_RETRIES, + baseDelay = DEFAULT_BASE_DELAY, + exponentialBackoff = DEFAULT_EXPONENTIAL_BACKOFF, + isRetryable = defaultIsRetryable, + } = options + + for (let attempt = 0; attempt < maxRetries; attempt++) { + try { + return await fn() + } catch (error: any) { + if (!isRetryable(error) || attempt === maxRetries - 1) { + throw error + } + + const delay = exponentialBackoff ? baseDelay * (attempt + 1) : baseDelay + await new Promise(resolve => setTimeout(resolve, delay)) + } + } + throw new Error('Retry failed') +} diff --git a/server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json b/server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json index 2120382d25..a5704fca16 100644 --- a/server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json +++ b/server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json @@ -516,6 +516,37 @@ ], "documentation": "

API to get code transformation status.

" }, + "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.

" + }, "GetUsageLimits": { "name": "GetUsageLimits", "http": { @@ -3126,6 +3157,24 @@ } } }, + "GetProfileRequest": { + "type": "structure", + "required": ["profileArn"], + "members": { + "profileArn": { + "shape": "ProfileArn" + } + } + }, + "GetProfileResponse": { + "type": "structure", + "required": ["profile"], + "members": { + "profile": { + "shape": "ProfileInfo" + } + } + }, "GetTaskAssistCodeGenerationRequest": { "type": "structure", "required": ["conversationId", "codeGenerationId"], @@ -3787,6 +3836,15 @@ "type": "long", "box": true }, + "MCPConfiguration": { + "type": "structure", + "required": ["toggle"], + "members": { + "toggle": { + "shape": "OptInFeatureToggle" + } + } + }, "MemoryEntry": { "type": "structure", "required": ["id", "memoryEntryString", "metadata"], @@ -3995,6 +4053,9 @@ }, "workspaceContext": { "shape": "WorkspaceContext" + }, + "mcpConfiguration": { + "shape": "MCPConfiguration" } } }, @@ -4189,6 +4250,30 @@ } } }, + "ProfileInfo": { + "type": "structure", + "required": ["arn"], + "members": { + "arn": { + "shape": "ProfileArn" + }, + "profileName": { + "shape": "ProfileName" + }, + "description": { + "shape": "ProfileDescription" + }, + "status": { + "shape": "ProfileStatus" + }, + "profileType": { + "shape": "ProfileType" + }, + "optInFeatures": { + "shape": "OptInFeatures" + } + } + }, "ProfileArn": { "type": "string", "max": 950, diff --git a/server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts b/server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts index 2189944d4f..c885612888 100644 --- a/server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts +++ b/server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts @@ -144,6 +144,14 @@ declare class CodeWhispererBearerTokenClient extends Service { * API to get code transformation status. */ getTransformationPlan(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetTransformationPlanResponse) => void): Request; + /** + * Get the requested CodeWhisperer profile. + */ + getProfile(params: CodeWhispererBearerTokenClient.Types.GetProfileRequest, callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetProfileResponse) => void): Request; + /** + * Get the requested CodeWhisperer profile. + */ + getProfile(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetProfileResponse) => void): Request; /** * API to get current usage limits */ @@ -961,6 +969,12 @@ declare namespace CodeWhispererBearerTokenClient { jobStatus?: CodeFixJobStatus; suggestedFix?: SuggestedFix; } + export interface GetProfileRequest { + profileArn: ProfileArn; + } + export interface GetProfileResponse { + profile: ProfileInfo; + } export interface GetTaskAssistCodeGenerationRequest { conversationId: ConversationId; codeGenerationId: CodeGenerationId; @@ -1189,6 +1203,9 @@ declare namespace CodeWhispererBearerTokenClient { nextToken?: String; } export type Long = number; + export interface MCPConfiguration { + toggle: OptInFeatureToggle; + } export interface MemoryEntry { /** * A unique identifier for a single memory entry @@ -1260,6 +1277,7 @@ declare namespace CodeWhispererBearerTokenClient { dashboardAnalytics?: DashboardAnalytics; notifications?: Notifications; workspaceContext?: WorkspaceContext; + mcpConfiguration?: MCPConfiguration; } export type OptOutPreference = "OPTIN"|"OPTOUT"|string; export type Origin = "CHATBOT"|"CONSOLE"|"DOCUMENTATION"|"MARKETING"|"MOBILE"|"SERVICE_INTERNAL"|"UNIFIED_SEARCH"|"UNKNOWN"|"MD"|"IDE"|"SAGE_MAKER"|"CLI"|"AI_EDITOR"|"OPENSEARCH_DASHBOARD"|"GITLAB"|string; @@ -1315,6 +1333,14 @@ declare namespace CodeWhispererBearerTokenClient { permissionUpdateRequired?: Boolean; applicationProperties?: ApplicationPropertiesList; } + export interface ProfileInfo { + arn: ProfileArn; + profileName?: ProfileName; + description?: ProfileDescription; + status?: ProfileStatus; + profileType?: ProfileType; + optInFeatures?: OptInFeatures; + } export type ProfileArn = string; export type ProfileDescription = string; export type ProfileList = Profile[]; diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/constants.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/constants.ts index 09fbb20436..1729c10c1d 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/constants.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/constants.ts @@ -9,6 +9,8 @@ export const RESPONSE_TIMEOUT_PARTIAL_MSG = 'Response processing timed out after export const LOADING_THRESHOLD_MS = 2000 export const CLIENT_TIMEOUT_MS = 245_000 export const RESPONSE_TIMEOUT_MS = 240_000 +export const SERVICE_MANAGER_TIMEOUT_MS = 10_000 //10 seconds +export const SERVICE_MANAGER_POLL_INTERVAL_MS = 100 // LLM Constants export const GENERATE_ASSISTANT_RESPONSE_INPUT_LIMIT = 500_000 diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.test.ts index 01984a310e..cb645eec43 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.test.ts @@ -240,4 +240,109 @@ describe('McpEventHandler error handling', () => { expect(result.header).to.not.be.undefined expect(result.header.title).to.equal('Edit MCP Server') }) + + describe('#getListMcpServersStatus', () => { + beforeEach(() => { + sinon.restore() + sinon.stub(mcpUtils, 'getGlobalAgentConfigPath').returns('/fake/home/.aws/amazonq/agents/default.json') + saveAgentConfigStub = sinon.stub(mcpUtils, 'saveAgentConfig').resolves() + }) + + it('returns admin disabled status when MCP state is false', async () => { + // Stub ProfileStatusMonitor.getMcpState to return false + const { ProfileStatusMonitor } = await import('./profileStatusMonitor') + sinon.stub(ProfileStatusMonitor, 'getMcpState').returns(false) + + loadStub = sinon.stub(mcpUtils, 'loadAgentConfig').resolves({ + servers: new Map(), + serverNameMapping: new Map(), + errors: new Map(), + agentConfig: { + name: 'test-agent', + version: '1.0.0', + description: 'Test agent', + mcpServers: {}, + tools: [], + allowedTools: [], + toolsSettings: {}, + includedFiles: [], + resources: [], + }, + }) + + await McpManager.init([], features) + const result = await eventHandler.onListMcpServers({}) + + expect(result.header.status).to.deep.equal({ + title: 'MCP functionality has been disabled by your administrator', + icon: 'info', + status: 'info', + }) + }) + + it('returns config error status when MCP state is not false but config errors exist', async () => { + // Stub ProfileStatusMonitor.getMcpState to return true + const { ProfileStatusMonitor } = await import('./profileStatusMonitor') + sinon.stub(ProfileStatusMonitor, 'getMcpState').returns(true) + + const mockErrors = new Map([['file1.json', 'Config error']]) + loadStub = sinon.stub(mcpUtils, 'loadAgentConfig').resolves({ + servers: new Map(), + serverNameMapping: new Map(), + errors: mockErrors, + agentConfig: { + name: 'test-agent', + version: '1.0.0', + description: 'Test agent', + mcpServers: {}, + tools: [], + allowedTools: [], + toolsSettings: {}, + includedFiles: [], + resources: [], + }, + }) + + await McpManager.init([], features) + sinon.stub(McpManager.instance, 'getConfigLoadErrors').returns('File: file1.json, Error: Config error') + + const result = await eventHandler.onListMcpServers({}) + + expect(result.header.status).to.deep.equal({ + title: 'File: file1.json, Error: Config error', + icon: 'cancel-circle', + status: 'error', + }) + }) + + it('returns undefined status when MCP state is not false and no config errors', async () => { + // Stub ProfileStatusMonitor.getMcpState to return true + const { ProfileStatusMonitor } = await import('./profileStatusMonitor') + sinon.stub(ProfileStatusMonitor, 'getMcpState').returns(true) + + loadStub = sinon.stub(mcpUtils, 'loadAgentConfig').resolves({ + servers: new Map(), + serverNameMapping: new Map(), + errors: new Map(), + agentConfig: { + name: 'test-agent', + version: '1.0.0', + description: 'Test agent', + mcpServers: {}, + tools: [], + allowedTools: [], + toolsSettings: {}, + includedFiles: [], + resources: [], + }, + }) + + await McpManager.init([], features) + sinon.stub(McpManager.instance, 'getConfigLoadErrors').returns(undefined) + + const result = await eventHandler.onListMcpServers({}) + + expect(result.header.status).to.be.undefined + }) + }) }) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts index 4596bc0646..bf8c91b533 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts @@ -30,6 +30,7 @@ import { } from './mcpTypes' import { TelemetryService } from '../../../../shared/telemetry/telemetryService' import { URI } from 'vscode-uri' +import { ProfileStatusMonitor } from './profileStatusMonitor' interface PermissionOption { label: string @@ -222,18 +223,63 @@ export class McpEventHandler { } // Return the result in the expected format + const mcpState = ProfileStatusMonitor.getMcpState() const header = { title: 'MCP Servers', - description: "Add MCP servers to extend Q's capabilities.", - // only show error on list mcp server page if unable to read mcp.json file - status: configLoadErrors - ? { title: configLoadErrors, icon: 'cancel-circle', status: 'error' as Status } - : undefined, + description: mcpState === false ? '' : "Add MCP servers to extend Q's capabilities.", + status: this.#getListMcpServersStatus(configLoadErrors, mcpState), + actions: this.#getListMcpServersActions(configLoadErrors, mcpState), } return { header, list: groups } } + /** + * Gets the status for the list MCP servers header + */ + #getListMcpServersStatus( + configLoadErrors: string | undefined, + mcpState: boolean | undefined + ): { title: string; icon: string; status: Status } | undefined { + if (mcpState === false) { + return { + title: 'MCP functionality has been disabled by your administrator', + icon: 'info', + status: 'info' as Status, + } + } + + if (configLoadErrors) { + return { title: configLoadErrors, icon: 'cancel-circle', status: 'error' as Status } + } + + return undefined + } + + /** + * Gets the actions for the list MCP servers header + */ + #getListMcpServersActions(configLoadErrors: string | undefined, mcpState: boolean | undefined) { + return mcpState !== false && (!configLoadErrors || configLoadErrors === '') + ? [ + { + id: 'add-new-mcp', + icon: 'plus', + status: 'clear', + text: 'Add new MCP server', + description: 'Add new MCP server', + }, + { + id: 'refresh-mcp-list', + icon: 'refresh', + status: 'clear', + text: 'Refresh MCP servers', + description: 'Refresh MCP servers', + }, + ] + : [] + } + /** * Handles MCP server click events */ diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts index 45ab34170c..7b212bad49 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts @@ -920,8 +920,8 @@ async function migrateConfigToAgent( ...existingAgentConfig, // Merge MCP servers, keeping existing ones if they exist mcpServers: { - ...existingAgentConfig.mcpServers, ...newAgentConfig.mcpServers, + ...existingAgentConfig.mcpServers, }, // Merge tools lists without duplicates tools: [...new Set([...existingAgentConfig.tools, ...newAgentConfig.tools])], diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.test.ts new file mode 100644 index 0000000000..8ee8454374 --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.test.ts @@ -0,0 +1,182 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. + * All Rights Reserved. SPDX-License-Identifier: Apache-2.0 + */ + +import { expect } from 'chai' +import * as sinon from 'sinon' +import { ProfileStatusMonitor } from './profileStatusMonitor' +import * as AmazonQTokenServiceManagerModule from '../../../../shared/amazonQServiceManager/AmazonQTokenServiceManager' + +describe('ProfileStatusMonitor', () => { + let profileStatusMonitor: ProfileStatusMonitor + let mockCredentialsProvider: any + let mockWorkspace: any + let mockLogging: any + let mockSdkInitializator: any + let mockOnMcpDisabled: sinon.SinonStub + let mockOnMcpEnabled: sinon.SinonStub + let clock: sinon.SinonFakeTimers + + beforeEach(() => { + clock = sinon.useFakeTimers() + + mockCredentialsProvider = { + hasCredentials: sinon.stub().returns(true), + } + + mockWorkspace = {} + + mockLogging = { + info: sinon.stub(), + debug: sinon.stub(), + } + + mockSdkInitializator = {} + mockOnMcpDisabled = sinon.stub() + mockOnMcpEnabled = sinon.stub() + + profileStatusMonitor = new ProfileStatusMonitor( + mockCredentialsProvider, + mockWorkspace, + mockLogging, + mockSdkInitializator, + mockOnMcpDisabled, + mockOnMcpEnabled + ) + }) + + afterEach(() => { + clock.restore() + sinon.restore() + profileStatusMonitor.stop() + }) + + describe('start', () => { + it('should start monitoring and log info message', () => { + profileStatusMonitor.start() + + expect( + mockLogging.info.calledWith('ProfileStatusMonitor started - checking MCP configuration every 24 hours') + ).to.be.true + }) + + it('should not start multiple times', () => { + profileStatusMonitor.start() + profileStatusMonitor.start() + + expect(mockLogging.info.callCount).to.equal(1) + }) + }) + + describe('stop', () => { + it('should stop monitoring and log info message', () => { + profileStatusMonitor.start() + profileStatusMonitor.stop() + + expect(mockLogging.info.calledWith('ProfileStatusMonitor stopped')).to.be.true + }) + }) + + describe('checkInitialState', () => { + it('should return true when no profile ARN is available', async () => { + sinon.stub(AmazonQTokenServiceManagerModule.AmazonQTokenServiceManager, 'getInstance').returns({ + getActiveProfileArn: () => undefined, + } as any) + + const result = await profileStatusMonitor.checkInitialState() + expect(result).to.be.true + }) + + it('should return true and log debug message on error', async () => { + // Stub the private isMcpEnabled method to throw an error + sinon.stub(profileStatusMonitor as any, 'isMcpEnabled').throws(new Error('Service manager not ready')) + + const result = await profileStatusMonitor.checkInitialState() + expect(result).to.be.true + expect(mockLogging.debug.calledWith(sinon.match('Initial MCP state check failed, defaulting to enabled'))) + .to.be.true + }) + }) + + describe('getMcpState', () => { + beforeEach(() => { + // Reset static state before each test + ;(ProfileStatusMonitor as any).lastMcpState = undefined + }) + + it('should return undefined initially', () => { + expect(ProfileStatusMonitor.getMcpState()).to.be.undefined + }) + + it('should return the last MCP state after it is set', () => { + // Access the private static property through reflection for testing + ;(ProfileStatusMonitor as any).lastMcpState = true + expect(ProfileStatusMonitor.getMcpState()).to.be.true + ;(ProfileStatusMonitor as any).lastMcpState = false + expect(ProfileStatusMonitor.getMcpState()).to.be.false + }) + + it('should be accessible across different instances', () => { + const monitor1 = new ProfileStatusMonitor( + mockCredentialsProvider, + mockWorkspace, + mockLogging, + mockSdkInitializator, + mockOnMcpDisabled, + mockOnMcpEnabled + ) + + const monitor2 = new ProfileStatusMonitor( + mockCredentialsProvider, + mockWorkspace, + mockLogging, + mockSdkInitializator, + mockOnMcpDisabled, + mockOnMcpEnabled + ) + + // Set state through static property + ;(ProfileStatusMonitor as any).lastMcpState = true + + // Should be accessible from both instances + expect(ProfileStatusMonitor.getMcpState()).to.be.true + }) + }) + + describe('static lastMcpState', () => { + beforeEach(() => { + // Reset static state before each test + ;(ProfileStatusMonitor as any).lastMcpState = undefined + }) + + it('should maintain state across multiple instances', () => { + const monitor1 = new ProfileStatusMonitor( + mockCredentialsProvider, + mockWorkspace, + mockLogging, + mockSdkInitializator, + mockOnMcpDisabled, + mockOnMcpEnabled + ) + + const monitor2 = new ProfileStatusMonitor( + mockCredentialsProvider, + mockWorkspace, + mockLogging, + mockSdkInitializator, + mockOnMcpDisabled, + mockOnMcpEnabled + ) + + // Initially undefined + expect(ProfileStatusMonitor.getMcpState()).to.be.undefined + + // Set through internal mechanism (simulating state change) + ;(ProfileStatusMonitor as any).lastMcpState = false + + // Both instances should see the same state + expect(ProfileStatusMonitor.getMcpState()).to.be.false + }) + }) +}) 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 new file mode 100644 index 0000000000..4a1cc705e4 --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts @@ -0,0 +1,129 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. + * All Rights Reserved. SPDX-License-Identifier: Apache-2.0 + */ + +import { + CredentialsProvider, + Logging, + SDKInitializator, + Workspace, +} from '@aws/language-server-runtimes/server-interface' +import { retryUtils } from '@aws/lsp-core' +import { CodeWhispererServiceToken } from '../../../../shared/codeWhispererService' +import { DEFAULT_AWS_Q_ENDPOINT_URL, DEFAULT_AWS_Q_REGION } from '../../../../shared/constants' +import { AmazonQTokenServiceManager } from '../../../../shared/amazonQServiceManager/AmazonQTokenServiceManager' + +export class ProfileStatusMonitor { + private intervalId?: NodeJS.Timeout + private readonly CHECK_INTERVAL = 24 * 60 * 60 * 1000 // 24 hours + private codeWhispererClient?: CodeWhispererServiceToken + private cachedProfileArn?: string + private static lastMcpState?: boolean + + constructor( + private credentialsProvider: CredentialsProvider, + private workspace: Workspace, + private logging: Logging, + private sdkInitializator: SDKInitializator, + private onMcpDisabled: () => void, + private onMcpEnabled?: () => void + ) {} + + async checkInitialState(): Promise { + try { + const isMcpEnabled = await this.isMcpEnabled() + return isMcpEnabled !== false // Return true if enabled or API failed + } catch (error) { + this.logging.debug(`Initial MCP state check failed, defaulting to enabled: ${error}`) + ProfileStatusMonitor.lastMcpState = true + return true + } + } + + start(): void { + if (this.intervalId) { + return + } + + this.intervalId = setInterval(() => { + void this.isMcpEnabled() + }, this.CHECK_INTERVAL) + + this.logging.info('ProfileStatusMonitor started - checking MCP configuration every 24 hours') + } + + stop(): void { + if (this.intervalId) { + clearInterval(this.intervalId) + this.intervalId = undefined + this.logging.info('ProfileStatusMonitor stopped') + } + } + + private async isMcpEnabled(): Promise { + try { + const profileArn = this.getProfileArn() + if (!profileArn) { + this.logging.debug('No profile ARN available for MCP configuration check') + ProfileStatusMonitor.lastMcpState = true // Default to enabled if no profile + return true + } + + if (!this.codeWhispererClient) { + this.codeWhispererClient = new CodeWhispererServiceToken( + this.credentialsProvider, + this.workspace, + this.logging, + process.env.CODEWHISPERER_REGION || DEFAULT_AWS_Q_REGION, + process.env.CODEWHISPERER_ENDPOINT || DEFAULT_AWS_Q_ENDPOINT_URL, + this.sdkInitializator + ) + this.codeWhispererClient.profileArn = profileArn + } + + const response = await retryUtils.retryWithBackoff(() => + this.codeWhispererClient!.getProfile({ profileArn }) + ) + const isMcpEnabled = response?.profile?.optInFeatures?.mcpConfiguration?.toggle === 'ON' + + if (ProfileStatusMonitor.lastMcpState !== isMcpEnabled) { + ProfileStatusMonitor.lastMcpState = isMcpEnabled + if (!isMcpEnabled) { + this.logging.info('MCP configuration disabled - removing tools') + this.onMcpDisabled() + } else if (isMcpEnabled && this.onMcpEnabled) { + this.logging.info('MCP configuration enabled - initializing tools') + this.onMcpEnabled() + } + } + + return isMcpEnabled + } catch (error) { + this.logging.debug(`MCP configuration check failed, defaulting to enabled: ${error}`) + ProfileStatusMonitor.lastMcpState = true + return true + } + } + + private getProfileArn(): string | undefined { + // Use cached value if available + if (this.cachedProfileArn) { + return this.cachedProfileArn + } + + try { + // Get profile ARN from service manager like in agenticChatController + const serviceManager = AmazonQTokenServiceManager.getInstance() + this.cachedProfileArn = serviceManager.getActiveProfileArn() + return this.cachedProfileArn + } catch (error) { + this.logging.debug(`Failed to get profile ARN: ${error}`) + } + return undefined + } + + static getMcpState(): boolean | undefined { + return ProfileStatusMonitor.lastMcpState + } +} 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 54eba4c852..fd692e6e8e 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 @@ -24,6 +24,9 @@ import { FsReplace, FsReplaceParams } from './fsReplace' import { CodeReviewUtils } from './qCodeAnalysis/codeReviewUtils' import { DEFAULT_AWS_Q_ENDPOINT_URL, DEFAULT_AWS_Q_REGION } from '../../../shared/constants' import { DisplayFindings } from './qCodeAnalysis/displayFindings' +import { ProfileStatusMonitor } from './mcp/profileStatusMonitor' +import { AmazonQTokenServiceManager } from '../../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { SERVICE_MANAGER_TIMEOUT_MS, SERVICE_MANAGER_POLL_INTERVAL_MS } from '../constants/constants' export const FsToolsServer: Server = ({ workspace, logging, agent, lsp }) => { const fsReadTool = new FsRead({ workspace, lsp, logging }) @@ -199,10 +202,45 @@ export const LspToolsServer: Server = ({ workspace, logging, lsp, agent }) => { return () => {} } -export const McpToolsServer: Server = ({ credentialsProvider, workspace, logging, lsp, agent, telemetry, runtime }) => { +export const McpToolsServer: Server = ({ + credentialsProvider, + workspace, + logging, + lsp, + agent, + telemetry, + runtime, + sdkInitializator, + chat, +}) => { const registered: Record = {} - const allNamespacedTools = new Set() + let profileStatusMonitor: ProfileStatusMonitor | undefined + + function removeAllMcpTools(): void { + logging.info('Removing all MCP tools due to admin configuration') + for (const [server, toolNames] of Object.entries(registered)) { + for (const name of toolNames) { + agent.removeTool(name) + allNamespacedTools.delete(name) + logging.info(`MCP: removed tool ${name}`) + } + registered[server] = [] + } + void McpManager.instance.close(true) //keep the instance but close all servers. + + try { + chat?.sendChatUpdate({ + tabId: 'mcpserver', + data: { + placeholderText: 'mcp-server-update', + messages: [], + }, + }) + } catch (error) { + logging.error(`Failed to send chatOptionsUpdate: ${error}`) + } + } function registerServerTools(server: string, defs: McpToolDefinition[]) { // 1) remove old tools @@ -257,20 +295,13 @@ export const McpToolsServer: Server = ({ credentialsProvider, workspace, logging } } - lsp.onInitialized(async () => { + async function initializeMcp() { try { - if (!enabledMCP(lsp.getClientInitializeParams())) { - logging.warn('MCP is currently not supported') - return - } - const wsUris = workspace.getAllWorkspaceFolders()?.map(f => f.uri) ?? [] - // Get agent paths const wsAgentPaths = getWorkspaceAgentConfigPaths(wsUris) const globalAgentPath = getGlobalAgentConfigPath(workspace.fs.getUserHomeDir()) const allAgentPaths = [...wsAgentPaths, globalAgentPath] - // Migrate config and persona files to agent config await migrateToAgentConfig(workspace, logging, agent) const mgr = await McpManager.init(allAgentPaths, { @@ -282,13 +313,9 @@ export const McpToolsServer: Server = ({ credentialsProvider, workspace, logging runtime, }) - // Clear tool name mapping before registering all tools to avoid conflicts from previous registrations McpManager.instance.clearToolNameMapping() const byServer: Record = {} - - logging.info(`enabled Tools: ${mgr.getEnabledTools().entries()}`) - // only register enabled tools for (const d of mgr.getEnabledTools()) { ;(byServer[d.serverName] ||= []).push(d) } @@ -300,11 +327,76 @@ export const McpToolsServer: Server = ({ credentialsProvider, workspace, logging registerServerTools(server, defs) }) } catch (e) { - console.warn('Caught error during MCP tool initialization; initialization may be incomplete:', e) + logging.error(`Failed to initialize MCP:' ${e}`) + } + } + + lsp.onInitialized(async () => { + try { + if (!enabledMCP(lsp.getClientInitializeParams())) { + logging.warn('MCP is currently not supported') + return + } + + if (sdkInitializator) { + profileStatusMonitor = new ProfileStatusMonitor( + credentialsProvider, + workspace, + logging, + sdkInitializator, + removeAllMcpTools, + async () => { + logging.info('MCP enabled by profile status monitor') + await initializeMcp() + } + ) + + // Wait for profile ARN to be available before checking MCP state + const checkAndInitialize = async () => { + const shouldInitialize = await profileStatusMonitor!.checkInitialState() + if (shouldInitialize) { + logging.info('MCP enabled, initializing immediately') + await initializeMcp() + } + profileStatusMonitor!.start() + } + + // Check if service manager is ready + try { + const serviceManager = AmazonQTokenServiceManager.getInstance() + if (serviceManager.getState() === 'INITIALIZED') { + await checkAndInitialize() + } else { + // Poll for service manager to be ready with 10s timeout + const startTime = Date.now() + const pollForReady = async () => { + if (serviceManager.getState() === 'INITIALIZED') { + await checkAndInitialize() + } else if (Date.now() - startTime < SERVICE_MANAGER_TIMEOUT_MS) { + setTimeout(pollForReady, SERVICE_MANAGER_POLL_INTERVAL_MS) + } else { + logging.warn('Service manager not ready after 10s, defaulting MCP to enabled') + await initializeMcp() + profileStatusMonitor!.start() + } + } + setTimeout(pollForReady, SERVICE_MANAGER_POLL_INTERVAL_MS) + } + } catch (error) { + // Service manager not initialized yet, default to enabled + logging.info('Service manager not ready, defaulting MCP to enabled') + await initializeMcp() + profileStatusMonitor!.start() + } + } + } catch (error) { + console.warn('Caught error during MCP tool initialization; initialization may be incomplete:', error) + logging.error(`Failed to initialize MCP in onInitialized: ${error}`) } }) return async () => { + profileStatusMonitor?.stop() await McpManager.instance.close() } } diff --git a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts index a162ff44e6..aafc0aaf4d 100644 --- a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts +++ b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts @@ -581,6 +581,13 @@ export class CodeWhispererServiceToken extends CodeWhispererServiceBase { return this.client.getCodeAnalysis(this.withProfileArn(request)).promise() } + /** + * @description Get profile details + */ + async getProfile(request: { profileArn: string }) { + return this.client.getProfile(request).promise() + } + /** * @description Once scan completed successfully, send a request to get list of all the findings for the given scan. */