Merge branch 'aws:main' into feature/mcp

Author: bywang56

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -459,6 +459,30 @@ export class AgenticChatController implements ChatHandlers {
     #getPendingToolUses(toolUses: Record<string, ToolUse & { stop: boolean }>): Array<ToolUse & { stop: boolean }> {
         return Object.values(toolUses).filter(toolUse => toolUse.stop)
     }
+    /**
+     * Creates a promise that does not resolve until the user accepts or rejects the tool usage.
+     * @param toolUseId
+     * @param toolUseName
+     * @param resultStream
+     * @param promptBlockId id of approval block. This allows us to overwrite the buttons with 'accepted' or 'rejected' text.
+     * @param session
+     */
+    async waitForToolApproval(
+        toolUseId: string,
+        toolUseName: string,
+        resultStream: AgenticChatResultStream,
+        promptBlockId: number,
+        session: ChatSessionService
+    ) {
+        const deferred = this.#createDeferred()
+        session.setDeferredToolExecution(toolUseId, deferred.resolve, deferred.reject)
+        this.#log(`Prompting for tool approval for tool: ${toolUseName}`)
+        await deferred.promise
+        if (toolUseName === 'executeBash') {
+            // Note: we want to overwrite the button block because it already exists in the stream.
+            await resultStream.overwriteResultBlock(this.#getUpdateBashConfirmResult(toolUseId, true), promptBlockId)
+        }
+    }
 
     /**
      * Processes tool uses by running the tools and collecting results
@@ -474,7 +498,6 @@ export class AgenticChatController implements ChatHandlers {
         for (const toolUse of toolUses) {
             if (!toolUse.name || !toolUse.toolUseId) continue
             this.#triggerContext.getToolUseLookup().set(toolUse.toolUseId, toolUse)
-            let needsConfirmation
 
             try {
                 const { explanation } = toolUse.input as unknown as ExplanatoryParams
@@ -506,9 +529,15 @@ export class AgenticChatController implements ChatHandlers {
                             toolUse.input as unknown as ExecuteBashParams
                         )
                         if (requiresAcceptance) {
-                            needsConfirmation = true
                             const confirmationResult = this.#processExecuteBashConfirmation(toolUse, warning)
-                            await chatResultStream.writeResultBlock(confirmationResult)
+                            const buttonBlockId = await chatResultStream.writeResultBlock(confirmationResult)
+                            await this.waitForToolApproval(
+                                toolUse.toolUseId,
+                                toolUse.name,
+                                chatResultStream,
+                                buttonBlockId,
+                                session
+                            )
                         }
                         break
                     default:
@@ -520,16 +549,6 @@ export class AgenticChatController implements ChatHandlers {
                         break
                 }
 
-                if (needsConfirmation) {
-                    const deferred = this.#createDeferred()
-                    session.setDeferredToolExecution(toolUse.toolUseId, deferred.resolve, deferred.reject)
-                    this.#log(`Prompting for tool approval for tool: ${toolUse.name}`)
-                    await deferred.promise
-                    if (toolUse.name === 'executeBash') {
-                        await chatResultStream.writeResultBlock(this.#getUpdateBashConfirmResult(toolUse, true))
-                    }
-                }
-
                 const result = await this.#features.agent.runTool(toolUse.name, toolUse.input, token)
                 let toolResultContent: ToolResultContentBlock
 
@@ -572,7 +591,9 @@ export class AgenticChatController implements ChatHandlers {
                 // If we did not approve a tool to be used or the user stopped the response, bubble this up to interrupt agentic loop
                 if (CancellationError.isUserCancelled(err) || err instanceof ToolApprovalException) {
                     if (err instanceof ToolApprovalException) {
-                        await chatResultStream.writeResultBlock(this.#getUpdateBashConfirmResult(toolUse, false))
+                        await chatResultStream.writeResultBlock(
+                            this.#getUpdateBashConfirmResult(toolUse.toolUseId, false)
+                        )
                     }
                     throw err
                 }
@@ -592,9 +613,9 @@ export class AgenticChatController implements ChatHandlers {
         return results
     }
 
-    #getUpdateBashConfirmResult(toolUse: ToolUse, isAccept: boolean): ChatResult {
+    #getUpdateBashConfirmResult(toolUseId: string, isAccept: boolean): ChatResult {
         return {
-            messageId: toolUse.toolUseId,
+            messageId: toolUseId,
             type: 'tool',
             body: '',
             header: {

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.test.ts

@@ -3,6 +3,7 @@ import { ChatResult } from '@aws/language-server-runtimes/protocol'
 import { AgenticChatResultStream } from './agenticChatResultStream'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
 
+// TODO: renable this test suite and update the following tests.
 xdescribe('agenticChatResponse', function () {
     let output: (ChatResult | string)[] = []
     const logging = new TestFeatures().logging
@@ -69,4 +70,17 @@ xdescribe('agenticChatResponse', function () {
 
         assert.ok(chatResultStream.getResultStreamWriter())
     })
+
+    it('allows blocks to overwritten on id', async function () {
+        const first = await chatResultStream.writeResultBlock({ body: 'first' })
+        const second = await chatResultStream.writeResultBlock({ body: 'second' })
+        await chatResultStream.writeResultBlock({ body: 'third' })
+
+        await chatResultStream.overwriteResultBlock({ body: 'fourth' }, first)
+        await chatResultStream.overwriteResultBlock({ body: 'fifth' }, second)
+
+        assert.deepStrictEqual(chatResultStream.getResult(), {
+            body: `fourth${AgenticChatResultStream.resultDelimiter}fifth${AgenticChatResultStream.resultDelimiter}third`,
+        })
+    })
 })

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts

@@ -97,9 +97,25 @@ export class AgenticChatResultStream {
             }, result)
     }
 
-    async writeResultBlock(result: ChatMessage) {
+    /**
+     * Add a block to the message block store and send it to the client.
+     * @param result the blockId associated with the block such that it can be overwritten later
+     * @returns
+     */
+    async writeResultBlock(result: ChatMessage): Promise<number> {
         this.#state.chatResultBlocks.push(result)
         await this.#sendProgress(this.getResult(result.messageId))
+        return this.#state.chatResultBlocks.length - 1
+    }
+
+    /**
+     * Overwrites a specific blockId and re-sends the resulting blocks to the client.
+     * @param result
+     * @param blockId
+     */
+    async overwriteResultBlock(result: ChatMessage, blockId: number) {
+        this.#state.chatResultBlocks[blockId] = result
+        await this.#sendProgress(this.getResult(result.messageId))
     }
 
     getResultStreamWriter(): ResultStreamWriter {