@tanstack/ai 0.47.2 → 0.48.0

package/src/activities/chat/stream/processor.ts

@@ -22,8 +22,15 @@ import {
   generateMessageId,
   uiMessageToModelMessages,
 } from '../messages.js'
-import { normalizeToolResult } from '../../../utilities/tool-result'
+import { runErrorEventToError } from '../../../utilities/errors'
 import { isProviderExecutedToolCall } from '../../../utilities/provider-executed'
+import {
+  mergeMetadata,
+  tanstackMetadata,
+} from '../../../utilities/merge-metadata'
+import { getChunkRunId } from '../../../utilities/chunk-ids'
+import type { AdapterYieldChunk } from '../../../utilities/adapter-yield-chunk'
+import { normalizeToolResult } from '../../../utilities/tool-result'
 import { defaultJSONParser } from './json-parser'
 import {
   appendStructuredOutputDelta,
@@ -55,6 +62,7 @@ import type {
   MessagePart,
   ModelMessage,
   StreamChunk,
+  ThinkingPart,
   ToolCall,
   ToolCallPart,
   ToolResultPart,
@@ -238,6 +246,7 @@ export class StreamProcessor {
    *
    * @param content - The message content (string or array of content parts)
    * @param id - Optional custom message ID (generated if not provided)
+   * @param metadata - Optional AG-UI metadata bag
    * @returns The created UIMessage
    *
    * @example
@@ -255,7 +264,11 @@ export class StreamProcessor {
    * processor.addUserMessage('Hello!', 'custom-id-123')
    * ```
    */
-  addUserMessage(content: string | Array<ContentPart>, id?: string): UIMessage {
+  addUserMessage(
+    content: string | Array<ContentPart>,
+    id?: string,
+    metadata?: UIMessage['metadata'],
+  ): UIMessage {
     // Convert content to message parts
     const parts: Array<MessagePart> =
       typeof content === 'string'
@@ -270,6 +283,7 @@ export class StreamProcessor {
       role: 'user',
       parts,
       createdAt: new Date(),
+      ...(metadata != null ? { metadata } : {}),
     }
 
     this.messages = [...this.messages, userMessage]
@@ -622,6 +636,12 @@ export class StreamProcessor {
         )
         break
 
+      case 'REASONING_ENCRYPTED_VALUE':
+        this.handleReasoningEncryptedValueEvent(
+          chunk as Extract<StreamChunk, { type: 'REASONING_ENCRYPTED_VALUE' }>,
+        )
+        break
+
       case 'TOOL_CALL_RESULT':
         this.handleToolCallResultEvent(
           chunk as Extract<StreamChunk, { type: 'TOOL_CALL_RESULT' }>,
@@ -783,6 +803,55 @@ export class StreamProcessor {
   // Event Handlers
   // ============================================
 
+  /**
+   * Merge event metadata onto a UIMessage. `tanstack` is deep-merged so a
+   * later delta does not wipe `tanstack.model`. High-frequency leftover
+   * keys (`content`, `args`) never stamp onto the message.
+   * Rebuilds `createdAt` when `tanstack.createdAt` is an ISO string.
+   */
+  private mergeMessageMetadata(messageId: string, incoming: unknown): void {
+    if (
+      incoming == null ||
+      typeof incoming !== 'object' ||
+      Array.isArray(incoming)
+    ) {
+      return
+    }
+    const message = this.messages.find((msg) => msg.id === messageId)
+    if (!message) return
+
+    const incomingRecord = incoming as NonNullable<UIMessage['metadata']>
+    const incomingTanstack = tanstackMetadata(incomingRecord)
+    const toMerge =
+      incomingTanstack != null &&
+      ('content' in incomingTanstack || 'args' in incomingTanstack)
+        ? {
+            ...incomingRecord,
+            tanstack: Object.fromEntries(
+              Object.entries(incomingTanstack).filter(
+                ([key]) => key !== 'content' && key !== 'args',
+              ),
+            ),
+          }
+        : incomingRecord
+    const metadata = mergeMetadata(message.metadata, toMerge)
+    const createdAtRaw = tanstackMetadata(incomingRecord)?.createdAt
+    const createdAt =
+      typeof createdAtRaw === 'string' ? new Date(createdAtRaw) : undefined
+    const createdAtValid =
+      createdAt !== undefined && !Number.isNaN(createdAt.getTime())
+    this.messages = this.messages.map((msg) =>
+      msg.id === messageId
+        ? {
+            ...msg,
+            ...(metadata !== undefined ? { metadata } : {}),
+            ...(createdAtValid ? { createdAt } : {}),
+          }
+        : msg,
+    )
+    this.emitMessagesChange()
+  }
+
   /**
    * Handle TEXT_MESSAGE_START event
    */
@@ -819,6 +888,15 @@ export class StreamProcessor {
         // Update activeMessageIds
         this.activeMessageIds.delete(pendingId)
         this.activeMessageIds.add(messageId)
+
+        // TOOL_CALL_ARGS/END route through toolCallToMessage. Keep those
+        // entries on the remapped id so later args still accumulate
+        // (interleaved text can arrive as a full START/CONTENT/END block).
+        for (const [toolCallId, mappedMessageId] of this.toolCallToMessage) {
+          if (mappedMessageId === pendingId) {
+            this.toolCallToMessage.set(toolCallId, messageId)
+          }
+        }
       }
 
       // Ensure state exists
@@ -827,6 +905,7 @@ export class StreamProcessor {
         this.activeMessageIds.add(messageId)
       }
 
+      this.mergeMessageMetadata(messageId, chunk.metadata)
       this.emitMessagesChange()
       return
     }
@@ -852,6 +931,7 @@ export class StreamProcessor {
           existingState.hasToolCallsSinceTextStart = false
         }
       }
+      this.mergeMessageMetadata(messageId, chunk.metadata)
       return
     }
 
@@ -867,6 +947,7 @@ export class StreamProcessor {
     this.createMessageState(messageId, uiRole)
     this.activeMessageIds.add(messageId)
 
+    this.mergeMessageMetadata(messageId, chunk.metadata)
     this.events.onStreamStart?.()
     this.emitMessagesChange()
   }
@@ -878,6 +959,7 @@ export class StreamProcessor {
     chunk: Extract<StreamChunk, { type: 'TEXT_MESSAGE_END' }>,
   ): void {
     const { messageId } = chunk
+    this.mergeMessageMetadata(messageId, chunk.metadata)
     const state = this.getMessageState(messageId)
     if (!state) return
     if (state.isComplete) return
@@ -886,9 +968,6 @@ export class StreamProcessor {
     if (state.currentSegmentText !== state.lastEmittedText) {
       this.emitTextUpdateForMessage(messageId)
     }
-
-    // Complete all tool calls for this message
-    this.completeAllToolCallsForMessage(messageId)
   }
 
   /**
@@ -914,8 +993,19 @@ export class StreamProcessor {
     // unreconstructable metadata, reconcile the normalized snapshot against
     // the pre-snapshot state; see `reconcileSnapshotToolCalls`.
     const prevMessages = this.messages
-    const normalized = chunk.messages.map(aguiSnapshotMessageToUIMessage)
-    this.messages = this.reconcileSnapshotToolCalls(normalized, prevMessages)
+    const prevById = new Map(prevMessages.map((msg) => [msg.id, msg]))
+    const normalized = this.mergeReasoningFanOut(
+      chunk.messages.map(aguiSnapshotMessageToUIMessage),
+    )
+    this.messages = this.reconcileSnapshotToolCalls(
+      normalized,
+      prevMessages,
+    ).map((msg) => {
+      if (msg.metadata != null) return msg
+      const prev = prevById.get(msg.id)
+      if (prev?.metadata == null) return msg
+      return { ...msg, metadata: prev.metadata }
+    })
     this.emitMessagesChange()
   }
 
@@ -942,6 +1032,54 @@ export class StreamProcessor {
    * matching `tool-call` (and prefer pre-snapshot complete/output when the
    * snapshot is poorer) so server tools keep the same UI shape as client tools.
    */
+  /**
+   * Wire order is reasoning fan-outs, then the assistant anchor.
+   * Snapshot conversion turns each reasoning row into its own assistant
+   * message. Fold leading thinking-only messages into the next real
+   * assistant. Do not fold into a tool-result-only message (`role: 'tool'`
+   * on the wire). `reconcileSnapshotToolCalls` anchors those results.
+   */
+  private mergeReasoningFanOut(messages: Array<UIMessage>): Array<UIMessage> {
+    const out: Array<UIMessage> = []
+    let pending: Array<UIMessage> = []
+    const thinkingParts = (msg: UIMessage) =>
+      msg.parts.filter((part): part is ThinkingPart => part.type === 'thinking')
+    const isThinkingOnly = (msg: UIMessage) =>
+      msg.role === 'assistant' &&
+      msg.parts.length > 0 &&
+      msg.parts.every((part) => part.type === 'thinking')
+    const isToolResultOnly = (msg: UIMessage) =>
+      msg.role === 'assistant' &&
+      msg.parts.length === 1 &&
+      msg.parts[0]?.type === 'tool-result'
+    const flushPending = () => {
+      out.push(...pending)
+      pending = []
+    }
+    for (const msg of messages) {
+      if (isThinkingOnly(msg)) {
+        pending.push(msg)
+        continue
+      }
+      if (
+        msg.role === 'assistant' &&
+        pending.length > 0 &&
+        !isToolResultOnly(msg)
+      ) {
+        out.push({
+          ...msg,
+          parts: [...pending.flatMap(thinkingParts), ...msg.parts],
+        })
+        pending = []
+        continue
+      }
+      flushPending()
+      out.push(msg)
+    }
+    flushPending()
+    return out
+  }
+
   private reconcileSnapshotToolCalls(
     snapshot: Array<UIMessage>,
     prevMessages: Array<UIMessage>,
@@ -1138,33 +1276,10 @@ export class StreamProcessor {
     chunk: Extract<StreamChunk, { type: 'TEXT_MESSAGE_CONTENT' }>,
   ): void {
     const { messageId, state } = this.ensureAssistantMessage(chunk.messageId)
-
-    // Content arriving means all current tool calls for this message are complete
-    this.completeAllToolCallsForMessage(messageId)
+    this.mergeMessageMetadata(messageId, chunk.metadata)
 
     if (this.structuredMessageIds.has(messageId)) {
-      // `chunk.delta` is incremental; `chunk.content` is sometimes cumulative
-      // (mirrors what the plain-text branch handles below). Reconcile against
-      // the existing raw buffer so adapters that emit cumulative content
-      // don't duplicate the JSON.
-      let delta = chunk.delta || ''
-      if (delta === '' && chunk.content !== undefined && chunk.content !== '') {
-        const existingRaw = (
-          this.messages
-            .find((m) => m.id === messageId)
-            ?.parts.find(
-              (p): p is Extract<MessagePart, { type: 'structured-output' }> =>
-                p.type === 'structured-output',
-            ) ?? { raw: '' }
-        ).raw
-        if (chunk.content.startsWith(existingRaw)) {
-          delta = chunk.content.slice(existingRaw.length)
-        } else if (existingRaw.startsWith(chunk.content)) {
-          delta = ''
-        } else {
-          delta = chunk.content
-        }
-      }
+      const delta = chunk.delta || ''
       if (delta !== '') {
         this.messages = appendStructuredOutputDelta(
           this.messages,
@@ -1198,31 +1313,15 @@ export class StreamProcessor {
     }
 
     const currentText = state.currentSegmentText
-    let nextText = currentText
-
-    // Prefer delta over content - delta is the incremental change
-    // Normalize to empty string to avoid "undefined" string concatenation
     const delta = chunk.delta || ''
-    if (delta !== '') {
-      nextText = currentText + delta
-    } else if (chunk.content !== undefined && chunk.content !== '') {
-      // Fallback: use content if delta is not provided
-      if (chunk.content.startsWith(currentText)) {
-        nextText = chunk.content
-      } else if (currentText.startsWith(chunk.content)) {
-        nextText = currentText
-      } else {
-        nextText = currentText + chunk.content
-      }
-    }
+    const nextText = delta !== '' ? currentText + delta : currentText
 
     // Calculate the delta for totalTextContent
     const textDelta = nextText.slice(currentText.length)
     state.currentSegmentText = nextText
     state.totalTextContent += textDelta
 
-    // Use delta for chunk strategy if available
-    const chunkPortion = chunk.delta || chunk.content || ''
+    const chunkPortion = chunk.delta || ''
     const shouldEmit = this.chunkStrategy.shouldEmit(
       chunkPortion,
       state.currentSegmentText,
@@ -1265,13 +1364,7 @@ export class StreamProcessor {
       // New tool call starting
       const initialState: ToolCallState = 'awaiting-input'
 
-      // `toolName` is a deprecated alias for `toolCallName` (see ToolCallStartEvent
-      // in types.ts). Accept either so chunks from older code paths or any
-      // adapter that only sets the deprecated field still produce a named part.
-      // The type marks both as required strings, but in practice some emitters
-      // only set one — fall back via the runtime value rather than the type.
-      const toolName =
-        (chunk as { toolCallName?: string }).toolCallName ?? chunk.toolName
+      const toolName = chunk.toolCallName
 
       // Capture provider metadata that arrived on TOOL_CALL_START so it
       // round-trips back through the assistant message on the next turn
@@ -1284,7 +1377,7 @@ export class StreamProcessor {
         arguments: '',
         state: initialState,
         parsedArguments: undefined,
-        index: chunk.index ?? state.toolCalls.size,
+        index: state.toolCalls.size,
         ...(chunkMetadata !== undefined && { metadata: chunkMetadata }),
       }
 
@@ -1373,17 +1466,15 @@ export class StreamProcessor {
   }
 
   /**
-   * Handle TOOL_CALL_END event — authoritative signal that a tool call's input is finalized.
+   * Handle TOOL_CALL_END event — arguments are finalized (input-complete).
+   * Tool output arrives on TOOL_CALL_RESULT, not on this event.
    *
-   * This event has a DUAL ROLE:
-   * - Without `result`: Signals arguments are done (from adapter). Transitions to input-complete.
-   * - With `result`: Signals tool was executed and result is available (from TextEngine).
-   *   Creates both output on the tool-call part AND a tool-result part.
+   * If TOOL_CALL_END carries parsed `input`, use it as the canonical arguments:
+   * back-fill the accumulated string when no TOOL_CALL_ARGS deltas were seen
+   * (adapters that deliver the whole input on END — e.g. Anthropic
+   * server_tool_use / web_search — issue #839) and override the rendered part's
+   * `input` with the canonical value.
    *
-   * If `input` is provided, it overrides the accumulated string parse as the
-   * canonical parsed arguments.
-   *
-   * @see docs/chat-architecture.md#tool-results-and-the-tool_call_end-dual-role — Full explanation
    * @see docs/chat-architecture.md#single-shot-tool-call-response — End-to-end flow
    */
   private handleToolCallEndEvent(
@@ -1395,15 +1486,22 @@ export class StreamProcessor {
     const msgState = this.getMessageState(messageId)
     if (!msgState) return
 
+    // The parsed input can ride on the spec `input` field or, for adapters
+    // that only stamp it into TanStack metadata, on `metadata.tanstack.input`.
+    const input =
+      chunk.input !== undefined
+        ? chunk.input
+        : (tanstackMetadata(chunk)?.input as unknown)
+
     // Transition the tool call to input-complete (the authoritative completion signal)
     const existingToolCall = msgState.toolCalls.get(chunk.toolCallId)
     if (existingToolCall && existingToolCall.state !== 'input-complete') {
-      // If TOOL_CALL_END provides parsed input and no TOOL_CALL_ARGS were
-      // received, back-fill the arguments string so the UIMessage ToolCallPart
-      // carries the correct value (defensive against adapters that skip ARGS).
-      if (chunk.input !== undefined && !existingToolCall.arguments) {
+      // Back-fill the arguments string from the parsed input when no
+      // TOOL_CALL_ARGS deltas were received, so completeToolCall's strict parse
+      // surfaces the correct value on the ToolCallPart.
+      if (input !== undefined && !existingToolCall.arguments) {
         try {
-          existingToolCall.arguments = JSON.stringify(chunk.input)
+          existingToolCall.arguments = JSON.stringify(input)
         } catch {
           // circular refs, BigInt, etc. — leave arguments empty rather than
           // aborting stream processing
@@ -1412,20 +1510,18 @@ export class StreamProcessor {
 
       const index = msgState.toolCallOrder.indexOf(chunk.toolCallId)
       this.completeToolCall(messageId, index, existingToolCall)
-      // If TOOL_CALL_END provides parsed input, use it as the canonical parsed
-      // arguments (overrides the accumulated string parse from completeToolCall)
-      // and refresh the rendered part's `input` so it reflects the canonical
-      // value rather than the possibly-divergent accumulated-args parse that
-      // completeToolCall wrote (e.g. an adapter that coerces values differently
-      // between the streamed args and the final structured input).
-      if (chunk.input !== undefined) {
-        existingToolCall.parsedArguments = chunk.input
+
+      // Canonicalize on the parsed input: overrides the accumulated-args parse
+      // that completeToolCall wrote (adapters may coerce values differently
+      // between streamed args and the final structured input).
+      if (input !== undefined) {
+        existingToolCall.parsedArguments = input
         this.messages = updateToolCallPart(this.messages, messageId, {
           id: existingToolCall.id,
           name: existingToolCall.name,
           arguments: existingToolCall.arguments,
           state: 'input-complete',
-          input: chunk.input,
+          input,
           ...(existingToolCall.metadata !== undefined && {
             metadata: existingToolCall.metadata,
           }),
@@ -1433,43 +1529,6 @@ export class StreamProcessor {
         this.emitMessagesChange()
       }
     }
-
-    // Update UIMessage if there's a result
-    if (chunk.result) {
-      // Step 1: Update the tool-call part's output field (for UI consistency
-      // with client tools — see GitHub issue #176)
-      let output: unknown
-      if (Array.isArray(chunk.result)) {
-        output = chunk.result
-      } else {
-        try {
-          output = JSON.parse(chunk.result)
-        } catch {
-          output = chunk.result
-        }
-      }
-      this.messages = updateToolCallWithOutput(
-        this.messages,
-        chunk.toolCallId,
-        output,
-        chunk.state === 'output-error' ? 'error' : undefined,
-      )
-
-      // Step 2: Create/update the tool-result part (for LLM conversation history)
-      const resultState: ToolResultState =
-        chunk.state === 'output-error' ? 'error' : 'complete'
-      this.messages = updateToolResultPart(
-        this.messages,
-        messageId,
-        chunk.toolCallId,
-        chunk.result,
-        resultState,
-        resultState === 'error'
-          ? this.extractToolResultError(output)
-          : undefined,
-      )
-      this.emitMessagesChange()
-    }
   }
 
   private extractToolResultError(output: unknown): string {
@@ -1494,9 +1553,27 @@ export class StreamProcessor {
   private handleToolCallResultEvent(
     chunk: Extract<StreamChunk, { type: 'TOOL_CALL_RESULT' }>,
   ): void {
-    const messageId = this.toolCallToMessage.get(chunk.toolCallId)
+    // A resume stream delivers TOOL_CALL_RESULT for a tool call that was
+    // started in a PRIOR run. A preceding MESSAGES_SNAPSHOT resets stream state
+    // (clearing `toolCallToMessage`), so fall back to locating the message that
+    // owns the tool-call part — matching `addToolResult`/`handleInterrupts`.
+    // Without this the result is dropped, the follow-up request omits the tool
+    // message, and the still-"pending" tool call re-interrupts (issue #532).
+    const messageId =
+      this.toolCallToMessage.get(chunk.toolCallId) ??
+      this.messages.find((m) =>
+        m.parts.some(
+          (p): p is ToolCallPart =>
+            p.type === 'tool-call' && p.id === chunk.toolCallId,
+        ),
+      )?.id
     if (!messageId) return
 
+    const extra = chunk as AdapterYieldChunk
+    const isOutputError =
+      extra.state === 'output-error' ||
+      tanstackMetadata(chunk)?.state === 'output-error'
+
     // Step 1: Update the tool-call part's output field
     let output: unknown
     try {
@@ -1508,12 +1585,11 @@ export class StreamProcessor {
       this.messages,
       chunk.toolCallId,
       output,
-      chunk.state === 'output-error' ? 'error' : undefined,
+      isOutputError ? 'error' : undefined,
     )
 
     // Step 2: Create/update the tool-result part
-    const resultState: ToolResultState =
-      chunk.state === 'output-error' ? 'error' : 'complete'
+    const resultState: ToolResultState = isOutputError ? 'error' : 'complete'
     this.messages = updateToolResultPart(
       this.messages,
       messageId,
@@ -1550,7 +1626,11 @@ export class StreamProcessor {
   private handleRunFinishedEvent(
     chunk: Extract<StreamChunk, { type: 'RUN_FINISHED' }>,
   ): void {
-    this.finishReason = chunk.finishReason ?? null
+    const extra = chunk as AdapterYieldChunk
+    this.finishReason =
+      extra.finishReason !== undefined
+        ? extra.finishReason
+        : (tanstackMetadata(chunk)?.finishReason ?? null)
     this.activeRuns.delete(chunk.runId)
 
     if (chunk.outcome?.type === 'interrupt') {
@@ -1643,10 +1723,7 @@ export class StreamProcessor {
     chunk: Extract<StreamChunk, { type: 'RUN_ERROR' }>,
   ): void {
     this.hasError = true
-    const runId =
-      'runId' in chunk && typeof chunk.runId === 'string'
-        ? chunk.runId
-        : undefined
+    const runId = getChunkRunId(chunk)
     if (runId) {
       this.activeRuns.delete(runId)
     } else {
@@ -1656,9 +1733,8 @@ export class StreamProcessor {
     // Prefer spec field `message`; fall back to deprecated `error.message`.
     // If neither is set, the chunk still carries debug context (provider
     // error codes, request ids, etc.) — log it so the failure isn't silent.
-    const errorMessage =
-      chunk.message || chunk.error?.message || 'An error occurred'
-    if (!chunk.message && !chunk.error?.message) {
+    const errorMessage = chunk.message || 'An error occurred'
+    if (!chunk.message) {
       console.error(
         '[StreamProcessor] RUN_ERROR with no message; original chunk:',
         chunk,
@@ -1681,28 +1757,20 @@ export class StreamProcessor {
     // the surfaced Error so consumers can recover the upstream detail that the
     // RUN_ERROR's `message` alone discards. Both are optional and added only
     // when present, keeping the Error backward compatible.
-    const error = new Error(errorMessage)
-    const code = chunk.code ?? chunk.error?.code
-    if (code !== undefined) {
-      Object.assign(error, { code })
-    }
-    if (chunk.rawEvent !== undefined) {
-      Object.assign(error, { rawEvent: chunk.rawEvent })
-    }
-    this.events.onError?.(error)
+    this.events.onError?.(runErrorEventToError(chunk))
   }
 
   /**
    * Handle STEP_STARTED event (for thinking/reasoning content).
    *
-   * Records the stepId so that subsequent STEP_FINISHED deltas accumulate
+   * Records the stepId so later REASONING_MESSAGE_CONTENT deltas accumulate
    * into their own ThinkingPart. Does not create a message — the message
-   * is lazily created when the first STEP_FINISHED content arrives.
+   * is lazily created when the first REASONING_MESSAGE_CONTENT arrives.
    */
   private handleStepStartedEvent(
     chunk: Extract<StreamChunk, { type: 'STEP_STARTED' }>,
   ): void {
-    const stepId = chunk.stepId ?? generateMessageId()
+    const stepId = chunk.stepName || generateMessageId()
     const activeId = this.getActiveAssistantMessageId()
     if (activeId) {
       const state = this.getMessageState(activeId)
@@ -1715,102 +1783,49 @@ export class StreamProcessor {
         // Clear any pending stepId from a prior STEP_STARTED that fired
         // before the assistant message existed. Now that we're tracking
         // the step directly on message state, the pending value is stale
-        // and must not leak into the next STEP_FINISHED (which would
-        // misattribute its delta to the stale step).
+        // and must not leak into the next REASONING_MESSAGE_CONTENT.
         this.pendingThinkingStepId = null
         return
       }
     }
 
-    // No active message yet — defer until ensureAssistantMessage in STEP_FINISHED
+    // No active message yet — defer until ensureAssistantMessage in
+    // REASONING_MESSAGE_CONTENT
     this.pendingThinkingStepId = stepId
   }
 
   /**
-   * Handle STEP_FINISHED event (for thinking/reasoning content).
-   *
-   * Accumulates delta into the current thinking step's content and updates
-   * the corresponding ThinkingPart in the UIMessage.
+   * Handle STEP_FINISHED event.
    *
-   * @see docs/chat-architecture.md#thinkingreasoning-content — Thinking flow
+   * Thinking *content* comes from REASONING_MESSAGE_CONTENT, not STEP_FINISHED.
+   * But some adapters (e.g. BytePlus thinking-summary) carry the provider
+   * signature blob ONLY on the STEP_FINISHED event, so still extract that here
+   * and attach it to the thinking step the reasoning events already built.
    */
   private handleStepFinishedEvent(
     chunk: Extract<StreamChunk, { type: 'STEP_FINISHED' }>,
   ): void {
+    const extra = chunk as AdapterYieldChunk
+    const signature = extra.signature
+    if (!signature) return
+
     const { messageId, state } = this.ensureAssistantMessage(
       this.getActiveAssistantMessageId() ?? undefined,
     )
+    const stepId = state.currentThinkingStepId ?? extra.stepId
+    if (!stepId) return
+    const thinking = state.thinkingSteps.get(stepId)
+    if (thinking === undefined) return
 
-    // During the transition period, adapters emit BOTH STEP_FINISHED and
-    // REASONING_MESSAGE_CONTENT with the same delta. If we've already processed
-    // REASONING_MESSAGE_CONTENT events for this message, skip the duplicate
-    // thinking content from STEP_FINISHED to avoid doubled content.
-    if (state.hasSeenReasoningEvents) {
-      if (chunk.signature) {
-        const stepId = state.currentThinkingStepId ?? chunk.stepId
-        if (!stepId) return
-        const thinking = state.thinkingSteps.get(stepId)
-        if (thinking !== undefined) {
-          state.thinkingStepSignatures.set(stepId, chunk.signature)
-          this.messages = updateThinkingPart(
-            this.messages,
-            messageId,
-            stepId,
-            thinking,
-            chunk.signature,
-          )
-          this.emitMessagesChange()
-        }
-      }
-      return
-    }
-
-    this.consumePendingThinkingStep(state)
-
-    const stepId =
-      state.currentThinkingStepId ?? chunk.stepId ?? generateMessageId()
-
-    // Auto-initialize if no prior STEP_STARTED (backward compat)
-    if (!state.thinkingSteps.has(stepId)) {
-      state.thinkingSteps.set(stepId, '')
-      state.thinkingStepOrder.push(stepId)
-      state.currentThinkingStepId = stepId
-    }
-
-    const previous = state.thinkingSteps.get(stepId) ?? ''
-    let nextThinking = previous
-
-    // Prefer delta over content
-    if (chunk.delta && chunk.delta !== '') {
-      nextThinking = previous + chunk.delta
-    } else if (chunk.content && chunk.content !== '') {
-      if (chunk.content.startsWith(previous)) {
-        nextThinking = chunk.content
-      } else if (previous.startsWith(chunk.content)) {
-        nextThinking = previous
-      } else {
-        nextThinking = previous + chunk.content
-      }
-    }
-
-    state.thinkingSteps.set(stepId, nextThinking)
-
-    if (chunk.signature) {
-      state.thinkingStepSignatures.set(stepId, chunk.signature)
-    }
-
-    // Update UIMessage
+    state.thinkingStepSignatures.set(stepId, signature)
     this.messages = updateThinkingPart(
       this.messages,
       messageId,
       stepId,
-      nextThinking,
-      state.thinkingStepSignatures.get(stepId),
+      thinking,
+      signature,
     )
     this.emitMessagesChange()
-
-    // Emit granular event
-    this.events.onThinkingUpdate?.(messageId, stepId, nextThinking)
   }
 
   /**
@@ -1853,6 +1868,74 @@ export class StreamProcessor {
     this.events.onThinkingUpdate?.(messageId, stepId, nextThinking)
   }
 
+  /**
+   * Attach a provider signature blob from REASONING_ENCRYPTED_VALUE.
+   * `subtype: 'message'` updates ThinkingPart.signature.
+   * `subtype: 'tool-call'` stores Gemini thoughtSignature on the tool-call part.
+   */
+  private handleReasoningEncryptedValueEvent(
+    chunk: Extract<StreamChunk, { type: 'REASONING_ENCRYPTED_VALUE' }>,
+  ): void {
+    const encryptedValue = chunk.encryptedValue
+    if (typeof encryptedValue !== 'string' || encryptedValue === '') return
+
+    if (chunk.subtype === 'tool-call') {
+      this.attachToolCallSignature(chunk.entityId, encryptedValue)
+      return
+    }
+
+    const { messageId, state } = this.ensureAssistantMessage(
+      this.getActiveAssistantMessageId() ?? undefined,
+    )
+    const stepId = state.currentThinkingStepId ?? chunk.entityId
+    state.thinkingStepSignatures.set(stepId, encryptedValue)
+    const content = state.thinkingSteps.get(stepId) ?? ''
+    if (!state.thinkingSteps.has(stepId)) {
+      state.thinkingSteps.set(stepId, content)
+      state.thinkingStepOrder.push(stepId)
+    }
+    this.messages = updateThinkingPart(
+      this.messages,
+      messageId,
+      stepId,
+      content,
+      encryptedValue,
+    )
+    this.emitMessagesChange()
+  }
+
+  private attachToolCallSignature(
+    toolCallId: string,
+    thoughtSignature: string,
+  ): void {
+    this.messages = this.messages.map((msg) => {
+      let changed = false
+      const parts = msg.parts.map((part) => {
+        if (part.type !== 'tool-call' || part.id !== toolCallId) return part
+        changed = true
+        return {
+          ...part,
+          metadata: {
+            ...(part.metadata != null && typeof part.metadata === 'object'
+              ? part.metadata
+              : {}),
+            thoughtSignature,
+          },
+        }
+      })
+      return changed ? { ...msg, parts } : msg
+    })
+    for (const state of this.messageStates.values()) {
+      const call = state.toolCalls.get(toolCallId)
+      if (!call) continue
+      call.metadata = {
+        ...(call.metadata ?? {}),
+        thoughtSignature,
+      }
+    }
+    this.emitMessagesChange()
+  }
+
   /**
    * Handle CUSTOM event.
    *
@@ -2026,22 +2109,10 @@ export class StreamProcessor {
    * Detect if an incoming content chunk represents a NEW text segment
    */
   private isNewTextSegment(
-    chunk: Extract<StreamChunk, { type: 'TEXT_MESSAGE_CONTENT' }>,
-    previous: string,
+    _chunk: Extract<StreamChunk, { type: 'TEXT_MESSAGE_CONTENT' }>,
+    _previous: string,
   ): boolean {
-    // Check if content is present (delta is always defined but may be empty string)
-    if (chunk.content !== undefined) {
-      if (chunk.content.length < previous.length) {
-        return true
-      }
-      if (
-        !chunk.content.startsWith(previous) &&
-        !previous.startsWith(chunk.content)
-      ) {
-        return true
-      }
-    }
-    return false
+    return true
   }
 
   /**
@@ -2087,8 +2158,17 @@ export class StreamProcessor {
     // counts as a completed tool call in getCompletedToolCalls()/getState().
     toolCall.state = 'input-complete'
 
-    // Try final parse
-    toolCall.parsedArguments = this.jsonParser.parse(toolCall.arguments)
+    // Only surface `input` from a strict parse. The streaming partial-JSON
+    // parser closes unterminated strings, so truncated arguments would become
+    // a plausible but wrong object (GitHub issue #1017). If parse fails,
+    // `input` stays unset and consumers use the raw `arguments` string.
+    let strictParseSucceeded = false
+    try {
+      toolCall.parsedArguments = JSON.parse(toolCall.arguments)
+      strictParseSucceeded = true
+    } catch {
+      toolCall.parsedArguments = undefined
+    }
 
     // Don't downgrade the rendered part of a call that already reached the
     // terminal 'error' state (e.g. an output-error TOOL_CALL_RESULT arrived
@@ -2106,19 +2186,13 @@ export class StreamProcessor {
     }
 
     // Update UIMessage. The arguments are complete now, so surface the parsed
-    // input on the part. For adapters that skip TOOL_CALL_ARGS the arguments
-    // string was back-filled from TOOL_CALL_END.input, so this parse matches
-    // the canonical input. If a TOOL_CALL_END.input diverges from the
-    // accumulated args, handleToolCallEndEvent re-updates the part with the
-    // canonical value after this call.
+    // input on the part from the accumulated TOOL_CALL_ARGS deltas.
     this.messages = updateToolCallPart(this.messages, messageId, {
       id: toolCall.id,
       name: toolCall.name,
       arguments: toolCall.arguments,
       state: 'input-complete',
-      ...(toolCall.parsedArguments !== undefined && {
-        input: toolCall.parsedArguments,
-      }),
+      ...(strictParseSucceeded && { input: toolCall.parsedArguments }),
       ...(toolCall.metadata !== undefined && { metadata: toolCall.metadata }),
     })
     this.emitMessagesChange()