fix(ai): preserve UI message IDs in model conversion

.changeset/fix-ui-message-id-roundtrip.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai': patch
+---
+
+Preserve stable UI message IDs when converting UI messages to model messages.

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

@@ -184,6 +184,7 @@ function buildUserOrToolMessage(uiMessage: UIMessage): ModelMessage {
   }
 
   return {
+    id: uiMessage.id,
     role: uiMessage.role as 'user' | 'assistant' | 'tool',
     content: collapseContentParts(contentParts),
   }
@@ -228,6 +229,8 @@ function isToolCallIncluded(part: ToolCallPart): boolean {
  * result is emitted as a tool message.
  */
 function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
+  // A single UI message can fan out into several model messages. Keep the
+  // shared UI id on each one so persistence can retain the original identity.
   const messageList: Array<ModelMessage> = []
   let current = createSegment()
   let pendingThinking: Array<{ content: string; signature?: string }> = []
@@ -244,6 +247,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
 
     if (hasContent || hasToolCalls) {
       messageList.push({
+        id: uiMessage.id,
         role: 'assistant',
         content,
         ...(hasToolCalls && { toolCalls: current.toolCalls }),
@@ -288,6 +292,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
           !emittedToolResultIds.has(part.toolCallId)
         ) {
           messageList.push({
+            id: uiMessage.id,
             role: 'tool',
             content: part.content,
             toolCallId: part.toolCallId,
@@ -347,6 +352,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
     // emit the concrete output regardless of approval metadata.
     if (part.output !== undefined && !emittedToolResultIds.has(part.id)) {
       messageList.push({
+        id: uiMessage.id,
         role: 'tool',
         content: normalizeToolResult(part.output),
         toolCallId: part.id,
@@ -363,6 +369,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
     ) {
       const approved = part.approval.approved
       messageList.push({
+        id: uiMessage.id,
         role: 'tool',
         content: JSON.stringify({
           approved,
@@ -380,6 +387,7 @@ function buildAssistantMessages(uiMessage: UIMessage): Array<ModelMessage> {
   // If no messages were produced (e.g., empty parts), emit a minimal assistant message
   if (messageList.length === 0) {
     messageList.push({
+      id: uiMessage.id,
       role: 'assistant',
       content: null,
     })

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

@@ -21,6 +21,7 @@ describe('Message Converters', () => {
 
       expect(result).toEqual([
         {
+          id: 'msg-1',
           role: 'user',
           content: 'Hello',
         },
@@ -41,6 +42,7 @@ describe('Message Converters', () => {
 
       expect(result).toEqual([
         {
+          id: 'msg-1',
           role: 'user',
           content: 'Hello world!',
         },
@@ -345,6 +347,74 @@ describe('Message Converters', () => {
       expect(result[1]?.content).toBe('{"temp": 72}')
     })
 
+    it('should preserve the UI message id on every generated model message', () => {
+      const uiMessage: UIMessage = {
+        id: 'assistant-1',
+        role: 'assistant',
+        parts: [
+          { type: 'text', content: 'Let me check.' },
+          {
+            type: 'tool-call',
+            id: 'tc-explicit',
+            name: 'getWeather',
+            arguments: '{}',
+            state: 'input-complete',
+          },
+          {
+            type: 'tool-result',
+            toolCallId: 'tc-explicit',
+            content: '{}',
+            state: 'complete',
+          },
+          { type: 'text', content: 'I found it.' },
+          {
+            type: 'tool-call',
+            id: 'tc-output',
+            name: 'getForecast',
+            arguments: '{}',
+            state: 'complete',
+            output: { temperature: 72 },
+          },
+          {
+            type: 'tool-call',
+            id: 'tc-approval',
+            name: 'deleteForecast',
+            arguments: '{}',
+            state: 'approval-responded',
+            approval: {
+              id: 'approval-tc-approval',
+              needsApproval: true,
+              approved: false,
+            },
+          },
+        ],
+      }
+
+      const result = uiMessageToModelMessages(uiMessage)
+
+      expect(result).toHaveLength(5)
+      expect(result.every((message) => message.id === uiMessage.id)).toBe(true)
+      expect(result.map((message) => message.role)).toEqual([
+        'assistant',
+        'tool',
+        'assistant',
+        'tool',
+        'tool',
+      ])
+    })
+
+    it('should preserve the UI message id on an empty assistant fallback', () => {
+      const uiMessage: UIMessage = {
+        id: 'assistant-empty',
+        role: 'assistant',
+        parts: [],
+      }
+
+      expect(uiMessageToModelMessages(uiMessage)).toEqual([
+        { id: uiMessage.id, role: 'assistant', content: null },
+      ])
+    })
+
     it('should preserve interleaving of text, tool calls, and tool results', () => {
       const uiMessage: UIMessage = {
         id: 'msg-1',
@@ -1132,7 +1202,7 @@ describe('Message Converters', () => {
 
       const result = convertMessagesToModelMessages(messages)
 
-      expect(result).toEqual([{ role: 'user', content: 'Hello' }])
+      expect(result).toEqual([{ id: 'msg-1', role: 'user', content: 'Hello' }])
     })
 
     it('should handle mixed UIMessage and ModelMessage array', () => {
@@ -1148,7 +1218,7 @@ describe('Message Converters', () => {
       const result = convertMessagesToModelMessages(messages)
 
       expect(result).toEqual([
-        { role: 'user', content: 'Hello' },
+        { id: 'msg-1', role: 'user', content: 'Hello' },
         { role: 'assistant', content: 'Hi there!' },
       ])
     })
@@ -1361,8 +1431,8 @@ describe('Message Converters', () => {
   describe('Round-trip symmetry: Model -> UI -> Model', () => {
     it('should round-trip simple text messages', () => {
       const original: Array<ModelMessage> = [
-        { role: 'user', content: 'Hello' },
-        { role: 'assistant', content: 'Hi there!' },
+        { id: 'user-1', role: 'user', content: 'Hello' },
+        { id: 'assistant-1', role: 'assistant', content: 'Hi there!' },
       ]
 
       const uiMessages = modelMessagesToUIMessages(original)
@@ -1374,6 +1444,7 @@ describe('Message Converters', () => {
     it('should round-trip assistant with toolCalls + tool result', () => {
       const original: Array<ModelMessage> = [
         {
+          id: 'assistant-1',
           role: 'assistant',
           content: 'Let me check.',
           toolCalls: [
@@ -1385,6 +1456,7 @@ describe('Message Converters', () => {
           ],
         },
         {
+          id: 'assistant-1',
           role: 'tool',
           content: '{"temp": 72}',
           toolCallId: 'tc-1',
@@ -1400,6 +1472,7 @@ describe('Message Converters', () => {
     it('should round-trip multimodal content array', () => {
       const original: Array<ModelMessage> = [
         {
+          id: 'user-1',
           role: 'user',
           content: [
             { type: 'text', content: 'What is this?' },
@@ -1419,8 +1492,9 @@ describe('Message Converters', () => {
 
     it('should round-trip multi-round tool conversation', () => {
       const original: Array<ModelMessage> = [
-        { role: 'user', content: 'Check guitars' },
+        { id: 'user-1', role: 'user', content: 'Check guitars' },
         {
+          id: 'assistant-1',
           role: 'assistant',
           content: 'Checking.',
           toolCalls: [
@@ -1431,8 +1505,14 @@ describe('Message Converters', () => {
             },
           ],
         },
-        { role: 'tool', content: '[{"id":7}]', toolCallId: 'tc-1' },
         {
+          id: 'assistant-1',
+          role: 'tool',
+          content: '[{"id":7}]',
+          toolCallId: 'tc-1',
+        },
+        {
+          id: 'assistant-2',
           role: 'assistant',
           content: 'Found one!',
           toolCalls: [
@@ -1444,11 +1524,16 @@ describe('Message Converters', () => {
           ],
         },
         {
+          id: 'assistant-2',
           role: 'tool',
           content: '{"recommended":true}',
           toolCallId: 'tc-2',
         },
-        { role: 'assistant', content: 'Here is my recommendation.' },
+        {
+          id: 'assistant-3',
+          role: 'assistant',
+          content: 'Here is my recommendation.',
+        },
       ]
 
       const uiMessages = modelMessagesToUIMessages(original)
@@ -1460,6 +1545,7 @@ describe('Message Converters', () => {
     it('should round-trip assistant with null content and toolCalls', () => {
       const original: Array<ModelMessage> = [
         {
+          id: 'assistant-1',
           role: 'assistant',
           content: null,
           toolCalls: [
@@ -1470,7 +1556,12 @@ describe('Message Converters', () => {
             },
           ],
         },
-        { role: 'tool', content: '{"temp":72}', toolCallId: 'tc-1' },
+        {
+          id: 'assistant-1',
+          role: 'tool',
+          content: '{"temp":72}',
+          toolCallId: 'tc-1',
+        },
       ]
 
       const uiMessages = modelMessagesToUIMessages(original)

packages/ai/tests/messages.test.ts

@@ -163,6 +163,7 @@ describe('convertMessagesToModelMessages — AG-UI dedup pre-pass', () => {
 
     expect(convertMessagesToModelMessages(messages)).toEqual([
       {
+        id: 'assistant-approval',
         role: 'assistant',
         content: null,
         toolCalls: [

testing/e2e/src/routeTree.gen.ts

@@ -44,6 +44,7 @@ import { Route as ApiOpenaiUsageDetailsRouteImport } from './routes/api.openai-u
 import { Route as ApiOpenaiShellSkillsWireRouteImport } from './routes/api.openai-shell-skills-wire'
 import { Route as ApiMultimodalToolResultWireRouteImport } from './routes/api.multimodal-tool-result-wire'
 import { Route as ApiMiddlewareTestRouteImport } from './routes/api.middleware-test'
+import { Route as ApiMessageIdsRouteImport } from './routes/api.message-ids'
 import { Route as ApiMcpTestRouteImport } from './routes/api.mcp-test'
 import { Route as ApiMcpStatusTestRouteImport } from './routes/api.mcp-status-test'
 import { Route as ApiMcpServerRouteImport } from './routes/api.mcp-server'
@@ -257,6 +258,11 @@ const ApiMiddlewareTestRoute = ApiMiddlewareTestRouteImport.update({
   path: '/api/middleware-test',
   getParentRoute: () => rootRouteImport,
 } as any)
+const ApiMessageIdsRoute = ApiMessageIdsRouteImport.update({
+  id: '/api/message-ids',
+  path: '/api/message-ids',
+  getParentRoute: () => rootRouteImport,
+} as any)
 const ApiMcpTestRoute = ApiMcpTestRouteImport.update({
   id: '/api/mcp-test',
   path: '/api/mcp-test',
@@ -454,6 +460,7 @@ export interface FileRoutesByFullPath {
   '/api/mcp-server': typeof ApiMcpServerRoute
   '/api/mcp-status-test': typeof ApiMcpStatusTestRoute
   '/api/mcp-test': typeof ApiMcpTestRoute
+  '/api/message-ids': typeof ApiMessageIdsRoute
   '/api/middleware-test': typeof ApiMiddlewareTestRoute
   '/api/multimodal-tool-result-wire': typeof ApiMultimodalToolResultWireRoute
   '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute
@@ -521,6 +528,7 @@ export interface FileRoutesByTo {
   '/api/mcp-server': typeof ApiMcpServerRoute
   '/api/mcp-status-test': typeof ApiMcpStatusTestRoute
   '/api/mcp-test': typeof ApiMcpTestRoute
+  '/api/message-ids': typeof ApiMessageIdsRoute
   '/api/middleware-test': typeof ApiMiddlewareTestRoute
   '/api/multimodal-tool-result-wire': typeof ApiMultimodalToolResultWireRoute
   '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute
@@ -589,6 +597,7 @@ export interface FileRoutesById {
   '/api/mcp-server': typeof ApiMcpServerRoute
   '/api/mcp-status-test': typeof ApiMcpStatusTestRoute
   '/api/mcp-test': typeof ApiMcpTestRoute
+  '/api/message-ids': typeof ApiMessageIdsRoute
   '/api/middleware-test': typeof ApiMiddlewareTestRoute
   '/api/multimodal-tool-result-wire': typeof ApiMultimodalToolResultWireRoute
   '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute
@@ -658,6 +667,7 @@ export interface FileRouteTypes {
     | '/api/mcp-server'
     | '/api/mcp-status-test'
     | '/api/mcp-test'
+    | '/api/message-ids'
     | '/api/middleware-test'
     | '/api/multimodal-tool-result-wire'
     | '/api/openai-shell-skills-wire'
@@ -725,6 +735,7 @@ export interface FileRouteTypes {
     | '/api/mcp-server'
     | '/api/mcp-status-test'
     | '/api/mcp-test'
+    | '/api/message-ids'
     | '/api/middleware-test'
     | '/api/multimodal-tool-result-wire'
     | '/api/openai-shell-skills-wire'
@@ -792,6 +803,7 @@ export interface FileRouteTypes {
     | '/api/mcp-server'
     | '/api/mcp-status-test'
     | '/api/mcp-test'
+    | '/api/message-ids'
     | '/api/middleware-test'
     | '/api/multimodal-tool-result-wire'
     | '/api/openai-shell-skills-wire'
@@ -860,6 +872,7 @@ export interface RootRouteChildren {
   ApiMcpServerRoute: typeof ApiMcpServerRoute
   ApiMcpStatusTestRoute: typeof ApiMcpStatusTestRoute
   ApiMcpTestRoute: typeof ApiMcpTestRoute
+  ApiMessageIdsRoute: typeof ApiMessageIdsRoute
   ApiMiddlewareTestRoute: typeof ApiMiddlewareTestRoute
   ApiMultimodalToolResultWireRoute: typeof ApiMultimodalToolResultWireRoute
   ApiOpenaiShellSkillsWireRoute: typeof ApiOpenaiShellSkillsWireRoute
@@ -1127,6 +1140,13 @@ declare module '@tanstack/react-router' {
       preLoaderRoute: typeof ApiMiddlewareTestRouteImport
       parentRoute: typeof rootRouteImport
     }
+    '/api/message-ids': {
+      id: '/api/message-ids'
+      path: '/api/message-ids'
+      fullPath: '/api/message-ids'
+      preLoaderRoute: typeof ApiMessageIdsRouteImport
+      parentRoute: typeof rootRouteImport
+    }
     '/api/mcp-test': {
       id: '/api/mcp-test'
       path: '/api/mcp-test'
@@ -1441,6 +1461,7 @@ const rootRouteChildren: RootRouteChildren = {
   ApiMcpServerRoute: ApiMcpServerRoute,
   ApiMcpStatusTestRoute: ApiMcpStatusTestRoute,
   ApiMcpTestRoute: ApiMcpTestRoute,
+  ApiMessageIdsRoute: ApiMessageIdsRoute,
   ApiMiddlewareTestRoute: ApiMiddlewareTestRoute,
   ApiMultimodalToolResultWireRoute: ApiMultimodalToolResultWireRoute,
   ApiOpenaiShellSkillsWireRoute: ApiOpenaiShellSkillsWireRoute,