diff --git a/.changeset/fix-1087-message-metadata.md b/.changeset/fix-1087-message-metadata.md new file mode 100644 index 000000000..57396e2e9 --- /dev/null +++ b/.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. diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 99df9d541..dbeeea292 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -265,6 +265,7 @@ interface RunStateEntry { * bubble in place. */ streamingMessageId?: string + streamingMessageCreatedAt?: Date } const runState = new WeakMap() @@ -450,6 +451,7 @@ function finishedTranscript( messages: ReadonlyArray, info: FinishInfo, messageId: string | undefined, + createdAt: Date | undefined, ): Array { 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( // 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( ...(snapshotState.streamingMessageId ? { id: snapshotState.streamingMessageId } : {}), + ...(snapshotState.streamingMessageCreatedAt + ? { createdAt: snapshotState.streamingMessageCreatedAt } + : {}), }, ]) } catch { @@ -1620,7 +1636,12 @@ export function withPersistence( // "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) diff --git a/packages/ai-persistence/tests/interrupts.test.ts b/packages/ai-persistence/tests/interrupts.test.ts index c848f58d3..ac46b31d0 100644 --- a/packages/ai-persistence/tests/interrupts.test.ts +++ b/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, tools: Array, + chunks: Array = 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') diff --git a/packages/ai-persistence/tests/with-persistence.test.ts b/packages/ai-persistence/tests/with-persistence.test.ts index 77d846e02..338431c92 100644 --- a/packages/ai-persistence/tests/with-persistence.test.ts +++ b/packages/ai-persistence/tests/with-persistence.test.ts @@ -1,6 +1,11 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { EventType, chat } from '@tanstack/ai' -import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { + AnyTextAdapter, + ModelMessage, + StreamChunk, + Tool, +} from '@tanstack/ai' import { memoryPersistence } from '../src/memory' import { withPersistence } from '../src/middleware' import { defineAIPersistence } from '../src/types' @@ -81,6 +86,25 @@ async function expectCollectRejects( await expect(collect(stream)).rejects.toThrow(pattern) } +function serverSearchTool(): Tool { + return { + name: 'search', + description: 'Search', + execute: () => ({ hits: [] }), + } +} + +function findAssistantToolCall( + messages: ReadonlyArray, + toolCallId: string, +) { + return messages.find( + (message) => + message.role === 'assistant' && + message.toolCalls?.some((call) => call.id === toolCallId), + ) +} + describe('withPersistence (state-only)', () => { it('completes the run and saves the transcript', async () => { const persistence = memoryPersistence() @@ -186,7 +210,12 @@ describe('withPersistence (state-only)', () => { // tagged with its stream messageId so a reload resumes the same bubble. expect(await persistence.stores.messages!.loadThread('t1')).toEqual([ { role: 'user', content: 'hi' }, - { role: 'assistant', content: 'Half a stor', id: 'm1' }, + expect.objectContaining({ + role: 'assistant', + content: 'Half a stor', + id: 'm1', + createdAt: expect.any(Date), + }), ]) }) @@ -220,8 +249,240 @@ describe('withPersistence (state-only)', () => { // `modelMessagesToUIMessages` reuses it and a reload can resume in place. expect(await persistence.stores.messages!.loadThread('t1')).toEqual([ { role: 'user', content: 'hi' }, - { role: 'assistant', content: 'hello', id: 'assistant-42' }, + expect.objectContaining({ + role: 'assistant', + content: 'hello', + id: 'assistant-42', + createdAt: expect.any(Date), + }), + ]) + }) + + it('persists a tool-call turn under the stream messageId', async () => { + const persistence = memoryPersistence() + const { adapter } = mockAdapter([ + [ + ev.runStarted(), + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'stream-assistant', + role: 'assistant', + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'search', + toolName: 'search', + parentMessageId: 'stream-assistant', + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: '{}', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'tool_calls', + timestamp: 1, + }, + ], + [ + ev.runStarted(), + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'stream-final', + role: 'assistant', + timestamp: 1, + }, + ev.text('done'), + ev.runFinished(), + ], + ]) + + await collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'search' }], + tools: [serverSearchTool()], + runId: 'r1', + threadId: 't1', + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + const thread = await persistence.stores.messages!.loadThread('t1') + const toolTurn = findAssistantToolCall(thread, 'call_1') + expect(toolTurn?.id).toBe('stream-assistant') + expect(toolTurn?.createdAt).toBeInstanceOf(Date) + expect(thread).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + role: 'assistant', + content: 'done', + id: 'stream-final', + }), + ]), + ) + }) + + it('stamps createdAt at TEXT_MESSAGE_START, not at iteration start', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')) + + try { + let call = 0 + const persistence = memoryPersistence() + const adapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': {}, + chatStream: () => { + call += 1 + return (async function* () { + if (call === 1) { + yield ev.runStarted() + await vi.advanceTimersByTimeAsync(5_000) + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: 'stream-assistant', + role: 'assistant', + timestamp: 1, + } satisfies StreamChunk + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'search', + toolName: 'search', + parentMessageId: 'stream-assistant', + timestamp: 1, + } satisfies StreamChunk + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: '{}', + timestamp: 1, + } satisfies StreamChunk + yield { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'tool_calls', + timestamp: 1, + } satisfies StreamChunk + return + } + yield ev.runStarted() + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: 'stream-final', + role: 'assistant', + timestamp: 1, + } satisfies StreamChunk + yield ev.text('done') + yield ev.runFinished() + })() + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } as unknown as AnyTextAdapter + + await collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'search' }], + tools: [serverSearchTool()], + runId: 'r1', + threadId: 't1', + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + const toolTurn = findAssistantToolCall( + await persistence.stores.messages!.loadThread('t1'), + 'call_1', + ) + expect(toolTurn?.createdAt).toEqual(new Date('2026-01-01T00:00:05.000Z')) + } finally { + vi.useRealTimers() + } + }) + + it('does not let structured-output TEXT_MESSAGE_START replace the agent-loop id', async () => { + const persistence = memoryPersistence() + const { adapter } = mockAdapter([ + [ + ev.runStarted(), + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'agent-tool', + role: 'assistant', + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'search', + toolName: 'search', + parentMessageId: 'agent-tool', + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: '{}', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'tool_calls', + timestamp: 1, + }, + ], + [ + ev.runStarted(), + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'agent-final', + role: 'assistant', + timestamp: 1, + }, + ev.text('hello'), + ev.runFinished(), + ], ]) + adapter.structuredOutput = async () => ({ + data: { name: 'Ada' }, + rawText: '{"name":"Ada"}', + }) + + await collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'extract' }], + tools: [serverSearchTool()], + runId: 'r1', + threadId: 't1', + stream: true, + outputSchema: { + type: 'object', + properties: { name: { type: 'string' } }, + }, + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + const thread = await persistence.stores.messages!.loadThread('t1') + const terminal = thread.find( + (message) => message.role === 'assistant' && message.content === 'hello', + ) + expect(terminal?.id).toBe('agent-final') }) it('records an interrupt and marks the run interrupted', async () => { diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index b884d5449..15fb215ba 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -729,6 +729,8 @@ class TextEngine< private streamStartTime = 0 private totalChunkCount = 0 private currentMessageId: string | null = null + private currentMessageCreatedAt: Date | null = null + private streamIdentityCaptured = false private accumulatedContent = '' private accumulatedThinking: Array<{ content: string; signature?: string }> = [] @@ -1268,6 +1270,8 @@ class TextEngine< private async beginIteration(): Promise { this.currentMessageId = this.createId('msg') + this.currentMessageCreatedAt = new Date() + this.streamIdentityCaptured = false this.accumulatedContent = '' this.accumulatedThinking = [] this.currentThinkingContent = '' @@ -1455,6 +1459,11 @@ class TextEngine< // eslint-disable-next-line @typescript-eslint/switch-exhaustiveness-check -- AG-UI EventType enum members vs string-literal case labels; default branch handles untraced events. switch (chunk.type) { // AG-UI Events + case 'TEXT_MESSAGE_START': + if (typeof chunk.messageId === 'string' && chunk.messageId !== '') { + this.captureStreamMessageIdentity(chunk.messageId) + } + break case 'TEXT_MESSAGE_CONTENT': this.handleTextMessageContentEvent(chunk) break @@ -1493,8 +1502,7 @@ class TextEngine< break default: - // RUN_STARTED, TEXT_MESSAGE_START, TEXT_MESSAGE_END, - // STATE_SNAPSHOT, STATE_DELTA, CUSTOM + // RUN_STARTED, TEXT_MESSAGE_END, STATE_SNAPSHOT, STATE_DELTA, CUSTOM // - no special handling needed in chat activity break } @@ -1513,7 +1521,22 @@ class TextEngine< this.middlewareCtx.accumulatedContent = this.accumulatedContent } + private captureStreamMessageIdentity(messageId: string): void { + this.currentMessageId = messageId + this.middlewareCtx.currentMessageId = messageId + if (!this.streamIdentityCaptured) { + this.currentMessageCreatedAt = new Date() + this.streamIdentityCaptured = true + } + } + private handleToolCallStartEvent(chunk: ToolCallStartEvent): void { + if ( + typeof chunk.parentMessageId === 'string' && + chunk.parentMessageId !== '' + ) { + this.captureStreamMessageIdentity(chunk.parentMessageId) + } this.toolCallManager.addToolCallStartEvent(chunk) } @@ -1954,6 +1977,8 @@ class TextEngine< role: 'assistant', content: this.accumulatedContent || null, toolCalls, + id: this.currentMessageId ?? undefined, + createdAt: this.currentMessageCreatedAt ?? undefined, ...(this.accumulatedThinking.length > 0 && { thinking: this.accumulatedThinking, }), diff --git a/packages/ai/tests/chat.test.ts b/packages/ai/tests/chat.test.ts index 24d7356d1..0e362875c 100644 --- a/packages/ai/tests/chat.test.ts +++ b/packages/ai/tests/chat.test.ts @@ -12,7 +12,7 @@ import { ev, serverTool, } from './test-utils' -import type { StreamChunk, Tool, UIMessage } from '../src/types' +import type { ModelMessage, StreamChunk, Tool, UIMessage } from '../src/types' import type { ChatMiddleware, ChatResumeToolState, @@ -3086,6 +3086,102 @@ describe('chat()', () => { expect(assistantToolMessage?.thinking).toEqual([ { content: 'Need inventory.', signature: 'sig-think-1' }, ]) + expect(assistantToolMessage?.id).toBeTruthy() + expect(assistantToolMessage?.createdAt).toBeInstanceOf(Date) + }) + + it('stamps the stream messageId on an assistant tool-call turn', async () => { + const toolSpy = vi.fn().mockReturnValue({ result: 'inventory' }) + + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart('stream-assistant'), + { + ...ev.toolStart('call_1', 'getInventory'), + parentMessageId: 'stream-assistant', + }, + ev.toolArgs('call_1', '{}'), + ev.runFinished('tool_calls'), + ], + [ + ev.runStarted(), + ev.textStart('stream-final'), + ev.textContent('Inventory loaded.'), + ev.textEnd('stream-final'), + ev.runFinished('stop'), + ], + ], + }) + + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Check inventory' }], + tools: [serverTool('getInventory', toolSpy)], + }) as AsyncIterable, + ) + + const assistantToolMessage = calls[1]!.messages.find( + (message: ModelMessage) => + message.role === 'assistant' && + message.toolCalls?.[0]?.id === 'call_1', + ) + + expect(assistantToolMessage?.id).toBe('stream-assistant') + expect(assistantToolMessage?.createdAt).toBeInstanceOf(Date) + }) + + it('stamps createdAt at TEXT_MESSAGE_START, not at iteration start', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')) + + try { + let call = 0 + const { adapter, calls } = createMockAdapter({ + chatStreamFn: async function* () { + call += 1 + if (call === 1) { + yield ev.runStarted() + await vi.advanceTimersByTimeAsync(5_000) + yield ev.textStart('stream-assistant') + yield { + ...ev.toolStart('call_1', 'getInventory'), + parentMessageId: 'stream-assistant', + } + yield ev.toolArgs('call_1', '{}') + yield ev.runFinished('tool_calls') + return + } + yield ev.runStarted() + yield ev.textStart('stream-final') + yield ev.textContent('done') + yield ev.textEnd('stream-final') + yield ev.runFinished('stop') + }, + }) + + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Check inventory' }], + tools: [serverTool('getInventory', () => ({ result: 'ok' }))], + }) as AsyncIterable, + ) + + const assistantToolMessage = calls[1]!.messages.find( + (message: ModelMessage) => + message.role === 'assistant' && + message.toolCalls?.[0]?.id === 'call_1', + ) + + expect(assistantToolMessage?.createdAt).toEqual( + new Date('2026-01-01T00:00:05.000Z'), + ) + } finally { + vi.useRealTimers() + } }) it('should execute tool calls that only provide the deprecated toolName field', async () => {