Skip to content

Commit 4445410

Browse files
DiegoGBrisaautofix-ci[bot]AlemTuzlak
authored
fix(ai): emit TOOL_CALL_START/ARGS during continuation re-executions (#372)
* test(ai): add tests for missing TOOL_CALL_START/ARGS during continuation Continuation re-executions (pending tool calls from message history) only emit TOOL_CALL_END, skipping TOOL_CALL_START and TOOL_CALL_ARGS. This causes clients to store tool calls with empty arguments, leading to infinite re-execution loops. Two tests added: - Single pending tool call emits full START → ARGS → END sequence - Batch of pending tool calls each emit full START → ARGS → END sequence * fix(ai): emit TOOL_CALL_START and TOOL_CALL_ARGS during continuation re-executions Pending tool calls resumed from message history only emit TOOL_CALL_END, causing clients to store empty arguments and triggering infinite re-execution loops. Emit the full START → ARGS → END sequence using the arguments from the original ToolCall objects. * ci: apply automated fixes * test(ai): add test for mixed pending batch with server and client tools * ci: apply automated fixes * test(e2e): add continuation re-execution tool call arguments tests Add 4 E2E tests verifying tool call arguments are preserved during continuation re-executions. Without TOOL_CALL_START/ARGS emission, clients store tool calls with empty {} arguments. Tests: - Single client tool args preserved after continuation - Sequential client tool args preserved across multiple continuations - Parallel client tool args preserved in batch continuation - Mixed server+client tool args preserved in sequence Also populate the `args` field on synthetic TOOL_CALL_ARGS chunks emitted during continuation, matching adapter-emitted chunks. --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Alem Tuzlak <t.zlak@hotmail.com>
1 parent e8204b2 commit 4445410

5 files changed

Lines changed: 466 additions & 0 deletions

File tree

‎.changeset/fix-continuation-chunk-emission.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@tanstack/ai': patch
3+
---
4+
5+
Emit TOOL_CALL_START and TOOL_CALL_ARGS for pending tool calls during continuation re-executions

‎packages/typescript/ai/src/activities/chat/index.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -731,6 +731,13 @@ class TextEngine<
731731
needsClientExecution: executionResult.needsClientExecution,
732732
})
733733

734+
// Build args lookup so buildToolResultChunks can emit TOOL_CALL_START +
735+
// TOOL_CALL_ARGS before TOOL_CALL_END during continuation re-executions.
736+
const argsMap = new Map<string, string>()
737+
for (const tc of pendingToolCalls) {
738+
argsMap.set(tc.id, tc.function.arguments)
739+
}
740+
734741
if (
735742
executionResult.needsApproval.length > 0 ||
736743
executionResult.needsClientExecution.length > 0
@@ -739,6 +746,7 @@ class TextEngine<
739746
for (const chunk of this.buildToolResultChunks(
740747
executionResult.results,
741748
finishEvent,
749+
argsMap,
742750
)) {
743751
yield chunk
744752
}
@@ -765,6 +773,7 @@ class TextEngine<
765773
const toolResultChunks = this.buildToolResultChunks(
766774
executionResult.results,
767775
finishEvent,
776+
argsMap,
768777
)
769778

770779
for (const chunk of toolResultChunks) {
@@ -1080,12 +1089,35 @@ class TextEngine<
10801089
private buildToolResultChunks(
10811090
results: Array<ToolResult>,
10821091
finishEvent: RunFinishedEvent,
1092+
argsMap?: Map<string, string>,
10831093
): Array<StreamChunk> {
10841094
const chunks: Array<StreamChunk> = []
10851095

10861096
for (const result of results) {
10871097
const content = JSON.stringify(result.result)
10881098

1099+
// Emit TOOL_CALL_START + TOOL_CALL_ARGS before TOOL_CALL_END so that
1100+
// the client can reconstruct the full tool call during continuations.
1101+
if (argsMap) {
1102+
chunks.push({
1103+
type: 'TOOL_CALL_START',
1104+
timestamp: Date.now(),
1105+
model: finishEvent.model,
1106+
toolCallId: result.toolCallId,
1107+
toolName: result.toolName,
1108+
})
1109+
1110+
const args = argsMap.get(result.toolCallId) ?? '{}'
1111+
chunks.push({
1112+
type: 'TOOL_CALL_ARGS',
1113+
timestamp: Date.now(),
1114+
model: finishEvent.model,
1115+
toolCallId: result.toolCallId,
1116+
delta: args,
1117+
args,
1118+
})
1119+
}
1120+
10891121
chunks.push({
10901122
type: 'TOOL_CALL_END',
10911123
timestamp: Date.now(),

‎packages/typescript/ai/tests/chat.test.ts‎

Lines changed: 226 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -656,6 +656,232 @@ describe('chat()', () => {
656656
expect(executeSpy).not.toHaveBeenCalled()
657657
expect(calls).toHaveLength(1)
658658
})
659+
660+
it('should emit TOOL_CALL_START and TOOL_CALL_ARGS before TOOL_CALL_END for pending tool calls', async () => {
661+
const executeSpy = vi.fn().mockReturnValue({ temp: 72 })
662+
663+
const { adapter } = createMockAdapter({
664+
iterations: [
665+
// After pending tool is executed, the engine calls the adapter for the next response
666+
[
667+
ev.runStarted(),
668+
ev.textStart(),
669+
ev.textContent('72F in NYC'),
670+
ev.textEnd(),
671+
ev.runFinished('stop'),
672+
],
673+
],
674+
})
675+
676+
const stream = chat({
677+
adapter,
678+
messages: [
679+
{ role: 'user', content: 'Weather?' },
680+
{
681+
role: 'assistant',
682+
content: 'Let me check.',
683+
toolCalls: [
684+
{
685+
id: 'call_1',
686+
type: 'function' as const,
687+
function: { name: 'getWeather', arguments: '{"city":"NYC"}' },
688+
},
689+
],
690+
},
691+
// No tool result message -> pending!
692+
],
693+
tools: [serverTool('getWeather', executeSpy)],
694+
})
695+
696+
const chunks = await collectChunks(stream as AsyncIterable<StreamChunk>)
697+
698+
// Tool should have been executed
699+
expect(executeSpy).toHaveBeenCalledTimes(1)
700+
701+
// The continuation re-execution should emit the full chunk sequence:
702+
// TOOL_CALL_START -> TOOL_CALL_ARGS -> TOOL_CALL_END
703+
// Without the fix, only TOOL_CALL_END is emitted, causing the client
704+
// to store the tool call with empty arguments {}.
705+
const toolStartChunks = chunks.filter(
706+
(c) =>
707+
c.type === 'TOOL_CALL_START' && (c as any).toolCallId === 'call_1',
708+
)
709+
expect(toolStartChunks).toHaveLength(1)
710+
expect((toolStartChunks[0] as any).toolName).toBe('getWeather')
711+
712+
const toolArgsChunks = chunks.filter(
713+
(c) =>
714+
c.type === 'TOOL_CALL_ARGS' && (c as any).toolCallId === 'call_1',
715+
)
716+
expect(toolArgsChunks).toHaveLength(1)
717+
expect((toolArgsChunks[0] as any).delta).toBe('{"city":"NYC"}')
718+
expect((toolArgsChunks[0] as any).args).toBe('{"city":"NYC"}')
719+
720+
const toolEndChunks = chunks.filter(
721+
(c) => c.type === 'TOOL_CALL_END' && (c as any).toolCallId === 'call_1',
722+
)
723+
expect(toolEndChunks).toHaveLength(1)
724+
725+
// Verify ordering: START before ARGS before END
726+
const startIdx = chunks.indexOf(toolStartChunks[0]!)
727+
const argsIdx = chunks.indexOf(toolArgsChunks[0]!)
728+
const endIdx = chunks.indexOf(toolEndChunks[0]!)
729+
expect(startIdx).toBeLessThan(argsIdx)
730+
expect(argsIdx).toBeLessThan(endIdx)
731+
})
732+
733+
it('should emit TOOL_CALL_START and TOOL_CALL_ARGS for each pending tool call in a batch', async () => {
734+
const weatherSpy = vi.fn().mockReturnValue({ temp: 72 })
735+
const timeSpy = vi.fn().mockReturnValue({ time: '3pm' })
736+
737+
const { adapter } = createMockAdapter({
738+
iterations: [
739+
[
740+
ev.runStarted(),
741+
ev.textStart(),
742+
ev.textContent('Done.'),
743+
ev.textEnd(),
744+
ev.runFinished('stop'),
745+
],
746+
],
747+
})
748+
749+
const stream = chat({
750+
adapter,
751+
messages: [
752+
{ role: 'user', content: 'Weather and time?' },
753+
{
754+
role: 'assistant',
755+
content: '',
756+
toolCalls: [
757+
{
758+
id: 'call_weather',
759+
type: 'function' as const,
760+
function: { name: 'getWeather', arguments: '{"city":"NYC"}' },
761+
},
762+
{
763+
id: 'call_time',
764+
type: 'function' as const,
765+
function: { name: 'getTime', arguments: '{"tz":"EST"}' },
766+
},
767+
],
768+
},
769+
// No tool results -> both pending
770+
],
771+
tools: [
772+
serverTool('getWeather', weatherSpy),
773+
serverTool('getTime', timeSpy),
774+
],
775+
})
776+
777+
const chunks = await collectChunks(stream as AsyncIterable<StreamChunk>)
778+
779+
// Both tools should have been executed
780+
expect(weatherSpy).toHaveBeenCalledTimes(1)
781+
expect(timeSpy).toHaveBeenCalledTimes(1)
782+
783+
// Each pending tool should get the full START -> ARGS -> END sequence
784+
for (const { id, name, args } of [
785+
{ id: 'call_weather', name: 'getWeather', args: '{"city":"NYC"}' },
786+
{ id: 'call_time', name: 'getTime', args: '{"tz":"EST"}' },
787+
]) {
788+
const starts = chunks.filter(
789+
(c) => c.type === 'TOOL_CALL_START' && (c as any).toolCallId === id,
790+
)
791+
expect(starts).toHaveLength(1)
792+
expect((starts[0] as any).toolName).toBe(name)
793+
794+
const argChunks = chunks.filter(
795+
(c) => c.type === 'TOOL_CALL_ARGS' && (c as any).toolCallId === id,
796+
)
797+
expect(argChunks).toHaveLength(1)
798+
expect((argChunks[0] as any).delta).toBe(args)
799+
800+
const ends = chunks.filter(
801+
(c) => c.type === 'TOOL_CALL_END' && (c as any).toolCallId === id,
802+
)
803+
expect(ends).toHaveLength(1)
804+
805+
// Verify ordering
806+
const startIdx = chunks.indexOf(starts[0]!)
807+
const argsIdx = chunks.indexOf(argChunks[0]!)
808+
const endIdx = chunks.indexOf(ends[0]!)
809+
expect(startIdx).toBeLessThan(argsIdx)
810+
expect(argsIdx).toBeLessThan(endIdx)
811+
}
812+
})
813+
814+
it('should emit TOOL_CALL_START and TOOL_CALL_ARGS for the server tool in a mixed pending batch', async () => {
815+
const weatherSpy = vi.fn().mockReturnValue({ temp: 72 })
816+
817+
const { adapter } = createMockAdapter({ iterations: [] })
818+
819+
const stream = chat({
820+
adapter,
821+
messages: [
822+
{ role: 'user', content: 'Weather and notify?' },
823+
{
824+
role: 'assistant',
825+
content: '',
826+
toolCalls: [
827+
{
828+
id: 'call_server',
829+
type: 'function' as const,
830+
function: { name: 'getWeather', arguments: '{"city":"NYC"}' },
831+
},
832+
{
833+
id: 'call_client',
834+
type: 'function' as const,
835+
function: {
836+
name: 'showNotification',
837+
arguments: '{"message":"done"}',
838+
},
839+
},
840+
],
841+
},
842+
// No tool results -> both pending
843+
],
844+
tools: [
845+
serverTool('getWeather', weatherSpy),
846+
clientTool('showNotification'),
847+
],
848+
})
849+
850+
const chunks = await collectChunks(stream as AsyncIterable<StreamChunk>)
851+
852+
// Server tool should have executed
853+
expect(weatherSpy).toHaveBeenCalledTimes(1)
854+
855+
// The executed server tool should get the full START -> ARGS -> END
856+
const starts = chunks.filter(
857+
(c) =>
858+
c.type === 'TOOL_CALL_START' &&
859+
(c as any).toolCallId === 'call_server',
860+
)
861+
expect(starts).toHaveLength(1)
862+
expect((starts[0] as any).toolName).toBe('getWeather')
863+
864+
const argChunks = chunks.filter(
865+
(c) =>
866+
c.type === 'TOOL_CALL_ARGS' &&
867+
(c as any).toolCallId === 'call_server',
868+
)
869+
expect(argChunks).toHaveLength(1)
870+
expect((argChunks[0] as any).delta).toBe('{"city":"NYC"}')
871+
872+
const ends = chunks.filter(
873+
(c) =>
874+
c.type === 'TOOL_CALL_END' && (c as any).toolCallId === 'call_server',
875+
)
876+
expect(ends).toHaveLength(1)
877+
878+
// Verify ordering
879+
const startIdx = chunks.indexOf(starts[0]!)
880+
const argsIdx = chunks.indexOf(argChunks[0]!)
881+
const endIdx = chunks.indexOf(ends[0]!)
882+
expect(startIdx).toBeLessThan(argsIdx)
883+
expect(argsIdx).toBeLessThan(endIdx)
884+
})
659885
})
660886

661887
// ==========================================================================

0 commit comments

Comments
 (0)