diff --git a/.changeset/preserve-multimodal-tool-results.md b/.changeset/preserve-multimodal-tool-results.md new file mode 100644 index 000000000..4a75e5898 --- /dev/null +++ b/.changeset/preserve-multimodal-tool-results.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +Preserve multimodal server tool results when they round-trip through the client. diff --git a/packages/ai/src/activities/chat/stream/message-updaters.ts b/packages/ai/src/activities/chat/stream/message-updaters.ts index 72648af11..11031424e 100644 --- a/packages/ai/src/activities/chat/stream/message-updaters.ts +++ b/packages/ai/src/activities/chat/stream/message-updaters.ts @@ -6,6 +6,7 @@ */ import { parsePartialJSON } from './json-parser' +import { isContentPartArray } from '../../../utilities/tool-result' import type { ContentPart, StructuredOutputPart, @@ -118,6 +119,18 @@ export function updateToolResultPart( state: ToolResultState, error?: string, ): Array { + let resolvedContent = content + if (typeof content === 'string') { + try { + const parsed = JSON.parse(content) + if (isContentPartArray(parsed)) { + resolvedContent = parsed + } + } catch { + // Keep non-JSON tool output as the original string. + } + } + return messages.map((msg) => { if (msg.id !== messageId) { return msg @@ -132,7 +145,7 @@ export function updateToolResultPart( const toolResultPart: ToolResultPart = { type: 'tool-result', toolCallId, - content, + content: resolvedContent, state, ...(error && { error }), } diff --git a/packages/ai/tests/multimodal-tool-roundtrip.test.ts b/packages/ai/tests/multimodal-tool-roundtrip.test.ts new file mode 100644 index 000000000..a6cfd092e --- /dev/null +++ b/packages/ai/tests/multimodal-tool-roundtrip.test.ts @@ -0,0 +1,70 @@ +import { describe, expect, it } from 'vitest' +import { StreamProcessor } from '../src/activities/chat/stream/processor' +import { uiMessagesToWire } from '../src/utilities/ag-ui-wire' +import type { ContentPart, StreamChunk } from '../src/types' + +function chunk( + type: T, + fields: Record, +): Extract { + return { type, timestamp: Date.now(), ...fields } as Extract< + StreamChunk, + { type: T } + > +} + +const MULTIMODAL_RESULT: Array = [ + { type: 'text', content: 'Layout rule 1.' }, + { + type: 'image', + source: { + type: 'data', + value: 'iVBORw0KGgoAAAANSUhEUg==', + mimeType: 'image/png', + }, + }, +] + +describe('multimodal TOOL_CALL_RESULT round trip', () => { + it('keeps structured content in the tool result metadata', () => { + const processor = new StreamProcessor() + const toolCallId = 'call-mm' + + processor.processChunk( + chunk('TOOL_CALL_START', { + toolCallId, + toolCallName: 'getLayoutRules', + parentMessageId: 'assistant-1', + }), + ) + processor.processChunk( + chunk('TOOL_CALL_ARGS', { toolCallId, delta: '{}' }), + ) + processor.processChunk(chunk('TOOL_CALL_END', { toolCallId })) + processor.processChunk( + chunk('TOOL_CALL_RESULT', { + toolCallId, + messageId: 'tool-result-1', + content: JSON.stringify(MULTIMODAL_RESULT), + }), + ) + + const toolResult = processor + .getMessages()[0]! + .parts.find((part) => part.type === 'tool-result') + + expect(toolResult?.content).toEqual(MULTIMODAL_RESULT) + + const wire = uiMessagesToWire(processor.getMessages()) + const toolMessage = wire.find((message) => message.role === 'tool') + const metadata = toolMessage?.metadata as + | { + tanstack?: { + toolResult?: { content?: Array } + } + } + | undefined + + expect(metadata?.tanstack?.toolResult?.content).toEqual(MULTIMODAL_RESULT) + }) +}) diff --git a/testing/e2e/tests/multimodal-tool-result-roundtrip.spec.ts b/testing/e2e/tests/multimodal-tool-result-roundtrip.spec.ts new file mode 100644 index 000000000..dd538e7f0 --- /dev/null +++ b/testing/e2e/tests/multimodal-tool-result-roundtrip.spec.ts @@ -0,0 +1,63 @@ +import { + StreamProcessor, + uiMessagesToWire, + type ContentPart, + type StreamChunk, +} from '@tanstack/ai' +import { test, expect } from './fixtures' + +function chunk( + type: T, + fields: Record, +): Extract { + return { type, timestamp: Date.now(), ...fields } as Extract< + StreamChunk, + { type: T } + > +} + +test('multimodal TOOL_CALL_RESULT keeps structured content on the wire', () => { + const content: Array = [ + { type: 'text', content: 'Layout rule 1.' }, + { + type: 'image', + source: { + type: 'data', + value: 'iVBORw0KGgoAAAANSUhEUg==', + mimeType: 'image/png', + }, + }, + ] + const processor = new StreamProcessor() + const toolCallId = 'call-mm' + + processor.processChunk( + chunk('TOOL_CALL_START', { + toolCallId, + toolCallName: 'getLayoutRules', + parentMessageId: 'assistant-1', + }), + ) + processor.processChunk(chunk('TOOL_CALL_ARGS', { toolCallId, delta: '{}' })) + processor.processChunk(chunk('TOOL_CALL_END', { toolCallId })) + processor.processChunk( + chunk('TOOL_CALL_RESULT', { + toolCallId, + messageId: 'tool-result-1', + content: JSON.stringify(content), + }), + ) + + const toolMessage = uiMessagesToWire(processor.getMessages()).find( + (message) => message.role === 'tool', + ) + const metadata = toolMessage?.metadata as + | { + tanstack?: { + toolResult?: { content?: Array } + } + } + | undefined + + expect(metadata?.tanstack?.toolResult?.content).toEqual(content) +})