diff --git a/.changeset/rejoin-keeps-replay-message.md b/.changeset/rejoin-keeps-replay-message.md new file mode 100644 index 0000000000..78aefdfcf4 --- /dev/null +++ b/.changeset/rejoin-keeps-replay-message.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-client': patch +--- + +Rejoining a run no longer loses the answer when the replay starts with reasoning. A reasoning start event that carries metadata, which `chat()` adds to every adapter event, already creates the replay's assistant message. The client used to drop that message at the first rebuild event. It now drops only the assistant message it hydrated before the replay. diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 3085d2f526..7bce935951 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -2155,6 +2155,12 @@ export class ChatClient< const connectTimer = setTimeout(() => { if (!attached) controller.abort() }, REJOIN_CONNECT_DEADLINE_MS) + // Captured before the replay: chunks before the first rebuild trigger + // (reasoning start with metadata) can already append the replay's own + // assistant message, which must survive the drop. + const hydrated = this.processor.getMessages().at(-1) + const hydratedAssistantId = + hydrated?.role === 'assistant' ? hydrated.id : undefined try { for await (const chunk of joinRun(runId, controller.signal)) { if (controller.signal.aborted) break @@ -2164,7 +2170,9 @@ export class ChatClient< } if (!rebuilt && rebuildsAssistantMessage(chunk)) { rebuilt = true - this.dropTrailingInFlightAssistant() + if (hydratedAssistantId) { + this.dropHydratedAssistant(hydratedAssistantId) + } } this.processIncomingChunk(chunk) } @@ -2222,17 +2230,15 @@ export class ChatClient< } /** - * Drop a hydrated, still-in-flight assistant turn so a resume replay can - * rebuild it cleanly. Only touches a trailing assistant message (the shape a - * reload-mid-stream leaves); a thread whose last turn is a user message (run - * never produced, or already settled) is left untouched. + * Drop the hydrated, still-in-flight assistant turn so a resume replay can + * rebuild it cleanly. `id` is the trailing assistant message from before the + * replay (the shape a reload-mid-stream leaves); a thread whose last turn was + * a user message (run never produced, or already settled) has none to drop. */ - private dropTrailingInFlightAssistant(): void { + private dropHydratedAssistant(id: string): void { const messages = this.processor.getMessages() - const last = messages[messages.length - 1] - if (last && last.role === 'assistant') { - this.processor.setMessages(messages.slice(0, -1)) - } + const kept = messages.filter((message) => message.id !== id) + if (kept.length < messages.length) this.processor.setMessages(kept) } private processIncomingChunk(chunk: StreamChunk): void { diff --git a/packages/ai-client/tests/chat-client-reasoning-rejoin.test.ts b/packages/ai-client/tests/chat-client-reasoning-rejoin.test.ts index e6ef5351fc..7d6e98217a 100644 --- a/packages/ai-client/tests/chat-client-reasoning-rejoin.test.ts +++ b/packages/ai-client/tests/chat-client-reasoning-rejoin.test.ts @@ -4,26 +4,60 @@ import { ChatClient } from '../src/chat-client' import type { ResumableConnectConnectionAdapter } from '../src/connection-adapters' import type { StreamChunk, UIMessage } from '@tanstack/ai/client' -const firstEvents: Array<{ +// chat() stamps this on every adapter chunk. On REASONING_MESSAGE_START it +// makes the processor create the assistant message before any rebuild trigger. +const callMetadata = { + tanstack: { + source: { provider: 'openai', api: 'responses', model: 'example-model' }, + }, +} + +const reasonings: Array<{ name: string chunks: Array - thinking?: string + thinking: { content: string; signature?: string } }> = [ { - name: 'reasoning content', + name: 'readable reasoning', chunks: [ + { type: EventType.REASONING_START, messageId: 'thinking', timestamp: 2 }, + { + type: EventType.REASONING_MESSAGE_START, + messageId: 'thinking', + role: 'reasoning', + timestamp: 2, + }, { type: EventType.REASONING_MESSAGE_CONTENT, messageId: 'thinking', delta: 'Inspect the files.', timestamp: 2, }, + { + type: EventType.REASONING_MESSAGE_END, + messageId: 'thinking', + timestamp: 2, + }, + { type: EventType.REASONING_END, messageId: 'thinking', timestamp: 2 }, ], - thinking: 'Inspect the files.', + thinking: { content: 'Inspect the files.' }, }, { - name: 'encrypted reasoning', + name: 'encrypted-only reasoning', chunks: [ + { type: EventType.REASONING_START, messageId: 'thinking', timestamp: 2 }, + { + type: EventType.REASONING_MESSAGE_START, + messageId: 'thinking', + role: 'reasoning', + timestamp: 2, + }, + { + type: EventType.REASONING_MESSAGE_END, + messageId: 'thinking', + timestamp: 2, + }, + { type: EventType.REASONING_END, messageId: 'thinking', timestamp: 2 }, { type: EventType.REASONING_ENCRYPTED_VALUE, subtype: 'message', @@ -32,13 +66,31 @@ const firstEvents: Array<{ timestamp: 2, }, ], - thinking: '', + thinking: { content: '', signature: 'signature' }, }, ] -it.each(firstEvents)( - 'keeps activity when replay starts with $name', - async ({ chunks, thinking }) => { +const hydratedTails: Array<{ name: string; tail: Array }> = [ + { name: 'a user turn', tail: [] }, + { + name: 'an in-flight assistant', + tail: [ + { + id: 'partial', + role: 'assistant', + parts: [{ type: 'thinking', content: 'Inspect' }], + }, + ], + }, +] + +const cases = reasonings.flatMap((reasoning) => + hydratedTails.map((hydrated) => ({ reasoning, hydrated })), +) + +it.each(cases)( + 'keeps the thread when replay starts with $reasoning.name and hydrate ends with $hydrated.name', + async ({ reasoning, hydrated }) => { const replay: Array = [ { type: EventType.RUN_STARTED, @@ -46,7 +98,7 @@ it.each(firstEvents)( threadId: 't1', timestamp: 1, }, - ...chunks, + ...reasoning.chunks, { type: EventType.TOOL_CALL_START, toolCallId: 'read', @@ -87,6 +139,12 @@ it.each(firstEvents)( timestamp: 10, }, ] + // chat() runs the tool itself, so the tool result is not an adapter chunk. + const stamped = replay.map((chunk) => + chunk.type === EventType.TOOL_CALL_RESULT + ? chunk + : { ...chunk, metadata: callMetadata }, + ) let sent: Array = [] const connection: ResumableConnectConnectionAdapter = { async hydrate() { @@ -97,18 +155,14 @@ it.each(firstEvents)( role: 'user', parts: [{ type: 'text', content: 'Build an app.' }], }, - { - id: 'partial', - role: 'assistant', - parts: [{ type: 'thinking', content: 'Inspect' }], - }, + ...hydrated.tail, ], activeRun: { runId: 'r1' }, interrupts: null, } }, async *joinRun() { - yield* replay + yield* stamped }, async *connect(messages) { sent = messages as Array @@ -122,30 +176,26 @@ it.each(firstEvents)( try { client.attach() await vi.waitFor(() => { - expect( - client.getMessages().flatMap((message) => message.parts), - ).toContainEqual( - expect.objectContaining({ type: 'text', content: 'Built the app.' }), - ) + expect(client.getMessages()).not.toHaveLength(0) + expect(client.getIsLoading()).toBe(false) }) await client.sendMessage('Change the title.') - const parts = sent.flatMap((message) => message.parts) - if (thinking !== undefined) { - expect(parts).toContainEqual( - expect.objectContaining({ type: 'thinking', content: thinking }), - ) - } - expect(parts).toContainEqual( + const assistants = sent.filter((message) => message.role === 'assistant') + expect(assistants).toHaveLength(1) + expect(assistants[0]!.id).not.toBe('partial') + expect(assistants[0]!.parts).toContainEqual( + expect.objectContaining({ type: 'thinking', ...reasoning.thinking }), + ) + expect(assistants[0]!.parts).toContainEqual( expect.objectContaining({ type: 'tool-call', id: 'read', output: 'file contents', }), ) - expect(parts).toContainEqual( + expect(assistants[0]!.parts).toContainEqual( expect.objectContaining({ type: 'text', content: 'Built the app.' }), ) - expect(sent.some((message) => message.id === 'partial')).toBe(false) } finally { client.dispose() } diff --git a/testing/e2e/README.md b/testing/e2e/README.md index e28cb81663..99442940a3 100644 --- a/testing/e2e/README.md +++ b/testing/e2e/README.md @@ -97,7 +97,7 @@ Notes: | -------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | | `tests/delivery-durability.spec.ts` | Transport layer: offset-tagged log, `Last-Event-ID` reconnect, second-tab join, SSE + NDJSON | | `tests/persistence-durability.spec.ts` | Client layer: a browser refresh restores the conversation and any pending interrupt | -| `tests/join-run-client-tool.spec.ts` | Mid-run reload: `joinRun` replay that ends on a client tool drains the continuation (issue #1058) | +| `tests/join-run-client-tool.spec.ts` | Mid-run reload: `joinRun` replay that ends on a client tool drains the continuation (issue #1058); a replay that starts with reasoning keeps its answer | | `tests/sandbox-durability.spec.ts` | Sandbox instances: a second run resumes the persisted sandbox | | `tests/durable-takeover.spec.ts` | Takeover with log alignment, detach-on-disconnect, out-of-band cancel in both bands, cancel-vs-disconnect divergence, and the superseded-driver epoch fence | diff --git a/testing/e2e/src/routes/api.join-run-client-tool.ts b/testing/e2e/src/routes/api.join-run-client-tool.ts index 81b243ca02..757938ed55 100644 --- a/testing/e2e/src/routes/api.join-run-client-tool.ts +++ b/testing/e2e/src/routes/api.join-run-client-tool.ts @@ -30,6 +30,7 @@ import type { const JOIN_OK = 'JOIN_OK the lookup is 42' const LATER_OK = 'LATER_OK the follow-up ran' +const REASON_OK = 'REASON_OK the answer after reasoning' const DISCONNECT_FALLBACK_MS = 3000 const SETTLE_AFTER_DISCONNECT_MS = 1500 @@ -51,6 +52,8 @@ const parentInterruptByThread = new Map< // User text for the #1429 scenario: the first run pauses on the client tool, // and the continuation run (not the first run) is held open. const CONTINUATION_SCENARIO = 'continue' +// User text for a held run that reasons before it answers. +const REASONING_SCENARIO = 'reason' async function* recordInterrupt( stream: AsyncIterable, @@ -154,7 +157,13 @@ const adapter: AnyTextAdapter = { ) const userText = lastUserText(options.messages) const reply = - userText === 'later' ? LATER_OK : hasToolResult ? JOIN_OK : null + userText === 'later' + ? LATER_OK + : userText === REASONING_SCENARIO + ? REASON_OK + : hasToolResult + ? JOIN_OK + : null yield { type: EventType.RUN_STARTED, @@ -165,6 +174,41 @@ const adapter: AnyTextAdapter = { } if (reply !== null) { + if (userText === REASONING_SCENARIO) { + const reasoningId = `${runId}-reasoning` + yield { + type: EventType.REASONING_START, + messageId: reasoningId, + model, + timestamp: Date.now(), + } + yield { + type: EventType.REASONING_MESSAGE_START, + messageId: reasoningId, + role: 'reasoning', + model, + timestamp: Date.now(), + } + yield { + type: EventType.REASONING_MESSAGE_CONTENT, + messageId: reasoningId, + delta: 'Check the order.', + model, + timestamp: Date.now(), + } + yield { + type: EventType.REASONING_MESSAGE_END, + messageId: reasoningId, + model, + timestamp: Date.now(), + } + yield { + type: EventType.REASONING_END, + messageId: reasoningId, + model, + timestamp: Date.now(), + } + } yield { type: EventType.TEXT_MESSAGE_START, messageId, diff --git a/testing/e2e/src/routes/join-run-client-tool.tsx b/testing/e2e/src/routes/join-run-client-tool.tsx index f8ac438b26..318e3df163 100644 --- a/testing/e2e/src/routes/join-run-client-tool.tsx +++ b/testing/e2e/src/routes/join-run-client-tool.tsx @@ -63,6 +63,14 @@ function JoinRunClientToolPage() { ), ) .join('') + const assistantThinking = messages + .filter((message) => message.role === 'assistant') + .flatMap((message) => + message.parts.flatMap((part) => + part.type === 'thinking' ? [part.content] : [], + ), + ) + .join('') const toolPending = messages.some((message) => message.parts.some((part) => { @@ -88,6 +96,7 @@ function JoinRunClientToolPage() {
{hydrated ?
: null}
{assistantText || 'none'}
+
{assistantThinking || 'none'}
{toolPending ? 'true' : 'false'}
diff --git a/testing/e2e/tests/join-run-client-tool.spec.ts b/testing/e2e/tests/join-run-client-tool.spec.ts index 4b29acb319..7f99c4bd8a 100644 --- a/testing/e2e/tests/join-run-client-tool.spec.ts +++ b/testing/e2e/tests/join-run-client-tool.spec.ts @@ -94,3 +94,30 @@ test.describe('joinRun client-tool continuation (issue #1058)', () => { }) }) }) + +test.describe('joinRun replay that starts with reasoning', () => { + test('a reload mid run keeps the thinking and the answer', async ({ + page, + }) => { + const threadId = `join-run-reasoning-${crypto.randomUUID()}` + await page.goto( + `/join-run-client-tool?threadId=${encodeURIComponent(threadId)}`, + ) + await expect(page.getByTestId('hydration-marker')).toBeAttached() + + await sendMessage(page, 'reason') + // The run is held open after its answer, so the remount sees `activeRun`. + await expect(page.getByTestId('assistant-text')).toContainText('REASON_OK') + + await page.reload() + await expect(page.getByTestId('hydration-marker')).toBeAttached() + + // Hydrate returns no messages, so both can only come from the replay. + await expect(page.getByTestId('assistant-text')).toContainText('REASON_OK') + await expect(page.getByTestId('assistant-thinking')).toHaveText( + 'Check the order.', + ) + await expect(page.getByTestId('assistant-message')).toHaveCount(1) + await expect(page.getByTestId('loading-indicator')).toHaveCount(0) + }) +})