Skip to content

Commit f1e6c68

Browse files
jguoamzjguoamz
committed
fix(chat): Fix ChatStream error caused by tool error
1 parent dfee9f7 commit f1e6c68

File tree

4 files changed

+117
-146
lines changed

4 files changed

+117
-146
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: 64 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -675,70 +675,7 @@ 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-
741-
private async processToolUseMessage(message: CustomFormActionMessage) {
678+
private async processToolUseMessage(message: CustomFormActionMessage, isToolAvailable: boolean = true) {
742679
const tabID = message.tabID
743680
if (!tabID) {
744681
return
@@ -756,59 +693,75 @@ 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 if (!isToolAvailable) {
715+
toolResults.push({
716+
content: [{ text: 'This tool is not an available tool in this mode' }],
717+
toolUseId: toolUseWithError.toolUse.toolUseId,
718+
status: ToolResultStatus.ERROR,
719+
})
720+
} else {
721+
const result = ToolUtils.tryFromToolUse(toolUse)
722+
if ('type' in result) {
723+
const tool: Tool = result
724+
725+
try {
726+
await ToolUtils.validate(tool)
727+
728+
const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
729+
requiresAcceptance: false,
730+
})
731+
const output = await ToolUtils.invoke(tool, chatStream)
732+
if (output.output.content.length > maxToolOutputCharacterLength) {
733+
throw Error(
734+
`Tool output exceeds maximum character limit of ${maxToolOutputCharacterLength}`
735+
)
736+
}
737+
738+
toolResults.push({
739+
content: [
740+
output.output.kind === OutputKind.Text
741+
? { text: output.output.content }
742+
: { json: output.output.content },
743+
],
744+
toolUseId: toolUse.toolUseId,
745+
status: ToolResultStatus.SUCCESS,
746+
})
747+
} catch (e: any) {
748+
toolResults.push({
749+
content: [{ text: e.message }],
750+
toolUseId: toolUse.toolUseId,
751+
status: ToolResultStatus.ERROR,
752+
})
784753
}
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-
})
754+
} else {
755+
const toolResult: ToolResult = result
756+
toolResults.push(toolResult)
801757
}
802-
} else {
803-
const toolResult: ToolResult = result
804-
toolResults.push(toolResult)
805-
}
806758

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-
)
759+
if (toolUse.name === ToolType.FsWrite) {
760+
await vscode.commands.executeCommand(
761+
'vscode.open',
762+
vscode.Uri.file((toolUse.input as unknown as FsWriteParams).path)
763+
)
764+
}
812765
}
813766

814767
await this.generateResponse(
@@ -880,7 +833,7 @@ export class ChatController {
880833
await this.rejectShellCommand(message)
881834
break
882835
case 'tool-unavailable':
883-
await this.processUnavailableToolUseMessage(message)
836+
await this.processToolUseMessage(message, false)
884837
break
885838
default:
886839
getLogger().warn(`Unhandled action: ${message.action.id}`)
@@ -920,19 +873,19 @@ export class ChatController {
920873
await fs.mkdir(resultArtifactsDir)
921874
const tempFilePath = path.join(
922875
resultArtifactsDir,
923-
`temp-${path.basename((session.toolUse?.input as unknown as FsWriteParams).path)}`
876+
`temp-${path.basename((session.toolUseWithError?.toolUse.input as unknown as FsWriteParams).path)}`
924877
)
925878

926879
// 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
880+
const filePath = (session.toolUseWithError?.toolUse.input as any).path ?? message.filePath
928881
const fileExists = await fs.existsFile(filePath)
929882
if (fileExists) {
930883
const fileContent = await fs.readFileText(filePath)
931884
await fs.writeFile(tempFilePath, fileContent)
932885
}
933886

934887
// 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)
888+
const clonedToolUse = structuredClone(session.toolUseWithError?.toolUse)
936889
if (!clonedToolUse) {
937890
return
938891
}

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

Lines changed: 42 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -240,49 +240,62 @@ 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

245244
const availableToolsNames = (session.pairProgrammingModeOn ? tools : noWriteTools).map(
246245
(item) => item.toolSpecification?.name
247246
)
248247
if (!availableToolsNames.includes(toolUse.name)) {
248+
session.setToolUseWithError({ toolUse, error: undefined })
249249
this.dispatcher.sendCustomFormActionMessage(
250250
new CustomFormActionMessage(tabID, {
251251
id: 'tool-unavailable',
252252
})
253253
)
254254
return
255255
}
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()
263-
}
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-
)
256+
let toolError = undefined
257+
try {
258+
const tool = ToolUtils.tryFromToolUse(toolUse)
259+
if ('type' in tool) {
260+
let changeList: Change[] | undefined = undefined
261+
if (tool.type === ToolType.FsWrite) {
262+
session.setShowDiffOnFileWrite(true)
263+
changeList = await tool.tool.getDiffChanges()
264+
}
265+
const validation = ToolUtils.requiresAcceptance(tool)
266+
const chatStream = new ChatStream(
267+
this,
268+
tabID,
269+
triggerID,
270+
toolUse,
271+
validation,
272+
changeList
273+
)
274+
await ToolUtils.queueDescription(tool, chatStream)
275+
276+
if (!validation.requiresAcceptance) {
277+
// 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.
278+
if (tool.type === ToolType.ExecuteBash) {
279+
this.dispatcher.sendCustomFormActionMessage(
280+
new CustomFormActionMessage(tabID, {
281+
id: 'run-shell-command',
282+
})
283+
)
284+
} else {
285+
this.dispatcher.sendCustomFormActionMessage(
286+
new CustomFormActionMessage(tabID, {
287+
id: 'generic-tool-execution',
288+
})
289+
)
290+
}
282291
}
292+
} else {
293+
toolError = new Error('Tool not found')
283294
}
284-
} else {
285-
// TODO: Handle the error
295+
} catch (error: any) {
296+
toolError = error
297+
} finally {
298+
session.setToolUseWithError({ toolUse, error: toolError })
286299
}
287300
} else if (cwChatEvent.toolUseEvent?.stop === undefined && toolUseInput !== '') {
288301
// This is for the case when writing tool is executed. The toolUseEvent is non stop but in toolUseInput is not empty. In this case we need show user the current spinner UI.

‎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)