@aws/lsp-codewhisperer 0.0.58 → 0.0.60

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_client_1 = require("@aws/codewhisperer-streaming-client");
+const codewhisperer_streaming_1 = require("@amzn/codewhisperer-streaming");
 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");
@@ -72,6 +72,7 @@ class AgenticChatController {
     #toolUseLatencies = [];
     #mcpEventHandler;
     #paidTierMode;
+    #origin;
     // latency metrics
     #llmRequestStartTime = 0;
     #toolCallLatencies = [];
@@ -106,6 +107,7 @@ class AgenticChatController {
         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);
+        this.#origin = (0, utils_2.getOriginFromClientInfo)(this.#features.lsp.getClientInitializeParams()?.clientInfo?.name);
     }
     async onExecuteCommand(params, _token) {
         this.#log(`onExecuteCommand: ${params.command}`);
@@ -362,11 +364,6 @@ class AgenticChatController {
                 // Then update UI to inform the user
                 await this.#showUndoAllIfRequired(chatResultStream, session);
                 await chatResultStream.updateOngoingProgressResult('Canceled');
-                await this.#getChatResultStream(params.partialResultToken).writeResultBlock({
-                    type: 'directive',
-                    messageId: 'stopped' + (0, uuid_1.v4)(),
-                    body: 'You stopped your current work, please provide additional examples or ask another question.',
-                });
                 // Finally, send telemetry/metrics
                 this.#telemetryController.emitInteractWithAgenticChat('StopChat', params.tabId, session.pairProgrammingMode, session.getConversationType());
                 await this.#telemetryController.emitAddMessageMetric(params.tabId, metric.metric, 'Cancelled');
@@ -419,8 +416,9 @@ 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_client_1.ChatTriggerType.MANUAL, this.#customizationArn, chatResultStream, profileArn, [], this.#getTools(session), additionalContext, session.modelId);
+        // Get profileArn from the service manager if available
+        const profileArn = this.#serviceManager?.getActiveProfileArn();
+        const requestInput = await this.#triggerContext.getChatParamsFromTrigger(params, triggerContext, codewhisperer_streaming_1.ChatTriggerType.MANUAL, this.#customizationArn, chatResultStream, profileArn, [], this.#getTools(session), additionalContext, session.modelId, this.#origin);
         return requestInput;
     }
     /**
@@ -443,17 +441,17 @@ class AgenticChatController {
                 this.#debug('Stopping agent loop - cancelled by user');
                 throw new lsp_core_1.CancellationError('user');
             }
+            this.truncateRequest(currentRequestInput);
             const currentMessage = currentRequestInput.conversationState?.currentMessage;
             const conversationId = conversationIdentifier ?? '';
             if (!currentMessage || !conversationId) {
                 this.#debug(`Warning: ${!currentMessage ? 'currentMessage' : ''}${!currentMessage && !conversationId ? ' and ' : ''}${!conversationId ? 'conversationIdentifier' : ''} is empty in agent loop iteration ${iterationCount}.`);
             }
-            const remainingCharacterBudget = this.truncateRequest(currentRequestInput);
             let messages = [];
             if (currentMessage) {
                 //  Get and process the messages from history DB to maintain invariants for service requests
                 try {
-                    messages = this.#chatHistoryDb.fixAndGetHistory(tabId, currentMessage, remainingCharacterBudget);
+                    messages = this.#chatHistoryDb.fixAndGetHistory(tabId, currentMessage);
                 }
                 catch (err) {
                     if (err instanceof chatDb_1.ToolResultValidationError) {
@@ -473,14 +471,14 @@ class AgenticChatController {
             this.#llmRequestStartTime = Date.now();
             // Phase 3: Request Execution
             // Note: these logs are very noisy, but contain information redacted on the backend.
-            this.#debug(`generateAssistantResponse Request: ${JSON.stringify(currentRequestInput, undefined, 2)}`);
-            const response = await session.generateAssistantResponse(currentRequestInput);
+            this.#debug(`generateAssistantResponse/SendMessage Request: ${JSON.stringify(currentRequestInput, undefined, 2)}`);
+            const response = await session.getChatResponse(currentRequestInput);
             if (response.$metadata.requestId) {
                 metric.mergeWith({
                     requestIds: [response.$metadata.requestId],
                 });
             }
-            this.#features.logging.info(`generateAssistantResponse ResponseMetadata: ${lsp_core_2.loggingUtils.formatObj(response.$metadata)}`);
+            this.#features.logging.info(`generateAssistantResponse/SendMessage ResponseMetadata: ${lsp_core_2.loggingUtils.formatObj(response.$metadata)}`);
             await chatResultStream.removeResultBlock(loadingMessageId);
             // Add the current user message to the history DB
             if (currentMessage && conversationIdentifier) {
@@ -502,7 +500,7 @@ class AgenticChatController {
             }
             shouldDisplayMessage = true;
             // Phase 4: Response Processing
-            const result = await this.#processGenerateAssistantResponseResponseWithTimeout(response, metric.mergeWith({
+            const result = await this.#processAgenticChatResponseWithTimeout(response, metric.mergeWith({
                 cwsprChatResponseCode: response.$metadata.httpStatusCode,
                 cwsprChatMessageId: response.$metadata.requestId,
             }), chatResultStream, session, documentReference);
@@ -565,7 +563,7 @@ class AgenticChatController {
             if (pendingToolUses.length === 0) {
                 this.recordChunk('agent_loop_done');
                 // No more tool uses, we're done
-                this.#telemetryController.emitAgencticLoop_InvokeLLM(response.$metadata.requestId, conversationId, 'AgenticChat', undefined, undefined, 'Succeeded', this.#features.runtime.serverInfo.version ?? '', [llmLatency], this.#toolCallLatencies, this.#timeToFirstChunk, this.#timeBetweenChunks, session.pairProgrammingMode);
+                this.#telemetryController.emitAgencticLoop_InvokeLLM(response.$metadata.requestId, conversationId, 'AgenticChat', undefined, undefined, 'Succeeded', this.#features.runtime.serverInfo.version ?? '', llmLatency, this.#toolCallLatencies, this.#timeToFirstChunk, this.#timeBetweenChunks, session.pairProgrammingMode);
                 finalResult = result;
                 break;
             }
@@ -586,16 +584,16 @@ class AgenticChatController {
                 metric.setDimension('requestIds', metric.metric.requestIds);
                 const toolNames = this.#toolUseLatencies.map(item => item.toolName);
                 const toolUseIds = this.#toolUseLatencies.map(item => item.toolUseId);
-                this.#telemetryController.emitAgencticLoop_InvokeLLM(response.$metadata.requestId, conversationId, 'AgenticChatWithToolUse', toolNames ?? undefined, toolUseIds ?? undefined, 'Succeeded', this.#features.runtime.serverInfo.version ?? '', [llmLatency], this.#toolCallLatencies, this.#timeToFirstChunk, this.#timeBetweenChunks, session.pairProgrammingMode);
+                this.#telemetryController.emitAgencticLoop_InvokeLLM(response.$metadata.requestId, conversationId, 'AgenticChatWithToolUse', toolNames ?? undefined, toolUseIds ?? undefined, 'Succeeded', this.#features.runtime.serverInfo.version ?? '', llmLatency, this.#toolCallLatencies, this.#timeToFirstChunk, this.#timeBetweenChunks, session.pairProgrammingMode);
             }
             else {
                 // Send an error card to UI?
                 toolResults = pendingToolUses.map(toolUse => ({
                     toolUseId: toolUse.toolUseId,
-                    status: codewhisperer_streaming_client_1.ToolResultStatus.ERROR,
+                    status: codewhisperer_streaming_1.ToolResultStatus.ERROR,
                     content: [{ text: result.error }],
                 }));
-                this.#telemetryController.emitAgencticLoop_InvokeLLM(response.$metadata.requestId, conversationId, 'AgenticChatWithToolUse', undefined, undefined, 'Failed', this.#features.runtime.serverInfo.version ?? '', [llmLatency], this.#toolCallLatencies, this.#timeToFirstChunk, this.#timeBetweenChunks, session.pairProgrammingMode);
+                this.#telemetryController.emitAgencticLoop_InvokeLLM(response.$metadata.requestId, conversationId, 'AgenticChatWithToolUse', undefined, undefined, 'Failed', this.#features.runtime.serverInfo.version ?? '', llmLatency, this.#toolCallLatencies, this.#timeToFirstChunk, this.#timeBetweenChunks, session.pairProgrammingMode);
                 if (result.error.startsWith('ToolUse input is invalid JSON:')) {
                     content =
                         'Your toolUse input is incomplete, try again. If the error happens consistently, break this task down into multiple tool uses with smaller input. Do not apologize.';
@@ -622,6 +620,7 @@ class AgenticChatController {
      * @param request
      */
     truncateRequest(request) {
+        // TODO: Confirm if this limit applies to SendMessage and rename this constant
         let remainingCharacterBudget = constants_2.generateAssistantResponseInputLimit;
         if (!request?.conversationState?.currentMessage?.userInputMessage) {
             return remainingCharacterBudget;
@@ -783,6 +782,15 @@ class AgenticChatController {
                     default:
                         // Get original server and tool names from the mapping
                         const originalNames = mcpManager_1.McpManager.instance.getOriginalToolNames(toolUse.name);
+                        // Remove explanation field from toolUse.input for MCP tools
+                        // many MCP servers do not support explanation field and it will break the tool if this is altered
+                        if (originalNames &&
+                            toolUse.input &&
+                            typeof toolUse.input === 'object' &&
+                            'explanation' in toolUse.input) {
+                            const { explanation, ...inputWithoutExplanation } = toolUse.input;
+                            toolUse.input = inputWithoutExplanation;
+                        }
                         if (originalNames) {
                             const { serverName, toolName } = originalNames;
                             const def = mcpManager_1.McpManager.instance
@@ -930,7 +938,7 @@ class AgenticChatController {
                     if (err instanceof lsp_core_1.CancellationError) {
                         results.push({
                             toolUseId: toolUse.toolUseId,
-                            status: codewhisperer_streaming_client_1.ToolResultStatus.ERROR,
+                            status: codewhisperer_streaming_1.ToolResultStatus.ERROR,
                             content: [{ text: 'Command stopped by user' }],
                         });
                         continue;
@@ -961,7 +969,7 @@ class AgenticChatController {
                                 },
                                 status: {
                                     status: 'error',
-                                    icon: 'error',
+                                    icon: 'cancel-circle',
                                     text: 'Error',
                                     description: customerFacingError,
                                 },
@@ -1010,7 +1018,7 @@ class AgenticChatController {
                 this.#log(`Error running tool ${toolUse.name}:`, errMsg);
                 results.push({
                     toolUseId: toolUse.toolUseId,
-                    status: codewhisperer_streaming_client_1.ToolResultStatus.ERROR,
+                    status: codewhisperer_streaming_1.ToolResultStatus.ERROR,
                     content: [{ json: { error: err instanceof Error ? err.message : 'Unknown error' } }],
                 });
             }
@@ -1018,7 +1026,7 @@ class AgenticChatController {
         return results;
     }
     #shouldSendBackErrorContent(toolResult) {
-        if (toolResult.status === codewhisperer_streaming_client_1.ToolResultStatus.ERROR) {
+        if (toolResult.status === codewhisperer_streaming_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
@@ -1406,7 +1414,7 @@ class AgenticChatController {
             }
         }
         // Determine if this is a built-in tool or MCP tool
-        const isStandardTool = toolName !== undefined && ['executeBash', 'fsWrite', 'fsRead', 'listDirectory'].includes(toolName);
+        const isStandardTool = toolName !== undefined && this.#features.agent.getBuiltInToolNames().includes(toolName);
         if (isStandardTool) {
             return {
                 type: 'tool',
@@ -1648,7 +1656,7 @@ class AgenticChatController {
         metric.setDimension('codewhispererCustomizationArn', this.#customizationArn);
         metric.setDimension('languageServerVersion', this.#features.runtime.serverInfo.version);
         metric.setDimension('enabled', session.pairProgrammingMode);
-        const profileArn = AmazonQTokenServiceManager_1.AmazonQTokenServiceManager.getInstance().getActiveProfileArn();
+        const profileArn = this.#serviceManager?.getActiveProfileArn();
         if (profileArn) {
             this.#telemetryService.updateProfileArn(profileArn);
         }
@@ -1772,6 +1780,19 @@ class AgenticChatController {
                     };
                     return emptyChatResult;
                 }
+                if (err.message === `I am experiencing high traffic, please try again shortly.`) {
+                    this.#features.chat.sendChatUpdate({
+                        tabId: tabId,
+                        data: { messages: [{ messageId: 'modelThrottled' }] },
+                    });
+                    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);
@@ -1793,7 +1814,7 @@ class AgenticChatController {
         let response;
         let requestInput;
         try {
-            requestInput = await this.#triggerContext.getChatParamsFromTrigger(params, triggerContext, codewhisperer_streaming_client_1.ChatTriggerType.INLINE_CHAT, this.#customizationArn);
+            requestInput = await this.#triggerContext.getChatParamsFromTrigger(params, triggerContext, codewhisperer_streaming_1.ChatTriggerType.INLINE_CHAT, this.#customizationArn);
             if (!this.#serviceManager) {
                 throw new Error('amazonQServiceManager is not initialized');
             }
@@ -1993,6 +2014,15 @@ class AgenticChatController {
             return new server_interface_1.ResponseError(protocol_2.ErrorCodes.InternalError, sessionResult.error);
         }
         session.modelId = modelId;
+        // Get the saved pair programming mode from the database or default to true if not found
+        const savedPairProgrammingMode = this.#chatHistoryDb.getPairProgrammingMode();
+        session.pairProgrammingMode = savedPairProgrammingMode !== undefined ? savedPairProgrammingMode : true;
+        // Update the client with the initial pair programming mode
+        this.#features.chat.chatOptionsUpdate({
+            tabId: params.tabId,
+            // Type assertion to support pairProgrammingMode
+            ...(session.pairProgrammingMode !== undefined ? { pairProgrammingMode: session.pairProgrammingMode } : {}),
+        });
         if (success && session) {
             // Set the logging object on the session
             session.setLogging(this.#features.logging);
@@ -2194,8 +2224,8 @@ class AgenticChatController {
         }
         else if (!mode) {
             // Note: intentionally async.
-            AmazonQTokenServiceManager_1.AmazonQTokenServiceManager.getInstance()
-                .getCodewhispererService()
+            this.#serviceManager
+                ?.getCodewhispererService()
                 .getSubscriptionStatus(true)
                 .then(o => {
                 this.#log(`setPaidTierMode: getSubscriptionStatus: ${o.status} ${o.encodedVerificationUrl}`);
@@ -2236,7 +2266,7 @@ class AgenticChatController {
      * @returns `undefined` on success, or error message on failure.
      */
     async onManageSubscription(tabId, awsAccountId) {
-        const client = AmazonQTokenServiceManager_1.AmazonQTokenServiceManager.getInstance().getCodewhispererService();
+        const client = this.#serviceManager?.getCodewhispererService();
         if (!awsAccountId) {
             // If no awsAccountId was provided:
             // 1. Check if the user is subscribed.
@@ -2245,7 +2275,7 @@ class AgenticChatController {
             //
             // Note: intentionally async.
             client
-                .getSubscriptionStatus()
+                ?.getSubscriptionStatus()
                 .then(o => {
                 this.#log(`onManageSubscription: getSubscriptionStatus: ${o.status} ${o.encodedVerificationUrl}`);
                 if (o.status !== 'none') {
@@ -2343,7 +2373,7 @@ class AgenticChatController {
             return;
         }
     }
-    async #processGenerateAssistantResponseResponseWithTimeout(response, metric, chatResultStream, session, contextList) {
+    async #processAgenticChatResponseWithTimeout(response, metric, chatResultStream, session, contextList) {
         const abortController = new AbortController();
         let timeoutId;
         const timeoutPromise = new Promise((_, reject) => {
@@ -2353,7 +2383,7 @@ class AgenticChatController {
             }, constants_2.responseTimeoutMs);
         });
         const streamWriter = chatResultStream.getResultStreamWriter();
-        const processResponsePromise = this.#processGenerateAssistantResponseResponse(response, metric, chatResultStream, streamWriter, session, contextList, abortController.signal);
+        const processResponsePromise = this.#processAgenticChatResponse(response, metric, chatResultStream, streamWriter, session, contextList, abortController.signal);
         try {
             const result = await Promise.race([processResponsePromise, timeoutPromise]);
             clearTimeout(timeoutId);
@@ -2418,7 +2448,7 @@ class AgenticChatController {
             }
         }
     }
-    async #processGenerateAssistantResponseResponse(response, metric, chatResultStream, streamWriter, session, contextList, abortSignal) {
+    async #processAgenticChatResponse(response, metric, chatResultStream, streamWriter, session, contextList, abortSignal) {
         const requestId = response.$metadata.requestId;
         const chatEventParser = new agenticChatEventParser_1.AgenticChatEventParser(requestId, metric, this.#features.logging);
         // Display context transparency list once at the beginning of response
@@ -2429,12 +2459,21 @@ class AgenticChatController {
         }
         const toolUseStartTimes = {};
         const toolUseLoadingTimeouts = {};
+        let chatEventStream = undefined;
+        if ('generateAssistantResponseResponse' in response) {
+            chatEventStream = response.generateAssistantResponseResponse;
+        }
+        else if ('sendMessageResponse' in response) {
+            chatEventStream = response.sendMessageResponse;
+        }
         let isEmptyResponse = true;
-        for await (const chatEvent of response.generateAssistantResponseResponse) {
+        for await (const chatEvent of chatEventStream) {
             isEmptyResponse = false;
             if (abortSignal?.aborted) {
                 throw new Error('Operation was aborted');
             }
+            // assistantResponseEvent is present in ChatResponseStream - used by both SendMessage and GenerateAssistantResponse
+            // https://code.amazon.com/packages/AWSVectorConsolasPlatformModel/blobs/mainline/--/model/types/conversation_types.smithy
             if (chatEvent.assistantResponseEvent) {
                 await this.#showUndoAllIfRequired(chatResultStream, session);
             }
@@ -2543,6 +2582,7 @@ class AgenticChatController {
         session.pairProgrammingMode = params.optionsValues['pair-programmer-mode'] === 'true';
         session.modelId = params.optionsValues['model-selection'];
         this.#chatHistoryDb.setModelId(session.modelId);
+        this.#chatHistoryDb.setPairProgrammingMode(session.pairProgrammingMode);
     }
     updateConfiguration = (newConfig) => {
         this.#customizationArn = newConfig.customizationArn;
@@ -2563,10 +2603,11 @@ class AgenticChatController {
         this.#paidTierMode = undefined;
     };
     #getTools(session) {
+        const builtInWriteTools = new Set(this.#features.agent.getBuiltInWriteToolNames());
         const allTools = this.#features.agent.getTools({ format: 'bedrock' });
         if (!(0, mcpUtils_1.enabledMCP)(this.#features.lsp.getClientInitializeParams())) {
             if (!session.pairProgrammingMode) {
-                return allTools.filter(tool => !['fsWrite', 'fsReplace', 'executeBash'].includes(tool.toolSpecification?.name || ''));
+                return allTools.filter(tool => !builtInWriteTools.has(tool.toolSpecification?.name || ''));
             }
             return allTools;
         }
@@ -2581,8 +2622,7 @@ class AgenticChatController {
             .getAllTools()
             .map(tool => (0, mcpUtils_1.createNamespacedToolName)(tool.serverName, tool.toolName, allNamespacedTools, tempMapping)));
         mcpManager_1.McpManager.instance.setToolNameMapping(tempMapping);
-        const writeToolNames = new Set(['fsWrite', 'fsReplace', 'executeBash']);
-        const restrictedToolNames = new Set([...mcpToolSpecNames, ...writeToolNames]);
+        const restrictedToolNames = new Set([...mcpToolSpecNames, ...builtInWriteTools]);
         const readOnlyTools = allTools.filter(tool => {
             const toolName = tool.toolSpecification.name;
             return !restrictedToolNames.has(toolName);