fix: preserve server-generated assistant message metadata

.changeset/fix-1087-message-metadata.md

@@ -0,0 +1,6 @@
+---
+'@tanstack/ai': patch
+'@tanstack/ai-persistence': patch
+---
+
+Preserve stable IDs and creation timestamps for server-generated assistant messages across persistence and hydration.

packages/ai-persistence/src/middleware.ts

@@ -265,6 +265,7 @@ interface RunStateEntry {
    * bubble in place.
    */
   streamingMessageId?: string
+  streamingMessageCreatedAt?: Date
 }
 
 const runState = new WeakMap<object, RunStateEntry>()
@@ -450,6 +451,7 @@ function finishedTranscript(
   messages: ReadonlyArray<ModelMessage>,
   info: FinishInfo,
   messageId: string | undefined,
+  createdAt: Date | undefined,
 ): Array<ModelMessage> {
   const transcript = [...messages]
   const last = transcript[transcript.length - 1]
@@ -464,6 +466,7 @@ function finishedTranscript(
       role: 'assistant',
       content: info.content,
       ...(messageId ? { id: messageId } : {}),
+      ...(createdAt ? { createdAt } : {}),
     })
   }
   return transcript
@@ -1536,11 +1539,21 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
       // regardless of snapshotStreaming — it's persisted onto the assistant
       // message so its identity survives hydrate and a reload resumes the same
       // bubble in place.
-      if (chunk.type === 'TEXT_MESSAGE_START') {
+      if (ctx.phase === 'modelStream') {
         const s = runState.get(ctx)
-        if (s) {
+        if (s && chunk.type === 'TEXT_MESSAGE_START') {
           s.streamingMessageId = chunk.messageId
+          s.streamingMessageCreatedAt = new Date()
           s.streamingText = ''
+        } else if (
+          s &&
+          chunk.type === 'TOOL_CALL_START' &&
+          typeof chunk.parentMessageId === 'string' &&
+          chunk.parentMessageId !== '' &&
+          s.streamingMessageId === undefined
+        ) {
+          s.streamingMessageId = chunk.parentMessageId
+          s.streamingMessageCreatedAt ??= new Date()
         }
       }
 
@@ -1571,6 +1584,9 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
                   ...(snapshotState.streamingMessageId
                     ? { id: snapshotState.streamingMessageId }
                     : {}),
+                  ...(snapshotState.streamingMessageCreatedAt
+                    ? { createdAt: snapshotState.streamingMessageCreatedAt }
+                    : {}),
                 },
               ])
             } catch {
@@ -1620,7 +1636,12 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
       // "finished" run whose transcript is missing the terminal turn.
       await messageStore.saveThread(
         ctx.threadId,
-        finishedTranscript(ctx.messages, info, state?.streamingMessageId),
+        finishedTranscript(
+          ctx.messages,
+          info,
+          state?.streamingMessageId,
+          state?.streamingMessageCreatedAt,
+        ),
       )
       await completeRun(runs, ctx.runId, info.usage)
       await commitPendingResumes(state, persistence.stores.interrupts)

packages/ai-persistence/tests/interrupts.test.ts

@@ -57,12 +57,13 @@ const runStarted = (): StreamChunk => ({
   timestamp: 1,
 })
 
-const toolStart = (): StreamChunk => ({
+const toolStart = (parentMessageId?: string): StreamChunk => ({
   type: EventType.TOOL_CALL_START,
   toolCallId: 'tool-call-1',
   toolCallName: 'clientSearch',
   toolName: 'clientSearch',
   timestamp: 1,
+  ...(parentMessageId ? { parentMessageId } : {}),
 })
 
 const toolArgs = (): StreamChunk => ({
@@ -105,8 +106,9 @@ const toolCallChunks = () => [
 async function persistClientToolTurn(
   persistence: ReturnType<typeof memoryPersistence>,
   tools: Array<Tool>,
+  chunks: Array<StreamChunk> = toolCallChunks(),
 ) {
-  const first = mockAdapter([toolCallChunks()])
+  const first = mockAdapter([chunks])
   await collect(
     chat({
       adapter: first.adapter,
@@ -174,6 +176,35 @@ describe('interrupt persistence', () => {
     ])
   })
 
+  it('keeps the stream messageId on an interrupted tool-call turn', async () => {
+    const persistence = memoryPersistence()
+    await persistClientToolTurn(
+      persistence,
+      [approvalClientTool('clientSearch')],
+      [
+        runStarted(),
+        {
+          type: EventType.TEXT_MESSAGE_START,
+          messageId: 'stream-assistant',
+          role: 'assistant',
+          timestamp: 1,
+        },
+        toolStart('stream-assistant'),
+        toolArgs(),
+        toolCallFinished(),
+      ],
+    )
+
+    const thread = await persistence.stores.messages!.loadThread('t1')
+    const toolTurn = thread.find(
+      (message) =>
+        message.role === 'assistant' &&
+        message.toolCalls?.some((call) => call.id === 'tool-call-1'),
+    )
+    expect(toolTurn?.id).toBe('stream-assistant')
+    expect(toolTurn?.createdAt).toBeInstanceOf(Date)
+  })
+
   it('does not persist duplicate records before terminal interrupt outcome', async () => {
     const persistence = memoryPersistence()
     const create = vi.spyOn(persistence.stores.interrupts!, 'create')