diff --git a/.changeset/tool-first-text-drop.md b/.changeset/tool-first-text-drop.md new file mode 100644 index 000000000..d332b6e54 --- /dev/null +++ b/.changeset/tool-first-text-drop.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +Fix `StreamProcessor` dropping the first `TEXT_MESSAGE_CONTENT` delta when a `TOOL_CALL_START` event's `parentMessageId` precedes that message's `TEXT_MESSAGE_START` — the normal AG-UI shape for "call a tool, then explain the result" as one assistant turn (#1247). diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 8b6f5bc37..a61e03364 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/packages/ai/src/activities/chat/stream/processor.ts @@ -900,9 +900,21 @@ export class StreamProcessor { } // Ensure state exists - if (!this.messageStates.has(messageId)) { - this.createMessageState(messageId, uiRole) + let pendingState = this.messageStates.get(messageId) + if (!pendingState) { + pendingState = this.createMessageState(messageId, uiRole) this.activeMessageIds.add(messageId) + } else if (pendingState.hasToolCallsSinceTextStart) { + // A tool call (e.g. TOOL_CALL_START with parentMessageId) marked + // this message before its "real" TEXT_MESSAGE_START arrived — same + // reset Case 2 performs, so the segment accumulator doesn't carry + // stale tool-call state into the text that follows. + if (pendingState.currentSegmentText !== pendingState.lastEmittedText) { + this.emitTextUpdateForMessage(messageId) + } + pendingState.currentSegmentText = '' + pendingState.lastEmittedText = '' + pendingState.hasToolCallsSinceTextStart = false } this.mergeMessageMetadata(messageId, chunk.metadata) diff --git a/packages/ai/tests/stream-processor.test.ts b/packages/ai/tests/stream-processor.test.ts index 61a37de65..eea81b70d 100644 --- a/packages/ai/tests/stream-processor.test.ts +++ b/packages/ai/tests/stream-processor.test.ts @@ -4109,6 +4109,38 @@ describe('StreamProcessor', () => { }, ]) }) + + it('should not drop the first TEXT_MESSAGE_CONTENT delta in a tool-first flow (#1247)', () => { + const processor = new StreamProcessor() + + processor.processChunk( + chunk(EventType.TOOL_CALL_START, { + toolCallId: 'tc-1', + toolCallName: 'lookupWeather', + toolName: 'lookupWeather', + parentMessageId: 'anthropic-msg-1', + }), + ) + processor.processChunk(ev.toolArgs('tc-1', '{"location":"Berlin"}')) + processor.processChunk(ev.toolEnd('tc-1', 'lookupWeather')) + + processor.processChunk( + chunk(EventType.TEXT_MESSAGE_START, { + messageId: 'anthropic-msg-1', + role: 'assistant' as const, + }), + ) + // Two deltas: the bug only surfaces once a second delta arrives after + // the tool-first message's real TEXT_MESSAGE_START. + processor.processChunk(ev.textContent('It is ', 'anthropic-msg-1')) + processor.processChunk(ev.textContent('sunny.', 'anthropic-msg-1')) + processor.processChunk(ev.textEnd('anthropic-msg-1')) + processor.finalizeStream() + + const messages = processor.getMessages() + const textParts = messages[0]?.parts.filter((p) => p.type === 'text') + expect(textParts).toEqual([{ type: 'text', content: 'It is sunny.' }]) + }) }) describe('double onStreamEnd guard', () => { diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 3acd134ba..70cd62134 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -11,6 +11,7 @@ import { Route as rootRouteImport } from './routes/__root' import { Route as WebsocketAdapterRouteImport } from './routes/websocket-adapter' import { Route as ToolsTestRouteImport } from './routes/tools-test' +import { Route as ToolFirstTextRouteImport } from './routes/tool-first-text' import { Route as PersistenceDurabilityRouteImport } from './routes/persistence-durability' import { Route as MiddlewareTestRouteImport } from './routes/middleware-test' import { Route as MarkdownCjkRouteImport } from './routes/markdown-cjk' @@ -34,6 +35,7 @@ import { Route as ApiVideoRouteImport } from './routes/api.video' import { Route as ApiTtsRouteImport } from './routes/api.tts' import { Route as ApiTranscriptionRouteImport } from './routes/api.transcription' import { Route as ApiToolsTestRouteImport } from './routes/api.tools-test' +import { Route as ApiToolFirstTextWireRouteImport } from './routes/api.tool-first-text-wire' import { Route as ApiToolCallLifecycleWireRouteImport } from './routes/api.tool-call-lifecycle-wire' import { Route as ApiSummarizeRouteImport } from './routes/api.summarize' import { Route as ApiSandboxToolHistoryRouteImport } from './routes/api.sandbox-tool-history' @@ -80,8 +82,8 @@ import { Route as ApiDurableTakeoverRouteImport } from './routes/api.durable-tak import { Route as ApiDurableDeliveryRouteImport } from './routes/api.durable-delivery' import { Route as ApiDevtoolsMemoryRouteImport } from './routes/api.devtools-memory' import { Route as ApiChatRouteImport } from './routes/api.chat' -import { Route as ApiByokChatRouteImport } from './routes/api.byok-chat' import { Route as ApiByteplusSeedance1080pWireRouteImport } from './routes/api.byteplus-seedance-1080p-wire' +import { Route as ApiByokChatRouteImport } from './routes/api.byok-chat' import { Route as ApiAudioRouteImport } from './routes/api.audio' import { Route as ApiArktypeToolWireRouteImport } from './routes/api.arktype-tool-wire' import { Route as ApiAnthropicStructuredUsageRouteImport } from './routes/api.anthropic-structured-usage' @@ -104,6 +106,11 @@ const ToolsTestRoute = ToolsTestRouteImport.update({ path: '/tools-test', getParentRoute: () => rootRouteImport, } as any) +const ToolFirstTextRoute = ToolFirstTextRouteImport.update({ + id: '/tool-first-text', + path: '/tool-first-text', + getParentRoute: () => rootRouteImport, +} as any) const PersistenceDurabilityRoute = PersistenceDurabilityRouteImport.update({ id: '/persistence-durability', path: '/persistence-durability', @@ -221,6 +228,11 @@ const ApiToolsTestRoute = ApiToolsTestRouteImport.update({ path: '/api/tools-test', getParentRoute: () => rootRouteImport, } as any) +const ApiToolFirstTextWireRoute = ApiToolFirstTextWireRouteImport.update({ + id: '/api/tool-first-text-wire', + path: '/api/tool-first-text-wire', + getParentRoute: () => rootRouteImport, +} as any) const ApiToolCallLifecycleWireRoute = ApiToolCallLifecycleWireRouteImport.update({ id: '/api/tool-call-lifecycle-wire', @@ -466,17 +478,17 @@ const ApiChatRoute = ApiChatRouteImport.update({ path: '/api/chat', getParentRoute: () => rootRouteImport, } as any) -const ApiByokChatRoute = ApiByokChatRouteImport.update({ - id: '/api/byok-chat', - path: '/api/byok-chat', - getParentRoute: () => rootRouteImport, -} as any) const ApiByteplusSeedance1080pWireRoute = ApiByteplusSeedance1080pWireRouteImport.update({ id: '/api/byteplus-seedance-1080p-wire', path: '/api/byteplus-seedance-1080p-wire', getParentRoute: () => rootRouteImport, } as any) +const ApiByokChatRoute = ApiByokChatRouteImport.update({ + id: '/api/byok-chat', + path: '/api/byok-chat', + getParentRoute: () => rootRouteImport, +} as any) const ApiAudioRoute = ApiAudioRouteImport.update({ id: '/api/audio', path: '/api/audio', @@ -553,6 +565,7 @@ export interface FileRoutesByFullPath { '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute '/persistence-durability': typeof PersistenceDurabilityRoute + '/tool-first-text': typeof ToolFirstTextRoute '/tools-test': typeof ToolsTestRoute '/websocket-adapter': typeof WebsocketAdapterRoute '/$provider/$feature': typeof ProviderFeatureRoute @@ -609,6 +622,7 @@ export interface FileRoutesByFullPath { '/api/sandbox-tool-history': typeof ApiSandboxToolHistoryRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute + '/api/tool-first-text-wire': typeof ApiToolFirstTextWireRoute '/api/tools-test': typeof ApiToolsTestRoute '/api/transcription': typeof ApiTranscriptionRouteWithChildren '/api/tts': typeof ApiTtsRouteWithChildren @@ -639,6 +653,7 @@ export interface FileRoutesByTo { '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute '/persistence-durability': typeof PersistenceDurabilityRoute + '/tool-first-text': typeof ToolFirstTextRoute '/tools-test': typeof ToolsTestRoute '/websocket-adapter': typeof WebsocketAdapterRoute '/$provider/$feature': typeof ProviderFeatureRoute @@ -695,6 +710,7 @@ export interface FileRoutesByTo { '/api/sandbox-tool-history': typeof ApiSandboxToolHistoryRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute + '/api/tool-first-text-wire': typeof ApiToolFirstTextWireRoute '/api/tools-test': typeof ApiToolsTestRoute '/api/transcription': typeof ApiTranscriptionRouteWithChildren '/api/tts': typeof ApiTtsRouteWithChildren @@ -726,6 +742,7 @@ export interface FileRoutesById { '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute '/persistence-durability': typeof PersistenceDurabilityRoute + '/tool-first-text': typeof ToolFirstTextRoute '/tools-test': typeof ToolsTestRoute '/websocket-adapter': typeof WebsocketAdapterRoute '/$provider/$feature': typeof ProviderFeatureRoute @@ -782,6 +799,7 @@ export interface FileRoutesById { '/api/sandbox-tool-history': typeof ApiSandboxToolHistoryRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute + '/api/tool-first-text-wire': typeof ApiToolFirstTextWireRoute '/api/tools-test': typeof ApiToolsTestRoute '/api/transcription': typeof ApiTranscriptionRouteWithChildren '/api/tts': typeof ApiTtsRouteWithChildren @@ -814,6 +832,7 @@ export interface FileRouteTypes { | '/markdown-cjk' | '/middleware-test' | '/persistence-durability' + | '/tool-first-text' | '/tools-test' | '/websocket-adapter' | '/$provider/$feature' @@ -870,6 +889,7 @@ export interface FileRouteTypes { | '/api/sandbox-tool-history' | '/api/summarize' | '/api/tool-call-lifecycle-wire' + | '/api/tool-first-text-wire' | '/api/tools-test' | '/api/transcription' | '/api/tts' @@ -900,6 +920,7 @@ export interface FileRouteTypes { | '/markdown-cjk' | '/middleware-test' | '/persistence-durability' + | '/tool-first-text' | '/tools-test' | '/websocket-adapter' | '/$provider/$feature' @@ -956,6 +977,7 @@ export interface FileRouteTypes { | '/api/sandbox-tool-history' | '/api/summarize' | '/api/tool-call-lifecycle-wire' + | '/api/tool-first-text-wire' | '/api/tools-test' | '/api/transcription' | '/api/tts' @@ -986,6 +1008,7 @@ export interface FileRouteTypes { | '/markdown-cjk' | '/middleware-test' | '/persistence-durability' + | '/tool-first-text' | '/tools-test' | '/websocket-adapter' | '/$provider/$feature' @@ -1042,6 +1065,7 @@ export interface FileRouteTypes { | '/api/sandbox-tool-history' | '/api/summarize' | '/api/tool-call-lifecycle-wire' + | '/api/tool-first-text-wire' | '/api/tools-test' | '/api/transcription' | '/api/tts' @@ -1073,6 +1097,7 @@ export interface RootRouteChildren { MarkdownCjkRoute: typeof MarkdownCjkRoute MiddlewareTestRoute: typeof MiddlewareTestRoute PersistenceDurabilityRoute: typeof PersistenceDurabilityRoute + ToolFirstTextRoute: typeof ToolFirstTextRoute ToolsTestRoute: typeof ToolsTestRoute WebsocketAdapterRoute: typeof WebsocketAdapterRoute ProviderFeatureRoute: typeof ProviderFeatureRoute @@ -1129,6 +1154,7 @@ export interface RootRouteChildren { ApiSandboxToolHistoryRoute: typeof ApiSandboxToolHistoryRoute ApiSummarizeRoute: typeof ApiSummarizeRoute ApiToolCallLifecycleWireRoute: typeof ApiToolCallLifecycleWireRoute + ApiToolFirstTextWireRoute: typeof ApiToolFirstTextWireRoute ApiToolsTestRoute: typeof ApiToolsTestRoute ApiTranscriptionRoute: typeof ApiTranscriptionRouteWithChildren ApiTtsRoute: typeof ApiTtsRouteWithChildren @@ -1152,6 +1178,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ToolsTestRouteImport parentRoute: typeof rootRouteImport } + '/tool-first-text': { + id: '/tool-first-text' + path: '/tool-first-text' + fullPath: '/tool-first-text' + preLoaderRoute: typeof ToolFirstTextRouteImport + parentRoute: typeof rootRouteImport + } '/persistence-durability': { id: '/persistence-durability' path: '/persistence-durability' @@ -1313,6 +1346,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiToolsTestRouteImport parentRoute: typeof rootRouteImport } + '/api/tool-first-text-wire': { + id: '/api/tool-first-text-wire' + path: '/api/tool-first-text-wire' + fullPath: '/api/tool-first-text-wire' + preLoaderRoute: typeof ApiToolFirstTextWireRouteImport + parentRoute: typeof rootRouteImport + } '/api/tool-call-lifecycle-wire': { id: '/api/tool-call-lifecycle-wire' path: '/api/tool-call-lifecycle-wire' @@ -1635,13 +1675,6 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiChatRouteImport parentRoute: typeof rootRouteImport } - '/api/byok-chat': { - id: '/api/byok-chat' - path: '/api/byok-chat' - fullPath: '/api/byok-chat' - preLoaderRoute: typeof ApiByokChatRouteImport - parentRoute: typeof rootRouteImport - } '/api/byteplus-seedance-1080p-wire': { id: '/api/byteplus-seedance-1080p-wire' path: '/api/byteplus-seedance-1080p-wire' @@ -1649,6 +1682,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiByteplusSeedance1080pWireRouteImport parentRoute: typeof rootRouteImport } + '/api/byok-chat': { + id: '/api/byok-chat' + path: '/api/byok-chat' + fullPath: '/api/byok-chat' + preLoaderRoute: typeof ApiByokChatRouteImport + parentRoute: typeof rootRouteImport + } '/api/audio': { id: '/api/audio' path: '/api/audio' @@ -1806,6 +1846,7 @@ const rootRouteChildren: RootRouteChildren = { MarkdownCjkRoute: MarkdownCjkRoute, MiddlewareTestRoute: MiddlewareTestRoute, PersistenceDurabilityRoute: PersistenceDurabilityRoute, + ToolFirstTextRoute: ToolFirstTextRoute, ToolsTestRoute: ToolsTestRoute, WebsocketAdapterRoute: WebsocketAdapterRoute, ProviderFeatureRoute: ProviderFeatureRoute, @@ -1862,6 +1903,7 @@ const rootRouteChildren: RootRouteChildren = { ApiSandboxToolHistoryRoute: ApiSandboxToolHistoryRoute, ApiSummarizeRoute: ApiSummarizeRoute, ApiToolCallLifecycleWireRoute: ApiToolCallLifecycleWireRoute, + ApiToolFirstTextWireRoute: ApiToolFirstTextWireRoute, ApiToolsTestRoute: ApiToolsTestRoute, ApiTranscriptionRoute: ApiTranscriptionRouteWithChildren, ApiTtsRoute: ApiTtsRouteWithChildren, diff --git a/testing/e2e/src/routes/api.tool-first-text-wire.ts b/testing/e2e/src/routes/api.tool-first-text-wire.ts new file mode 100644 index 000000000..c03a1fbe9 --- /dev/null +++ b/testing/e2e/src/routes/api.tool-first-text-wire.ts @@ -0,0 +1,88 @@ +import { createFileRoute } from '@tanstack/react-router' +import { toServerSentEventsResponse } from '@tanstack/ai' +import type { StreamChunk } from '@tanstack/ai' + +/** + * Wire-format regression for issue #1247. + * + * A provider-free harness run where `TOOL_CALL_START` carries a + * `parentMessageId` that has not had a `TEXT_MESSAGE_START` yet — the normal + * AG-UI shape for "call a tool, then explain the result" as one assistant + * turn. The two `TEXT_MESSAGE_CONTENT` deltas below are the assertion: the + * bug only surfaces once a *second* delta arrives after the tool-first + * message's real `TEXT_MESSAGE_START`. + */ +function toolFirstRun( + threadId: string, + runId: string, +): AsyncIterable { + const messageId = 'msg-1' + const toolCallId = 'call-1' + return (async function* () { + yield { type: 'RUN_STARTED', threadId, runId, timestamp: Date.now() } + yield { + type: 'TOOL_CALL_START', + toolCallId, + toolCallName: 'lookupWeather', + parentMessageId: messageId, + timestamp: Date.now(), + } + yield { + type: 'TOOL_CALL_ARGS', + toolCallId, + delta: '{}', + timestamp: Date.now(), + } + yield { type: 'TOOL_CALL_END', toolCallId, timestamp: Date.now() } + yield { + type: 'TOOL_CALL_RESULT', + messageId: 'tool-1', + toolCallId, + role: 'tool', + content: '{"ok":true}', + timestamp: Date.now(), + } + yield { + type: 'TEXT_MESSAGE_START', + messageId, + role: 'assistant', + timestamp: Date.now(), + } + yield { + type: 'TEXT_MESSAGE_CONTENT', + messageId, + delta: 'Hello, ', + timestamp: Date.now(), + } + yield { + type: 'TEXT_MESSAGE_CONTENT', + messageId, + delta: 'world.', + timestamp: Date.now(), + } + yield { type: 'TEXT_MESSAGE_END', messageId, timestamp: Date.now() } + yield { + type: 'RUN_FINISHED', + threadId, + runId, + timestamp: Date.now(), + outcome: { type: 'success' }, + } + })() as AsyncIterable +} + +export const Route = createFileRoute('/api/tool-first-text-wire')({ + server: { + handlers: { + POST: async ({ request }) => { + const body: unknown = await request.json() + const threadId = + typeof body === 'object' && body !== null && 'threadId' in body + ? String((body as Record).threadId) + : 'thread-1' + const runId = `run-${threadId}` + return toServerSentEventsResponse(toolFirstRun(threadId, runId)) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/tool-first-text.tsx b/testing/e2e/src/routes/tool-first-text.tsx new file mode 100644 index 000000000..80b9a211b --- /dev/null +++ b/testing/e2e/src/routes/tool-first-text.tsx @@ -0,0 +1,41 @@ +import { useEffect } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' + +/** + * Harness page for issue #1247: a tool call whose `parentMessageId` precedes + * that message's `TEXT_MESSAGE_START`. `/api/tool-first-text-wire` streams + * that shape once; this page renders the resulting assistant text so the + * spec can assert the first delta was not dropped. + */ +function ToolFirstTextPage() { + const { messages, sendMessage } = useChat({ + threadId: 'tool-first-text-1', + connection: fetchServerSentEvents('/api/tool-first-text-wire'), + }) + + const assistantText = messages + .filter((message) => message.role === 'assistant') + .flatMap((message) => + message.parts.flatMap((part) => + part.type === 'text' ? [part.content] : [], + ), + ) + .join('') + + useEffect(() => { + void sendMessage('go') + // Fire the single run once on mount; the harness route ignores the content. + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []) + + return ( +
+
{assistantText}
+
+ ) +} + +export const Route = createFileRoute('/tool-first-text')({ + component: ToolFirstTextPage, +}) diff --git a/testing/e2e/tests/tool-first-text.spec.ts b/testing/e2e/tests/tool-first-text.spec.ts new file mode 100644 index 000000000..e856615b3 --- /dev/null +++ b/testing/e2e/tests/tool-first-text.spec.ts @@ -0,0 +1,19 @@ +import { expect, test } from '@playwright/test' + +/** + * `StreamProcessor` must not drop the first `TEXT_MESSAGE_CONTENT` delta when + * a tool call's `parentMessageId` precedes that message's real + * `TEXT_MESSAGE_START` — the normal AG-UI shape for "call a tool, then + * explain the result" as one assistant turn. + */ +test.describe('tool-first text (#1247)', () => { + test('does not drop the first text delta after a tool-first message', async ({ + page, + }) => { + await page.goto('/tool-first-text') + + await expect(page.getByTestId('assistant-text')).toHaveText( + 'Hello, world.', + ) + }) +})