Skip to content

Commit eddfbbd

Browse files
shaked-frametombeckenhamclaude
authored
fix: preserve anthropic assistant ids in tool-first streams (#480)
* fix: preserve anthropic assistant ids in tool-first streams * fix(adapters): set AG-UI parentMessageId on tool-first tool calls Bind TOOL_CALL_START to the stream's stable assistant message id via AG-UI `parentMessageId` in every text adapter, so a tool call that streams before any text no longer forces the assistant message id to change mid-stream (which destabilises UIMessage.id and can remount the message subtree in useChat). Fixes #477. Extends #480 (Anthropic only) to: - @tanstack/openai-base (Responses + Chat Completions) - @tanstack/ai-openrouter (Responses + Chat Completions) - @tanstack/ai-gemini (text + experimental text-interactions) - @tanstack/ai-ollama Adds a tool-first regression per adapter asserting the TOOL_CALL_START parentMessageId equals the TEXT_MESSAGE_START messageId. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 31de22b commit eddfbbd

18 files changed

Lines changed: 631 additions & 0 deletions

File tree

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
---
2+
'@tanstack/ai-anthropic': patch
3+
'@tanstack/openai-base': patch
4+
'@tanstack/ai-openrouter': patch
5+
'@tanstack/ai-gemini': patch
6+
'@tanstack/ai-ollama': patch
7+
---
8+
9+
Bind tool calls to the assistant message in tool-first streams by setting AG-UI's
10+
`parentMessageId` on `TOOL_CALL_START`.
11+
12+
When a provider streams a tool call **before** any text, the `StreamProcessor` had no
13+
active assistant message to attach it to, so it created one under a temporary local id.
14+
The later `TEXT_MESSAGE_START` then carried the real provider message id, forcing a
15+
mid-stream id change — which destabilizes `UIMessage.id` and can remount the message
16+
subtree in `useChat` (React list keys, etc.). See #477.
17+
18+
Every text adapter generates one stable assistant message id per stream and already uses
19+
it for `TEXT_MESSAGE_START`; they now also emit it as `parentMessageId` on
20+
`TOOL_CALL_START`. The processor reads `chunk.parentMessageId` (`?? active assistant id`)
21+
so the message is created with the correct id immediately and the subsequent
22+
`TEXT_MESSAGE_START` matches — no rename, no remount.
23+
24+
Fixed across all adapters that emit `TOOL_CALL_START` (Anthropic, OpenAI Responses +
25+
Chat Completions via `@tanstack/openai-base`, OpenRouter, Gemini including the
26+
experimental text-interactions adapter, and Ollama).

packages/ai-anthropic/src/adapters/text.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -999,6 +999,7 @@ export class AnthropicTextAdapter<
999999
toolCallId: existing.id,
10001000
toolCallName: existing.name,
10011001
toolName: existing.name,
1002+
parentMessageId: messageId,
10021003
model,
10031004
timestamp: Date.now(),
10041005
index: currentToolIndex,
@@ -1053,6 +1054,7 @@ export class AnthropicTextAdapter<
10531054
toolCallId: existing.id,
10541055
toolCallName: existing.name,
10551056
toolName: existing.name,
1057+
parentMessageId: messageId,
10561058
model,
10571059
timestamp: Date.now(),
10581060
index: currentToolIndex,

packages/ai-anthropic/tests/anthropic-adapter.test.ts

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1385,6 +1385,73 @@ describe('Anthropic stream processing', () => {
13851385
type: 'RUN_FINISHED',
13861386
})
13871387
})
1388+
1389+
it('emits parentMessageId on tool-first tool call chunks', async () => {
1390+
const mockStream = (async function* () {
1391+
yield {
1392+
type: 'content_block_start',
1393+
index: 0,
1394+
content_block: {
1395+
type: 'tool_use',
1396+
id: 'toolu_weather',
1397+
name: 'lookup_weather',
1398+
input: {},
1399+
},
1400+
}
1401+
yield {
1402+
type: 'content_block_delta',
1403+
index: 0,
1404+
delta: {
1405+
type: 'input_json_delta',
1406+
partial_json: '{"location":"Berlin"}',
1407+
},
1408+
}
1409+
yield { type: 'content_block_stop', index: 0 }
1410+
yield {
1411+
type: 'content_block_start',
1412+
index: 1,
1413+
content_block: { type: 'text', text: '' },
1414+
}
1415+
yield {
1416+
type: 'content_block_delta',
1417+
index: 1,
1418+
delta: { type: 'text_delta', text: 'It is sunny.' },
1419+
}
1420+
yield { type: 'content_block_stop', index: 1 }
1421+
yield {
1422+
type: 'message_delta',
1423+
delta: { stop_reason: 'end_turn' },
1424+
usage: { output_tokens: 7 },
1425+
}
1426+
yield { type: 'message_stop' }
1427+
})()
1428+
1429+
mocks.betaMessagesCreate.mockResolvedValueOnce(mockStream)
1430+
1431+
const adapter = createAdapter('claude-3-7-sonnet')
1432+
1433+
const chunks: StreamChunk[] = []
1434+
for await (const chunk of chat({
1435+
adapter,
1436+
messages: [{ role: 'user', content: 'What is the weather in Berlin?' }],
1437+
tools: [weatherTool],
1438+
})) {
1439+
chunks.push(chunk)
1440+
}
1441+
1442+
const textStart = chunks.find(
1443+
(chunk): chunk is Extract<StreamChunk, { type: 'TEXT_MESSAGE_START' }> =>
1444+
chunk.type === 'TEXT_MESSAGE_START',
1445+
)
1446+
const toolStart = chunks.find(
1447+
(chunk): chunk is Extract<StreamChunk, { type: 'TOOL_CALL_START' }> =>
1448+
chunk.type === 'TOOL_CALL_START',
1449+
)
1450+
1451+
expect(textStart).toBeDefined()
1452+
expect(toolStart).toBeDefined()
1453+
expect(toolStart?.parentMessageId).toBe(textStart?.messageId)
1454+
})
13881455
})
13891456

13901457
describe('Anthropic adapter error handling', () => {

packages/ai-gemini/src/adapters/text.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -446,6 +446,7 @@ export class GeminiTextAdapter<
446446
toolCallId,
447447
toolCallName: toolCallData.name,
448448
toolName: toolCallData.name,
449+
parentMessageId: messageId,
449450
model,
450451
timestamp: Date.now(),
451452
index: toolCallData.index,
@@ -524,6 +525,7 @@ export class GeminiTextAdapter<
524525
toolCallId,
525526
toolCallName: functionCall.name || '',
526527
toolName: functionCall.name || '',
528+
parentMessageId: messageId,
527529
model,
528530
timestamp: Date.now(),
529531
index: nextToolIndex - 1,

packages/ai-gemini/src/experimental/text-interactions/adapter.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1064,6 +1064,7 @@ async function* translateInteractionEvents(
10641064
toolCallId,
10651065
toolCallName: state.name,
10661066
toolName: state.name,
1067+
parentMessageId: messageId,
10671068
model,
10681069
timestamp,
10691070
index: state.index,

packages/ai-gemini/tests/gemini-adapter.test.ts

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,67 @@ describe('GeminiAdapter through AI', () => {
409409
})
410410
})
411411

412+
it('emits parentMessageId on tool-first tool calls matching the assistant message id', async () => {
413+
// A functionCall part arrives before any text. parentMessageId must bind the
414+
// tool call to the same assistant message id the eventual TEXT_MESSAGE_START
415+
// uses so the message id stays stable mid-stream (#477).
416+
const streamChunks = [
417+
{
418+
candidates: [
419+
{
420+
content: {
421+
parts: [
422+
{
423+
functionCall: {
424+
name: 'lookup_weather',
425+
args: { location: 'Berlin' },
426+
},
427+
},
428+
],
429+
},
430+
},
431+
],
432+
},
433+
{
434+
candidates: [
435+
{
436+
content: { parts: [{ text: 'It is sunny.' }] },
437+
finishReason: 'STOP',
438+
},
439+
],
440+
usageMetadata: {
441+
promptTokenCount: 4,
442+
candidatesTokenCount: 7,
443+
totalTokenCount: 11,
444+
},
445+
},
446+
]
447+
448+
mocks.generateContentStreamSpy.mockResolvedValue(createStream(streamChunks))
449+
450+
const adapter = createTextAdapter()
451+
const received: StreamChunk[] = []
452+
for await (const chunk of chat({
453+
adapter,
454+
messages: [{ role: 'user', content: 'What is the weather in Berlin?' }],
455+
tools: [weatherTool],
456+
})) {
457+
received.push(chunk)
458+
}
459+
460+
const textStart = received.find((c) => c.type === 'TEXT_MESSAGE_START')
461+
const toolStart = received.find((c) => c.type === 'TOOL_CALL_START')
462+
463+
expect(textStart?.type).toBe('TEXT_MESSAGE_START')
464+
expect(toolStart?.type).toBe('TOOL_CALL_START')
465+
if (
466+
textStart?.type === 'TEXT_MESSAGE_START' &&
467+
toolStart?.type === 'TOOL_CALL_START'
468+
) {
469+
expect(toolStart.parentMessageId).toBe(textStart.messageId)
470+
}
471+
})
472+
412473
it('merges consecutive user messages when tool results precede a follow-up user message', async () => {
413474
const streamChunks = [
414475
{

packages/ai-gemini/tests/text-interactions-adapter.test.ts

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -448,6 +448,77 @@ describe('GeminiTextInteractionsAdapter', () => {
448448
expect(finished.finishReason).toBe('tool_calls')
449449
})
450450

451+
it('emits parentMessageId on tool-first tool calls matching the assistant message id', async () => {
452+
// The function_call content arrives before any text. parentMessageId must
453+
// bind the tool call to the same assistant message id the eventual
454+
// TEXT_MESSAGE_START uses so the message id stays stable mid-stream (#477).
455+
mocks.interactionsCreateSpy.mockResolvedValue(
456+
mkStream([
457+
{
458+
event_type: 'interaction.start',
459+
interaction: { id: 'int_tf', status: 'in_progress' },
460+
},
461+
{
462+
event_type: 'content.start',
463+
index: 0,
464+
content: { type: 'function_call' },
465+
},
466+
{
467+
event_type: 'content.delta',
468+
index: 0,
469+
delta: {
470+
type: 'function_call',
471+
id: 'call_tf',
472+
name: 'lookup_weather',
473+
arguments: { location: 'Berlin' },
474+
},
475+
},
476+
{ event_type: 'content.stop', index: 0 },
477+
{
478+
event_type: 'content.start',
479+
index: 1,
480+
content: { type: 'text', text: '' },
481+
},
482+
{
483+
event_type: 'content.delta',
484+
index: 1,
485+
delta: { type: 'text', text: 'It is sunny.' },
486+
},
487+
{ event_type: 'content.stop', index: 1 },
488+
{
489+
event_type: 'interaction.complete',
490+
interaction: { id: 'int_tf', status: 'completed' },
491+
},
492+
]),
493+
)
494+
495+
const weatherTool: Tool = {
496+
name: 'lookup_weather',
497+
description: 'Return the weather for a location',
498+
}
499+
500+
const adapter = createAdapter()
501+
const chunks = await collectChunks(
502+
chat({
503+
adapter,
504+
messages: [{ role: 'user', content: 'Weather in Berlin?' }],
505+
tools: [weatherTool],
506+
}),
507+
)
508+
509+
const textStart = chunks.find((c) => c.type === 'TEXT_MESSAGE_START')
510+
const toolStart = chunks.find((c) => c.type === 'TOOL_CALL_START')
511+
512+
expect(textStart?.type).toBe('TEXT_MESSAGE_START')
513+
expect(toolStart?.type).toBe('TOOL_CALL_START')
514+
if (
515+
textStart?.type === 'TEXT_MESSAGE_START' &&
516+
toolStart?.type === 'TOOL_CALL_START'
517+
) {
518+
expect(toolStart.parentMessageId).toBe(textStart.messageId)
519+
}
520+
})
521+
451522
it('serializes tool results as function_result content blocks', async () => {
452523
mocks.interactionsCreateSpy.mockResolvedValue(
453524
mkStream([

packages/ai-ollama/src/adapters/text.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,7 @@ export class OllamaTextAdapter<TModel extends string> extends BaseTextAdapter<
226226
toolCallId,
227227
toolCallName: actualToolCall.function.name || '',
228228
toolName: actualToolCall.function.name || '',
229+
parentMessageId: messageId,
229230
model: chunk.model,
230231
timestamp: Date.now(),
231232
index: actualToolCall.function.index,

packages/ai-ollama/tests/text-adapter.test.ts

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,60 @@ describe('OllamaTextAdapter.chatStream (tool calls)', () => {
184184
expect(startChunk!.toolCallId).toBe('tc-123')
185185
})
186186

187+
it('emits parentMessageId on tool-first tool calls matching the assistant message id', async () => {
188+
// The tool call arrives before any text content. parentMessageId must bind
189+
// the tool call to the same assistant message id the eventual
190+
// TEXT_MESSAGE_START uses so the message id stays stable mid-stream (#477).
191+
chatMock.mockResolvedValueOnce(
192+
asyncIterable([
193+
{
194+
message: {
195+
role: 'assistant',
196+
content: '',
197+
tool_calls: [
198+
{
199+
id: 'tc-tf',
200+
function: { name: 'search', arguments: { q: 'cats' } },
201+
},
202+
],
203+
},
204+
done: false,
205+
},
206+
{
207+
message: { role: 'assistant', content: 'Found some cats.' },
208+
done: false,
209+
},
210+
{
211+
message: { role: 'assistant', content: '' },
212+
done: true,
213+
done_reason: 'stop',
214+
},
215+
]),
216+
)
217+
218+
const adapter = createOllamaChat('llama3.2')
219+
const chunks = await collectStream(
220+
adapter.chatStream({
221+
logger: testLogger,
222+
model: 'llama3.2',
223+
messages: [{ role: 'user', content: 'find cats' }],
224+
tools: [searchTool],
225+
}),
226+
)
227+
228+
const textStart = chunks.find((c) => c.type === 'TEXT_MESSAGE_START')
229+
const toolStart = chunks.find((c) => c.type === 'TOOL_CALL_START')
230+
231+
expect(textStart?.type).toBe('TEXT_MESSAGE_START')
232+
expect(toolStart?.type).toBe('TOOL_CALL_START')
233+
if (
234+
textStart?.type === 'TEXT_MESSAGE_START' &&
235+
toolStart?.type === 'TOOL_CALL_START'
236+
) {
237+
expect(toolStart.parentMessageId).toBe(textStart.messageId)
238+
}
239+
})
240+
187241
it('synthesises a tool-call id when Ollama omits one', async () => {
188242
chatMock.mockResolvedValueOnce(
189243
asyncIterable([

packages/ai-openrouter/src/adapters/responses-text.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1173,6 +1173,7 @@ export class OpenRouterResponsesTextAdapter<
11731173
toolCallId: item.id,
11741174
toolCallName: metadata.name,
11751175
toolName: metadata.name,
1176+
parentMessageId: aguiState.messageId,
11761177
model: model || options.model,
11771178
timestamp: Date.now(),
11781179
index: chunk.outputIndex ?? 0,
@@ -1286,6 +1287,7 @@ export class OpenRouterResponsesTextAdapter<
12861287
toolCallId: item.id,
12871288
toolCallName: metadata.name,
12881289
toolName: metadata.name,
1290+
parentMessageId: aguiState.messageId,
12891291
model: model || options.model,
12901292
timestamp: Date.now(),
12911293
index: metadata.index,
@@ -1361,6 +1363,7 @@ export class OpenRouterResponsesTextAdapter<
13611363
toolCallId: item.id,
13621364
toolCallName: metadata.name,
13631365
toolName: metadata.name,
1366+
parentMessageId: aguiState.messageId,
13641367
model: model || options.model,
13651368
timestamp: Date.now(),
13661369
index: metadata.index,

0 commit comments

Comments
 (0)