feat: add text based tool updates for agentic-chat (#984)

Hweinstock · web-flow · commit 12dc8d767be4 · 2025-04-17T16:11:37.000-07:00
* feat: add bash tools to agentic-chat

* refactor: move tool processing out of main function

* feat: add tool results log after running

* feat: add responseStream obj to abstract streaming to client

* feat: add simple markdown ui

* refactor: clean up implementation of response stream and add tests

* refactor: clean up tests

* refactor: remove unused intermediate results

* test: update controller tests to ignore tool based updates

* refactor: rename response to chatresultstream

* refactor: clean up some loose ends

* fix: avoid exposing executeBash in current state

* docs: add comment for future steps

* docs: add note about formatting
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
@@ -439,7 +439,7 @@ describe('AgenticChatController', () => {
             assert.ok(secondCallArgs.conversationState?.history[1].assistantResponseMessage)
 
             // Verify the final result
-            assert.deepStrictEqual(chatResult, expectedCompleteChatResult)
+            assertChatResultsMatch(chatResult, expectedCompleteChatResult)
         })
 
         it('propagates tool execution errors to the model in toolResults', async () => {
@@ -574,7 +574,7 @@ describe('AgenticChatController', () => {
             }
 
             // Verify the final result includes both messages
-            assert.deepStrictEqual(chatResult, expectedErrorChatResult)
+            assertChatResultsMatch(chatResult, expectedErrorChatResult)
         })
 
         it('handles multiple iterations of tool uses with proper history updates', async () => {
@@ -764,7 +764,7 @@ describe('AgenticChatController', () => {
             )
 
             // Verify the final result
-            assert.deepStrictEqual(chatResult, expectedCompleteChatResult)
+            assertChatResultsMatch(chatResult, expectedCompleteChatResult)
         })
 
         it('returns help message if it is a help follow up action', async () => {
@@ -1655,3 +1655,16 @@ ${' '.repeat(8)}}
         sinon.assert.calledOnce(tabBarActionStub)
     })
 })
+
+// The body may include text-based progress updates from tool invocations.
+// We want to ignore these in the tests.
+function assertChatResultsMatch(actual: any, expected: ChatResult) {
+    if (actual?.body && expected?.body) {
+        assert.ok(
+            actual.body.endsWith(expected.body),
+            `Body should end with "${expected.body}"\nActual: "${actual.body}"`
+        )
+    }
+
+    assert.deepStrictEqual({ ...actual, body: undefined }, { ...expected, body: undefined })
+}
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts
@@ -69,6 +69,8 @@ import {
     ChatResultWithMetadata as AgenticChatResultWithMetadata,
 } from './agenticChatEventParser'
 import { ChatSessionService } from '../chat/chatSessionService'
+import { AgenticChatResultStream } from './agenticChatResultStream'
+import { executeToolMessage, toolErrorMessage, toolResultMessage } from './textFormatting'
 
 type ChatHandlers = Omit<
     LspHandlers<Chat>,
@@ -122,10 +124,21 @@ export class AgenticChatController implements ChatHandlers {
         return this.#tabBarController.onConversationClick(params)
     }
 
+    async #sendProgressToClient(chunk: ChatResult | string, partialResultToken?: string | number) {
+        if (!isNullish(partialResultToken)) {
+            await this.#features.lsp.sendProgress(chatRequestType, partialResultToken, chunk)
+        }
+    }
+
+    #getChatResultStream(partialResultToken?: string | number): AgenticChatResultStream {
+        return new AgenticChatResultStream(async (chunk: ChatResult | string) =>
+            this.#sendProgressToClient(chunk, partialResultToken)
+        )
+    }
+
     async onChatPrompt(params: ChatParams, token: CancellationToken): Promise<ChatResult | ResponseError<ChatResult>> {
         // Phase 1: Initial Setup - This happens only once
         const maybeDefaultResponse = getDefaultChatResponse(params.prompt.prompt)
-
         if (maybeDefaultResponse) {
             return maybeDefaultResponse
         }
@@ -151,7 +164,7 @@ export class AgenticChatController implements ChatHandlers {
         })
 
         const conversationIdentifier = session?.conversationId ?? 'New conversation'
-
+        const chatResultStream = this.#getChatResultStream(params.partialResultToken)
         try {
             // Get the initial request input
             const initialRequestInput = await this.#prepareRequestInput(params, session, triggerContext)
@@ -161,7 +174,7 @@ export class AgenticChatController implements ChatHandlers {
                 initialRequestInput,
                 session,
                 metric,
-                params.partialResultToken,
+                chatResultStream,
                 conversationIdentifier,
                 token
             )
@@ -173,7 +186,8 @@ export class AgenticChatController implements ChatHandlers {
                 params,
                 metric,
                 triggerContext,
-                isNewConversation
+                isNewConversation,
+                chatResultStream
             )
         } catch (err) {
             return this.#handleRequestError(err, params.tabId, metric)
@@ -214,15 +228,14 @@ export class AgenticChatController implements ChatHandlers {
         initialRequestInput: GenerateAssistantResponseCommandInput,
         session: ChatSessionService,
         metric: Metric<CombinedConversationEvent>,
-        partialResultToken?: string | number,
+        chatResultStream: AgenticChatResultStream,
         conversationIdentifier?: string,
         token?: CancellationToken
     ): Promise<Result<AgenticChatResultWithMetadata, string>> {
         let currentRequestInput = { ...initialRequestInput }
         let finalResult: Result<AgenticChatResultWithMetadata, string> | null = null
         let iterationCount = 0
         const maxIterations = 10 // Safety limit to prevent infinite loops
-
         metric.recordStart()
 
         while (iterationCount < maxIterations) {
@@ -247,7 +260,7 @@ export class AgenticChatController implements ChatHandlers {
                     cwsprChatResponseCode: response.$metadata.httpStatusCode,
                     cwsprChatMessageId: response.$metadata.requestId,
                 }),
-                partialResultToken
+                chatResultStream
             )
 
             // Store the conversation ID from the first response
@@ -267,7 +280,7 @@ export class AgenticChatController implements ChatHandlers {
             const currentMessage = currentRequestInput.conversationState?.currentMessage
 
             // Process tool uses and update the request input for the next iteration
-            const toolResults = await this.#processToolUses(pendingToolUses)
+            const toolResults = await this.#processToolUses(pendingToolUses, chatResultStream)
             currentRequestInput = this.#updateRequestInputWithToolResults(currentRequestInput, toolResults)
 
             if (!currentRequestInput.conversationState!.history) {
@@ -318,14 +331,17 @@ export class AgenticChatController implements ChatHandlers {
     /**
      * Processes tool uses by running the tools and collecting results
      */
-    async #processToolUses(toolUses: Array<ToolUse & { stop: boolean }>): Promise<ToolResult[]> {
+    async #processToolUses(
+        toolUses: Array<ToolUse & { stop: boolean }>,
+        chatResultStream: AgenticChatResultStream
+    ): Promise<ToolResult[]> {
         const results: ToolResult[] = []
 
         for (const toolUse of toolUses) {
             if (!toolUse.name || !toolUse.toolUseId) continue
 
             try {
-                this.#debug(`Running tool ${toolUse.name} with input:`, JSON.stringify(toolUse.input))
+                await chatResultStream.writeResultBlock({ body: `${executeToolMessage(toolUse)}` })
 
                 const result = await this.#features.agent.runTool(toolUse.name, toolUse.input)
                 let toolResultContent: ToolResultContentBlock
@@ -343,10 +359,13 @@ export class AgenticChatController implements ChatHandlers {
                     status: 'success',
                     content: [toolResultContent],
                 })
-
-                this.#debug(`Tool ${toolUse.name} completed with result:`, JSON.stringify(result))
+                await chatResultStream.writeResultBlock({ body: toolResultMessage(toolUse, result) })
             } catch (err) {
-                this.#log(`Error running tool ${toolUse.name}:`, err instanceof Error ? err.message : 'unknown error')
+                const errMsg = err instanceof Error ? err.message : 'unknown error'
+                await chatResultStream.writeResultBlock({
+                    body: toolErrorMessage(toolUse, errMsg),
+                })
+                this.#log(`Error running tool ${toolUse.name}:`, errMsg)
                 results.push({
                     toolUseId: toolUse.toolUseId,
                     status: 'error',
@@ -394,12 +413,12 @@ export class AgenticChatController implements ChatHandlers {
         params: ChatParams,
         metric: Metric<CombinedConversationEvent>,
         triggerContext: TriggerContext,
-        isNewConversation: boolean
+        isNewConversation: boolean,
+        chatResultStream: AgenticChatResultStream
     ): Promise<ChatResult | ResponseError<ChatResult>> {
         if (!result.success) {
             return new ResponseError<ChatResult>(LSPErrorCodes.RequestFailed, result.error)
         }
-
         const conversationId = session.conversationId
         this.#debug('Final session conversation id:', conversationId || '')
 
@@ -452,7 +471,7 @@ export class AgenticChatController implements ChatHandlers {
             })
         }
 
-        return result.data.chatResult
+        return chatResultStream.getResult()
     }
 
     /**
@@ -792,11 +811,11 @@ export class AgenticChatController implements ChatHandlers {
     async #processGenerateAssistantResponseResponse(
         response: GenerateAssistantResponseCommandOutput,
         metric: Metric<AddMessageEvent>,
-        partialResultToken?: string | number
+        chatResultStream: AgenticChatResultStream
     ): Promise<Result<AgenticChatResultWithMetadata, string>> {
         const requestId = response.$metadata.requestId!
         const chatEventParser = new AgenticChatEventParser(requestId, metric)
-
+        const streamWriter = chatResultStream.getResultStreamWriter()
         for await (const chatEvent of response.generateAssistantResponseResponse!) {
             const result = chatEventParser.processPartialEvent(chatEvent)
 
@@ -805,10 +824,9 @@ export class AgenticChatController implements ChatHandlers {
                 return result
             }
 
-            if (!isNullish(partialResultToken)) {
-                await this.#features.lsp.sendProgress(chatRequestType, partialResultToken, result.data.chatResult)
-            }
+            await streamWriter.write(result.data.chatResult)
         }
+        await streamWriter.close()
 
         metric.mergeWith({
             cwsprChatFullResponseLatency: metric.getTimeElapsed(),
@@ -837,9 +855,7 @@ export class AgenticChatController implements ChatHandlers {
                 return result
             }
 
-            if (!isNullish(partialResultToken)) {
-                await this.#features.lsp.sendProgress(chatRequestType, partialResultToken, result.data.chatResult)
-            }
+            await this.#sendProgressToClient(result.data.chatResult, partialResultToken)
         }
 
         return chatEventParser.getResult()
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.test.ts
@@ -0,0 +1,72 @@
+import * as assert from 'assert'
+import { ChatResult } from '@aws/language-server-runtimes/protocol'
+import { AgenticChatResultStream } from './agenticChatResultStream'
+import { TestFeatures } from '@aws/language-server-runtimes/testing'
+
+describe('agenticChatResponse', function () {
+    let output: (ChatResult | string)[] = []
+    const logging = new TestFeatures().logging
+    const sendProgress = async (s: ChatResult | string) => void output.push(s)
+    let chatResultStream: AgenticChatResultStream
+
+    beforeEach(function () {
+        output = []
+        chatResultStream = new AgenticChatResultStream(sendProgress)
+    })
+
+    it('combines all previous results on write', async function () {
+        await chatResultStream.writeResultBlock({ body: 'first' })
+        await chatResultStream.writeResultBlock({ body: 'second' })
+        await chatResultStream.writeResultBlock({ body: 'third' })
+
+        assert.deepStrictEqual(chatResultStream.getResult(), {
+            body: `first${AgenticChatResultStream.resultDelimiter}second${AgenticChatResultStream.resultDelimiter}third`,
+        })
+    })
+
+    it('inherits properties from the last result', async function () {
+        await chatResultStream.writeResultBlock({ body: 'first', messageId: '1' })
+        await chatResultStream.writeResultBlock({ body: 'second', messageId: '2' })
+        await chatResultStream.writeResultBlock({ body: 'third', messageId: '3' })
+
+        assert.deepStrictEqual(chatResultStream.getResult(), {
+            body: `first${AgenticChatResultStream.resultDelimiter}second${AgenticChatResultStream.resultDelimiter}third`,
+            messageId: '3',
+        })
+    })
+
+    it('streams the results to the chat', async function () {
+        const writer = chatResultStream.getResultStreamWriter()
+        await writer.write({ body: 'f' })
+        await writer.write({ body: 'fi' })
+        await writer.write({ body: 'firs' })
+        await writer.write({ body: 'first' })
+
+        await writer.close()
+
+        assert.deepStrictEqual(chatResultStream.getResult(), { body: 'first' })
+    })
+
+    it('combines results properly', async function () {
+        const writer = chatResultStream.getResultStreamWriter()
+        await writer.write({ body: 'f', messageId: '1' })
+        await writer.write({ body: 'fi', messageId: '1' })
+        await writer.write({ body: 'firs', messageId: '1' })
+        await writer.write({ body: 'first', messageId: '1' })
+
+        await writer.close()
+
+        assert.deepStrictEqual(chatResultStream.getResult(), {
+            messageId: '1',
+            body: 'first',
+        })
+    })
+
+    it('throws error if multiple stream writers are initialized', async function () {
+        const writer = chatResultStream.getResultStreamWriter()
+        assert.throws(() => chatResultStream.getResultStreamWriter(), Error)
+        await writer.close()
+
+        assert.ok(chatResultStream.getResultStreamWriter())
+    })
+})
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatResultStream.ts
@@ -0,0 +1,64 @@
+import { ChatResult } from '@aws/language-server-runtimes/protocol'
+
+interface ResultStreamWriter {
+    write(chunk: ChatResult): Promise<void>
+    close(): Promise<void>
+}
+
+/**
+ * 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.
+ */
+export class AgenticChatResultStream {
+    static readonly resultDelimiter = '\n\n'
+    #state = {
+        chatResultBlocks: [] as ChatResult[],
+        isLocked: false,
+    }
+    readonly #sendProgress: (newChatResult: ChatResult | string) => Promise<void>
+
+    constructor(sendProgress: (newChatResult: ChatResult | string) => Promise<void>) {
+        this.#sendProgress = sendProgress
+    }
+
+    getResult(): ChatResult {
+        return this.#joinResults(this.#state.chatResultBlocks)
+    }
+
+    #joinResults(chatResults: ChatResult[]): ChatResult {
+        // TODO: if we add ui elements to ChatResult in the response, we need to be more aware of how we combine them.
+        return chatResults.reduceRight((acc, c) => ({
+            ...acc,
+            body: c.body + AgenticChatResultStream.resultDelimiter + acc.body,
+        }))
+    }
+
+    async writeResultBlock(result: ChatResult) {
+        this.#state.chatResultBlocks.push(result)
+        await this.#sendProgress(this.getResult())
+    }
+
+    getResultStreamWriter(): ResultStreamWriter {
+        // Note: if write calls are not awaited, stream can be out-of-order.
+        if (this.#state.isLocked) {
+            throw new Error('AgenticChatResultStream: already locked')
+        }
+        this.#state.isLocked = true
+        let lastResult: ChatResult | undefined
+
+        return {
+            write: async (intermediateChatResult: ChatResult) => {
+                const combinedResult = this.#joinResults([...this.#state.chatResultBlocks, intermediateChatResult])
+                lastResult = intermediateChatResult
+                return await this.#sendProgress(combinedResult)
+            },
+            close: async () => {
+                if (lastResult) {
+                    this.#state.chatResultBlocks.push(lastResult)
+                }
+                this.#state.isLocked = false
+            },
+        }
+    }
+}
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/textFormatting.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/textFormatting.ts
