fix(ai-client): keep the replay's answer when a rejoin starts with reasoning

.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.

packages/ai-client/src/chat-client.ts

@@ -2192,6 +2192,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
@@ -2201,7 +2207,9 @@ export class ChatClient<
           }
           if (!rebuilt && rebuildsAssistantMessage(chunk)) {
             rebuilt = true
-            this.dropTrailingInFlightAssistant()
+            if (hydratedAssistantId) {
+              this.dropHydratedAssistant(hydratedAssistantId)
+            }
           }
           this.processIncomingChunk(chunk)
         }
@@ -2259,17 +2267,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 {

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<StreamChunk>
-  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,21 +66,39 @@ 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<UIMessage> }> = [
+  { 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<StreamChunk> = [
       {
         type: EventType.RUN_STARTED,
         runId: 'r1',
         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<UIMessage> = []
     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<UIMessage>
@@ -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()
     }

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/use-chat-hidden-view.spec.ts`   | Client layer: a chat that `<Activity>` hides mid-reply loads the thread and joins the run when it shows again                                               |
 | `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 |

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<AdapterYieldChunk>,
@@ -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,

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() {
     <div data-testid="join-run-client-tool-page" style={{ padding: 16 }}>
       {hydrated ? <div data-testid="hydration-marker" /> : null}
       <div data-testid="assistant-text">{assistantText || 'none'}</div>
+      <div data-testid="assistant-thinking">{assistantThinking || 'none'}</div>
       <div data-testid="client-tool-pending">
         {toolPending ? 'true' : 'false'}
       </div>

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)
+  })
+})