11import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
22import { getInitialSessionState } from '@codebuff/common/types/session-state'
33import { promptSuccess } from '@codebuff/common/util/error'
4+ import { jsonToolResult } from '@codebuff/common/util/messages'
45import { beforeEach , describe , expect , it } from 'bun:test'
56
67import { mockFileContext } from './test-utils'
@@ -12,6 +13,10 @@ import type {
1213 AgentRuntimeScopedDeps ,
1314} from '@codebuff/common/types/contracts/agent-runtime'
1415import type { StreamChunk } from '@codebuff/common/types/contracts/llm'
16+ import type {
17+ AssistantMessage ,
18+ ToolMessage ,
19+ } from '@codebuff/common/types/messages/codebuff-message'
1520import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
1621
1722describe ( 'tool validation error handling' , ( ) => {
@@ -225,4 +230,127 @@ describe('tool validation error handling', () => {
225230 )
226231 expect ( errorEvents . length ) . toBe ( 0 )
227232 } )
233+
234+ it ( 'should preserve tool_call/tool_result ordering when custom tool setup is async' , async ( ) => {
235+ const toolName = 'delayed_custom_tool'
236+ const agentWithCustomTool : AgentTemplate = {
237+ ...testAgentTemplate ,
238+ toolNames : [ toolName , 'end_turn' ] ,
239+ }
240+
241+ const delayedToolCallChunk : StreamChunk = {
242+ type : 'tool-call' ,
243+ toolName,
244+ toolCallId : 'delayed-custom-tool-call-id' ,
245+ input : {
246+ query : 'test' ,
247+ } ,
248+ }
249+
250+ async function * mockStream ( ) {
251+ yield delayedToolCallChunk
252+ return promptSuccess ( 'mock-message-id' )
253+ }
254+
255+ const fileContextWithCustomTool = {
256+ ...mockFileContext ,
257+ customToolDefinitions : {
258+ [ toolName ] : {
259+ inputSchema : {
260+ type : 'object' ,
261+ properties : {
262+ query : { type : 'string' } ,
263+ } ,
264+ required : [ 'query' ] ,
265+ additionalProperties : false ,
266+ } ,
267+ endsAgentStep : false ,
268+ description : 'A delayed custom tool for ordering tests' ,
269+ } ,
270+ } ,
271+ }
272+
273+ const sessionState = getInitialSessionState ( fileContextWithCustomTool )
274+ const agentState = sessionState . mainAgentState
275+
276+ agentRuntimeImpl . requestMcpToolData = async ( ) => {
277+ // Force an async gap so tool_call emission happens after stream completion.
278+ await new Promise ( ( resolve ) => setTimeout ( resolve , 20 ) )
279+ return [ ]
280+ }
281+ agentRuntimeImpl . requestToolCall = async ( ) => ( {
282+ output : jsonToolResult ( { ok : true } ) ,
283+ } )
284+
285+ await processStream ( {
286+ ...agentRuntimeImpl ,
287+ agentContext : { } ,
288+ agentState,
289+ agentStepId : 'test-step-id' ,
290+ agentTemplate : agentWithCustomTool ,
291+ ancestorRunIds : [ ] ,
292+ clientSessionId : 'test-session' ,
293+ fileContext : fileContextWithCustomTool ,
294+ fingerprintId : 'test-fingerprint' ,
295+ fullResponse : '' ,
296+ localAgentTemplates : { 'test-agent' : agentWithCustomTool } ,
297+ messages : [ ] ,
298+ prompt : 'test prompt' ,
299+ repoId : undefined ,
300+ repoUrl : undefined ,
301+ runId : 'test-run-id' ,
302+ signal : new AbortController ( ) . signal ,
303+ stream : mockStream ( ) ,
304+ system : 'test system' ,
305+ tools : { } ,
306+ userId : 'test-user' ,
307+ userInputId : 'test-input-id' ,
308+ onCostCalculated : async ( ) => { } ,
309+ onResponseChunk : ( ) => { } ,
310+ } )
311+
312+ const assistantToolCallMessages = agentState . messageHistory . filter (
313+ ( m ) : m is AssistantMessage =>
314+ m . role === 'assistant' &&
315+ m . content . some ( ( c ) => c . type === 'tool-call' && c . toolName === toolName ) ,
316+ )
317+ const toolMessages = agentState . messageHistory . filter (
318+ ( m ) : m is ToolMessage => m . role === 'tool' && m . toolName === toolName ,
319+ )
320+
321+ expect ( assistantToolCallMessages . length ) . toBe ( 1 )
322+ expect ( toolMessages . length ) . toBe ( 1 )
323+
324+ const assistantToolCallPart = assistantToolCallMessages [ 0 ] . content . find (
325+ (
326+ c ,
327+ ) : c is Extract < AssistantMessage [ 'content' ] [ number ] , { type : 'tool-call' } > =>
328+ c . type === 'tool-call' && c . toolName === toolName ,
329+ )
330+ expect ( assistantToolCallPart ) . toBeDefined ( )
331+ expect ( toolMessages [ 0 ] . toolCallId ) . toBe ( assistantToolCallPart ! . toolCallId )
332+
333+ const assistantIndex = agentState . messageHistory . indexOf (
334+ assistantToolCallMessages [ 0 ] ,
335+ )
336+ const toolResultIndex = agentState . messageHistory . indexOf ( toolMessages [ 0 ] )
337+ expect ( assistantIndex ) . toBeGreaterThanOrEqual ( 0 )
338+ expect ( toolResultIndex ) . toBeGreaterThan ( assistantIndex )
339+
340+ const assistantToolCallIds = new Set (
341+ agentState . messageHistory . flatMap ( ( message ) => {
342+ if ( message . role !== 'assistant' ) {
343+ return [ ]
344+ }
345+ return message . content . flatMap ( ( part ) =>
346+ part . type === 'tool-call' ? [ part . toolCallId ] : [ ] ,
347+ )
348+ } ) ,
349+ )
350+ const orphanToolResults = agentState . messageHistory . filter (
351+ ( message ) : message is ToolMessage =>
352+ message . role === 'tool' && ! assistantToolCallIds . has ( message . toolCallId ) ,
353+ )
354+ expect ( orphanToolResults . length ) . toBe ( 0 )
355+ } )
228356} )
0 commit comments