fix(ai): preserve metadata on AG-UI TOOL_CALL_RESULT

.changeset/wacky-gifts-show.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai': patch
+---
+
+Preserve metadata from live AG-UI TOOL_CALL_RESULT events on ToolResultPart.

packages/ai/src/activities/chat/stream/message-updaters.ts

@@ -6,6 +6,10 @@
  */
 
 import { isRedactedThinkingId } from '../../../utilities/reasoning-encrypted-value'
+import {
+  mergeMetadata,
+  tanstackMetadata,
+} from '../../../utilities/merge-metadata'
 import { parsePartialJSON } from './json-parser'
 import type {
   ContentPart,
@@ -109,6 +113,23 @@ export function updateToolCallPart(
   })
 }
 
+/** Drop a control outcome left on metadata after the result itself succeeded. */
+export function withoutStaleToolResultOutcome(
+  metadata: Record<string, unknown> | undefined,
+): Record<string, unknown> | undefined {
+  if (metadata == null) return undefined
+  const tanstack = tanstackMetadata({ metadata })
+  if (tanstack?.toolResultOutcome === undefined) return metadata
+  const { toolResultOutcome: _staleOutcome, ...restTanstack } = tanstack
+  const nextMetadata = { ...metadata }
+  if (Object.keys(restTanstack).length > 0) {
+    nextMetadata.tanstack = restTanstack
+  } else {
+    delete nextMetadata.tanstack
+  }
+  return Object.keys(nextMetadata).length > 0 ? nextMetadata : undefined
+}
+
 /**
  * Update or add a tool result part to a message.
  */
@@ -120,6 +141,7 @@ export function updateToolResultPart(
   state: ToolResultState,
   error?: string,
   outcome?: ToolResultOutcome,
+  metadata?: Record<string, unknown>,
 ): Array<UIMessage> {
   return messages.map((msg) => {
     if (msg.id !== messageId) {
@@ -132,13 +154,35 @@ export function updateToolResultPart(
         p.type === 'tool-result' && p.toolCallId === toolCallId,
     )
 
+    const existingResultPart =
+      resultPartIndex >= 0
+        ? (parts[resultPartIndex] as ToolResultPart)
+        : undefined
+    // Preserve existing metadata when a later result update does not include
+    // metadata, and merge partial metadata (including metadata.tanstack).
+    let mergedMetadata = mergeMetadata(existingResultPart?.metadata, metadata)
+    const hasIncomingToolResultOutcome =
+      tanstackMetadata({ metadata })?.toolResultOutcome !== undefined
+
+    // A successful retry supersedes a previous denial/cancellation. Do not
+    // carry the old control outcome into the new result, or snapshot hydration
+    // will incorrectly restore the successful result as an error.
+    if (
+      state === 'complete' &&
+      outcome === undefined &&
+      !hasIncomingToolResultOutcome
+    ) {
+      mergedMetadata = withoutStaleToolResultOutcome(mergedMetadata)
+    }
+
     const toolResultPart: ToolResultPart = {
       type: 'tool-result',
       toolCallId,
       content,
       state,
       ...(outcome !== undefined && { outcome }),
       ...(error && { error }),
+      ...(mergedMetadata !== undefined && { metadata: mergedMetadata }),
     }
 
     if (resultPartIndex >= 0) {

packages/ai/src/activities/chat/stream/processor.ts

@@ -58,6 +58,7 @@ import {
   updateToolCallPart,
   updateToolCallWithOutput,
   updateToolResultPart,
+  withoutStaleToolResultOutcome,
 } from './message-updaters'
 import { ImmediateStrategy } from './strategies'
 import { INTERRUPT_BINDING_METADATA_KEY } from '../../../interrupt-resume'
@@ -2095,11 +2096,16 @@ export class StreamProcessor {
     // restore metadata the snapshot cannot re-emit. Duplicate ids resolve
     // last-write-wins: the same tool call can appear in multiple messages
     // across reconnects, and the most recent part carries the freshest state.
+    // Snapshot rows may omit metadata from previously streamed tool results.
+    // Index both part types by id so we can restore only the missing metadata.
     const prevToolCalls = new Map<string, ToolCallPart>()
+    const prevToolResults = new Map<string, ToolResultPart>()
     for (const msg of prevMessages) {
       for (const part of msg.parts) {
         if (part.type === 'tool-call') {
           prevToolCalls.set(part.id, part)
+        } else if (part.type === 'tool-result') {
+          prevToolResults.set(part.toolCallId, part)
         }
       }
     }
@@ -2116,20 +2122,36 @@ export class StreamProcessor {
 
     const reconciled: Array<UIMessage> = []
     for (const msg of snapshot) {
-      const toolResultParts = msg.parts.filter(
+      let restoredToolResultMetadata = false
+      const snapshotParts = msg.parts.map((part) => {
+        if (part.type !== 'tool-result' || part.metadata != null) return part
+        const previous = prevToolResults.get(part.toolCallId)
+        if (previous?.metadata == null) return part
+        const metadata =
+          part.state === 'complete' && part.outcome === undefined
+            ? withoutStaleToolResultOutcome(previous.metadata)
+            : previous.metadata
+        if (metadata == null) return part
+        restoredToolResultMetadata = true
+        return { ...part, metadata }
+      })
+      const snapshotMessage = restoredToolResultMetadata
+        ? { ...msg, parts: snapshotParts }
+        : msg
+      const toolResultParts = snapshotParts.filter(
         (part): part is ToolResultPart => part.type === 'tool-result',
       )
       const toolResultPart =
         msg.role === 'assistant' &&
         toolResultParts.length === 1 &&
-        msg.parts.every(
+        snapshotParts.every(
           (part) => part.type === 'tool-result' || part.type === 'ui-resource',
         )
           ? toolResultParts[0]
           : undefined
 
       if (!toolResultPart) {
-        reconciled.push(msg)
+        reconciled.push(snapshotMessage)
         continue
       }
 
@@ -2157,7 +2179,7 @@ export class StreamProcessor {
             `[StreamProcessor] MESSAGES_SNAPSHOT contains a tool-result for "${toolResultPart.toolCallId}" but no matching tool-call exists in the snapshot, and there is no assistant message to anchor into; addToolResult("${toolResultPart.toolCallId}") will not be able to locate this call`,
           )
         }
-        reconciled.push(msg)
+        reconciled.push(snapshotMessage)
         continue
       }
 
@@ -2183,7 +2205,7 @@ export class StreamProcessor {
           )
         }
       }
-      parts.push(...msg.parts)
+      parts.push(...snapshotParts)
       // Replace rather than push into `target.parts`: a snapshot message that
       // arrived already carrying `parts` (TanStack server echoing UIMessages)
       // shares its array with the incoming chunk, and mutating it in place
@@ -2647,6 +2669,7 @@ export class StreamProcessor {
       resultState,
       resultState === 'error' ? toolResultErrorText(output) : undefined,
       toolResultOutcome,
+      chunk.metadata,
     )
     this.emitMessagesChange()
   }

packages/ai/tests/message-updaters.test.ts

@@ -386,6 +386,28 @@ describe('message-updaters', () => {
       })
     })
 
+    it('should include metadata on a new tool result', () => {
+      const messages = [createMessage('msg-1')]
+      const metadata = { displayType: 'account-assets', accountId: 'A001' }
+
+      const result = updateToolResultPart(
+        messages,
+        'msg-1',
+        'call-1',
+        '{"totalAssets":1000000}',
+        'complete',
+        undefined,
+        undefined,
+        metadata,
+      )
+
+      expect(result[0]?.parts[0]).toMatchObject({
+        type: 'tool-result',
+        toolCallId: 'call-1',
+        metadata,
+      })
+    })
+
     it('should update an existing tool result by toolCallId', () => {
       const messages = [
         createMessage('msg-1', 'assistant', [
@@ -414,6 +436,150 @@ describe('message-updaters', () => {
       })
     })
 
+    it('should preserve existing metadata when an update omits metadata', () => {
+      const metadata = { displayType: 'account-assets' }
+      const messages = [
+        createMessage('msg-1', 'assistant', [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            content: '{"temperature":',
+            state: 'streaming',
+            metadata,
+          },
+        ]),
+      ]
+
+      const result = updateToolResultPart(
+        messages,
+        'msg-1',
+        'call-1',
+        '{"temperature":20}',
+        'complete',
+      )
+
+      expect(result[0]?.parts[0]).toMatchObject({
+        content: '{"temperature":20}',
+        state: 'complete',
+        metadata,
+      })
+    })
+
+    it('should clear stale tool result outcome metadata on a successful update', () => {
+      const messages = [
+        createMessage('msg-1', 'assistant', [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            content: '{"error":"User declined"}',
+            state: 'error',
+            metadata: {
+              displayType: 'account-assets',
+              accountId: 'A001',
+              tanstack: {
+                toolResultOutcome: 'denied',
+                retained: true,
+              },
+            },
+          },
+        ]),
+      ]
+
+      const result = updateToolResultPart(
+        messages,
+        'msg-1',
+        'call-1',
+        '{"ok":true}',
+        'complete',
+      )
+
+      expect(result[0]?.parts[0]).toEqual({
+        type: 'tool-result',
+        toolCallId: 'call-1',
+        content: '{"ok":true}',
+        state: 'complete',
+        metadata: {
+          displayType: 'account-assets',
+          accountId: 'A001',
+          tanstack: { retained: true },
+        },
+      })
+    })
+
+    it('should preserve outcome metadata explicitly supplied by an update', () => {
+      const incomingMetadata = {
+        displayType: 'account-assets',
+        tanstack: { toolResultOutcome: 'cancelled' },
+      }
+      const messages = [
+        createMessage('msg-1', 'assistant', [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            content: '{"error":"User declined"}',
+            state: 'error',
+            metadata: { tanstack: { toolResultOutcome: 'denied' } },
+          },
+        ]),
+      ]
+
+      const result = updateToolResultPart(
+        messages,
+        'msg-1',
+        'call-1',
+        '{"ok":true}',
+        'complete',
+        undefined,
+        undefined,
+        incomingMetadata,
+      )
+
+      expect(result[0]?.parts[0]).toMatchObject({
+        content: '{"ok":true}',
+        state: 'complete',
+        metadata: incomingMetadata,
+      })
+    })
+
+    it('should merge incoming metadata with existing metadata', () => {
+      const messages = [
+        createMessage('msg-1', 'assistant', [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            content: '{"temperature":',
+            state: 'streaming',
+            metadata: {
+              displayType: 'account-assets',
+              tanstack: { existing: true, state: 'streaming' },
+            },
+          },
+        ]),
+      ]
+
+      const result = updateToolResultPart(
+        messages,
+        'msg-1',
+        'call-1',
+        '{"temperature":20}',
+        'complete',
+        undefined,
+        undefined,
+        {
+          accountId: 'A001',
+          tanstack: { state: 'complete' },
+        },
+      )
+
+      expect(result[0]?.parts[0]).toMatchObject({
+        metadata: {
+          displayType: 'account-assets',
+          accountId: 'A001',
+          tanstack: { existing: true, state: 'complete' },
+        },
+      })
+    })
+
     it('should include error when provided', () => {
       const messages = [createMessage('msg-1')]
       const result = updateToolResultPart(