@aws/lsp-codewhisperer 0.0.48 → 0.0.50

package/out/language-server/agenticChat/agenticChatController.js

@@ -7,7 +7,7 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.AgenticChatController = void 0;
 const crypto = require("crypto");
 const path = require("path");
-const codewhisperer_streaming_1 = require("@amzn/codewhisperer-streaming");
+const codewhisperer_streaming_client_1 = require("@aws/codewhisperer-streaming-client");
 const protocol_1 = require("@aws/language-server-runtimes/protocol");
 const protocol_2 = require("@aws/language-server-runtimes/protocol");
 const server_interface_1 = require("@aws/language-server-runtimes/server-interface");
@@ -52,6 +52,7 @@ const mcpUtils_1 = require("./tools/mcp/mcpUtils");
 const mcpManager_1 = require("./tools/mcp/mcpManager");
 const mcpTool_1 = require("./tools/mcp/mcpTool");
 const paidTier_1 = require("../paidTier/paidTier");
+const util_1 = require("./tools/chatDb/util");
 class AgenticChatController {
     #features;
     #chatSessionManagementService;
@@ -89,8 +90,8 @@ class AgenticChatController {
         this.#telemetryService = telemetryService;
         this.#serviceManager = serviceManager;
         this.#chatHistoryDb = new chatDb_1.ChatDatabase(features);
-        this.#tabBarController = new tabBarController_1.TabBarController(features, this.#chatHistoryDb, telemetryService);
-        this.#additionalContextProvider = new addtionalContextProvider_1.AdditionalContextProvider(features.workspace);
+        this.#tabBarController = new tabBarController_1.TabBarController(features, this.#chatHistoryDb, telemetryService, (tabId) => this.sendPinnedContext(tabId));
+        this.#additionalContextProvider = new addtionalContextProvider_1.AdditionalContextProvider(features, this.#chatHistoryDb);
         this.#contextCommandsProvider = new contextCommandsProvider_1.ContextCommandsProvider(this.#features.logging, this.#features.chat, this.#features.workspace, this.#features.lsp);
         this.#mcpEventHandler = new mcpEventHandler_1.McpEventHandler(features, telemetryService);
     }
@@ -247,6 +248,23 @@ class AgenticChatController {
         }
     }
     async onCreatePrompt(params) {
+        if (params.isRule) {
+            let workspaceFolders = lsp_core_1.workspaceUtils.getWorkspaceFolderPaths(this.#features.workspace);
+            let workspaceRulesDirectory = path.join(workspaceFolders[0], '.amazonq', 'rules');
+            if (workspaceFolders.length > 0) {
+                const newFilePath = (0, contextUtils_1.getNewRuleFilePath)(params.promptName, workspaceRulesDirectory);
+                const newFileContent = '';
+                try {
+                    await this.#features.workspace.fs.mkdir(workspaceRulesDirectory, { recursive: true });
+                    await this.#features.workspace.fs.writeFile(newFilePath, newFileContent, { mode: 0o600 });
+                    await this.#features.lsp.window.showDocument({ uri: vscode_uri_1.URI.file(newFilePath).toString() });
+                }
+                catch (e) {
+                    this.#features.logging.warn(`Error creating rule file: ${e}`);
+                }
+                return;
+            }
+        }
         const newFilePath = (0, contextUtils_1.getNewPromptFilePath)(params.promptName);
         const newFileContent = '';
         try {
@@ -271,6 +289,12 @@ class AgenticChatController {
     async onConversationClick(params) {
         return this.#tabBarController.onConversationClick(params);
     }
+    async onRuleClick(params) {
+        return this.#additionalContextProvider.onRuleClick(params);
+    }
+    async onListRules(params) {
+        return this.#additionalContextProvider.onListRules(params);
+    }
     async onListMcpServers(params) {
         return this.#mcpEventHandler.onListMcpServers(params);
     }
@@ -336,7 +360,7 @@ class AgenticChatController {
                 await this.#telemetryController.emitAddMessageMetric(params.tabId, metric.metric, 'Cancelled');
             });
             session.setConversationType('AgenticChat');
-            const additionalContext = await this.#additionalContextProvider.getAdditionalContext(triggerContext, params.context);
+            const additionalContext = await this.#additionalContextProvider.getAdditionalContext(triggerContext, params.tabId, params.context);
             if (additionalContext.length) {
                 triggerContext.documentReference =
                     this.#additionalContextProvider.getFileListFromContext(additionalContext);
@@ -384,7 +408,7 @@ class AgenticChatController {
     async #prepareRequestInput(params, session, triggerContext, additionalContext, chatResultStream) {
         this.#debug('Preparing request input');
         const profileArn = AmazonQTokenServiceManager_1.AmazonQTokenServiceManager.getInstance().getActiveProfileArn();
-        const requestInput = await this.#triggerContext.getChatParamsFromTrigger(params, triggerContext, codewhisperer_streaming_1.ChatTriggerType.MANUAL, this.#customizationArn, chatResultStream, profileArn, [], this.#getTools(session), additionalContext, session.modelId);
+        const requestInput = await this.#triggerContext.getChatParamsFromTrigger(params, triggerContext, codewhisperer_streaming_client_1.ChatTriggerType.MANUAL, this.#customizationArn, chatResultStream, profileArn, [], this.#getTools(session), additionalContext, session.modelId);
         return requestInput;
     }
     /**
@@ -410,19 +434,24 @@ class AgenticChatController {
                 this.#debug(`Warning: ${!currentMessage ? 'currentMessage' : ''}${!currentMessage && !conversationId ? ' and ' : ''}${!conversationId ? 'conversationIdentifier' : ''} is empty in agent loop iteration ${iterationCount}.`);
             }
             const remainingCharacterBudget = this.truncateRequest(currentRequestInput);
-            //  Fix the history to maintain invariants
+            let messages = [];
             if (currentMessage) {
-                const isHistoryValid = this.#chatHistoryDb.fixAndValidateHistory(tabId, currentMessage, conversationIdentifier ?? '', remainingCharacterBudget);
-                if (!isHistoryValid) {
-                    this.#features.logging.warn('Skipping request due to invalid tool result/tool use relationship');
-                    break;
+                //  Get and process the messages from history DB to maintain invariants for service requests
+                try {
+                    messages = this.#chatHistoryDb.fixAndGetHistory(tabId, currentMessage, remainingCharacterBudget);
+                }
+                catch (err) {
+                    if (err instanceof chatDb_1.ToolResultValidationError) {
+                        this.#features.logging.warn(`Tool validation error: ${err.message}`);
+                        break;
+                    }
                 }
             }
-            //  Retrieve the history from DB; Do not include chatHistory for requests going to Mynah Backend
+            //  Do not include chatHistory for requests going to Mynah Backend
             currentRequestInput.conversationState.history = currentRequestInput.conversationState?.currentMessage
                 ?.userInputMessage?.userIntent
                 ? []
-                : this.#chatHistoryDb.getMessages(tabId);
+                : messages.map(msg => (0, util_1.messageToStreamingMessage)(msg));
             // Add loading message before making the request
             const loadingMessageId = `loading-${(0, uuid_1.v4)()}`;
             await chatResultStream.writeResultBlock({ ...constants_1.loadingMessage, messageId: loadingMessageId });
@@ -437,9 +466,10 @@ class AgenticChatController {
             }
             this.#features.logging.info(`generateAssistantResponse ResponseMetadata: ${lsp_core_2.loggingUtils.formatObj(response.$metadata)}`);
             await chatResultStream.removeResultBlock(loadingMessageId);
-            //  Add the current user message to the history DB
+            // Add the current user message to the history DB
             if (currentMessage && conversationIdentifier) {
                 if (this.#isPromptCanceled(token, session, promptId)) {
+                    // Only skip adding message to history, continue executing to avoid unexpected stop for the conversation
                     this.#debug('Skipping adding user message to history - cancelled by user');
                 }
                 else {
@@ -450,6 +480,7 @@ class AgenticChatController {
                         origin: currentMessage.userInputMessage?.origin,
                         userInputMessageContext: currentMessage.userInputMessage?.userInputMessageContext,
                         shouldDisplayMessage: shouldDisplayMessage,
+                        timestamp: new Date(),
                     });
                 }
             }
@@ -463,11 +494,16 @@ class AgenticChatController {
             // and we want to auto-retry
             if (!result.success && result.error.startsWith(constants_2.responseTimeoutPartialMsg)) {
                 const content = 'You took too long to respond - try to split up the work into smaller steps. Do not apologize.';
-                if (!this.#isPromptCanceled(token, session, promptId)) {
+                if (this.#isPromptCanceled(token, session, promptId)) {
+                    // Only skip adding message to the history DB, continue executing to avoid unexpected stop for the conversation
+                    this.#debug('Skipping adding messages to history - cancelled by user');
+                }
+                else {
                     this.#chatHistoryDb.addMessage(tabId, 'cwc', conversationIdentifier ?? '', {
                         body: 'Response timed out - message took too long to generate',
                         type: 'answer',
                         shouldDisplayMessage: false,
+                        timestamp: new Date(),
                     });
                 }
                 currentRequestInput = this.#updateRequestInputWithToolResults(currentRequestInput, [], content);
@@ -476,10 +512,11 @@ class AgenticChatController {
                 await chatResultStream.updateOngoingProgressResult('Error');
                 continue;
             }
-            //  Add the current assistantResponse message to the history DB
+            // Add the current assistantResponse message to the history DB
             if (result.data?.chatResult.body !== undefined) {
                 if (this.#isPromptCanceled(token, session, promptId)) {
-                    this.#debug('Skipping adding assistant message to history - cancelled by user');
+                    // Only skip adding message to the history DB, continue executing to avoid unexpected stop for the conversation
+                    this.#debug('Skipping adding messages to history - cancelled by user');
                 }
                 else {
                     this.#chatHistoryDb.addMessage(tabId, 'cwc', conversationIdentifier ?? '', {
@@ -498,6 +535,7 @@ class AgenticChatController {
                             input: result.data.toolUses[k].input,
                         })),
                         shouldDisplayMessage: shouldDisplayMessage,
+                        timestamp: new Date(),
                     });
                 }
             }
@@ -534,7 +572,7 @@ class AgenticChatController {
                 // Send an error card to UI?
                 toolResults = pendingToolUses.map(toolUse => ({
                     toolUseId: toolUse.toolUseId,
-                    status: codewhisperer_streaming_1.ToolResultStatus.ERROR,
+                    status: codewhisperer_streaming_client_1.ToolResultStatus.ERROR,
                     content: [{ text: result.error }],
                 }));
                 this.#telemetryController.emitAgencticLoop_InvokeLLM(response.$metadata.requestId, conversationId, 'AgenticChatWithToolUse', undefined, undefined, 'Failed', this.#features.runtime.serverInfo.version ?? '', undefined, session.pairProgrammingMode);
@@ -627,13 +665,13 @@ class AgenticChatController {
      * @param promptBlockId id of approval block. This allows us to overwrite the buttons with 'accepted' or 'rejected' text.
      * @param session
      */
-    async waitForToolApproval(toolUse, resultStream, promptBlockId, session) {
+    async waitForToolApproval(toolUse, resultStream, promptBlockId, session, toolName) {
         const deferred = this.#createDeferred();
         session.setDeferredToolExecution(toolUse.toolUseId, deferred.resolve, deferred.reject);
-        this.#log(`Prompting for tool approval for tool: ${toolUse.name}`);
+        this.#log(`Prompting for tool approval for tool: ${toolName ?? toolUse.name}`);
         await deferred.promise;
         // Note: we want to overwrite the button block because it already exists in the stream.
-        await resultStream.overwriteResultBlock(this.#getUpdateToolConfirmResult(toolUse, true), promptBlockId);
+        await resultStream.overwriteResultBlock(this.#getUpdateToolConfirmResult(toolUse, true, toolName), promptBlockId);
     }
     /**
      * Processes tool uses by running the tools and collecting results
@@ -709,7 +747,7 @@ class AgenticChatController {
                                 this.#telemetryController.emitInteractWithAgenticChat('GeneratedCommand', tabId, session.pairProgrammingMode, session.getConversationType());
                             }
                             if (requiresAcceptance) {
-                                await this.waitForToolApproval(toolUse, chatResultStream, cachedButtonBlockId, session);
+                                await this.waitForToolApproval(toolUse, chatResultStream, cachedButtonBlockId, session, toolUse.name);
                             }
                             if (isExecuteBash) {
                                 this.#telemetryController.emitInteractWithAgenticChat('RunCommand', tabId, session.pairProgrammingMode, session.getConversationType());
@@ -736,7 +774,7 @@ class AgenticChatController {
                                     const confirmation = this.#processToolConfirmation(toolUse, requiresAcceptance, warning, undefined, toolName // Pass the original tool name here
                                     );
                                     cachedButtonBlockId = await chatResultStream.writeResultBlock(confirmation);
-                                    await this.waitForToolApproval(toolUse, chatResultStream, cachedButtonBlockId, session);
+                                    await this.waitForToolApproval(toolUse, chatResultStream, cachedButtonBlockId, session, toolName);
                                 }
                                 // Store the blockId in the session for later use
                                 if (toolUse.toolUseId) {
@@ -848,7 +886,7 @@ class AgenticChatController {
                 if (this.isUserAction(err, token)) {
                     // Handle ToolApprovalException for any tool
                     if (err instanceof toolShared_1.ToolApprovalException && cachedButtonBlockId) {
-                        await chatResultStream.overwriteResultBlock(this.#getUpdateToolConfirmResult(toolUse, false), cachedButtonBlockId);
+                        await chatResultStream.overwriteResultBlock(this.#getUpdateToolConfirmResult(toolUse, false, toolUse.name), cachedButtonBlockId);
                         if (err.shouldShowMessage) {
                             await chatResultStream.writeResultBlock({
                                 type: 'answer',
@@ -864,7 +902,7 @@ class AgenticChatController {
                     if (err instanceof lsp_core_1.CancellationError) {
                         results.push({
                             toolUseId: toolUse.toolUseId,
-                            status: codewhisperer_streaming_1.ToolResultStatus.ERROR,
+                            status: codewhisperer_streaming_client_1.ToolResultStatus.ERROR,
                             content: [{ text: 'Command stopped by user' }],
                         });
                         continue;
@@ -940,7 +978,7 @@ class AgenticChatController {
                 this.#log(`Error running tool ${toolUse.name}:`, errMsg);
                 results.push({
                     toolUseId: toolUse.toolUseId,
-                    status: codewhisperer_streaming_1.ToolResultStatus.ERROR,
+                    status: codewhisperer_streaming_client_1.ToolResultStatus.ERROR,
                     content: [{ json: { error: err instanceof Error ? err.message : 'Unknown error' } }],
                 });
             }
@@ -948,7 +986,7 @@ class AgenticChatController {
         return results;
     }
     #shouldSendBackErrorContent(toolResult) {
-        if (toolResult.status === codewhisperer_streaming_1.ToolResultStatus.ERROR) {
+        if (toolResult.status === codewhisperer_streaming_client_1.ToolResultStatus.ERROR) {
             for (const content of toolResult.content ?? []) {
                 if (content.json && JSON.stringify(content.json).includes(constants_2.outputLimitExceedsPartialMsg)) {
                     // do not send the content response back for this case to avoid unnecessary messages
@@ -1117,8 +1155,8 @@ class AgenticChatController {
      * @param toolType Optional tool type for specialized handling
      * @returns ChatResult with appropriate confirmation UI
      */
-    #getUpdateToolConfirmResult(toolUse, isAccept, toolType) {
-        const toolName = toolType || toolUse.name;
+    #getUpdateToolConfirmResult(toolUse, isAccept, originalToolName, toolType) {
+        const toolName = originalToolName ?? (toolType || toolUse.name);
         // Handle bash commands with special formatting
         if (toolName === 'executeBash') {
             return {
@@ -1169,7 +1207,7 @@ class AgenticChatController {
                 body = `File search ${isAccept ? 'allowed' : 'rejected'}: \`${searchPath}\``;
                 break;
             default:
-                // Default tool (not MCP)
+                // Default tool (not only MCP)
                 return {
                     type: 'tool',
                     messageId: toolUse.toolUseId,
@@ -1177,7 +1215,7 @@ class AgenticChatController {
                         content: {
                             header: {
                                 icon: 'tools',
-                                body: `${toolUse.name}`,
+                                body: `${toolName}`,
                                 status: {
                                     status: isAccept ? 'success' : 'error',
                                     icon: isAccept ? 'ok' : 'cancel',
@@ -1579,7 +1617,8 @@ class AgenticChatController {
                 cwsprChatHasContextList: triggerContext.documentReference?.filePaths?.length ? true : false,
                 cwsprChatFolderContextCount: triggerContext.contextInfo.contextCount.folderContextCount,
                 cwsprChatFileContextCount: triggerContext.contextInfo.contextCount.fileContextCount,
-                cwsprChatRuleContextCount: triggerContext.contextInfo.contextCount.ruleContextCount,
+                cwsprChatRuleContextCount: triggerContext.contextInfo.contextCount.activeRuleContextCount,
+                cwsprChatTotalRuleContextCount: triggerContext.contextInfo.contextCount.totalRuleContextCount,
                 cwsprChatPromptContextCount: triggerContext.contextInfo.contextCount.promptContextCount,
                 cwsprChatFileContextLength: triggerContext.contextInfo.contextLength.fileContextLength,
                 cwsprChatRuleContextLength: triggerContext.contextInfo.contextLength.ruleContextLength,
@@ -1587,6 +1626,10 @@ class AgenticChatController {
                 cwsprChatCodeContextCount: triggerContext.contextInfo.contextCount.codeContextCount,
                 cwsprChatCodeContextLength: triggerContext.contextInfo.contextLength.codeContextLength,
                 cwsprChatFocusFileContextLength: triggerContext.text?.length,
+                cwsprChatPinnedCodeContextCount: triggerContext.contextInfo.pinnedContextCount.codeContextCount,
+                cwsprChatPinnedFileContextCount: triggerContext.contextInfo.pinnedContextCount.fileContextCount,
+                cwsprChatPinnedFolderContextCount: triggerContext.contextInfo.pinnedContextCount.folderContextCount,
+                cwsprChatPinnedPromptContextCount: triggerContext.contextInfo.pinnedContextCount.promptContextCount,
             });
         }
         await this.#telemetryController.emitAddMessageMetric(params.tabId, metric.metric, 'Succeeded');
@@ -1661,6 +1704,21 @@ class AgenticChatController {
                 buttons: [],
             };
             if (err.code === 'QModelResponse') {
+                // special case for throttling where we show error card instead of chat message
+                if (err.message ===
+                    `The model you selected is temporarily unavailable. Please switch to a different model and try again.`) {
+                    this.#features.chat.sendChatUpdate({
+                        tabId: tabId,
+                        data: { messages: [{ messageId: 'modelUnavailable' }] },
+                    });
+                    const emptyChatResult = {
+                        type: 'answer',
+                        body: '',
+                        messageId: errorMessageId,
+                        buttons: [],
+                    };
+                    return emptyChatResult;
+                }
                 return responseData;
             }
             return new server_interface_1.ResponseError(server_interface_1.LSPErrorCodes.RequestFailed, err.message, responseData);
@@ -1682,7 +1740,7 @@ class AgenticChatController {
         let response;
         let requestInput;
         try {
-            requestInput = await this.#triggerContext.getChatParamsFromTrigger(params, triggerContext, codewhisperer_streaming_1.ChatTriggerType.INLINE_CHAT, this.#customizationArn);
+            requestInput = await this.#triggerContext.getChatParamsFromTrigger(params, triggerContext, codewhisperer_streaming_client_1.ChatTriggerType.INLINE_CHAT, this.#customizationArn);
             if (!this.#serviceManager) {
                 throw new Error('amazonQServiceManager is not initialized');
             }
@@ -1715,6 +1773,11 @@ class AgenticChatController {
     async onInlineChatResult(params) {
         await this.#telemetryService.emitInlineChatResultLog(params);
     }
+    async onActiveEditorChanged(params) {
+        if (this.#telemetryController.activeTabId) {
+            this.sendPinnedContext(this.#telemetryController.activeTabId);
+        }
+    }
     async onCodeInsertToCursorPosition(params) {
         // Implementation based on https://github.com/aws/aws-toolkit-vscode/blob/1814cc84228d4bf20270574c5980b91b227f31cf/packages/core/src/amazonq/commons/controllers/contentController.ts#L38
         if (!params.textDocument || !params.cursorPosition || !params.code) {
@@ -1861,6 +1924,9 @@ class AgenticChatController {
         // In that case, we use the default modelId.
         const modelId = this.#chatHistoryDb.getModelId() ?? constants_2.defaultModelId;
         this.#features.chat.chatOptionsUpdate({ modelId: modelId, tabId: params.tabId });
+        if (!params.restoredTab) {
+            this.sendPinnedContext(params.tabId);
+        }
         const sessionResult = this.#chatSessionManagementService.createSession(params.tabId);
         const { data: session, success } = sessionResult;
         if (!success) {
@@ -1879,12 +1945,16 @@ class AgenticChatController {
             data: {},
         });
         this.#telemetryController.activeTabId = params.tabId;
+        this.sendPinnedContext(params.tabId);
         this.#telemetryController.emitConversationMetric({
             name: types_1.ChatTelemetryEventName.EnterFocusConversation,
             data: {},
         });
         this.setPaidTierMode(params.tabId);
     }
+    sendPinnedContext(tabId) {
+        this.#additionalContextProvider.sendPinnedContext(tabId);
+    }
     onTabRemove(params) {
         if (this.#telemetryController.activeTabId === params.tabId) {
             this.#telemetryController.emitConversationMetric({
@@ -1923,21 +1993,16 @@ class AgenticChatController {
                     messageId: (0, uuid_1.v4)(),
                     body: constants_1.HELP_MESSAGE,
                 };
-            // "Manage Subscription" (paid-tier user), or "Upgrade Q" (free-tier user)
-            case quickActions_1.QuickAction.Manage:
-                this.#telemetryController.emitChatMetric({
-                    name: types_1.ChatTelemetryEventName.RunCommand,
-                    data: {
-                        cwsprChatCommandType: params.quickAction,
-                        cwsprChatCommandName: '/manage',
-                    },
-                });
-                void this.onManageSubscription(params.tabId);
-                return {};
             default:
                 return {};
         }
     }
+    onPinnedContextAdd(params) {
+        this.#additionalContextProvider.onPinnedContextAdd(params);
+    }
+    onPinnedContextRemove(params) {
+        this.#additionalContextProvider.onPinnedContextRemove(params);
+    }
     async onTabBarAction(params) {
         return this.#tabBarController.onTabBarAction(params);
     }