diff --git a/.changeset/wacky-gifts-show.md b/.changeset/wacky-gifts-show.md new file mode 100644 index 0000000000..71d0887086 --- /dev/null +++ b/.changeset/wacky-gifts-show.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +Preserve metadata from live AG-UI TOOL_CALL_RESULT events on ToolResultPart. diff --git a/packages/ai/src/activities/chat/stream/message-updaters.ts b/packages/ai/src/activities/chat/stream/message-updaters.ts index b7e936beab..edf4edb58c 100644 --- a/packages/ai/src/activities/chat/stream/message-updaters.ts +++ b/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 | undefined, +): Record | 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, ): Array { return messages.map((msg) => { if (msg.id !== messageId) { @@ -132,6 +154,27 @@ 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, @@ -139,6 +182,7 @@ export function updateToolResultPart( state, ...(outcome !== undefined && { outcome }), ...(error && { error }), + ...(mergedMetadata !== undefined && { metadata: mergedMetadata }), } if (resultPartIndex >= 0) { diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 917406236f..a552040067 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/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() + const prevToolResults = new Map() 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 = [] 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() } diff --git a/packages/ai/tests/message-updaters.test.ts b/packages/ai/tests/message-updaters.test.ts index fc74c4e347..50bbc8700b 100644 --- a/packages/ai/tests/message-updaters.test.ts +++ b/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( diff --git a/packages/ai/tests/stream-processor.test.ts b/packages/ai/tests/stream-processor.test.ts index 40e04e1541..1a2e8a76b1 100644 --- a/packages/ai/tests/stream-processor.test.ts +++ b/packages/ai/tests/stream-processor.test.ts @@ -176,6 +176,25 @@ function spyEvents(): MockedEvents { // ============================================================================ describe('StreamProcessor', () => { + it('preserves metadata on snapshot tool messages', () => { + const metadata = { + displayType: 'account-assets', + accountId: 'A001', + } + const message = aguiSnapshotMessageToUIMessage({ + id: 'tool-1', + role: 'tool', + toolCallId: 'call-1', + content: '{"totalAssets":1000000}', + metadata, + } as unknown as AGUIMessage) + + const toolResultPart = message.parts.find( + (part): part is ToolResultPart => part.type === 'tool-result', + ) + expect(toolResultPart?.metadata).toEqual(metadata) + }) + it('normalizes nested dates on snapshot existing tool-result parts', () => { const message = aguiSnapshotMessageToUIMessage({ id: 'a1', @@ -3874,6 +3893,140 @@ describe('StreamProcessor', () => { ]) }) + it('preserves streamed tool-result metadata when a snapshot omits it', () => { + const processor = new StreamProcessor() + const streamedMetadata = { + displayType: 'account-assets', + accountId: 'A001', + } + + processor.processChunk(ev.runStarted()) + processor.processChunk(ev.textStart('a1')) + processor.processChunk(ev.textContent('Checking.', 'a1')) + processor.processChunk(ev.toolStart('tc-1', 'getWeather')) + processor.processChunk(ev.toolEnd('tc-1', 'getWeather')) + processor.processChunk( + chunk(EventType.TOOL_CALL_RESULT, { + messageId: 't1', + toolCallId: 'tc-1', + content: '{"temp":72}', + role: 'tool', + metadata: streamedMetadata, + }), + ) + processor.processChunk(ev.runFinished('tool_calls')) + + const snapshotMessages = (metadata?: Record) => [ + { + id: 'a1', + role: 'assistant', + content: 'Checking.', + toolCalls: [ + { + id: 'tc-1', + type: 'function', + function: { name: 'getWeather', arguments: '{"loc":"Berlin"}' }, + }, + ], + }, + { + id: 't1', + role: 'tool', + content: '{"temp":72}', + toolCallId: 'tc-1', + ...(metadata !== undefined && { metadata }), + }, + ] + + processor.processChunk({ + type: EventType.MESSAGES_SNAPSHOT, + messages: snapshotMessages(), + timestamp: Date.now(), + } as unknown as StreamChunk) + + const getToolResult = () => + processor + .getMessages() + .flatMap((message) => message.parts) + .find( + (part): part is ToolResultPart => + part.type === 'tool-result' && part.toolCallId === 'tc-1', + ) + + expect(getToolResult()?.metadata).toEqual(streamedMetadata) + + const snapshotMetadata = { + displayType: 'snapshot-account-assets', + accountId: 'A002', + } + processor.processChunk({ + type: EventType.MESSAGES_SNAPSHOT, + messages: snapshotMessages(snapshotMetadata), + timestamp: Date.now(), + } as unknown as StreamChunk) + + expect(getToolResult()?.metadata).toEqual(snapshotMetadata) + }) + + it('drops a stale denial when a snapshot success omits metadata', () => { + const processor = new StreamProcessor() + + processor.processChunk(ev.runStarted()) + processor.processChunk(ev.textStart('a1')) + processor.processChunk(ev.toolStart('tc-1', 'getWeather')) + processor.processChunk( + chunk(EventType.TOOL_CALL_RESULT, { + messageId: 't1', + toolCallId: 'tc-1', + content: '{"error":"User declined tool execution"}', + role: 'tool', + metadata: { + displayType: 'account-assets', + tanstack: { toolResultOutcome: 'denied', retained: true }, + }, + }), + ) + processor.processChunk(ev.runFinished('tool_calls')) + processor.processChunk({ + type: EventType.MESSAGES_SNAPSHOT, + messages: [ + { + id: 'a1', + role: 'assistant', + content: 'Checking.', + toolCalls: [ + { + id: 'tc-1', + type: 'function', + function: { name: 'getWeather', arguments: '{}' }, + }, + ], + }, + { + id: 't1', + role: 'tool', + content: '{"temp":72}', + toolCallId: 'tc-1', + }, + ], + timestamp: Date.now(), + } as unknown as StreamChunk) + + const toolResult = processor + .getMessages() + .flatMap((message) => message.parts) + .find( + (part): part is ToolResultPart => + part.type === 'tool-result' && part.toolCallId === 'tc-1', + ) + expect(toolResult?.state).toBe('complete') + expect(toolResult?.outcome).toBeUndefined() + expect(toolResult?.metadata).toEqual({ + displayType: 'account-assets', + tanstack: { retained: true }, + }) + }) + it('anchors a tool message that also has a valid UI resource', () => { const processor = new StreamProcessor() const uiResource = { @@ -6365,6 +6518,10 @@ describe('StreamProcessor', () => { it('should create tool-result part and update tool-call output', () => { const events = spyEvents() const processor = new StreamProcessor({ events }) + const metadata = { + displayType: 'account-assets', + accountId: 'A001', + } processor.processChunk(ev.runStarted()) processor.processChunk(ev.textStart()) @@ -6380,6 +6537,7 @@ describe('StreamProcessor', () => { toolCallId: 'tc-1', content: '{"temp": 72}', role: 'tool', + metadata, }), ) processor.processChunk(ev.runFinished('tool_calls')) @@ -6398,6 +6556,7 @@ describe('StreamProcessor', () => { expect(toolResultPart.toolCallId).toBe('tc-1') expect(toolResultPart.content).toBe('{"temp": 72}') expect(toolResultPart.state).toBe('complete') + expect(toolResultPart.metadata).toEqual(metadata) }) it('should mark output-error tool results as errored message parts', () => {