Skip to content

Commit f0a7578

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

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
} else if (cwChatEvent.toolUseEvent?.stop === undefined && toolUseInput !== '') {
288295
// 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)