Skip to content

Commit adcbf5d

Browse files
mikemikimikemikemikimike
committed
fix(ai): preserve UI message ids in model conversion
1 parent 7d92296 commit adcbf5d

7 files changed

Lines changed: 205 additions & 8 deletions

File tree

‎.changeset/fix-ui-message-id-roundtrip.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+
Preserve stable UI message IDs when converting UI messages to model messages.

‎packages/ai/src/activities/chat/messages.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,7 @@ function buildUserOrToolMessage(uiMessage: UIMessage): ModelMessage {
184184
}
185185

186186
return {
187+
id: uiMessage.id,
187188
role: uiMessage.role as 'user' | 'assistant' | 'tool',
188189
content: collapseContentParts(contentParts),
189190
}
@@ -228,6 +229,8 @@ function isToolCallIncluded(part: ToolCallPart): boolean {
228229
* result is emitted as a tool message.
229230
*/
230231
function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
232+
// A single UI message can fan out into several model messages. Keep the
233+
// shared UI id on each one so persistence can retain the original identity.
231234
const messageList: Array<ModelMessage> = []
232235
let current = createSegment()
233236
let pendingThinking: Array<{ content: string; signature?: string }> = []
@@ -244,6 +247,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
244247

245248
if (hasContent || hasToolCalls) {
246249
messageList.push({
250+
id: uiMessage.id,
247251
role: 'assistant',
248252
content,
249253
...(hasToolCalls && { toolCalls: current.toolCalls }),
@@ -288,6 +292,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
288292
!emittedToolResultIds.has(part.toolCallId)
289293
) {
290294
messageList.push({
295+
id: uiMessage.id,
291296
role: 'tool',
292297
content: part.content,
293298
toolCallId: part.toolCallId,
@@ -347,6 +352,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
347352
// emit the concrete output regardless of approval metadata.
348353
if (part.output !== undefined && !emittedToolResultIds.has(part.id)) {
349354
messageList.push({
355+
id: uiMessage.id,
350356
role: 'tool',
351357
content: normalizeToolResult(part.output),
352358
toolCallId: part.id,
@@ -363,6 +369,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
363369
) {
364370
const approved = part.approval.approved
365371
messageList.push({
372+
id: uiMessage.id,
366373
role: 'tool',
367374
content: JSON.stringify({
368375
approved,
@@ -380,6 +387,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
380387
// If no messages were produced (e.g., empty parts), emit a minimal assistant message
381388
if (messageList.length === 0) {
382389
messageList.push({
390+
id: uiMessage.id,
383391
role: 'assistant',
384392
content: null,
385393
})

‎packages/ai/tests/message-converters.test.ts‎

Lines changed: 87 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ describe('Message Converters', () => {
2121

2222
expect(result).toEqual([
2323
{
24+
id: 'msg-1',
2425
role: 'user',
2526
content: 'Hello',
2627
},
@@ -41,6 +42,7 @@ describe('Message Converters', () => {
4142

4243
expect(result).toEqual([
4344
{
45+
id: 'msg-1',
4446
role: 'user',
4547
content: 'Hello world!',
4648
},
@@ -345,6 +347,62 @@ describe('Message Converters', () => {
345347
expect(result[1]?.content).toBe('{"temp": 72}')
346348
})
347349

350+
it('should preserve the UI message id on every generated model message', () => {
351+
const uiMessage: UIMessage = {
352+
id: 'assistant-1',
353+
role: 'assistant',
354+
parts: [
355+
{ type: 'text', content: 'Let me check.' },
356+
{
357+
type: 'tool-call',
358+
id: 'tc-explicit',
359+
name: 'getWeather',
360+
arguments: '{}',
361+
state: 'input-complete',
362+
},
363+
{
364+
type: 'tool-result',
365+
toolCallId: 'tc-explicit',
366+
content: '{}',
367+
state: 'complete',
368+
},
369+
{ type: 'text', content: 'I found it.' },
370+
{
371+
type: 'tool-call',
372+
id: 'tc-output',
373+
name: 'getForecast',
374+
arguments: '{}',
375+
state: 'complete',
376+
output: { temperature: 72 },
377+
},
378+
{
379+
type: 'tool-call',
380+
id: 'tc-approval',
381+
name: 'deleteForecast',
382+
arguments: '{}',
383+
state: 'approval-responded',
384+
approval: {
385+
id: 'approval-tc-approval',
386+
needsApproval: true,
387+
approved: false,
388+
},
389+
},
390+
],
391+
}
392+
393+
const result = uiMessageToModelMessages(uiMessage)
394+
395+
expect(result).toHaveLength(5)
396+
expect(result.every((message) => message.id === uiMessage.id)).toBe(true)
397+
expect(result.map((message) => message.role)).toEqual([
398+
'assistant',
399+
'tool',
400+
'assistant',
401+
'tool',
402+
'tool',
403+
])
404+
})
405+
348406
it('should preserve interleaving of text, tool calls, and tool results', () => {
349407
const uiMessage: UIMessage = {
350408
id: 'msg-1',
@@ -1132,7 +1190,7 @@ describe('Message Converters', () => {
11321190

11331191
const result = convertMessagesToModelMessages(messages)
11341192

1135-
expect(result).toEqual([{ role: 'user', content: 'Hello' }])
1193+
expect(result).toEqual([{ id: 'msg-1', role: 'user', content: 'Hello' }])
11361194
})
11371195

11381196
it('should handle mixed UIMessage and ModelMessage array', () => {
@@ -1148,7 +1206,7 @@ describe('Message Converters', () => {
11481206
const result = convertMessagesToModelMessages(messages)
11491207

11501208
expect(result).toEqual([
1151-
{ role: 'user', content: 'Hello' },
1209+
{ id: 'msg-1', role: 'user', content: 'Hello' },
11521210
{ role: 'assistant', content: 'Hi there!' },
11531211
])
11541212
})
@@ -1361,8 +1419,8 @@ describe('Message Converters', () => {
13611419
describe('Round-trip symmetry: Model -> UI -> Model', () => {
13621420
it('should round-trip simple text messages', () => {
13631421
const original: Array<ModelMessage> = [
1364-
{ role: 'user', content: 'Hello' },
1365-
{ role: 'assistant', content: 'Hi there!' },
1422+
{ id: 'user-1', role: 'user', content: 'Hello' },
1423+
{ id: 'assistant-1', role: 'assistant', content: 'Hi there!' },
13661424
]
13671425

13681426
const uiMessages = modelMessagesToUIMessages(original)
@@ -1374,6 +1432,7 @@ describe('Message Converters', () => {
13741432
it('should round-trip assistant with toolCalls + tool result', () => {
13751433
const original: Array<ModelMessage> = [
13761434
{
1435+
id: 'assistant-1',
13771436
role: 'assistant',
13781437
content: 'Let me check.',
13791438
toolCalls: [
@@ -1385,6 +1444,7 @@ describe('Message Converters', () => {
13851444
],
13861445
},
13871446
{
1447+
id: 'assistant-1',
13881448
role: 'tool',
13891449
content: '{"temp": 72}',
13901450
toolCallId: 'tc-1',
@@ -1400,6 +1460,7 @@ describe('Message Converters', () => {
14001460
it('should round-trip multimodal content array', () => {
14011461
const original: Array<ModelMessage> = [
14021462
{
1463+
id: 'user-1',
14031464
role: 'user',
14041465
content: [
14051466
{ type: 'text', content: 'What is this?' },
@@ -1419,8 +1480,9 @@ describe('Message Converters', () => {
14191480

14201481
it('should round-trip multi-round tool conversation', () => {
14211482
const original: Array<ModelMessage> = [
1422-
{ role: 'user', content: 'Check guitars' },
1483+
{ id: 'user-1', role: 'user', content: 'Check guitars' },
14231484
{
1485+
id: 'assistant-1',
14241486
role: 'assistant',
14251487
content: 'Checking.',
14261488
toolCalls: [
@@ -1431,8 +1493,14 @@ describe('Message Converters', () => {
14311493
},
14321494
],
14331495
},
1434-
{ role: 'tool', content: '[{"id":7}]', toolCallId: 'tc-1' },
14351496
{
1497+
id: 'assistant-1',
1498+
role: 'tool',
1499+
content: '[{"id":7}]',
1500+
toolCallId: 'tc-1',
1501+
},
1502+
{
1503+
id: 'assistant-2',
14361504
role: 'assistant',
14371505
content: 'Found one!',
14381506
toolCalls: [
@@ -1444,11 +1512,16 @@ describe('Message Converters', () => {
14441512
],
14451513
},
14461514
{
1515+
id: 'assistant-2',
14471516
role: 'tool',
14481517
content: '{"recommended":true}',
14491518
toolCallId: 'tc-2',
14501519
},
1451-
{ role: 'assistant', content: 'Here is my recommendation.' },
1520+
{
1521+
id: 'assistant-3',
1522+
role: 'assistant',
1523+
content: 'Here is my recommendation.',
1524+
},
14521525
]
14531526

14541527
const uiMessages = modelMessagesToUIMessages(original)
@@ -1460,6 +1533,7 @@ describe('Message Converters', () => {
14601533
it('should round-trip assistant with null content and toolCalls', () => {
14611534
const original: Array<ModelMessage> = [
14621535
{
1536+
id: 'assistant-1',
14631537
role: 'assistant',
14641538
content: null,
14651539
toolCalls: [
@@ -1470,7 +1544,12 @@ describe('Message Converters', () => {
14701544
},
14711545
],
14721546
},
1473-
{ role: 'tool', content: '{"temp":72}', toolCallId: 'tc-1' },
1547+
{
1548+
id: 'assistant-1',
1549+
role: 'tool',
1550+
content: '{"temp":72}',
1551+
toolCallId: 'tc-1',
1552+
},
14741553
]
14751554

14761555
const uiMessages = modelMessagesToUIMessages(original)

‎packages/ai/tests/messages.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,7 @@ describe('convertMessagesToModelMessages — AG-UI dedup pre-pass', () => {
163163

164164
expect(convertMessagesToModelMessages(messages)).toEqual([
165165
{
166+
id: 'assistant-approval',
166167
role: 'assistant',
167168
content: null,
168169
toolCalls: [

‎testing/e2e/src/routeTree.gen.ts‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import { Route as ApiOpenaiUsageDetailsRouteImport } from './routes/api.openai-u
4444
import { Route as ApiOpenaiShellSkillsWireRouteImport } from './routes/api.openai-shell-skills-wire'
4545
import { Route as ApiMultimodalToolResultWireRouteImport } from './routes/api.multimodal-tool-result-wire'
4646
import { Route as ApiMiddlewareTestRouteImport } from './routes/api.middleware-test'
47+
import { Route as ApiMessageIdsRouteImport } from './routes/api.message-ids'
4748
import { Route as ApiMcpTestRouteImport } from './routes/api.mcp-test'
4849
import { Route as ApiMcpStatusTestRouteImport } from './routes/api.mcp-status-test'
4950
import { Route as ApiMcpServerRouteImport } from './routes/api.mcp-server'
@@ -257,6 +258,11 @@ const ApiMiddlewareTestRoute = ApiMiddlewareTestRouteImport.update({
257258
path: '/api/middleware-test',
258259
getParentRoute: () => rootRouteImport,
259260
} as any)
261+
const ApiMessageIdsRoute = ApiMessageIdsRouteImport.update({
262+
id: '/api/message-ids',
263+
path: '/api/message-ids',
264+
getParentRoute: () => rootRouteImport,
265+
} as any)
260266
const ApiMcpTestRoute = ApiMcpTestRouteImport.update({
261267
id: '/api/mcp-test',
262268
path: '/api/mcp-test',
@@ -454,6 +460,7 @@ export interface FileRoutesByFullPath {
454460
'/api/mcp-server': typeof ApiMcpServerRoute
455461
'/api/mcp-status-test': typeof ApiMcpStatusTestRoute
456462
'/api/mcp-test': typeof ApiMcpTestRoute
463+
'/api/message-ids': typeof ApiMessageIdsRoute
457464
'/api/middleware-test': typeof ApiMiddlewareTestRoute
458465
'/api/multimodal-tool-result-wire': typeof ApiMultimodalToolResultWireRoute
459466
'/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute
@@ -521,6 +528,7 @@ export interface FileRoutesByTo {
521528
'/api/mcp-server': typeof ApiMcpServerRoute
522529
'/api/mcp-status-test': typeof ApiMcpStatusTestRoute
523530
'/api/mcp-test': typeof ApiMcpTestRoute
531+
'/api/message-ids': typeof ApiMessageIdsRoute
524532
'/api/middleware-test': typeof ApiMiddlewareTestRoute
525533
'/api/multimodal-tool-result-wire': typeof ApiMultimodalToolResultWireRoute
526534
'/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute
@@ -589,6 +597,7 @@ export interface FileRoutesById {
589597
'/api/mcp-server': typeof ApiMcpServerRoute
590598
'/api/mcp-status-test': typeof ApiMcpStatusTestRoute
591599
'/api/mcp-test': typeof ApiMcpTestRoute
600+
'/api/message-ids': typeof ApiMessageIdsRoute
592601
'/api/middleware-test': typeof ApiMiddlewareTestRoute
593602
'/api/multimodal-tool-result-wire': typeof ApiMultimodalToolResultWireRoute
594603
'/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute
@@ -658,6 +667,7 @@ export interface FileRouteTypes {
658667
| '/api/mcp-server'
659668
| '/api/mcp-status-test'
660669
| '/api/mcp-test'
670+
| '/api/message-ids'
661671
| '/api/middleware-test'
662672
| '/api/multimodal-tool-result-wire'
663673
| '/api/openai-shell-skills-wire'
@@ -725,6 +735,7 @@ export interface FileRouteTypes {
725735
| '/api/mcp-server'
726736
| '/api/mcp-status-test'
727737
| '/api/mcp-test'
738+
| '/api/message-ids'
728739
| '/api/middleware-test'
729740
| '/api/multimodal-tool-result-wire'
730741
| '/api/openai-shell-skills-wire'
@@ -792,6 +803,7 @@ export interface FileRouteTypes {
792803
| '/api/mcp-server'
793804
| '/api/mcp-status-test'
794805
| '/api/mcp-test'
806+
| '/api/message-ids'
795807
| '/api/middleware-test'
796808
| '/api/multimodal-tool-result-wire'
797809
| '/api/openai-shell-skills-wire'
@@ -860,6 +872,7 @@ export interface RootRouteChildren {
860872
ApiMcpServerRoute: typeof ApiMcpServerRoute
861873
ApiMcpStatusTestRoute: typeof ApiMcpStatusTestRoute
862874
ApiMcpTestRoute: typeof ApiMcpTestRoute
875+
ApiMessageIdsRoute: typeof ApiMessageIdsRoute
863876
ApiMiddlewareTestRoute: typeof ApiMiddlewareTestRoute
864877
ApiMultimodalToolResultWireRoute: typeof ApiMultimodalToolResultWireRoute
865878
ApiOpenaiShellSkillsWireRoute: typeof ApiOpenaiShellSkillsWireRoute
@@ -1127,6 +1140,13 @@ declare module '@tanstack/react-router' {
11271140
preLoaderRoute: typeof ApiMiddlewareTestRouteImport
11281141
parentRoute: typeof rootRouteImport
11291142
}
1143+
'/api/message-ids': {
1144+
id: '/api/message-ids'
1145+
path: '/api/message-ids'
1146+
fullPath: '/api/message-ids'
1147+
preLoaderRoute: typeof ApiMessageIdsRouteImport
1148+
parentRoute: typeof rootRouteImport
1149+
}
11301150
'/api/mcp-test': {
11311151
id: '/api/mcp-test'
11321152
path: '/api/mcp-test'
@@ -1441,6 +1461,7 @@ const rootRouteChildren: RootRouteChildren = {
14411461
ApiMcpServerRoute: ApiMcpServerRoute,
14421462
ApiMcpStatusTestRoute: ApiMcpStatusTestRoute,
14431463
ApiMcpTestRoute: ApiMcpTestRoute,
1464+
ApiMessageIdsRoute: ApiMessageIdsRoute,
14441465
ApiMiddlewareTestRoute: ApiMiddlewareTestRoute,
14451466
ApiMultimodalToolResultWireRoute: ApiMultimodalToolResultWireRoute,
14461467
ApiOpenaiShellSkillsWireRoute: ApiOpenaiShellSkillsWireRoute,

‎testing/e2e/src/routes/api.message-ids.ts‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import { createFileRoute } from '@tanstack/react-router'
2+
import { convertMessagesToModelMessages } from '@tanstack/ai'
3+
import type { ModelMessage, UIMessage } from '@tanstack/ai'
4+
5+
/**
6+
* Provider-free harness for the UIMessage -> ModelMessage identity contract.
7+
* The route keeps this regression at the public server boundary without
8+
* starting a provider request or requiring an API key.
9+
*/
10+
export const Route = createFileRoute('/api/message-ids')({
11+
server: {
12+
handlers: {
13+
POST: async ({ request }) => {
14+
const body = (await request.json()) as {
15+
messages: Array<UIMessage | ModelMessage>
16+
}
17+
18+
return Response.json(convertMessagesToModelMessages(body.messages))
19+
},
20+
},
21+
},
22+
})

0 commit comments

Comments
 (0)