Skip to content

Commit 8a20b10

Browse files
jguoamzjguoamz
authored
fix(chat): Fix ChatStream error caused by tool error (aws#6954)
## Problem - Sometimes we receive error in sendAIResponse where `await ToolUtils.queueDescription(tool, chatStream)` throws error and we don't handle it currently - Minor refactoring to remove `tool-unavailble` custom action and `processUnavailableToolUseMessage` to re-use `processToolUseMessage` instead ## Solution - Store the error along with toolUse in chatSession so that we can handle it correctly in `processToolUseMessage` --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent 9ce08d2 commit 8a20b10

File tree

4 files changed

+112
-156
lines changed

4 files changed

+112
-156
lines changed

‎packages/core/src/codewhispererChat/clients/chat/v0/chat.ts‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,11 @@ import { createQDeveloperStreamingClient } from '../../../../shared/clients/qDev
1616
import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWrittenCodeTracker'
1717
import { PromptMessage } from '../../../controllers/chat/model'
1818

19+
export type ToolUseWithError = {
20+
toolUse: ToolUse
21+
error: Error | undefined
22+
}
23+
1924
export class ChatSession {
2025
private sessionId?: string
2126
/**
@@ -24,7 +29,7 @@ export class ChatSession {
2429
* _context = Additional context to be passed to the LLM for generating the response
2530
*/
2631
private _readFiles: string[] = []
27-
private _toolUse: ToolUse | undefined
32+
private _toolUseWithError: ToolUseWithError | undefined
2833
private _showDiffOnFileWrite: boolean = false
2934
private _context: PromptMessage['context']
3035
private _pairProgrammingModeOn: boolean = true
@@ -49,12 +54,12 @@ export class ChatSession {
4954
this._pairProgrammingModeOn = pairProgrammingModeOn
5055
}
5156

52-
public get toolUse(): ToolUse | undefined {
53-
return this._toolUse
57+
public get toolUseWithError(): ToolUseWithError | undefined {
58+
return this._toolUseWithError
5459
}
5560

56-
public setToolUse(toolUse: ToolUse | undefined) {
57-
this._toolUse = toolUse
61+
public setToolUseWithError(toolUseWithError: ToolUseWithError | undefined) {
62+
this._toolUseWithError = toolUseWithError
5863
}
5964

6065
public get context(): PromptMessage['context'] {

‎packages/core/src/codewhispererChat/controllers/chat/controller.ts‎

Lines changed: 56 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -675,69 +675,6 @@ export class ChatController {
675675
telemetry.ui_click.emit({ elementId: 'amazonq_createSavedPrompt' })
676676
}
677677

678-
private async processUnavailableToolUseMessage(message: CustomFormActionMessage) {
679-
const tabID = message.tabID
680-
if (!tabID) {
681-
return
682-
}
683-
this.editorContextExtractor
684-
.extractContextForTrigger('ChatMessage')
685-
.then(async (context) => {
686-
const triggerID = randomUUID()
687-
this.triggerEventsStorage.addTriggerEvent({
688-
id: triggerID,
689-
tabID: message.tabID,
690-
message: undefined,
691-
type: 'chat_message',
692-
context,
693-
})
694-
const session = this.sessionStorage.getSession(tabID)
695-
const toolUse = session.toolUse
696-
if (!toolUse || !toolUse.input) {
697-
return
698-
}
699-
session.setToolUse(undefined)
700-
701-
const toolResults: ToolResult[] = []
702-
703-
toolResults.push({
704-
content: [{ text: 'This tool is not an available tool in this mode' }],
705-
toolUseId: toolUse.toolUseId,
706-
status: ToolResultStatus.ERROR,
707-
})
708-
709-
await this.generateResponse(
710-
{
711-
message: '',
712-
trigger: ChatTriggerType.ChatMessage,
713-
query: undefined,
714-
codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,
715-
fileText: context?.focusAreaContext?.extendedCodeBlock ?? '',
716-
fileLanguage: context?.activeFileContext?.fileLanguage,
717-
filePath: context?.activeFileContext?.filePath,
718-
matchPolicy: context?.activeFileContext?.matchPolicy,
719-
codeQuery: context?.focusAreaContext?.names,
720-
userIntent: undefined,
721-
customization: getSelectedCustomization(),
722-
toolResults: toolResults,
723-
origin: Origin.IDE,
724-
context: session.context ?? [],
725-
relevantTextDocuments: [],
726-
additionalContents: [],
727-
documentReferences: [],
728-
useRelevantDocuments: false,
729-
contextLengths: {
730-
...defaultContextLengths,
731-
},
732-
},
733-
triggerID
734-
)
735-
})
736-
.catch((e) => {
737-
this.processException(e, tabID)
738-
})
739-
}
740-
741678
private async processToolUseMessage(message: CustomFormActionMessage) {
742679
const tabID = message.tabID
743680
if (!tabID) {
@@ -756,59 +693,69 @@ export class ChatController {
756693
})
757694
this.messenger.sendAsyncEventProgress(tabID, true, '')
758695
const session = this.sessionStorage.getSession(tabID)
759-
const toolUse = session.toolUse
760-
if (!toolUse || !toolUse.input) {
696+
const toolUseWithError = session.toolUseWithError
697+
if (!toolUseWithError || !toolUseWithError.toolUse || !toolUseWithError.toolUse.input) {
761698
// Turn off AgentLoop flag if there's no tool use
762699
this.sessionStorage.setAgentLoopInProgress(tabID, false)
763700
return
764701
}
765-
session.setToolUse(undefined)
702+
session.setToolUseWithError(undefined)
766703

704+
const toolUse = toolUseWithError.toolUse
705+
const toolUseError = toolUseWithError.error
767706
const toolResults: ToolResult[] = []
768707

769-
const result = ToolUtils.tryFromToolUse(toolUse)
770-
if ('type' in result) {
771-
const tool: Tool = result
772-
773-
try {
774-
await ToolUtils.validate(tool)
775-
776-
const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
777-
requiresAcceptance: false,
778-
})
779-
const output = await ToolUtils.invoke(tool, chatStream)
780-
if (output.output.content.length > maxToolOutputCharacterLength) {
781-
throw Error(
782-
`Tool output exceeds maximum character limit of ${maxToolOutputCharacterLength}`
783-
)
708+
if (toolUseError) {
709+
toolResults.push({
710+
content: [{ text: toolUseError.message }],
711+
toolUseId: toolUse.toolUseId,
712+
status: ToolResultStatus.ERROR,
713+
})
714+
} else {
715+
const result = ToolUtils.tryFromToolUse(toolUse)
716+
if ('type' in result) {
717+
const tool: Tool = result
718+
719+
try {
720+
await ToolUtils.validate(tool)
721+
722+
const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
723+
requiresAcceptance: false,
724+
})
725+
const output = await ToolUtils.invoke(tool, chatStream)
726+
if (output.output.content.length > maxToolOutputCharacterLength) {
727+
throw Error(
728+
`Tool output exceeds maximum character limit of ${maxToolOutputCharacterLength}`
729+
)
730+
}
731+
732+
toolResults.push({
733+
content: [
734+
output.output.kind === OutputKind.Text
735+
? { text: output.output.content }
736+
: { json: output.output.content },
737+
],
738+
toolUseId: toolUse.toolUseId,
739+
status: ToolResultStatus.SUCCESS,
740+
})
741+
} catch (e: any) {
742+
toolResults.push({
743+
content: [{ text: e.message }],
744+
toolUseId: toolUse.toolUseId,
745+
status: ToolResultStatus.ERROR,
746+
})
784747
}
785-
786-
toolResults.push({
787-
content: [
788-
output.output.kind === OutputKind.Text
789-
? { text: output.output.content }
790-
: { json: output.output.content },
791-
],
792-
toolUseId: toolUse.toolUseId,
793-
status: ToolResultStatus.SUCCESS,
794-
})
795-
} catch (e: any) {
796-
toolResults.push({
797-
content: [{ text: e.message }],
798-
toolUseId: toolUse.toolUseId,
799-
status: ToolResultStatus.ERROR,
800-
})
748+
} else {
749+
const toolResult: ToolResult = result
750+
toolResults.push(toolResult)
801751
}
802-
} else {
803-
const toolResult: ToolResult = result
804-
toolResults.push(toolResult)
805-
}
806752

807-
if (toolUse.name === ToolType.FsWrite) {
808-
await vscode.commands.executeCommand(
809-
'vscode.open',
810-
vscode.Uri.file((toolUse.input as unknown as FsWriteParams).path)
811-
)
753+
if (toolUse.name === ToolType.FsWrite) {
754+
await vscode.commands.executeCommand(
755+
'vscode.open',
756+
vscode.Uri.file((toolUse.input as unknown as FsWriteParams).path)
757+
)
758+
}
812759
}
813760

814761
await this.generateResponse(
@@ -879,9 +826,6 @@ export class ChatController {
879826
case 'reject-shell-command':
880827
await this.rejectShellCommand(message)
881828
break
882-
case 'tool-unavailable':
883-
await this.processUnavailableToolUseMessage(message)
884-
break
885829
default:
886830
getLogger().warn(`Unhandled action: ${message.action.id}`)
887831
}
@@ -920,19 +864,19 @@ export class ChatController {
920864
await fs.mkdir(resultArtifactsDir)
921865
const tempFilePath = path.join(
922866
resultArtifactsDir,
923-
`temp-${path.basename((session.toolUse?.input as unknown as FsWriteParams).path)}`
867+
`temp-${path.basename((session.toolUseWithError?.toolUse.input as unknown as FsWriteParams).path)}`
924868
)
925869

926870
// If we have existing filePath copy file content from existing file to temporary file.
927-
const filePath = (session.toolUse?.input as any).path ?? message.filePath
871+
const filePath = (session.toolUseWithError?.toolUse.input as any).path ?? message.filePath
928872
const fileExists = await fs.existsFile(filePath)
929873
if (fileExists) {
930874
const fileContent = await fs.readFileText(filePath)
931875
await fs.writeFile(tempFilePath, fileContent)
932876
}
933877

934878
// Create a deep clone of the toolUse object and pass this toolUse to FsWrite tool execution to get the modified temporary file.
935-
const clonedToolUse = structuredClone(session.toolUse)
879+
const clonedToolUse = structuredClone(session.toolUseWithError?.toolUse)
936880
if (!clonedToolUse) {
937881
return
938882
}

‎packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts‎

Lines changed: 45 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -240,49 +240,56 @@ export class Messenger {
240240
toolUse.input = JSON.parse(toolUseInput)
241241
toolUse.toolUseId = cwChatEvent.toolUseEvent.toolUseId ?? ''
242242
toolUse.name = cwChatEvent.toolUseEvent.name ?? ''
243-
session.setToolUse(toolUse)
244243

245-
const availableToolsNames = (session.pairProgrammingModeOn ? tools : noWriteTools).map(
246-
(item) => item.toolSpecification?.name
247-
)
248-
if (!availableToolsNames.includes(toolUse.name)) {
249-
this.dispatcher.sendCustomFormActionMessage(
250-
new CustomFormActionMessage(tabID, {
251-
id: 'tool-unavailable',
252-
})
244+
let toolError = undefined
245+
try {
246+
const availableToolsNames = (session.pairProgrammingModeOn ? tools : noWriteTools).map(
247+
(item) => item.toolSpecification?.name
253248
)
254-
return
255-
}
256-
257-
const tool = ToolUtils.tryFromToolUse(toolUse)
258-
if ('type' in tool) {
259-
let changeList: Change[] | undefined = undefined
260-
if (tool.type === ToolType.FsWrite) {
261-
session.setShowDiffOnFileWrite(true)
262-
changeList = await tool.tool.getDiffChanges()
249+
if (!availableToolsNames.includes(toolUse.name)) {
250+
throw new Error(`Tool ${toolUse.name} is not available in the current mode`)
263251
}
264-
const validation = ToolUtils.requiresAcceptance(tool)
265-
const chatStream = new ChatStream(this, tabID, triggerID, toolUse, validation, changeList)
266-
await ToolUtils.queueDescription(tool, chatStream)
267-
268-
if (!validation.requiresAcceptance) {
269-
// Need separate id for read tool and safe bash command execution as 'run-shell-command' id is required to state in cwChatConnector.ts which will impact generic tool execution.
270-
if (tool.type === ToolType.ExecuteBash) {
271-
this.dispatcher.sendCustomFormActionMessage(
272-
new CustomFormActionMessage(tabID, {
273-
id: 'run-shell-command',
274-
})
275-
)
276-
} else {
277-
this.dispatcher.sendCustomFormActionMessage(
278-
new CustomFormActionMessage(tabID, {
279-
id: 'generic-tool-execution',
280-
})
281-
)
252+
const tool = ToolUtils.tryFromToolUse(toolUse)
253+
if ('type' in tool) {
254+
let changeList: Change[] | undefined = undefined
255+
if (tool.type === ToolType.FsWrite) {
256+
session.setShowDiffOnFileWrite(true)
257+
changeList = await tool.tool.getDiffChanges()
258+
}
259+
const validation = ToolUtils.requiresAcceptance(tool)
260+
const chatStream = new ChatStream(
261+
this,
262+
tabID,
263+
triggerID,
264+
toolUse,
265+
validation,
266+
changeList
267+
)
268+
await ToolUtils.queueDescription(tool, chatStream)
269+
270+
if (!validation.requiresAcceptance) {
271+
// Need separate id for read tool and safe bash command execution as 'run-shell-command' id is required to state in cwChatConnector.ts which will impact generic tool execution.
272+
if (tool.type === ToolType.ExecuteBash) {
273+
this.dispatcher.sendCustomFormActionMessage(
274+
new CustomFormActionMessage(tabID, {
275+
id: 'run-shell-command',
276+
})
277+
)
278+
} else {
279+
this.dispatcher.sendCustomFormActionMessage(
280+
new CustomFormActionMessage(tabID, {
281+
id: 'generic-tool-execution',
282+
})
283+
)
284+
}
282285
}
286+
} else {
287+
toolError = new Error('Tool not found')
283288
}
284-
} else {
285-
// TODO: Handle the error
289+
} catch (error: any) {
290+
toolError = error
291+
} finally {
292+
session.setToolUseWithError({ toolUse, error: toolError })
286293
}
287294
// TODO: Add a spinner component for fsWrite, previous implementation is causing lag in mynah UX.
288295
}

‎packages/core/src/shared/utilities/messageUtil.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ export interface MessageErrorInfo {
1313
}
1414

1515
export function extractErrorInfo(error: any): MessageErrorInfo {
16-
let errorMessage = 'Error reading chat stream.'
16+
let errorMessage = 'Error reading chat response stream: ' + error.message
1717
let statusCode = undefined
1818
let requestId = undefined
1919

0 commit comments

Comments
 (0)