From 301f2c59b17eaf1f6ce6305166efd78dd14b4856 Mon Sep 17 00:00:00 2001 From: Tai Lai Date: Tue, 12 Aug 2025 15:14:04 -0700 Subject: [PATCH 1/3] feat(amazonq): read tool ui revamp (#2113) * feat(amazonq): read tool ui revamp * feat(amazonq): read tool message revamp (#2049) * feat(amazonq): read tool message revamp * fix tests * feat: file search ui (#2078) * feat: file search ui * fix tests * fix integration tests * remove unnecessary type check * fix: use quotes instead of backticks * fix header update issue * fix integration test --- chat-client/src/client/mynahUi.ts | 16 +- .../src/tests/agenticChatInteg.test.ts | 13 +- .../agenticChat/agenticChatController.test.ts | 6 +- .../agenticChat/agenticChatController.ts | 188 ++++++++++++------ .../agenticChat/agenticChatResultStream.ts | 63 ++---- .../agenticChat/tools/fileSearch.ts | 4 + 6 files changed, 163 insertions(+), 127 deletions(-) diff --git a/chat-client/src/client/mynahUi.ts b/chat-client/src/client/mynahUi.ts index 3b42249331..527dea06f3 100644 --- a/chat-client/src/client/mynahUi.ts +++ b/chat-client/src/client/mynahUi.ts @@ -1353,10 +1353,15 @@ export const createMynahUi = ( fileTreeTitle: '', hideFileCount: true, details: toDetailsWithoutIcon(header.fileList.details), + renderAsPills: + !header.fileList.details || + (Object.values(header.fileList.details).every(detail => !detail.changes) && + (!header.buttons || !header.buttons.some(button => button.id === 'undo-changes')) && + !header.status?.icon), } } if (!isPartialResult) { - if (processedHeader) { + if (processedHeader && !message.header?.status) { processedHeader.status = undefined } } @@ -1369,7 +1374,8 @@ export const createMynahUi = ( processedHeader.buttons !== null && processedHeader.buttons.length > 0) || processedHeader.status !== undefined || - processedHeader.icon !== undefined) + processedHeader.icon !== undefined || + processedHeader.fileList !== undefined) const padding = message.type === 'tool' ? (fileList ? true : message.messageId?.endsWith('_permission')) : undefined @@ -1380,8 +1386,10 @@ export const createMynahUi = ( // Adding this conditional check to show the stop message in the center. const contentHorizontalAlignment: ChatItem['contentHorizontalAlignment'] = undefined - // If message.header?.status?.text is Stopped or Rejected or Ignored or Completed etc.. card should be in disabled state. - const shouldMute = message.header?.status?.text !== undefined && message.header?.status?.text !== 'Completed' + // If message.header?.status?.text is Stopped or Rejected or Ignored etc.. card should be in disabled state. + const shouldMute = + message.header?.status?.text !== undefined && + ['Stopped', 'Rejected', 'Ignored', 'Failed', 'Error'].includes(message.header.status.text) return { body: message.body, diff --git a/integration-tests/q-agentic-chat-server/src/tests/agenticChatInteg.test.ts b/integration-tests/q-agentic-chat-server/src/tests/agenticChatInteg.test.ts index 897ee02071..c5a8046bd5 100644 --- a/integration-tests/q-agentic-chat-server/src/tests/agenticChatInteg.test.ts +++ b/integration-tests/q-agentic-chat-server/src/tests/agenticChatInteg.test.ts @@ -169,11 +169,11 @@ describe('Q Agentic Chat Server Integration Tests', async () => { expect(decryptedResult.additionalMessages).to.be.an('array') const fsReadMessage = decryptedResult.additionalMessages?.find( - msg => msg.type === 'tool' && msg.fileList?.rootFolderTitle === '1 file read' + msg => msg.type === 'tool' && msg.header?.body === '1 file read' ) expect(fsReadMessage).to.exist const expectedPath = path.join(rootPath, 'test.py') - const actualPaths = fsReadMessage?.fileList?.filePaths?.map(normalizePath) || [] + const actualPaths = fsReadMessage?.header?.fileList?.filePaths?.map(normalizePath) || [] expect(actualPaths).to.include.members([normalizePath(expectedPath)]) expect(fsReadMessage?.messageId?.startsWith('tooluse_')).to.be.true }) @@ -191,10 +191,10 @@ describe('Q Agentic Chat Server Integration Tests', async () => { expect(decryptedResult.additionalMessages).to.be.an('array') const listDirectoryMessage = decryptedResult.additionalMessages?.find( - msg => msg.type === 'tool' && msg.fileList?.rootFolderTitle === '1 directory listed' + msg => msg.type === 'tool' && msg.header?.body === '1 directory listed' ) expect(listDirectoryMessage).to.exist - const actualPaths = listDirectoryMessage?.fileList?.filePaths?.map(normalizePath) || [] + const actualPaths = listDirectoryMessage?.header?.fileList?.filePaths?.map(normalizePath) || [] expect(actualPaths).to.include.members([normalizePath(rootPath)]) expect(listDirectoryMessage?.messageId?.startsWith('tooluse_')).to.be.true }) @@ -371,11 +371,12 @@ describe('Q Agentic Chat Server Integration Tests', async () => { expect(decryptedResult.additionalMessages).to.be.an('array') const fileSearchMessage = decryptedResult.additionalMessages?.find( - msg => msg.type === 'tool' && msg.fileList?.rootFolderTitle === '1 directory searched' + msg => msg.type === 'tool' && msg.header?.body === 'Searched for "test" in ' ) expect(fileSearchMessage).to.exist expect(fileSearchMessage?.messageId?.startsWith('tooluse_')).to.be.true - const actualPaths = fileSearchMessage?.fileList?.filePaths?.map(normalizePath) || [] + expect(fileSearchMessage?.header?.status?.text).to.equal('3 results found') + const actualPaths = fileSearchMessage?.header?.fileList?.filePaths?.map(normalizePath) || [] expect(actualPaths).to.include.members([normalizePath(rootPath)]) }) }) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts index 6c8af7a272..616848d0b5 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts @@ -451,7 +451,7 @@ describe('AgenticChatController', () => { assert.deepStrictEqual(chatResult, { additionalMessages: [], - body: '\n\nHello World!', + body: '\nHello World!', messageId: 'mock-message-id', buttons: [], codeReference: [], @@ -1150,7 +1150,7 @@ describe('AgenticChatController', () => { sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length + 1) // response length + 1 loading messages assert.deepStrictEqual(chatResult, { additionalMessages: [], - body: '\n\nHello World!', + body: '\nHello World!', messageId: 'mock-message-id', codeReference: [], buttons: [], @@ -1169,7 +1169,7 @@ describe('AgenticChatController', () => { sinon.assert.callCount(testFeatures.lsp.sendProgress, mockChatResponseList.length + 1) // response length + 1 loading message assert.deepStrictEqual(chatResult, { additionalMessages: [], - body: '\n\nHello World!', + body: '\nHello World!', messageId: 'mock-message-id', buttons: [], codeReference: [], diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts index 72217fae07..09bf965819 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts @@ -169,7 +169,7 @@ import { ExecuteBash, ExecuteBashParams } from './tools/executeBash' import { ExplanatoryParams, InvokeOutput, ToolApprovalException } from './tools/toolShared' import { validatePathBasic, validatePathExists, validatePaths as validatePathsSync } from './utils/pathValidation' import { GrepSearch, SanitizedRipgrepOutput } from './tools/grepSearch' -import { FileSearch, FileSearchParams } from './tools/fileSearch' +import { FileSearch, FileSearchParams, isFileSearchParams } from './tools/fileSearch' import { FsReplace, FsReplaceParams } from './tools/fsReplace' import { loggingUtils, timeoutUtils } from '@aws/lsp-core' import { diffLines } from 'diff' @@ -1695,8 +1695,7 @@ export class AgenticChatController implements ChatHandlers { // remove progress UI await chatResultStream.removeResultBlockAndUpdateUI(progressPrefix + toolUse.toolUseId) - // fsRead and listDirectory write to an existing card and could show nothing in the current position - if (![FS_WRITE, FS_REPLACE, FS_READ, LIST_DIRECTORY].includes(toolUse.name)) { + if (![FS_WRITE, FS_REPLACE].includes(toolUse.name)) { await this.#showUndoAllIfRequired(chatResultStream, session) } // fsWrite can take a long time, so we render fsWrite Explanatory upon partial streaming responses. @@ -1911,10 +1910,19 @@ export class AgenticChatController implements ChatHandlers { switch (toolUse.name) { case FS_READ: case LIST_DIRECTORY: + const readToolResult = await this.#processReadTool(toolUse, chatResultStream) + if (readToolResult) { + await chatResultStream.writeResultBlock(readToolResult) + } + break case FILE_SEARCH: - const initialListDirResult = this.#processReadOrListOrSearch(toolUse, chatResultStream) - if (initialListDirResult) { - await chatResultStream.writeResultBlock(initialListDirResult) + if (isFileSearchParams(toolUse.input)) { + await this.#processFileSearchTool( + toolUse.input, + toolUse.toolUseId, + result, + chatResultStream + ) } break // no need to write tool result for listDir,fsRead,fileSearch into chat stream @@ -2315,7 +2323,6 @@ export class AgenticChatController implements ChatHandlers { } const toolMsgId = toolUse.toolUseId! - const chatMsgId = chatResultStream.getResult().messageId let headerEmitted = false const initialHeader: ChatMessage['header'] = { @@ -2353,13 +2360,6 @@ export class AgenticChatController implements ChatHandlers { header: completedHeader, }) - await chatResultStream.writeResultBlock({ - type: 'answer', - messageId: chatMsgId, - body: '', - header: undefined, - }) - this.#stoppedToolUses.add(toolMsgId) }, }) @@ -2877,70 +2877,135 @@ export class AgenticChatController implements ChatHandlers { } } - #processReadOrListOrSearch(toolUse: ToolUse, chatResultStream: AgenticChatResultStream): ChatMessage | undefined { - let messageIdToUpdate = toolUse.toolUseId! - const currentId = chatResultStream.getMessageIdToUpdateForTool(toolUse.name!) + async #processFileSearchTool( + toolInput: FileSearchParams, + toolUseId: string, + result: InvokeOutput, + chatResultStream: AgenticChatResultStream + ): Promise { + if (typeof result.output.content !== 'string') return - if (currentId) { - messageIdToUpdate = currentId - } else { - chatResultStream.setMessageIdToUpdateForTool(toolUse.name!, messageIdToUpdate) + const { queryName, path: inputPath } = toolInput + const resultCount = result.output.content + .split('\n') + .filter(line => line.trim().startsWith('[F]') || line.trim().startsWith('[D]')).length + + const chatMessage: ChatMessage = { + type: 'tool', + messageId: toolUseId, + header: { + body: `Searched for "${queryName}" in `, + icon: 'search', + status: { + text: `${resultCount} result${resultCount !== 1 ? 's' : ''} found`, + }, + fileList: { + filePaths: [inputPath], + details: { + [inputPath]: { + description: inputPath, + visibleName: path.basename(inputPath), + clickable: false, + }, + }, + }, + }, } - let currentPaths = [] + await chatResultStream.writeResultBlock(chatMessage) + } + + async #processReadTool( + toolUse: ToolUse, + chatResultStream: AgenticChatResultStream + ): Promise { + let currentPaths: string[] = [] if (toolUse.name === FS_READ) { - currentPaths = (toolUse.input as unknown as FsReadParams)?.paths + currentPaths = (toolUse.input as unknown as FsReadParams)?.paths || [] + } else if (toolUse.name === LIST_DIRECTORY) { + const singlePath = (toolUse.input as unknown as ListDirectoryParams)?.path + if (singlePath) { + currentPaths = [singlePath] + } + } else if (toolUse.name === FILE_SEARCH) { + const queryName = (toolUse.input as unknown as FileSearchParams)?.queryName + if (queryName) { + currentPaths = [queryName] + } } else { - currentPaths.push((toolUse.input as unknown as ListDirectoryParams | FileSearchParams)?.path) + return } - if (!currentPaths) return + if (currentPaths.length === 0) return - for (const currentPath of currentPaths) { - const existingPaths = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths || [] - // Check if path already exists in the list - const isPathAlreadyProcessed = existingPaths.some(path => path.relativeFilePath === currentPath) - if (!isPathAlreadyProcessed) { - const currentFileDetail = { - relativeFilePath: currentPath, - lineRanges: [{ first: -1, second: -1 }], - } - chatResultStream.addMessageOperation(messageIdToUpdate, toolUse.name!, [ - ...existingPaths, - currentFileDetail, - ]) + // Check if the last message is the same tool type + const lastMessage = chatResultStream.getLastMessage() + const isSameToolType = + lastMessage?.type === 'tool' && lastMessage.header?.icon === this.#toolToIcon(toolUse.name) + + let allPaths = currentPaths + + if (isSameToolType && lastMessage.messageId) { + // Combine with existing paths and overwrite the last message + const existingPaths = lastMessage.header?.fileList?.filePaths || [] + allPaths = [...existingPaths, ...currentPaths] + + const blockId = chatResultStream.getMessageBlockId(lastMessage.messageId) + if (blockId !== undefined) { + // Create the updated message with combined paths + const updatedMessage = this.#createFileListToolMessage(toolUse, allPaths, lastMessage.messageId) + // Overwrite the existing block + await chatResultStream.overwriteResultBlock(updatedMessage, blockId) + return undefined // Don't return a message since we already wrote it } } + + // Create new message with current paths + return this.#createFileListToolMessage(toolUse, allPaths, toolUse.toolUseId!) + } + + #createFileListToolMessage(toolUse: ToolUse, filePaths: string[], messageId: string): ChatMessage { + const itemCount = filePaths.length let title: string - const itemCount = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths.length - const filePathsPushed = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths ?? [] - if (!itemCount) { + if (itemCount === 0) { title = 'Gathering context' } else { title = toolUse.name === FS_READ ? `${itemCount} file${itemCount > 1 ? 's' : ''} read` - : toolUse.name === FILE_SEARCH - ? `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} searched` - : `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} listed` + : toolUse.name === LIST_DIRECTORY + ? `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} listed` + : '' } const details: Record = {} - for (const item of filePathsPushed) { - details[item.relativeFilePath] = { - lineRanges: item.lineRanges, - description: item.relativeFilePath, + for (const filePath of filePaths) { + details[filePath] = { + description: filePath, + visibleName: path.basename(filePath), + clickable: toolUse.name === FS_READ, } } - - const fileList: FileList = { - rootFolderTitle: title, - filePaths: filePathsPushed.map(item => item.relativeFilePath), - details, - } return { type: 'tool', - fileList, - messageId: messageIdToUpdate, - body: '', + header: { + body: title, + icon: this.#toolToIcon(toolUse.name), + fileList: { + filePaths, + details, + }, + }, + messageId, + } + } + + #toolToIcon(toolName: string | undefined): string | undefined { + switch (toolName) { + case FS_READ: + return 'eye' + case LIST_DIRECTORY: + return 'check-list' + default: + return undefined } } @@ -2956,14 +3021,7 @@ export class AgenticChatController implements ChatHandlers { return undefined } - let messageIdToUpdate = toolUse.toolUseId! - const currentId = chatResultStream.getMessageIdToUpdateForTool(toolUse.name!) - - if (currentId) { - messageIdToUpdate = currentId - } else { - chatResultStream.setMessageIdToUpdateForTool(toolUse.name!, messageIdToUpdate) - } + const messageIdToUpdate = toolUse.toolUseId! // Extract search results from the tool output const output = result.output.content as SanitizedRipgrepOutput diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts index 70b3452361..195e1e880b 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts @@ -1,4 +1,4 @@ -import { ChatResult, FileDetails, ChatMessage } from '@aws/language-server-runtimes/protocol' +import { ChatResult, ChatMessage } from '@aws/language-server-runtimes/protocol' import { randomUUID } from 'crypto' export interface ResultStreamWriter { @@ -32,33 +32,20 @@ export interface ResultStreamWriter { close(): Promise } +export const progressPrefix = 'progress_' + /** * This class wraps around lsp.sendProgress to provide a more helpful interface for streaming a ChatResult to the client. * ChatResults are grouped into blocks that can be written directly, or streamed in. * In the final message, blocks are seperated by resultDelimiter defined below. */ - -interface FileDetailsWithPath extends FileDetails { - relativeFilePath: string -} - -type OperationType = 'read' | 'write' | 'listDir' - -export const progressPrefix = 'progress_' - -interface FileOperation { - type: OperationType - filePaths: FileDetailsWithPath[] -} export class AgenticChatResultStream { - static readonly resultDelimiter = '\n\n' + static readonly resultDelimiter = '\n' #state = { chatResultBlocks: [] as ChatMessage[], isLocked: false, uuid: randomUUID(), messageId: undefined as string | undefined, - messageIdToUpdateForTool: new Map(), - messageOperations: new Map(), } readonly #sendProgress: (newChatResult: ChatResult | string) => Promise @@ -70,33 +57,6 @@ export class AgenticChatResultStream { return this.#joinResults(this.#state.chatResultBlocks, only) } - setMessageIdToUpdateForTool(toolName: string, messageId: string) { - this.#state.messageIdToUpdateForTool.set(toolName as OperationType, messageId) - } - - getMessageIdToUpdateForTool(toolName: string): string | undefined { - return this.#state.messageIdToUpdateForTool.get(toolName as OperationType) - } - - /** - * Adds a file operation for a specific message - * @param messageId The ID of the message - * @param type The type of operation ('fsRead' or 'listDirectory' or 'fsWrite') - * @param filePaths Array of FileDetailsWithPath involved in the operation - */ - addMessageOperation(messageId: string, type: string, filePaths: FileDetailsWithPath[]) { - this.#state.messageOperations.set(messageId, { type: type as OperationType, filePaths }) - } - - /** - * Gets the file operation details for a specific message - * @param messageId The ID of the message - * @returns The file operation details or undefined if not found - */ - getMessageOperation(messageId: string): FileOperation | undefined { - return this.#state.messageOperations.get(messageId) - } - #joinResults(chatResults: ChatMessage[], only?: string): ChatResult { const result: ChatResult = { body: '', @@ -111,9 +71,9 @@ export class AgenticChatResultStream { return { ...acc, buttons: [...(acc.buttons ?? []), ...(c.buttons ?? [])], - body: acc.body + AgenticChatResultStream.resultDelimiter + c.body, - ...(c.contextList && { contextList: c.contextList }), - header: Object.prototype.hasOwnProperty.call(c, 'header') ? c.header : acc.header, + body: acc.body + (c.body ? AgenticChatResultStream.resultDelimiter + c.body : ''), + ...(c.contextList && c.type !== 'tool' && { contextList: c.contextList }), + header: c.header !== undefined ? c.header : acc.header, codeReference: [...(acc.codeReference ?? []), ...(c.codeReference ?? [])], } } else if (acc.additionalMessages!.some(am => am.messageId === c.messageId)) { @@ -127,9 +87,10 @@ export class AgenticChatResultStream { : am.buttons, body: am.messageId === c.messageId - ? am.body + AgenticChatResultStream.resultDelimiter + c.body + ? am.body + (c.body ? AgenticChatResultStream.resultDelimiter + c.body : '') : am.body, ...(am.messageId === c.messageId && + c.type !== 'tool' && (c.contextList || acc.contextList) && { contextList: { filePaths: [ @@ -161,7 +122,7 @@ export class AgenticChatResultStream { }, }, }), - header: Object.prototype.hasOwnProperty.call(c, 'header') ? c.header : am.header, + ...(am.messageId === c.messageId && c.header !== undefined && { header: c.header }), })), } } else { @@ -246,6 +207,10 @@ export class AgenticChatResultStream { return undefined } + getLastMessage(): ChatMessage { + return this.#state.chatResultBlocks[this.#state.chatResultBlocks.length - 1] + } + getResultStreamWriter(): ResultStreamWriter { // Note: if write calls are not awaited, stream can be out-of-order. if (this.#state.isLocked) { diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fileSearch.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fileSearch.ts index fb6486996e..37d11afe4f 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fileSearch.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fileSearch.ts @@ -158,3 +158,7 @@ export class FileSearch { } as const } } + +export function isFileSearchParams(input: any): input is FileSearchParams { + return input && typeof input.path === 'string' && typeof input.queryName === 'string' +} From 96f4a236138154a67c55803fd99710f469884edf Mon Sep 17 00:00:00 2001 From: Dung Dong Date: Tue, 12 Aug 2025 17:03:29 -0700 Subject: [PATCH 2/3] fix(amazonq): persist mcp configs in agent json on start-up (#2112) --- .../src/language-server/agenticChat/tools/mcp/mcpUtils.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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])], From 0947079f99aefc103eacb4d74e8175c51ea62f4c Mon Sep 17 00:00:00 2001 From: invictus <149003065+ashishrp-aws@users.noreply.github.com> Date: Wed, 13 Aug 2025 10:28:53 -0700 Subject: [PATCH 3/3] feat(amazonq): added mcp admin level configuration with GetProfile (#2000) * feat(amazonq): added mcp admin level configuration with GetProfile * feat(amazonq): added UX message for mcp admin control * test: add unit tests for ProfileStatusMonitor static functionality * test: add comprehensive unit tests for MCP admin control features * fix: fix to wait for serviceManager is initialzied to initialize mcp managers * fix: fix for unit test failures * fix: fix for UI changes * fix(amazonq): fix to to rename mcp enabled function and max time limit of 10 seconds * fix: fix to add async initialization for mcp manager * fix: fix to move action buttons to mcpEventHandler for listMcpServers * fix: added try and catch block for mcp initialization * fix: fix for merge conflicts * fix: fix for test failure * fix: remove the unnecessary feature flag * fix: fix for mynah test failure * fix: fix to retry function to common util * fix: fix to add retryUtils --- chat-client/src/client/mcpMynahUi.test.ts | 6 +- chat-client/src/client/mcpMynahUi.ts | 19 +- core/aws-lsp-core/src/index.ts | 1 + core/aws-lsp-core/src/util/retryUtils.test.ts | 120 ++++++++++++ core/aws-lsp-core/src/util/retryUtils.ts | 77 ++++++++ .../client/token/bearer-token-service.json | 85 ++++++++ .../token/codewhispererbearertokenclient.d.ts | 26 +++ .../agenticChat/constants/constants.ts | 2 + .../tools/mcp/mcpEventHandler.test.ts | 105 ++++++++++ .../agenticChat/tools/mcp/mcpEventHandler.ts | 56 +++++- .../tools/mcp/profileStatusMonitor.test.ts | 182 ++++++++++++++++++ .../tools/mcp/profileStatusMonitor.ts | 129 +++++++++++++ .../agenticChat/tools/toolServer.ts | 122 ++++++++++-- .../src/shared/codeWhispererService.ts | 7 + 14 files changed, 899 insertions(+), 38 deletions(-) create mode 100644 core/aws-lsp-core/src/util/retryUtils.test.ts create mode 100644 core/aws-lsp-core/src/util/retryUtils.ts create mode 100644 server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.test.ts create mode 100644 server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts 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/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. */