From 6fa7ac09b825883e6d02edff02743e65d9f6b206 Mon Sep 17 00:00:00 2001 From: hugo Date: Thu, 26 Feb 2026 12:19:09 +0100 Subject: [PATCH] feat: replay of took calls for a2ui elements --- src/renderer/a2ui/surfaceAssociation.ts | 72 ++++++ src/renderer/a2ui/useA2UISurface.ts | 16 ++ .../components/ChatPanel/ChatPanel.tsx | 8 +- tests/renderer/a2ui/surfaceReplay.test.ts | 213 ++++++++++++++++++ 4 files changed, 307 insertions(+), 2 deletions(-) create mode 100644 tests/renderer/a2ui/surfaceReplay.test.ts diff --git a/src/renderer/a2ui/surfaceAssociation.ts b/src/renderer/a2ui/surfaceAssociation.ts index 8efd676..5e30462 100644 --- a/src/renderer/a2ui/surfaceAssociation.ts +++ b/src/renderer/a2ui/surfaceAssociation.ts @@ -3,9 +3,14 @@ * * Computes the turn index for a message based on its position * in the message array, enabling inline surface rendering. + * + * Also provides replay logic to reconstruct A2UI surfaces from + * persisted tool calls when reloading a saved conversation. */ import type { ChatMessage } from '../../main/shared/electronApi'; +import type { A2UIServerMessage } from '../../main/a2ui/types'; +import { isRenderTool, generateFromToolCall } from '../../main/a2ui/generator'; /** * Compute the turn index for a message at a given position. @@ -25,3 +30,70 @@ export function computeTurnIndex(messages: ChatMessage[], currentIndex: number): } return userCount - 1; } + +interface StoredToolCall { + name: string; + args?: Record; +} + +/** + * Replay A2UI surfaces from persisted chat messages. + * + * Scans assistant messages for render tool calls stored in their + * `toolCalls` JSON field and regenerates the A2UI server messages + * needed to reconstruct the surfaces in the manager. + * + * @param conversationId - The conversation ID for surface creation + * @param messages - The full ordered message history from the database + * @returns Array of A2UI server messages to feed into A2UISurfaceManager + */ +export function replaySurfacesFromMessages( + conversationId: string, + messages: ChatMessage[], +): A2UIServerMessage[] { + const allA2UIMessages: A2UIServerMessage[] = []; + + for (let i = 0; i < messages.length; i++) { + const message = messages[i]; + if (message.role !== 'assistant' || !message.toolCalls) { + continue; + } + + let toolCalls: StoredToolCall[]; + try { + toolCalls = JSON.parse(message.toolCalls) as StoredToolCall[]; + } catch { + continue; + } + + if (!Array.isArray(toolCalls)) { + continue; + } + + const turnIndex = computeTurnIndex(messages, i); + + for (const toolCall of toolCalls) { + if (!isRenderTool(toolCall.name)) { + continue; + } + + const a2uiMessages = generateFromToolCall( + conversationId, + toolCall.name, + toolCall.args ?? {}, + ); + + if (a2uiMessages) { + // Inject turnIndex into createSurface metadata, matching live behavior + for (const msg of a2uiMessages) { + if (msg.type === 'createSurface') { + msg.metadata = { ...msg.metadata, turnIndex }; + } + } + allA2UIMessages.push(...a2uiMessages); + } + } + } + + return allA2UIMessages; +} diff --git a/src/renderer/a2ui/useA2UISurface.ts b/src/renderer/a2ui/useA2UISurface.ts index 3189396..f23f47e 100644 --- a/src/renderer/a2ui/useA2UISurface.ts +++ b/src/renderer/a2ui/useA2UISurface.ts @@ -7,7 +7,9 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { A2UISurfaceManager } from './A2UISurfaceManager'; +import { replaySurfacesFromMessages } from './surfaceAssociation'; import type { A2UIResolvedComponent, A2UIServerMessage, A2UIClientAction } from '../../main/a2ui/types'; +import type { ChatMessage } from '../../main/shared/electronApi'; interface UseA2UISurfaceInput { conversationId: string | null; @@ -37,6 +39,8 @@ interface UseA2UISurfaceResult { getDataModel: (surfaceId: string) => Record; /** Clear all surfaces for the conversation */ clearSurfaces: () => void; + /** Replay surfaces from persisted chat messages */ + replayFromMessages: (messages: ChatMessage[]) => void; } export function useA2UISurface(input: UseA2UISurfaceInput): UseA2UISurfaceResult { @@ -162,6 +166,17 @@ export function useA2UISurface(input: UseA2UISurfaceInput): UseA2UISurfaceResult } }, [conversationId]); + const replayFromMessages = useCallback((msgs: ChatMessage[]) => { + if (!conversationId) return; + const manager = managerRef.current; + // Only replay if no surfaces exist yet (avoid duplicates on re-render) + if (manager.getSurfaceIds(conversationId).length > 0) return; + const a2uiMessages = replaySurfacesFromMessages(conversationId, msgs); + for (const msg of a2uiMessages) { + manager.processMessage(msg); + } + }, [conversationId]); + return { surfaces, surfacesByTurn, @@ -172,5 +187,6 @@ export function useA2UISurface(input: UseA2UISurfaceInput): UseA2UISurfaceResult updateLocalData, getDataModel, clearSurfaces, + replayFromMessages, }; } diff --git a/src/renderer/components/ChatPanel/ChatPanel.tsx b/src/renderer/components/ChatPanel/ChatPanel.tsx index ece7cb4..e768d9b 100644 --- a/src/renderer/components/ChatPanel/ChatPanel.tsx +++ b/src/renderer/components/ChatPanel/ChatPanel.tsx @@ -66,6 +66,7 @@ export const ChatPanel: React.FC = ({ conversationId }) => { dismissSurface, dispatchAction, updateLocalData, + replayFromMessages, } = useA2UISurface({ conversationId }); // Current turn index for associating streaming surfaces @@ -102,12 +103,15 @@ export const ChatPanel: React.FC = ({ conversationId }) => { ]); if (conv) setConversation(conv); - if (msgs) setMessages(msgs); + if (msgs) { + setMessages(msgs); + replayFromMessages(msgs); + } if (modelsResult?.models) setAvailableModels(modelsResult.models); } catch (error) { console.error('Failed to load chat data:', error); } - }, [conversationId]); + }, [conversationId, replayFromMessages]); useEffect(() => { checkReady(); diff --git a/tests/renderer/a2ui/surfaceReplay.test.ts b/tests/renderer/a2ui/surfaceReplay.test.ts new file mode 100644 index 0000000..60eb03f --- /dev/null +++ b/tests/renderer/a2ui/surfaceReplay.test.ts @@ -0,0 +1,213 @@ +import { describe, expect, it } from 'vitest'; +import { replaySurfacesFromMessages } from '../../../src/renderer/a2ui/surfaceAssociation'; +import type { ChatMessage } from '../../../src/main/shared/electronApi'; + +function msg( + role: ChatMessage['role'], + overrides: Partial = {}, +): ChatMessage { + return { + id: `msg-${Math.random()}`, + conversationId: 'conv-1', + role, + content: '', + createdAt: new Date().toISOString(), + ...overrides, + }; +} + +describe('replaySurfacesFromMessages', () => { + it('returns empty array when there are no messages', () => { + const result = replaySurfacesFromMessages('conv-1', []); + expect(result).toEqual([]); + }); + + it('returns empty array when no assistant messages have toolCalls', () => { + const messages = [ + msg('user'), + msg('assistant', { content: 'Hello!' }), + ]; + const result = replaySurfacesFromMessages('conv-1', messages); + expect(result).toEqual([]); + }); + + it('returns empty array when assistant has non-render tool calls only', () => { + const messages = [ + msg('user'), + msg('assistant', { + toolCalls: JSON.stringify([ + { name: 'search_posts', args: { query: 'test' } }, + ]), + }), + ]; + const result = replaySurfacesFromMessages('conv-1', messages); + expect(result).toEqual([]); + }); + + it('replays a single render_chart tool call', () => { + const messages = [ + msg('user'), + msg('assistant', { + toolCalls: JSON.stringify([ + { + name: 'render_chart', + args: { + chartType: 'bar', + title: 'Test Chart', + series: [{ label: 'A', value: 10 }], + }, + }, + ]), + }), + ]; + const result = replaySurfacesFromMessages('conv-1', messages); + + expect(result.length).toBeGreaterThanOrEqual(2); // createSurface + updateComponents + updateDataModel + + // Verify createSurface has correct turnIndex + const createMsg = result.find((m) => m.type === 'createSurface'); + expect(createMsg).toBeDefined(); + expect(createMsg!.type).toBe('createSurface'); + if (createMsg!.type === 'createSurface') { + expect(createMsg!.conversationId).toBe('conv-1'); + expect(createMsg!.metadata?.turnIndex).toBe(0); + } + + // Verify chart component was generated + const updateMsg = result.find((m) => m.type === 'updateComponents'); + expect(updateMsg).toBeDefined(); + if (updateMsg!.type === 'updateComponents') { + expect(updateMsg!.components[0].type).toBe('chart'); + expect(updateMsg!.components[0].properties.chartType).toBe('bar'); + } + }); + + it('assigns correct turnIndex based on message position', () => { + const messages = [ + msg('user'), + msg('assistant', { content: 'turn 0' }), + msg('user'), + msg('assistant', { + toolCalls: JSON.stringify([ + { + name: 'render_metric', + args: { label: 'Total', value: '42' }, + }, + ]), + }), + ]; + const result = replaySurfacesFromMessages('conv-1', messages); + + const createMsg = result.find((m) => m.type === 'createSurface'); + expect(createMsg).toBeDefined(); + if (createMsg!.type === 'createSurface') { + expect(createMsg!.metadata?.turnIndex).toBe(1); + } + }); + + it('replays multiple render tools from the same assistant message', () => { + const messages = [ + msg('user'), + msg('assistant', { + toolCalls: JSON.stringify([ + { + name: 'render_chart', + args: { chartType: 'bar', series: [{ label: 'A', value: 1 }] }, + }, + { + name: 'render_metric', + args: { label: 'Count', value: '5' }, + }, + ]), + }), + ]; + const result = replaySurfacesFromMessages('conv-1', messages); + + // Should have two createSurface messages (one per render tool) + const creates = result.filter((m) => m.type === 'createSurface'); + expect(creates).toHaveLength(2); + + // Both should have the same turnIndex + for (const c of creates) { + if (c.type === 'createSurface') { + expect(c.metadata?.turnIndex).toBe(0); + } + } + }); + + it('replays render tools across multiple turns', () => { + const messages = [ + msg('user'), + msg('assistant', { + toolCalls: JSON.stringify([ + { + name: 'render_chart', + args: { chartType: 'bar', series: [{ label: 'A', value: 1 }] }, + }, + ]), + }), + msg('user'), + msg('assistant', { + toolCalls: JSON.stringify([ + { + name: 'render_table', + args: { columns: ['Name'], rows: [['X']] }, + }, + ]), + }), + ]; + const result = replaySurfacesFromMessages('conv-1', messages); + + const creates = result.filter((m) => m.type === 'createSurface'); + expect(creates).toHaveLength(2); + + if (creates[0].type === 'createSurface' && creates[1].type === 'createSurface') { + expect(creates[0].metadata?.turnIndex).toBe(0); + expect(creates[1].metadata?.turnIndex).toBe(1); + } + }); + + it('mixes data tools and render tools, only replaying render tools', () => { + const messages = [ + msg('user'), + msg('assistant', { + toolCalls: JSON.stringify([ + { name: 'search_posts', args: { query: 'hello' } }, + { + name: 'render_chart', + args: { chartType: 'pie', series: [{ label: 'Yes', value: 3 }] }, + }, + { name: 'list_tags', args: {} }, + ]), + }), + ]; + const result = replaySurfacesFromMessages('conv-1', messages); + + const creates = result.filter((m) => m.type === 'createSurface'); + expect(creates).toHaveLength(1); + }); + + it('handles malformed toolCalls JSON gracefully', () => { + const messages = [ + msg('user'), + msg('assistant', { toolCalls: 'not-valid-json{{{' }), + ]; + const result = replaySurfacesFromMessages('conv-1', messages); + expect(result).toEqual([]); + }); + + it('handles toolCalls with missing args gracefully', () => { + const messages = [ + msg('user'), + msg('assistant', { + toolCalls: JSON.stringify([ + { name: 'render_chart' }, + ]), + }), + ]; + // Should not throw + const result = replaySurfacesFromMessages('conv-1', messages); + // We expect it to try to generate (generator handles validation) + expect(result.length).toBeGreaterThanOrEqual(0); + }); +});