diff --git a/.changeset/fix-empty-text-message-start-id.md b/.changeset/fix-empty-text-message-start-id.md new file mode 100644 index 000000000..950847214 --- /dev/null +++ b/.changeset/fix-empty-text-message-start-id.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-persistence': patch +--- + +Ignore empty `TEXT_MESSAGE_START` message IDs so tool-call `parentMessageId` can be used. diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index dbeeea292..91cabab81 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -1542,7 +1542,14 @@ export function withPersistence( if (ctx.phase === 'modelStream') { const s = runState.get(ctx) if (s && chunk.type === 'TEXT_MESSAGE_START') { - s.streamingMessageId = chunk.messageId + // An empty/malformed messageId means "no identity" (matching the + // engine's convention), leaving room for the TOOL_CALL_START + // parentMessageId fallback below — but the per-turn accumulator + // still resets so snapshots never mix text across turns. + s.streamingMessageId = + typeof chunk.messageId === 'string' && chunk.messageId !== '' + ? chunk.messageId + : undefined s.streamingMessageCreatedAt = new Date() s.streamingText = '' } else if ( diff --git a/packages/ai-persistence/tests/with-persistence.test.ts b/packages/ai-persistence/tests/with-persistence.test.ts index 338431c92..b70a6d751 100644 --- a/packages/ai-persistence/tests/with-persistence.test.ts +++ b/packages/ai-persistence/tests/with-persistence.test.ts @@ -219,6 +219,147 @@ describe('withPersistence (state-only)', () => { ]) }) + it('does not let an empty TEXT_MESSAGE_START id replace parentMessageId', async () => { + const persistence = memoryPersistence() + const adapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': {}, + chatStream: () => + (async function* () { + yield ev.runStarted() + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: '', + timestamp: 1, + } + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'search', + toolName: 'search', + parentMessageId: 'stream-assistant', + timestamp: 1, + } + yield ev.text('Half a stor') + throw new Error('crash mid-stream') + })(), + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } as unknown as AnyTextAdapter + + await expect( + collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + middleware: [ + withPersistence(persistence, { snapshotStreaming: true }), + ], + }) as AsyncIterable, + ), + ).rejects.toThrow('crash mid-stream') + + expect(await persistence.stores.messages!.loadThread('t1')).toEqual([ + { role: 'user', content: 'hi' }, + expect.objectContaining({ + role: 'assistant', + content: 'Half a stor', + id: 'stream-assistant', + createdAt: expect.any(Date), + }), + ]) + }) + + it('resets streaming state on an empty-id TEXT_MESSAGE_START between turns', async () => { + const persistence = memoryPersistence() + let call = 0 + const adapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': {}, + chatStream: () => { + call++ + if (call === 1) { + return (async function* () { + yield ev.runStarted() + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: '', + timestamp: 1, + } + yield ev.text('Let me search.') + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'search', + toolName: 'search', + parentMessageId: 'assistant-turn-1', + timestamp: 1, + } + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: '{}', + timestamp: 1, + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'tool_calls', + timestamp: 1, + } + })() + } + return (async function* () { + yield ev.runStarted() + yield { + type: EventType.TEXT_MESSAGE_START, + messageId: '', + timestamp: 1, + } + yield ev.text('The answer is 42.') + throw new Error('crash mid-stream') + })() + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } as unknown as AnyTextAdapter + + await expect( + collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'search' }], + tools: [serverSearchTool()], + runId: 'r1', + threadId: 't1', + middleware: [ + withPersistence(persistence, { + snapshotStreaming: true, + snapshotIntervalMs: 0, + }), + ], + }) as AsyncIterable, + ), + ).rejects.toThrow('crash mid-stream') + + // The empty-id start on turn 2 still resets the per-turn accumulator: the + // crash-window snapshot holds only turn-2 text, and it must not inherit + // turn 1's tool-call id — two persisted messages may never share an id. + const thread = await persistence.stores.messages!.loadThread('t1') + expect(findAssistantToolCall(thread, 'call_1')?.id).toBe('assistant-turn-1') + const terminal = thread.at(-1) + expect(terminal).toMatchObject({ + role: 'assistant', + content: 'The answer is 42.', + }) + expect(terminal).not.toHaveProperty('id') + }) + it('stamps the terminal assistant turn with its stream messageId', async () => { const persistence = memoryPersistence() const { adapter } = mockAdapter([