@@ -16,11 +16,27 @@ import type {
1616} from './types'
1717import type { ChatMessage } from '@shared/types/core/chat-message'
1818import { nanoid } from 'nanoid'
19- import type { ToolOutputGuard } from './toolOutputGuard'
19+ import type { ToolBatchOutputFitItem , ToolOutputGuard } from './toolOutputGuard'
2020import { buildTerminalErrorBlocks } from './messageStore'
2121
2222type PermissionType = 'read' | 'write' | 'all' | 'command'
2323
24+ type ExtractedSearchPayload = ReturnType < typeof extractSearchPayload >
25+
26+ type StagedToolResult = {
27+ toolCallId : string
28+ toolName : string
29+ toolArgs : string
30+ responseText : string
31+ isError : boolean
32+ offloadPath ?: string
33+ searchPayload : ExtractedSearchPayload
34+ rtkApplied ?: boolean
35+ rtkMode ?: 'rewrite' | 'direct' | 'bypass'
36+ rtkFallbackReason ?: string
37+ postHookKind : 'success' | 'failure'
38+ }
39+
2440type PermissionRequestLike = {
2541 toolName ?: string
2642 serverName ?: string
@@ -189,6 +205,90 @@ function updateToolCallBlock(
189205 }
190206}
191207
208+ function persistToolExecutionState ( io : IoParams , state : StreamState ) : void {
209+ if ( ! state . dirty ) {
210+ return
211+ }
212+
213+ flushBlocksToRenderer ( io , state . blocks )
214+ io . messageStore . updateAssistantContent ( io . messageId , state . blocks )
215+ state . dirty = false
216+ }
217+
218+ function applyFinalizedToolResults ( params : {
219+ stagedResults : StagedToolResult [ ]
220+ fittedResults : ToolBatchOutputFitItem [ ]
221+ conversation : ChatMessage [ ]
222+ state : StreamState
223+ io : IoParams
224+ hooks ?: ProcessHooks
225+ appendToConversation : boolean
226+ } ) : void {
227+ const { stagedResults, fittedResults, conversation, state, io, hooks, appendToConversation } =
228+ params
229+
230+ for ( let index = 0 ; index < stagedResults . length ; index += 1 ) {
231+ const stagedResult = stagedResults [ index ]
232+ const fittedResult = fittedResults [ index ]
233+ if ( ! fittedResult ) {
234+ continue
235+ }
236+
237+ if ( appendToConversation ) {
238+ conversation . push ( {
239+ role : 'tool' ,
240+ tool_call_id : fittedResult . toolCallId ,
241+ content : fittedResult . contextResponseText
242+ } )
243+ }
244+
245+ if ( ! fittedResult . downgraded && stagedResult . searchPayload ) {
246+ state . blocks . push ( stagedResult . searchPayload . block )
247+ for ( const result of stagedResult . searchPayload . results ) {
248+ io . messageStore . addSearchResult ( {
249+ sessionId : io . sessionId ,
250+ messageId : io . messageId ,
251+ searchId : result . searchId ,
252+ rank : typeof result . rank === 'number' ? result . rank : null ,
253+ result
254+ } )
255+ }
256+ }
257+
258+ updateToolCallBlock (
259+ state . blocks ,
260+ fittedResult . toolCallId ,
261+ fittedResult . responseText ,
262+ fittedResult . isError ,
263+ fittedResult . downgraded
264+ ? undefined
265+ : {
266+ rtkApplied : stagedResult . rtkApplied ,
267+ rtkMode : stagedResult . rtkMode ,
268+ rtkFallbackReason : stagedResult . rtkFallbackReason
269+ }
270+ )
271+
272+ if ( fittedResult . isError ) {
273+ hooks ?. onPostToolUseFailure ?.( {
274+ callId : stagedResult . toolCallId ,
275+ name : stagedResult . toolName ,
276+ params : stagedResult . toolArgs ,
277+ error : fittedResult . responseText
278+ } )
279+ } else if ( stagedResult . postHookKind === 'success' ) {
280+ hooks ?. onPostToolUse ?.( {
281+ callId : stagedResult . toolCallId ,
282+ name : stagedResult . toolName ,
283+ params : stagedResult . toolArgs ,
284+ response : fittedResult . responseText
285+ } )
286+ }
287+ }
288+
289+ state . dirty = true
290+ }
291+
192292function isPermissionType ( value : unknown ) : value is PermissionType {
193293 return value === 'read' || value === 'write' || value === 'all' || value === 'command'
194294}
@@ -450,6 +550,7 @@ export async function executeTools(
450550
451551 let executed = 0
452552 const pendingInteractions : PendingToolInteraction [ ] = [ ]
553+ const stagedResults : StagedToolResult [ ] = [ ]
453554
454555 for ( const tc of state . completedToolCalls ) {
455556 if ( io . abortSignal . aborted ) break
@@ -486,8 +587,7 @@ export async function executeTools(
486587 updateToolCallBlock ( state . blocks , tc . id , errorText , true )
487588 state . dirty = true
488589 executed += 1
489- flushBlocksToRenderer ( io , state . blocks )
490- io . messageStore . updateAssistantContent ( io . messageId , state . blocks )
590+ persistToolExecutionState ( io , state )
491591 continue
492592 }
493593
@@ -584,100 +684,83 @@ export async function executeTools(
584684 toolContext . name ,
585685 toolContext . serverName
586686 )
587- if ( searchPayload ) {
588- state . blocks . push ( searchPayload . block )
589- for ( const result of searchPayload . results ) {
590- io . messageStore . addSearchResult ( {
591- sessionId : io . sessionId ,
592- messageId : io . messageId ,
593- searchId : result . searchId ,
594- rank : typeof result . rank === 'number' ? result . rank : null ,
595- result
596- } )
597- }
598- }
599687
600688 const responseText = toolResponseToText ( toolRawData . content )
601- const guardedResult = await toolOutputGuard . guardToolOutput ( {
689+ const preparedResult = await toolOutputGuard . prepareToolOutput ( {
602690 sessionId : io . sessionId ,
603691 toolCallId : tc . id ,
604692 toolName : toolContext . name ,
605- rawContent : responseText ,
606- conversationMessages : conversation ,
607- toolDefinitions : tools ,
608- contextLength,
609- maxTokens
693+ rawContent : responseText
610694 } )
695+ const stagedResponseText =
696+ preparedResult . kind === 'tool_error' ? preparedResult . message : preparedResult . content
697+ const stagedIsError = preparedResult . kind === 'tool_error' || toolRawData . isError === true
611698
612- if ( guardedResult . kind === 'terminal_error' ) {
613- updateToolCallBlock ( state . blocks , tc . id , guardedResult . message , true )
614- hooks ?. onPostToolUseFailure ?.( {
615- callId : tc . id ,
616- name : tc . name ,
617- params : tc . arguments ,
618- error : guardedResult . message
619- } )
620- state . dirty = true
621- executed += 1
622- flushBlocksToRenderer ( io , state . blocks )
623- io . messageStore . updateAssistantContent ( io . messageId , state . blocks )
624- return {
625- executed,
626- pendingInteractions,
627- terminalError : guardedResult . message
628- }
629- }
630-
631- const isToolError = guardedResult . kind === 'tool_error' || toolRawData . isError === true
632- const toolMessageContent =
633- guardedResult . kind === 'tool_error' ? guardedResult . message : guardedResult . content
634- conversation . push ( {
635- role : 'tool' ,
636- tool_call_id : tc . id ,
637- content : toolMessageContent
638- } )
639- updateToolCallBlock ( state . blocks , tc . id , toolMessageContent , isToolError , {
699+ stagedResults . push ( {
700+ toolCallId : tc . id ,
701+ toolName : tc . name ,
702+ toolArgs : tc . arguments ,
703+ responseText : stagedResponseText ,
704+ isError : stagedIsError ,
705+ offloadPath : preparedResult . kind === 'ok' ? preparedResult . offloadPath : undefined ,
706+ searchPayload,
640707 rtkApplied : toolRawData . rtkApplied ,
641708 rtkMode : toolRawData . rtkMode ,
642- rtkFallbackReason : toolRawData . rtkFallbackReason
709+ rtkFallbackReason : toolRawData . rtkFallbackReason ,
710+ postHookKind : stagedIsError ? 'failure' : 'success'
643711 } )
644- if ( isToolError ) {
645- hooks ?. onPostToolUseFailure ?.( {
646- callId : tc . id ,
647- name : tc . name ,
648- params : tc . arguments ,
649- error : toolMessageContent
650- } )
651- } else {
652- hooks ?. onPostToolUse ?.( {
653- callId : tc . id ,
654- name : tc . name ,
655- params : tc . arguments ,
656- response : toolMessageContent
657- } )
658- }
712+ executed += 1
659713 } catch ( err ) {
660714 const errorText = err instanceof Error ? err . message : String ( err )
661- conversation . push ( {
662- role : 'tool' ,
663- tool_call_id : tc . id ,
664- content : `Error: ${ errorText } `
665- } )
666- updateToolCallBlock ( state . blocks , tc . id , `Error: ${ errorText } ` , true )
667- hooks ?. onPostToolUseFailure ?.( {
668- callId : tc . id ,
669- name : tc . name ,
670- params : tc . arguments ,
671- error : `Error: ${ errorText } `
715+ stagedResults . push ( {
716+ toolCallId : tc . id ,
717+ toolName : tc . name ,
718+ toolArgs : tc . arguments ,
719+ responseText : `Error: ${ errorText } ` ,
720+ isError : true ,
721+ searchPayload : null ,
722+ postHookKind : 'failure'
672723 } )
724+ executed += 1
673725 }
726+ }
727+
728+ if ( stagedResults . length > 0 ) {
729+ const fittedResults = await toolOutputGuard . fitToolBatchOutputs ( {
730+ conversationMessages : conversation ,
731+ results : stagedResults . map ( ( result ) => ( {
732+ toolCallId : result . toolCallId ,
733+ toolName : result . toolName ,
734+ responseText : result . responseText ,
735+ isError : result . isError ,
736+ offloadPath : result . offloadPath
737+ } ) ) ,
738+ toolDefinitions : tools ,
739+ contextLength,
740+ maxTokens
741+ } )
674742
675- state . dirty = true
676- executed += 1
677- flushBlocksToRenderer ( io , state . blocks )
678- io . messageStore . updateAssistantContent ( io . messageId , state . blocks )
743+ applyFinalizedToolResults ( {
744+ stagedResults,
745+ fittedResults : fittedResults . results ,
746+ conversation,
747+ state,
748+ io,
749+ hooks,
750+ appendToConversation : fittedResults . kind === 'ok'
751+ } )
752+ persistToolExecutionState ( io , state )
753+
754+ if ( fittedResults . kind === 'terminal_error' ) {
755+ return {
756+ executed,
757+ pendingInteractions,
758+ terminalError : fittedResults . message
759+ }
760+ }
679761 }
680762
763+ persistToolExecutionState ( io , state )
681764 return { executed, pendingInteractions }
682765}
683766
0 commit comments