diff --git a/chat-client/package.json b/chat-client/package.json index ef3ba87421..161cb50a0b 100644 --- a/chat-client/package.json +++ b/chat-client/package.json @@ -27,7 +27,7 @@ "@aws/chat-client-ui-types": "^0.1.56", "@aws/language-server-runtimes": "^0.2.123", "@aws/language-server-runtimes-types": "^0.1.50", - "@aws/mynah-ui": "^4.36.2" + "@aws/mynah-ui": "^4.36.4" }, "devDependencies": { "@types/jsdom": "^21.1.6", 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..efc0467d2f 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/package-lock.json b/package-lock.json index 9bcaeb0f84..6b55b375a3 100644 --- a/package-lock.json +++ b/package-lock.json @@ -257,7 +257,7 @@ "@aws/chat-client-ui-types": "^0.1.56", "@aws/language-server-runtimes": "^0.2.123", "@aws/language-server-runtimes-types": "^0.1.50", - "@aws/mynah-ui": "^4.36.2" + "@aws/mynah-ui": "^4.36.4" }, "devDependencies": { "@types/jsdom": "^21.1.6", @@ -4204,9 +4204,9 @@ "link": true }, "node_modules/@aws/mynah-ui": { - "version": "4.36.2", - "resolved": "https://registry.npmjs.org/@aws/mynah-ui/-/mynah-ui-4.36.2.tgz", - "integrity": "sha512-3ibfK2CTj7dlFFdgTIE1DdEyDpy+P3hdP/Fmlx76T9GGSYiGHqwunDSi59L1P61Kj46WADBrQ52mLUQ6FR8Rzg==", + "version": "4.36.4", + "resolved": "https://registry.npmjs.org/@aws/mynah-ui/-/mynah-ui-4.36.4.tgz", + "integrity": "sha512-vGW4wlNindpr2Ep9x3iuKbrZTXe5KrE8vWpg15DjkN3qK42KMuMEQ67Pqtfgl5EseNYC1ukZm4HIQIMmt+vevA==", "hasInstallScript": true, "license": "Apache License 2.0", "dependencies": { 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 48e9e2adec..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 @@ -67,6 +67,7 @@ import { McpManager } from './tools/mcp/mcpManager' import { AgenticChatResultStream } from './agenticChatResultStream' import { AgenticChatError } from './errors' import * as sharedUtils from '../../shared/utils' +import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager' describe('AgenticChatController', () => { let mcpInstanceStub: sinon.SinonStub @@ -450,7 +451,7 @@ describe('AgenticChatController', () => { assert.deepStrictEqual(chatResult, { additionalMessages: [], - body: '\n\nHello World!', + body: '\nHello World!', messageId: 'mock-message-id', buttons: [], codeReference: [], @@ -475,6 +476,15 @@ describe('AgenticChatController', () => { assert.strictEqual(typeof session.conversationId, 'string') }) + it('invokes IdleWorkspaceManager recordActivityTimestamp', async () => { + const recordActivityTimestampStub = sinon.stub(IdleWorkspaceManager, 'recordActivityTimestamp') + + await chatController.onChatPrompt({ tabId: mockTabId, prompt: { prompt: 'Hello' } }, mockCancellationToken) + + sinon.assert.calledOnce(recordActivityTimestampStub) + recordActivityTimestampStub.restore() + }) + it('includes chat history from the database in the request input', async () => { // Mock chat history const mockHistory = [ @@ -1140,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: [], @@ -1159,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 60bba5885a..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' @@ -230,6 +230,7 @@ import { ActiveUserTracker } from '../../shared/activeUserTracker' import { UserContext } from '../../client/token/codewhispererbearertokenclient' import { CodeWhispererServiceToken } from '../../shared/codeWhispererService' import { DisplayFindings } from './tools/qCodeAnalysis/displayFindings' +import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager' type ChatHandlers = Omit< LspHandlers, @@ -722,6 +723,8 @@ export class AgenticChatController implements ChatHandlers { // Phase 1: Initial Setup - This happens only once params.prompt.prompt = sanitizeInput(params.prompt.prompt || '') + IdleWorkspaceManager.recordActivityTimestamp() + const maybeDefaultResponse = !params.prompt.command && getDefaultChatResponse(params.prompt.prompt) if (maybeDefaultResponse) { return maybeDefaultResponse @@ -1692,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. @@ -1908,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 @@ -2312,7 +2323,6 @@ export class AgenticChatController implements ChatHandlers { } const toolMsgId = toolUse.toolUseId! - const chatMsgId = chatResultStream.getResult().messageId let headerEmitted = false const initialHeader: ChatMessage['header'] = { @@ -2350,13 +2360,6 @@ export class AgenticChatController implements ChatHandlers { header: completedHeader, }) - await chatResultStream.writeResultBlock({ - type: 'answer', - messageId: chatMsgId, - body: '', - header: undefined, - }) - this.#stoppedToolUses.add(toolMsgId) }, }) @@ -2874,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 + + const { queryName, path: inputPath } = toolInput + const resultCount = result.output.content + .split('\n') + .filter(line => line.trim().startsWith('[F]') || line.trim().startsWith('[D]')).length - if (currentId) { - messageIdToUpdate = currentId - } else { - chatResultStream.setMessageIdToUpdateForTool(toolUse.name!, messageIdToUpdate) + 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 } } @@ -2953,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 @@ -3285,6 +3346,9 @@ export class AgenticChatController implements ChatHandlers { const metric = new Metric({ cwsprChatConversationType: 'Chat', }) + + IdleWorkspaceManager.recordActivityTimestamp() + const triggerContext = await this.#getInlineChatTriggerContext(params) let response: ChatCommandOutput 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..5fb5c39bab 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,7 +87,7 @@ 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.contextList || acc.contextList) && { @@ -161,7 +121,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 +206,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' +} diff --git a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.test.ts b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.test.ts index 1636c17b84..29390248f7 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.test.ts @@ -61,6 +61,7 @@ import { INVALID_TOKEN } from '../../shared/constants' import { AmazonQError } from '../../shared/amazonQServiceManager/errors' import * as path from 'path' import { CONTEXT_CHARACTERS_LIMIT } from './constants' +import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager' const updateConfiguration = async ( features: TestFeatures, @@ -770,6 +771,22 @@ describe('CodeWhisperer Server', () => { assert.rejects(promise, ResponseError) }) + it('invokes IdleWorkspaceManager recordActivityTimestamp', async () => { + const recordActivityTimestampStub = sinon.stub(IdleWorkspaceManager, 'recordActivityTimestamp') + + await features.doInlineCompletionWithReferences( + { + textDocument: { uri: SOME_FILE.uri }, + position: { line: 0, character: 0 }, + context: { triggerKind: InlineCompletionTriggerKind.Invoked }, + }, + CancellationToken.None + ) + + sinon.assert.calledOnce(recordActivityTimestampStub) + recordActivityTimestampStub.restore() + }) + describe('Supplemental Context', () => { it('should send supplemental context when using token authentication', async () => { const test_service = sinon.createStubInstance( diff --git a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts index f2b1a2c43d..56cda299f8 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts @@ -55,7 +55,9 @@ import { } from './telemetry' import { DocumentChangedListener } from './documentChangedListener' import { EditCompletionHandler } from './editCompletionHandler' -import { EMPTY_RESULT } from './constants' +import { EMPTY_RESULT, ABAP_EXTENSIONS } from './constants' +import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager' +import { URI } from 'vscode-uri' const mergeSuggestionsWithRightContext = ( rightFileContext: string, @@ -141,6 +143,8 @@ export const CodewhispererServerFactory = // 2. it is not designed to handle concurrent changes to these state variables. // when one handler is at the API call stage, it has not yet update the session state // but another request can start, causing the state to be incorrect. + IdleWorkspaceManager.recordActivityTimestamp() + if (isOnInlineCompletionHandlerInProgress) { logging.log(`Skip concurrent inline completion`) return EMPTY_RESULT @@ -158,7 +162,7 @@ export const CodewhispererServerFactory = if (cursorTracker) { cursorTracker.trackPosition(params.textDocument.uri, params.position) } - const textDocument = await workspace.getTextDocument(params.textDocument.uri) + const textDocument = await getTextDocument(params.textDocument.uri, workspace, logging) const codeWhispererService = amazonQServiceManager.getCodewhispererService() if (params.partialResultToken && currentSession) { @@ -498,7 +502,6 @@ export const CodewhispererServerFactory = partialResultToken: suggestionResponse.responseContext.nextToken, } } else { - session.hasEditsPending = suggestionResponse.responseContext.nextToken ? true : false return { items: suggestionResponse.suggestions .map(suggestion => { @@ -687,14 +690,8 @@ export const CodewhispererServerFactory = if (firstCompletionDisplayLatency) emitPerceivedLatencyTelemetry(telemetry, session) // Always emit user trigger decision at session close - // Close session unless Edit suggestion was accepted with more pending - const shouldKeepSessionOpen = - session.suggestionType === SuggestionType.EDIT && isAccepted && session.hasEditsPending - - if (!shouldKeepSessionOpen) { - completionSessionManager.closeSession(session) - } - const streakLength = editsEnabled ? completionSessionManager.getAndUpdateStreakLength(isAccepted) : 0 + sessionManager.closeSession(session) + const streakLength = editsEnabled ? sessionManager.getAndUpdateStreakLength(isAccepted) : 0 await emitUserTriggerDecisionTelemetry( telemetry, telemetryService, @@ -722,10 +719,9 @@ export const CodewhispererServerFactory = userWrittenCodeTracker.customizationArn = customizationArn } logging.debug(`CodePercentageTracker customizationArn updated to ${customizationArn}`) - /* - The flag enableTelemetryEventsToDestination is set to true temporarily. It's value will be determined through destination - configuration post all events migration to STE. It'll be replaced by qConfig['enableTelemetryEventsToDestination'] === true - */ + + // The flag enableTelemetryEventsToDestination is set to true temporarily. It's value will be determined through destination + // configuration post all events migration to STE. It'll be replaced by qConfig['enableTelemetryEventsToDestination'] === true // const enableTelemetryEventsToDestination = true // telemetryService.updateEnableTelemetryEventsToDestination(enableTelemetryEventsToDestination) telemetryService.updateOptOutPreference(optOutTelemetryPreference) @@ -902,3 +898,27 @@ export const CodewhispererServerFactory = export const CodeWhispererServerIAM = CodewhispererServerFactory(getOrThrowBaseIAMServiceManager) export const CodeWhispererServerToken = CodewhispererServerFactory(getOrThrowBaseTokenServiceManager) + +const getLanguageIdFromUri = (uri: string, logging?: any): string => { + try { + const extension = uri.split('.').pop()?.toLowerCase() + return ABAP_EXTENSIONS.has(extension || '') ? 'abap' : '' + } catch (err) { + logging?.log(`Error parsing URI to determine language: ${uri}: ${err}`) + return '' + } +} + +const getTextDocument = async (uri: string, workspace: any, logging: any): Promise => { + let textDocument = await workspace.getTextDocument(uri) + if (!textDocument) { + try { + const content = await workspace.fs.readFile(URI.parse(uri).fsPath) + const languageId = getLanguageIdFromUri(uri) + textDocument = TextDocument.create(uri, languageId, 0, content) + } catch (err) { + logging.log(`Unable to load from ${uri}: ${err}`) + } + } + return textDocument +} diff --git a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/constants.ts b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/constants.ts index 49a33b35de..c5810924b6 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/constants.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/constants.ts @@ -4,3 +4,28 @@ export const CONTEXT_CHARACTERS_LIMIT = 10240 export const EMPTY_RESULT = { sessionId: '', items: [] } export const EDIT_DEBOUNCE_INTERVAL_MS = 500 export const EDIT_STALE_RETRY_COUNT = 3 +// ABAP ADT extensions commonly used with Eclipse +export const ABAP_EXTENSIONS = new Set([ + 'asprog', + 'aclass', + 'asinc', + 'aint', + 'assrvds', + 'asbdef', + 'asddls', + 'astablds', + 'astabldt', + 'amdp', + 'apack', + 'asrv', + 'aobj', + 'aexit', + 'abdef', + 'acinc', + 'asfugr', + 'apfugr', + 'asfunc', + 'asfinc', + 'apfunc', + 'apfinc', +]) diff --git a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/editCompletionHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/editCompletionHandler.ts index 5a6f8644b1..3c72e368f4 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/editCompletionHandler.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/editCompletionHandler.ts @@ -120,16 +120,35 @@ export class EditCompletionHandler { } if (params.partialResultToken && currentSession) { + // Close ACTIVE session. We shouldn't record Discard trigger decision for trigger with nextToken. + if (currentSession && currentSession.state === 'ACTIVE') { + this.sessionManager.discardSession(currentSession) + } + + const newSession = this.sessionManager.createSession({ + document: textDocument, + startPosition: params.position, + triggerType: 'AutoTrigger', + language: currentSession.language, + requestContext: currentSession.requestContext, + autoTriggerType: undefined, + triggerCharacter: '', + classifierResult: undefined, + classifierThreshold: undefined, + credentialStartUrl: currentSession.credentialStartUrl, + supplementalMetadata: currentSession.supplementalMetadata, + customizationArn: currentSession.customizationArn, + }) // subsequent paginated requests for current session try { const suggestionResponse = await this.codeWhispererService.generateSuggestions({ - ...currentSession.requestContext, + ...newSession.requestContext, nextToken: `${params.partialResultToken}`, }) return await this.processSuggestionResponse( suggestionResponse, - currentSession, - false, + newSession, + true, params.context.selectedCompletionInfo?.range ) } catch (error) { diff --git a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/session/sessionManager.ts b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/session/sessionManager.ts index 8ac71737cf..cb873a2920 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/session/sessionManager.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/session/sessionManager.ts @@ -80,7 +80,6 @@ export class CodeWhispererSession { includeImportsWithSuggestions?: boolean codewhispererSuggestionImportCount: number = 0 suggestionType?: string - hasEditsPending?: boolean = false // Track the most recent itemId for paginated Edit suggestions constructor(data: SessionData) { diff --git a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/telemetry.ts b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/telemetry.ts index b990ccfceb..d53d141a2b 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/telemetry.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/telemetry.ts @@ -147,13 +147,7 @@ export const emitUserTriggerDecisionTelemetry = async ( streakLength ) - // Mark telemetry as complete unless Edit suggestion was accepted with more pending - const hasPendingEditTelemetry = - session.suggestionType === SuggestionType.EDIT && session.acceptedSuggestionId && session.hasEditsPending - - if (!hasPendingEditTelemetry) { - session.reportedUserDecision = true - } + session.reportedUserDecision = true } export const emitAggregatedUserTriggerDecisionTelemetry = ( diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/IdleWorkspaceManager.test.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/IdleWorkspaceManager.test.ts new file mode 100644 index 0000000000..a5d81526c9 --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/IdleWorkspaceManager.test.ts @@ -0,0 +1,75 @@ +import { IdleWorkspaceManager } from './IdleWorkspaceManager' +import { WorkspaceFolderManager } from './workspaceFolderManager' +import sinon, { stubInterface, StubbedInstance } from 'ts-sinon' + +describe('IdleWorkspaceManager', () => { + let clock: sinon.SinonFakeTimers + let mockWorkspaceFolderManager: StubbedInstance + + beforeEach(() => { + clock = sinon.useFakeTimers() + mockWorkspaceFolderManager = stubInterface() + sinon.stub(WorkspaceFolderManager, 'getInstance').returns(mockWorkspaceFolderManager) + sinon.stub(console, 'error') + }) + + afterEach(() => { + clock.restore() + sinon.restore() + }) + + describe('isSessionIdle', () => { + it('should return false when session is not idle', () => { + IdleWorkspaceManager.recordActivityTimestamp() + + const result = IdleWorkspaceManager.isSessionIdle() + + expect(result).toBe(false) + }) + + it('should return true when session exceeds idle threshold', () => { + IdleWorkspaceManager.recordActivityTimestamp() + clock.tick(31 * 60 * 1000) // 31 minutes + + const result = IdleWorkspaceManager.isSessionIdle() + + expect(result).toBe(true) + }) + }) + + describe('recordActivityTimestamp', () => { + it('should update activity timestamp', async () => { + IdleWorkspaceManager.recordActivityTimestamp() + + expect(IdleWorkspaceManager.isSessionIdle()).toBe(false) + }) + + it('should not trigger workspace check when session was not idle', async () => { + mockWorkspaceFolderManager.isContinuousMonitoringStopped.returns(false) + + IdleWorkspaceManager.recordActivityTimestamp() + + sinon.assert.notCalled(mockWorkspaceFolderManager.checkRemoteWorkspaceStatusAndReact) + }) + + it('should trigger workspace check when session was idle and monitoring is active', async () => { + // Make session idle first + clock.tick(31 * 60 * 1000) + mockWorkspaceFolderManager.isContinuousMonitoringStopped.returns(false) + mockWorkspaceFolderManager.checkRemoteWorkspaceStatusAndReact.resolves() + + IdleWorkspaceManager.recordActivityTimestamp() + + sinon.assert.calledOnce(mockWorkspaceFolderManager.checkRemoteWorkspaceStatusAndReact) + }) + + it('should not trigger workspace check when session was idle but monitoring is stopped', async () => { + clock.tick(31 * 60 * 1000) + mockWorkspaceFolderManager.isContinuousMonitoringStopped.returns(true) + + IdleWorkspaceManager.recordActivityTimestamp() + + sinon.assert.notCalled(mockWorkspaceFolderManager.checkRemoteWorkspaceStatusAndReact) + }) + }) +}) diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/IdleWorkspaceManager.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/IdleWorkspaceManager.ts new file mode 100644 index 0000000000..5a8359ccac --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/IdleWorkspaceManager.ts @@ -0,0 +1,38 @@ +import { WorkspaceFolderManager } from './workspaceFolderManager' + +export class IdleWorkspaceManager { + private static readonly idleThreshold = 30 * 60 * 1000 // 30 minutes + private static lastActivityTimestamp = 0 // treat session as idle as the start + + private constructor() {} + + /** + * Records activity timestamp and triggers workspace status check if session was idle. + * + * When transitioning from idle to active, proactively checks remote workspace status + * (if continuous monitoring is enabled) without blocking the current operation. + */ + public static recordActivityTimestamp(): void { + try { + const wasSessionIdle = IdleWorkspaceManager.isSessionIdle() + IdleWorkspaceManager.lastActivityTimestamp = Date.now() + + const workspaceFolderManager = WorkspaceFolderManager.getInstance() + if (workspaceFolderManager && wasSessionIdle && !workspaceFolderManager.isContinuousMonitoringStopped()) { + // Proactively check the remote workspace status instead of waiting for the next scheduled check + // Fire and forget - don't await to avoid blocking + workspaceFolderManager.checkRemoteWorkspaceStatusAndReact().catch(err => { + // ignore errors + }) + } + } catch (err) { + // ignore errors + } + } + + public static isSessionIdle(): boolean { + const currentTime = Date.now() + const timeSinceLastActivity = currentTime - IdleWorkspaceManager.lastActivityTimestamp + return timeSinceLastActivity > IdleWorkspaceManager.idleThreshold + } +} diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts index 8a4a9755cb..ed387ca2ca 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts @@ -191,7 +191,8 @@ export class DependencyDiscoverer { } } - public resetFromDisposal(): void { + public disposeAndReset(): void { + this.dispose() this.sharedState.isDisposed = false this.sharedState.dependencyUploadedSizeSum = 0 } diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts index 09336ffb2f..c536c1087d 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts @@ -221,7 +221,6 @@ export const WorkspaceContextServer = (): Server => features => { isLoggedInUsingBearerToken(credentialsProvider) && abTestingEnabled && !workspaceFolderManager.getOptOutStatus() && - !workspaceFolderManager.getServiceQuotaExceededStatus() && workspaceIdentifier ) } @@ -303,17 +302,15 @@ export const WorkspaceContextServer = (): Server => features => { await evaluateABTesting() isWorkflowInitialized = true - workspaceFolderManager.resetAdminOptOutAndServiceQuotaStatus() + workspaceFolderManager.resetAdminOptOutStatus() if (!isUserEligibleForWorkspaceContext()) { return } fileUploadJobManager.startFileUploadJobConsumer() dependencyEventBundler.startDependencyEventBundler() - await Promise.all([ - workspaceFolderManager.initializeWorkspaceStatusMonitor(), - workspaceFolderManager.processNewWorkspaceFolders(workspaceFolders), - ]) + + workspaceFolderManager.initializeWorkspaceStatusMonitor() logging.log(`Workspace context workflow initialized`) } else if (!isLoggedIn) { if (isWorkflowInitialized) { diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts index 0b4ec53f25..7ab596a930 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts @@ -6,8 +6,8 @@ import { DependencyDiscoverer } from './dependency/dependencyDiscoverer' import { WorkspaceFolder } from 'vscode-languageserver-protocol' import { ArtifactManager } from './artifactManager' import { CodeWhispererServiceToken } from '../../shared/codeWhispererService' -import { CreateWorkspaceResponse } from '../../client/token/codewhispererbearertokenclient' -import { AWSError } from 'aws-sdk' +import { ListWorkspaceMetadataResponse } from '../../client/token/codewhispererbearertokenclient' +import { IdleWorkspaceManager } from './IdleWorkspaceManager' describe('WorkspaceFolderManager', () => { let mockServiceManager: StubbedInstance @@ -33,8 +33,8 @@ describe('WorkspaceFolderManager', () => { sinon.restore() }) - describe('getServiceQuotaExceededStatus', () => { - it('should return true when service quota is exceeded', async () => { + describe('checkRemoteWorkspaceStatusAndReact', () => { + it('should check and react when IDE session is not idle', async () => { // Setup const workspaceFolders: WorkspaceFolder[] = [ { @@ -43,17 +43,20 @@ describe('WorkspaceFolderManager', () => { }, ] - // Mock the createWorkspace method to throw a ServiceQuotaExceededException - const mockError: AWSError = { - name: 'ServiceQuotaExceededException', - message: 'You have too many active running workspaces.', - code: 'ServiceQuotaExceededException', - time: new Date(), - retryable: false, - statusCode: 400, + // Mock IdleSessionManager to return false (not idle) + sinon.stub(IdleWorkspaceManager, 'isSessionIdle').returns(false) + + // Mock successful response + const mockResponse: ListWorkspaceMetadataResponse = { + workspaces: [ + { + workspaceId: 'test-workspace-id', + workspaceStatus: 'CREATED', + }, + ], } - mockCodeWhispererService.createWorkspace.rejects(mockError) + mockCodeWhispererService.listWorkspaceMetadata.resolves(mockResponse as any) // Create the WorkspaceFolderManager instance using the static createInstance method workspaceFolderManager = WorkspaceFolderManager.createInstance( @@ -66,23 +69,24 @@ describe('WorkspaceFolderManager', () => { 'test-workspace-identifier' ) - // Spy on clearAllWorkspaceResources and related methods - const clearAllWorkspaceResourcesSpy = sinon.stub( + // Spy on resetWebSocketClient + const resetWebSocketClientSpy = sinon.stub(workspaceFolderManager as any, 'resetWebSocketClient') + + // Spy on handleWorkspaceCreatedState + const handleWorkspaceCreatedStateSpy = sinon.stub( workspaceFolderManager as any, - 'clearAllWorkspaceResources' + 'handleWorkspaceCreatedState' ) - // Act - trigger the createNewWorkspace method which sets isServiceQuotaExceeded - await (workspaceFolderManager as any).createNewWorkspace() + // Act - trigger the checkRemoteWorkspaceStatusAndReact method + await workspaceFolderManager.checkRemoteWorkspaceStatusAndReact() - // Assert - expect(workspaceFolderManager.getServiceQuotaExceededStatus()).toBe(true) - - // Verify that clearAllWorkspaceResources was called - sinon.assert.calledOnce(clearAllWorkspaceResourcesSpy) + // Verify that resetWebSocketClient was called once + sinon.assert.notCalled(resetWebSocketClientSpy) + sinon.assert.calledOnce(handleWorkspaceCreatedStateSpy) }) - it('should return false when service quota is not exceeded', async () => { + it('should skip checking and reacting when IDE session is idle', async () => { // Setup const workspaceFolders: WorkspaceFolder[] = [ { @@ -91,15 +95,20 @@ describe('WorkspaceFolderManager', () => { }, ] + // Mock IdleSessionManager to return true (idle) + sinon.stub(IdleWorkspaceManager, 'isSessionIdle').returns(true) + // Mock successful response - const mockResponse: CreateWorkspaceResponse = { - workspace: { - workspaceId: 'test-workspace-id', - workspaceStatus: 'RUNNING', - }, + const mockResponse: ListWorkspaceMetadataResponse = { + workspaces: [ + { + workspaceId: 'test-workspace-id', + workspaceStatus: 'CREATED', + }, + ], } - mockCodeWhispererService.createWorkspace.resolves(mockResponse as any) + mockCodeWhispererService.listWorkspaceMetadata.resolves(mockResponse as any) // Create the WorkspaceFolderManager instance using the static createInstance method workspaceFolderManager = WorkspaceFolderManager.createInstance( @@ -112,20 +121,18 @@ describe('WorkspaceFolderManager', () => { 'test-workspace-identifier' ) - // Spy on clearAllWorkspaceResources - const clearAllWorkspaceResourcesSpy = sinon.stub( - workspaceFolderManager as any, - 'clearAllWorkspaceResources' - ) + // Spy on resetWebSocketClient + const resetWebSocketClientSpy = sinon.stub(workspaceFolderManager as any, 'resetWebSocketClient') - // Act - trigger the createNewWorkspace method - await (workspaceFolderManager as any).createNewWorkspace() + // Act - trigger the checkRemoteWorkspaceStatusAndReact method + await workspaceFolderManager.checkRemoteWorkspaceStatusAndReact() - // Assert - expect(workspaceFolderManager.getServiceQuotaExceededStatus()).toBe(false) - - // Verify that clearAllWorkspaceResources was not called - sinon.assert.notCalled(clearAllWorkspaceResourcesSpy) + // Verify that resetWebSocketClient was called once + sinon.assert.calledOnce(resetWebSocketClientSpy) + sinon.assert.calledWith( + mockLogging.log, + sinon.match(/Session is idle, skipping remote workspace status check/) + ) }) }) }) diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts index 00048cd62f..99fc9c4628 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts @@ -20,6 +20,7 @@ import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/A import { URI } from 'vscode-uri' import path = require('path') import { isAwsError } from '../../shared/utils' +import { IdleWorkspaceManager } from './IdleWorkspaceManager' interface WorkspaceState { remoteWorkspaceState: WorkspaceStatus @@ -54,8 +55,8 @@ export class WorkspaceFolderManager { private optOutMonitorInterval: NodeJS.Timeout | undefined private messageQueueConsumerInterval: NodeJS.Timeout | undefined private isOptedOut: boolean = false - // Tracks if the user has reached their maximum allowed remote workspaces quota - private isServiceQuotaExceeded: boolean = false + private isCheckingRemoteWorkspaceStatus: boolean = false + private isArtifactUploadedToRemoteWorkspace: boolean = false static createInstance( serviceManager: AmazonQTokenServiceManager, @@ -138,13 +139,8 @@ export class WorkspaceFolderManager { return this.isOptedOut } - getServiceQuotaExceededStatus(): boolean { - return this.isServiceQuotaExceeded - } - - resetAdminOptOutAndServiceQuotaStatus(): void { + resetAdminOptOutStatus(): void { this.isOptedOut = false - this.isServiceQuotaExceeded = false } getWorkspaceState(): WorkspaceState { @@ -324,14 +320,13 @@ export class WorkspaceFolderManager { this.workspaceState.webSocketClient = webSocketClient } - async initializeWorkspaceStatusMonitor() { + initializeWorkspaceStatusMonitor() { this.logging.log(`Initializing workspace status check for workspace [${this.workspaceIdentifier}]`) // Reset workspace ID to force operations to wait for new remote workspace information this.resetRemoteWorkspaceId() - this.artifactManager.resetFromDisposal() - this.dependencyDiscoverer.resetFromDisposal() + this.isArtifactUploadedToRemoteWorkspace = false // Set up message queue consumer if (this.messageQueueConsumerInterval === undefined) { @@ -350,12 +345,8 @@ export class WorkspaceFolderManager { }, this.MESSAGE_PUBLISH_INTERVAL) } - // Perform a one-time checkRemoteWorkspaceStatusAndReact first - // Pass skipUploads as true since it would be handled by processNewWorkspaceFolders - await this.checkRemoteWorkspaceStatusAndReact(true) - // Set up continuous monitoring which periodically invokes checkRemoteWorkspaceStatusAndReact - if (!this.isOptedOut && !this.isServiceQuotaExceeded && this.continuousMonitorInterval === undefined) { + if (!this.isOptedOut && this.continuousMonitorInterval === undefined) { this.logging.log(`Starting continuous monitor for workspace [${this.workspaceIdentifier}]`) this.continuousMonitorInterval = setInterval(async () => { try { @@ -427,61 +418,92 @@ export class WorkspaceFolderManager { }) } - private async checkRemoteWorkspaceStatusAndReact(skipUploads: boolean = false) { - if (this.workspaceFolders.length === 0) { - this.logging.log(`No workspace folders added, skipping workspace status check`) + public async checkRemoteWorkspaceStatusAndReact() { + if (this.isCheckingRemoteWorkspaceStatus) { + // Skip checking remote workspace if a previous check is still in progress return } + this.isCheckingRemoteWorkspaceStatus = true + try { + if (IdleWorkspaceManager.isSessionIdle()) { + this.resetWebSocketClient() + this.logging.log('Session is idle, skipping remote workspace status check') + return + } - this.logging.log(`Checking remote workspace status for workspace [${this.workspaceIdentifier}]`) - const { metadata, optOut, error } = await this.listWorkspaceMetadata(this.workspaceIdentifier) + if (this.workspaceFolders.length === 0) { + this.logging.log(`No workspace folders added, skipping workspace status check`) + return + } - if (optOut) { - this.logging.log('User opted out, clearing all resources and starting opt-out monitor') - this.isOptedOut = true - this.clearAllWorkspaceResources() - this.startOptOutMonitor() - return - } + this.logging.log(`Checking remote workspace status for workspace [${this.workspaceIdentifier}]`) + const { metadata, optOut, error } = await this.listWorkspaceMetadata(this.workspaceIdentifier) - if (error) { - // Do not do anything if we received an exception but not caused by optOut - return - } + if (optOut) { + this.logging.log('User opted out, clearing all resources and starting opt-out monitor') + this.isOptedOut = true + this.clearAllWorkspaceResources() + this.startOptOutMonitor() + return + } - if (!metadata) { - // Workspace no longer exists, Recreate it. - this.resetRemoteWorkspaceId() // workspaceId would change if remote record is gone - await this.handleWorkspaceCreatedState(skipUploads) - return - } + if (error) { + // Do not do anything if we received an exception but not caused by optOut + return + } - this.workspaceState.remoteWorkspaceState = metadata.workspaceStatus - if (this.workspaceState.workspaceId === undefined) { - this.setRemoteWorkspaceId(metadata.workspaceId) - } + if (!metadata) { + // Workspace no longer exists, Recreate it. + this.resetRemoteWorkspaceId() // workspaceId would change if remote record is gone + await this.handleWorkspaceCreatedState() + return + } - switch (metadata.workspaceStatus) { - case 'READY': - // Check if connection exists - const client = this.workspaceState.webSocketClient - if (!client || !client.isConnected()) { - this.logging.log( - `Workspace is ready but no connection exists or connection lost. Re-establishing connection...` - ) - await this.establishConnection(metadata) - } - break - case 'PENDING': - // Schedule an initial connection when pending - await this.waitForInitialConnection() - break - case 'CREATED': - // Workspace has no environment, Recreate it. - await this.handleWorkspaceCreatedState(skipUploads) - break - default: - this.logging.warn(`Unknown workspace status: ${metadata.workspaceStatus}`) + this.workspaceState.remoteWorkspaceState = metadata.workspaceStatus + if (this.workspaceState.workspaceId === undefined) { + this.setRemoteWorkspaceId(metadata.workspaceId) + } + + switch (metadata.workspaceStatus) { + case 'READY': + // Check if connection exists + const client = this.workspaceState.webSocketClient + if (!client || !client.isConnected()) { + this.logging.log( + `Workspace is ready but no connection exists or connection lost. Re-establishing connection...` + ) + let uploadArtifactsPromise: Promise | undefined + if (!this.isArtifactUploadedToRemoteWorkspace) { + uploadArtifactsPromise = this.uploadAllArtifactsToRemoteWorkspace() + } + await this.establishConnection(metadata) + if (uploadArtifactsPromise) { + await uploadArtifactsPromise + } + } + break + case 'PENDING': + // Schedule an initial connection when pending + let uploadArtifactsPromise: Promise | undefined + if (!this.isArtifactUploadedToRemoteWorkspace) { + uploadArtifactsPromise = this.uploadAllArtifactsToRemoteWorkspace() + } + await this.waitForInitialConnection() + if (uploadArtifactsPromise) { + await uploadArtifactsPromise + } + break + case 'CREATED': + // Workspace has no environment, Recreate it. + await this.handleWorkspaceCreatedState() + break + default: + this.logging.warn(`Unknown workspace status: ${metadata.workspaceStatus}`) + } + } catch (error) { + this.logging.error(`Error checking remote workspace status: ${error}`) + } finally { + this.isCheckingRemoteWorkspaceStatus = false } } @@ -515,9 +537,7 @@ export class WorkspaceFolderManager { ) clearInterval(intervalId) this.optOutMonitorInterval = undefined - this.initializeWorkspaceStatusMonitor().catch(error => { - this.logging.error(`Error while initializing workspace status monitoring: ${error}`) - }) + this.initializeWorkspaceStatusMonitor() this.processNewWorkspaceFolders(this.workspaceFolders).catch(error => { this.logging.error(`Error while processing workspace folders: ${error}`) }) @@ -530,25 +550,18 @@ export class WorkspaceFolderManager { } } - private async handleWorkspaceCreatedState(skipUploads: boolean = false): Promise { + private async handleWorkspaceCreatedState(): Promise { this.logging.log(`No READY / PENDING remote workspace found, creating a new one`) // If remote state is CREATED, call create API to create a new workspace - if (this.workspaceState.webSocketClient) { - this.workspaceState.webSocketClient.destroyClient() - this.workspaceState.webSocketClient = undefined - } + this.resetWebSocketClient() const initialResult = await this.createNewWorkspace() // If creation succeeds, establish connection if (initialResult.response) { this.logging.log(`Workspace [${this.workspaceIdentifier}] created successfully, establishing connection`) + const uploadArtifactsPromise = this.uploadAllArtifactsToRemoteWorkspace() await this.waitForInitialConnection() - if (!skipUploads) { - await this.syncSourceCodesToS3(this.workspaceFolders) - this.dependencyDiscoverer.reSyncDependenciesToS3(this.workspaceFolders).catch(e => { - this.logging.warn(`Error during re-syncing dependencies: ${e}`) - }) - } + await uploadArtifactsPromise return } @@ -570,13 +583,27 @@ export class WorkspaceFolderManager { } this.logging.log(`Retry succeeded for workspace creation, establishing connection`) + const uploadArtifactsPromise = this.uploadAllArtifactsToRemoteWorkspace() await this.waitForInitialConnection() - if (!skipUploads) { - await this.syncSourceCodesToS3(this.workspaceFolders) - this.dependencyDiscoverer.reSyncDependenciesToS3(this.workspaceFolders).catch(e => { - this.logging.warn(`Error during re-syncing dependencies: ${e}`) - }) - } + await uploadArtifactsPromise + } + + private async uploadAllArtifactsToRemoteWorkspace() { + // initialize source codes + this.artifactManager.resetFromDisposal() + await this.syncSourceCodesToS3(this.workspaceFolders) + + // initialize dependencies + this.dependencyDiscoverer.disposeAndReset() + this.dependencyDiscoverer.searchDependencies(this.workspaceFolders).catch(e => { + this.logging.warn(`Error during dependency discovery: ${e}`) + }) + + this.isArtifactUploadedToRemoteWorkspace = true + } + + public isContinuousMonitoringStopped(): boolean { + return this.continuousMonitorInterval === undefined } private stopContinuousMonitoring() { @@ -602,15 +629,15 @@ export class WorkspaceFolderManager { } } - private async createNewWorkspace() { - const createWorkspaceResult = await this.createWorkspace(this.workspaceIdentifier) - - this.isServiceQuotaExceeded = createWorkspaceResult.isServiceQuotaExceeded - if (this.isServiceQuotaExceeded) { - // Stop continuous monitor and all actions - this.clearAllWorkspaceResources() + private resetWebSocketClient() { + if (this.workspaceState.webSocketClient) { + this.workspaceState.webSocketClient.destroyClient() + this.workspaceState.webSocketClient = undefined } + } + private async createNewWorkspace() { + const createWorkspaceResult = await this.createWorkspace(this.workspaceIdentifier) const workspaceDetails = createWorkspaceResult.response if (!workspaceDetails) { this.logging.warn(`Failed to create remote workspace for [${this.workspaceIdentifier}]`)