feat: replay of took calls for a2ui elements

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<string, unknown>;
+}
+
+/**
+ * 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;
+}

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<string, unknown>;
   /** 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,
   };
 }

src/renderer/components/ChatPanel/ChatPanel.tsx

@@ -66,6 +66,7 @@ export const ChatPanel: React.FC<ChatPanelProps> = ({ conversationId }) => {
     dismissSurface,
     dispatchAction,
     updateLocalData,
+    replayFromMessages,
   } = useA2UISurface({ conversationId });
 
   // Current turn index for associating streaming surfaces
@@ -102,12 +103,15 @@ export const ChatPanel: React.FC<ChatPanelProps> = ({ 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();

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> = {},
+): 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);
+  });
+});