feat: finally a good working state

.claude/settings.local.json

@@ -10,7 +10,9 @@
       "WebFetch(domain:github.com)",
       "WebFetch(domain:www.npmjs.com)",
       "WebFetch(domain:a2ui-sdk.js.org)",
-      "WebFetch(domain:www.copilotkit.ai)"
+      "WebFetch(domain:www.copilotkit.ai)",
+      "Bash(grep -l \"A2UIRenderer\\\\|useA2UISurface\\\\|a2ui-surface\" /Users/gb/Projects/bDS/src/renderer/components/**/*.tsx)",
+      "Bash(npm test)"
     ]
   }
 }

src/main/engine/OpenCodeManager.ts

@@ -293,10 +293,16 @@ export class OpenCodeManager {
       let fullResponse = '';
       const toolCallsCollected: Array<{ name: string; args: unknown }> = [];
 
+      // Compute turn index for surface-to-message association
+      const turnIndex = dbMessages.filter(m => m.role === 'user').length - 1;
+
       // Wrap onA2UIMessage emission for render tools
       const emitA2UIMessages = (messages: A2UIServerMessage[]) => {
         if (onA2UIMessage) {
           for (const msg of messages) {
+            if (msg.type === 'createSurface') {
+              msg.metadata = { ...msg.metadata, turnIndex };
+            }
             onA2UIMessage(msg);
           }
         }

src/renderer/a2ui/InlineSurface.css

@@ -0,0 +1,69 @@
+.inline-surface {
+  margin: 8px 0;
+  border-radius: 8px;
+  border: 1px solid var(--vscode-editorGroup-border, var(--vscode-panel-border));
+  overflow: hidden;
+}
+
+.inline-surface.collapsed {
+  display: flex;
+  align-items: center;
+  gap: 8px;
+  padding: 8px 12px;
+  cursor: pointer;
+  background: var(--vscode-textBlockQuote-background, rgba(127, 127, 127, 0.1));
+  transition: background 0.15s;
+}
+
+.inline-surface.collapsed:hover {
+  background: var(--vscode-list-hoverBackground);
+}
+
+.inline-surface-icon {
+  font-size: 16px;
+  flex-shrink: 0;
+}
+
+.inline-surface-title {
+  flex: 1;
+  font-size: 13px;
+  font-weight: 500;
+  color: var(--vscode-foreground);
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.inline-surface-header {
+  display: flex;
+  align-items: center;
+  gap: 8px;
+  padding: 8px 12px;
+  border-bottom: 1px solid var(--vscode-editorGroup-border, var(--vscode-panel-border));
+  background: var(--vscode-textBlockQuote-background, rgba(127, 127, 127, 0.1));
+}
+
+.inline-surface-expand,
+.inline-surface-collapse,
+.inline-surface-dismiss {
+  background: transparent;
+  border: none;
+  color: var(--vscode-descriptionForeground);
+  cursor: pointer;
+  padding: 2px 6px;
+  font-size: 14px;
+  border-radius: 4px;
+  line-height: 1;
+}
+
+.inline-surface-expand:hover,
+.inline-surface-collapse:hover,
+.inline-surface-dismiss:hover {
+  background: var(--vscode-toolbar-hoverBackground);
+  color: var(--vscode-foreground);
+}
+
+.inline-surface.expanded > .a2ui-surface {
+  border-top: none;
+  padding-top: 0;
+}

src/renderer/a2ui/InlineSurface.tsx

@@ -0,0 +1,153 @@
+/**
+ * InlineSurface component.
+ *
+ * Wraps A2UIRenderer with expand/collapse and dismiss controls.
+ * Renders inline within the chat transcript, anchored to the
+ * assistant message turn that created the surface.
+ */
+
+import React, { useEffect, useState } from 'react';
+import type { A2UIResolvedComponent, A2UIClientAction } from '../../main/a2ui/types';
+import { A2UIRenderer } from './A2UIRenderer';
+import './InlineSurface.css';
+
+interface InlineSurfaceProps {
+  surfaceId: string;
+  tree: A2UIResolvedComponent[];
+  isExpanded: boolean;
+  onDismiss?: (surfaceId: string) => void;
+  onAction?: (action: A2UIClientAction) => void;
+  onDataChange?: (surfaceId: string, path: string, value: unknown) => void;
+}
+
+/**
+ * Derive a display title from the surface's component tree.
+ * Tries the root component's `title` or `label` property,
+ * then falls back to the capitalized component type.
+ */
+export function deriveSurfaceTitle(tree: A2UIResolvedComponent[]): string {
+  if (tree.length === 0) {
+    return 'Surface';
+  }
+  const root = tree[0];
+  const title = root.properties?.title as string | undefined;
+  if (title) {
+    return title;
+  }
+  const label = root.properties?.label as string | undefined;
+  if (label) {
+    return label;
+  }
+  return root.type.charAt(0).toUpperCase() + root.type.slice(1);
+}
+
+/**
+ * Get an icon character for the surface based on the root component type.
+ */
+export function getSurfaceIcon(tree: A2UIResolvedComponent[]): string {
+  if (tree.length === 0) {
+    return '\u25A0';
+  }
+  const type = tree[0].type;
+  const icons: Record<string, string> = {
+    chart: '\u{1F4CA}',
+    table: '\u{1F4CB}',
+    form: '\u{1F4DD}',
+    card: '\u{1F4C4}',
+    metric: '\u{1F4CF}',
+    list: '\u{1F4CB}',
+    tabs: '\u{1F4C2}',
+  };
+  return icons[type] || '\u25A0';
+}
+
+export const InlineSurface: React.FC<InlineSurfaceProps> = ({
+  surfaceId,
+  tree,
+  isExpanded: defaultExpanded,
+  onDismiss,
+  onAction,
+  onDataChange,
+}) => {
+  const [expanded, setExpanded] = useState(defaultExpanded);
+
+  // Auto-collapse/expand when the parent changes which surface is latest
+  useEffect(() => {
+    setExpanded(defaultExpanded);
+  }, [defaultExpanded]);
+
+  const surfaceTitle = deriveSurfaceTitle(tree);
+  const surfaceIcon = getSurfaceIcon(tree);
+
+  if (!expanded) {
+    return (
+      <div
+        className="inline-surface collapsed"
+        onClick={() => setExpanded(true)}
+        role="button"
+        tabIndex={0}
+        onKeyDown={(e) => {
+          if (e.key === 'Enter' || e.key === ' ') {
+            setExpanded(true);
+          }
+        }}
+      >
+        <span className="inline-surface-icon">{surfaceIcon}</span>
+        <span className="inline-surface-title">{surfaceTitle}</span>
+        <button
+          className="inline-surface-expand"
+          onClick={(e) => {
+            e.stopPropagation();
+            setExpanded(true);
+          }}
+          aria-label="Expand surface"
+        >
+          {'\u25B6'}
+        </button>
+        {onDismiss && (
+          <button
+            className="inline-surface-dismiss"
+            onClick={(e) => {
+              e.stopPropagation();
+              onDismiss(surfaceId);
+            }}
+            aria-label="Dismiss surface"
+          >
+            {'\u2715'}
+          </button>
+        )}
+      </div>
+    );
+  }
+
+  return (
+    <div className="inline-surface expanded">
+      <div className="inline-surface-header">
+        <span className="inline-surface-icon">{surfaceIcon}</span>
+        <span className="inline-surface-title">{surfaceTitle}</span>
+        <button
+          className="inline-surface-collapse"
+          onClick={() => setExpanded(false)}
+          aria-label="Collapse surface"
+        >
+          {'\u25BC'}
+        </button>
+        {onDismiss && (
+          <button
+            className="inline-surface-dismiss"
+            onClick={() => onDismiss(surfaceId)}
+            aria-label="Dismiss surface"
+          >
+            {'\u2715'}
+          </button>
+        )}
+      </div>
+      <A2UIRenderer
+        surfaceId={surfaceId}
+        tree={tree}
+        onAction={onAction!}
+        onDataChange={onDataChange}
+      />
+    </div>
+  );
+};

src/renderer/a2ui/surfaceAssociation.ts

@@ -0,0 +1,27 @@
+/**
+ * Surface-to-message association utilities.
+ *
+ * Computes the turn index for a message based on its position
+ * in the message array, enabling inline surface rendering.
+ */
+
+import type { ChatMessage } from '../../main/shared/electronApi';
+
+/**
+ * Compute the turn index for a message at a given position.
+ *
+ * Turn index is defined as the 0-based count of user messages
+ * seen up to and including the given position, minus 1.
+ * System and tool messages do not affect the count.
+ *
+ * Returns -1 if no user message has been seen at or before the index.
+ */
+export function computeTurnIndex(messages: ChatMessage[], currentIndex: number): number {
+  let userCount = 0;
+  for (let i = 0; i <= currentIndex; i++) {
+    if (messages[i].role === 'user') {
+      userCount++;
+    }
+  }
+  return userCount - 1;
+}

src/renderer/a2ui/useA2UISurface.ts

@@ -13,9 +13,22 @@ interface UseA2UISurfaceInput {
   conversationId: string | null;
 }
 
+export interface SurfaceEntry {
+  surfaceId: string;
+  tree: A2UIResolvedComponent[];
+}
+
 interface UseA2UISurfaceResult {
   /** All active surface trees for this conversation */
-  surfaces: Array<{ surfaceId: string; tree: A2UIResolvedComponent[] }>;
+  surfaces: SurfaceEntry[];
+  /** Surfaces grouped by the turn index that created them */
+  surfacesByTurn: Map<number, SurfaceEntry[]>;
+  /** The surfaceId of the most recently created surface */
+  latestSurfaceId: string | null;
+  /** Set of surface IDs that the user has dismissed */
+  dismissedSurfaceIds: Set<string>;
+  /** Dismiss a surface by ID */
+  dismissSurface: (surfaceId: string) => void;
   /** Dispatch an action back to the main process */
   dispatchAction: (action: A2UIClientAction) => void;
   /** Update a local data binding (for form inputs) */
@@ -30,6 +43,7 @@ export function useA2UISurface(input: UseA2UISurfaceInput): UseA2UISurfaceResult
   const { conversationId } = input;
   const managerRef = useRef<A2UISurfaceManager>(new A2UISurfaceManager());
   const [renderTick, setRenderTick] = useState(0);
+  const [dismissedSurfaceIds, setDismissedSurfaceIds] = useState<Set<string>>(new Set());
 
   // Subscribe to surface changes
   useEffect(() => {
@@ -58,8 +72,9 @@ export function useA2UISurface(input: UseA2UISurfaceInput): UseA2UISurfaceResult
     };
   }, [conversationId]);
 
-  // Clear surfaces when conversation changes
+  // Clear surfaces and dismissed set when conversation changes
   useEffect(() => {
+    setDismissedSurfaceIds(new Set());
     return () => {
       if (conversationId) {
         managerRef.current.clearConversation(conversationId);
@@ -83,6 +98,52 @@ export function useA2UISurface(input: UseA2UISurfaceInput): UseA2UISurfaceResult
     }));
   }, [conversationId, renderTick]);
 
+  const surfacesByTurn = useMemo(() => {
+    void renderTick;
+
+    const map = new Map<number, SurfaceEntry[]>();
+    if (!conversationId) {
+      return map;
+    }
+
+    const manager = managerRef.current;
+    const surfaceIds = manager.getSurfaceIds(conversationId);
+
+    for (const surfaceId of surfaceIds) {
+      const surface = manager.getSurface(surfaceId);
+      const turnIndex = (surface?.metadata?.turnIndex as number) ?? -1;
+      const entry: SurfaceEntry = { surfaceId, tree: manager.resolveTree(surfaceId) };
+
+      const existing = map.get(turnIndex);
+      if (existing) {
+        existing.push(entry);
+      } else {
+        map.set(turnIndex, [entry]);
+      }
+    }
+
+    return map;
+  }, [conversationId, renderTick]);
+
+  const latestSurfaceId = useMemo(() => {
+    void renderTick;
+
+    if (!conversationId) {
+      return null;
+    }
+
+    const ids = managerRef.current.getSurfaceIds(conversationId);
+    return ids.length > 0 ? ids[ids.length - 1] : null;
+  }, [conversationId, renderTick]);
+
+  const dismissSurface = useCallback((surfaceId: string) => {
+    setDismissedSurfaceIds((prev) => {
+      const next = new Set(prev);
+      next.add(surfaceId);
+      return next;
+    });
+  }, []);
+
   const dispatchAction = useCallback((action: A2UIClientAction) => {
     window.electronAPI?.chat.dispatchA2UIAction?.(action);
   }, []);
@@ -103,6 +164,10 @@ export function useA2UISurface(input: UseA2UISurfaceInput): UseA2UISurfaceResult
 
   return {
     surfaces,
+    surfacesByTurn,
+    latestSurfaceId,
+    dismissedSurfaceIds,
+    dismissSurface,
     dispatchAction,
     updateLocalData,
     getDataModel,

src/renderer/components/AssistantSidebar/AssistantSidebar.tsx

@@ -8,7 +8,6 @@ import { getChatSurfaceMode } from '../../navigation/chatSurfaceMode';
 import { useChatMessageSender } from '../../navigation/useChatMessageSender';
 import { useChatSurfaceState } from '../../navigation/useChatSurfaceState';
 import { useA2UISurface } from '../../a2ui/useA2UISurface';
-import { A2UIRenderer } from '../../a2ui/A2UIRenderer';
 import { ChatTranscript } from '../ChatSurface';
 import { useI18n } from '../../i18n';
 import '../../styles/chatSurface.css';
@@ -55,7 +54,19 @@ export const AssistantSidebar: React.FC = () => {
   } = useChatSurfaceState();
 
   // A2UI surface rendering
-  const { surfaces, dispatchAction, updateLocalData } = useA2UISurface({ conversationId });
+  const {
+    surfacesByTurn,
+    latestSurfaceId,
+    dismissedSurfaceIds,
+    dismissSurface,
+    dispatchAction,
+    updateLocalData,
+  } = useA2UISurface({ conversationId });
+
+  // Current turn index for associating streaming surfaces
+  const currentTurnIndex = useMemo(() => {
+    return messages.filter(m => m.role === 'user').length - 1;
+  }, [messages]);
 
   const activeTab = useMemo(() => tabs.find((tab) => tab.id === activeTabId) ?? null, [tabs, activeTabId]);
 
@@ -266,19 +277,16 @@ export const AssistantSidebar: React.FC = () => {
             assistantRoleLabel={tr('chat.role.assistant')}
             userRoleLabel={tr('chat.role.you')}
             showToolMarkers={surfaceMode.showToolMarkers}
+            surfacesByTurn={surfacesByTurn}
+            latestSurfaceId={latestSurfaceId}
+            dismissedSurfaceIds={dismissedSurfaceIds}
+            onSurfaceDismiss={dismissSurface}
+            onSurfaceAction={dispatchAction}
+            onSurfaceDataChange={updateLocalData}
+            currentTurnIndex={currentTurnIndex}
           />
         </div>
       )}
-
-      {surfaces.map((surface) => (
-        <A2UIRenderer
-          key={surface.surfaceId}
-          surfaceId={surface.surfaceId}
-          tree={surface.tree}
-          onAction={dispatchAction}
-          onDataChange={updateLocalData}
-        />
-      ))}
     </div>
   );
 };

src/renderer/components/ChatPanel/ChatPanel.tsx

@@ -1,11 +1,10 @@
-import React, { useState, useEffect, useRef, useCallback } from 'react';
+import React, { useState, useEffect, useRef, useCallback, useMemo } from 'react';
 import type { ChatConversation, ChatModel } from '../../types/electron';
 import { useChatMessageSender } from '../../navigation/useChatMessageSender';
 import { useChatSurfaceState } from '../../navigation/useChatSurfaceState';
 import { getChatSurfaceMode } from '../../navigation/chatSurfaceMode';
 import { dispatchAssistantAction } from '../../navigation/assistantActionDispatcher';
 import { useA2UISurface } from '../../a2ui/useA2UISurface';
-import { A2UIRenderer } from '../../a2ui/A2UIRenderer';
 import { useAppStore } from '../../store';
 import { ChatTranscript } from '../ChatSurface';
 import { useI18n } from '../../i18n';
@@ -60,7 +59,19 @@ export const ChatPanel: React.FC<ChatPanelProps> = ({ conversationId }) => {
   } = useChatSurfaceState();
 
   // A2UI surface rendering
-  const { surfaces, dispatchAction, updateLocalData } = useA2UISurface({ conversationId });
+  const {
+    surfacesByTurn,
+    latestSurfaceId,
+    dismissedSurfaceIds,
+    dismissSurface,
+    dispatchAction,
+    updateLocalData,
+  } = useA2UISurface({ conversationId });
+
+  // Current turn index for associating streaming surfaces
+  const currentTurnIndex = useMemo(() => {
+    return messages.filter(m => m.role === 'user').length - 1;
+  }, [messages]);
 
   // Scroll to bottom when messages change
   const scrollToBottom = useCallback(() => {
@@ -368,18 +379,15 @@ export const ChatPanel: React.FC<ChatPanelProps> = ({ conversationId }) => {
           userRoleLabel={tr('chat.role.you')}
           showToolMarkers={surfaceMode.showToolMarkers}
           endRef={messagesEndRef}
+          surfacesByTurn={surfacesByTurn}
+          latestSurfaceId={latestSurfaceId}
+          dismissedSurfaceIds={dismissedSurfaceIds}
+          onSurfaceDismiss={dismissSurface}
+          onSurfaceAction={dispatchAction}
+          onSurfaceDataChange={updateLocalData}
+          currentTurnIndex={currentTurnIndex}
         />
 
-        {surfaces.map((surface) => (
-          <A2UIRenderer
-            key={surface.surfaceId}
-            surfaceId={surface.surfaceId}
-            tree={surface.tree}
-            onAction={dispatchAction}
-            onDataChange={updateLocalData}
-          />
-        ))}
-
         {actionError && <p className="chat-surface-error">{actionError}</p>}
       </div>
 

src/renderer/components/ChatSurface/ChatTranscript.tsx

@@ -2,6 +2,10 @@ import React from 'react';
 import Markdown from 'marked-react';
 import type { ChatMessage } from '../../types/electron';
 import type { ChatToolEvent } from '../../navigation/useChatSurfaceState';
+import type { A2UIResolvedComponent, A2UIClientAction } from '../../../main/a2ui/types';
+import { InlineSurface } from '../../a2ui/InlineSurface';
+import type { SurfaceEntry } from '../../a2ui/useA2UISurface';
+import { computeTurnIndex } from '../../a2ui/surfaceAssociation';
 
 interface ChatTranscriptProps {
   messages: ChatMessage[];
@@ -12,6 +16,20 @@ interface ChatTranscriptProps {
   userRoleLabel: string;
   showToolMarkers?: boolean;
   endRef?: React.RefObject<HTMLDivElement | null>;
+  /** Surfaces grouped by the turn index that created them */
+  surfacesByTurn?: Map<number, SurfaceEntry[]>;
+  /** The surfaceId of the most recently created surface */
+  latestSurfaceId?: string | null;
+  /** Set of surface IDs the user has dismissed */
+  dismissedSurfaceIds?: Set<string>;
+  /** Callback to dismiss a surface */
+  onSurfaceDismiss?: (surfaceId: string) => void;
+  /** Callback for surface actions */
+  onSurfaceAction?: (action: A2UIClientAction) => void;
+  /** Callback for surface data changes */
+  onSurfaceDataChange?: (surfaceId: string, path: string, value: unknown) => void;
+  /** The current streaming turn index */
+  currentTurnIndex?: number;
 }
 
 export const ChatTranscript: React.FC<ChatTranscriptProps> = ({
@@ -23,6 +41,13 @@ export const ChatTranscript: React.FC<ChatTranscriptProps> = ({
   userRoleLabel,
   showToolMarkers = true,
   endRef,
+  surfacesByTurn,
+  latestSurfaceId,
+  dismissedSurfaceIds,
+  onSurfaceDismiss,
+  onSurfaceAction,
+  onSurfaceDataChange,
+  currentTurnIndex,
 }) => {
   const renderToolMarkers = (events: ChatToolEvent[]) => {
     if (events.length === 0) {
@@ -65,7 +90,34 @@ export const ChatTranscript: React.FC<ChatTranscriptProps> = ({
     );
   };
 
-  const renderMessage = (message: ChatMessage) => {
+  const renderInlineSurfaces = (turnIndex: number) => {
+    if (!surfacesByTurn?.has(turnIndex)) {
+      return null;
+    }
+
+    const turnSurfaces = surfacesByTurn.get(turnIndex)!;
+    const visibleSurfaces = turnSurfaces.filter(
+      (s) => !dismissedSurfaceIds?.has(s.surfaceId),
+    );
+
+    if (visibleSurfaces.length === 0) {
+      return null;
+    }
+
+    return visibleSurfaces.map((surface) => (
+      <InlineSurface
+        key={surface.surfaceId}
+        surfaceId={surface.surfaceId}
+        tree={surface.tree}
+        isExpanded={surface.surfaceId === latestSurfaceId}
+        onDismiss={onSurfaceDismiss}
+        onAction={onSurfaceAction}
+        onDataChange={onSurfaceDataChange}
+      />
+    ));
+  };
+
+  const renderMessage = (message: ChatMessage, messageIndex: number) => {
     if (message.role === 'system' || message.role === 'tool') {
       return null;
     }
@@ -80,7 +132,7 @@ export const ChatTranscript: React.FC<ChatTranscriptProps> = ({
       }
     }
 
-    return (
+    const messageEl = (
       <div key={message.id} className={`chat-message ${message.role}`}>
         <div className="chat-message-avatar">
           {message.role === 'user' ? '\u{1F464}' : '\u{1F916}'}
@@ -113,28 +165,47 @@ export const ChatTranscript: React.FC<ChatTranscriptProps> = ({
         </div>
       </div>
     );
+
+    // After an assistant message, render any inline surfaces for this turn
+    if (message.role === 'assistant' && surfacesByTurn) {
+      const turnIndex = computeTurnIndex(messages, messageIndex);
+      const inlineSurfaces = renderInlineSurfaces(turnIndex);
+      if (inlineSurfaces) {
+        return (
+          <React.Fragment key={message.id}>
+            {messageEl}
+            {inlineSurfaces}
+          </React.Fragment>
+        );
+      }
+    }
+
+    return messageEl;
   };
 
   return (
     <>
-      {messages.map(renderMessage)}
+      {messages.map((message, index) => renderMessage(message, index))}
 
       {isStreaming && (streamingContent || (showToolMarkers && toolEvents.length > 0)) && (
-        <div className="chat-message assistant streaming">
+        <>
-          <div className="chat-message-avatar">{'\u{1F916}'}</div>
+          <div className="chat-message assistant streaming">
-          <div className="chat-message-content">
+            <div className="chat-message-avatar">{'\u{1F916}'}</div>
-            <div className="chat-message-header">
+            <div className="chat-message-content">
-              <span className="chat-message-role">{assistantRoleLabel}</span>
+              <div className="chat-message-header">
-              <span className="streaming-indicator">{'\u25CF'}</span>
+                <span className="chat-message-role">{assistantRoleLabel}</span>
-            </div>
+                <span className="streaming-indicator">{'\u25CF'}</span>
-            {showToolMarkers ? renderToolMarkers(toolEvents) : null}
-            {streamingContent && (
-              <div className="chat-message-text">
-                <Markdown gfm>{streamingContent}</Markdown>
               </div>
-            )}
+              {showToolMarkers ? renderToolMarkers(toolEvents) : null}
+              {streamingContent && (
+                <div className="chat-message-text">
+                  <Markdown gfm>{streamingContent}</Markdown>
+                </div>
+              )}
+            </div>
           </div>
-        </div>
+          {currentTurnIndex !== undefined && renderInlineSurfaces(currentTurnIndex)}
+        </>
       )}
 
       {isStreaming && !streamingContent && (!showToolMarkers || toolEvents.length === 0) && (

tests/engine/a2ui/surfaceManager.test.ts

@@ -34,6 +34,21 @@ describe('A2UISurfaceManager', () => {
       expect(surface!.dataModel).toEqual({});
     });
 
+    it('preserves metadata including turnIndex on createSurface', () => {
+      const manager = new A2UISurfaceManager();
+
+      manager.processMessage({
+        type: 'createSurface',
+        surfaceId: 'surface-1',
+        conversationId: 'conv-1',
+        metadata: { turnIndex: 3 },
+      });
+
+      const surface = manager.getSurface('surface-1');
+      expect(surface).toBeDefined();
+      expect(surface!.metadata).toEqual({ turnIndex: 3 });
+    });
+
     it('notifies listeners on surface creation', () => {
       const manager = new A2UISurfaceManager();
       const listener = vi.fn();

tests/renderer/a2ui/InlineSurface.test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, it } from 'vitest';
+import { deriveSurfaceTitle, getSurfaceIcon } from '../../../src/renderer/a2ui/InlineSurface';
+import type { A2UIResolvedComponent } from '../../../src/main/a2ui/types';
+
+function makeTree(type: string, properties: Record<string, unknown> = {}): A2UIResolvedComponent[] {
+  return [{
+    id: `${type}-1`,
+    type: type as A2UIResolvedComponent['type'],
+    properties,
+    children: [],
+  }];
+}
+
+describe('deriveSurfaceTitle', () => {
+  it('extracts title from root component properties', () => {
+    expect(deriveSurfaceTitle(makeTree('chart', { title: 'Post Views' }))).toBe('Post Views');
+  });
+
+  it('extracts label from root component properties when no title', () => {
+    expect(deriveSurfaceTitle(makeTree('metric', { label: 'Total Posts' }))).toBe('Total Posts');
+  });
+
+  it('falls back to capitalized component type when no title or label', () => {
+    expect(deriveSurfaceTitle(makeTree('chart'))).toBe('Chart');
+  });
+
+  it('returns "Surface" for an empty tree', () => {
+    expect(deriveSurfaceTitle([])).toBe('Surface');
+  });
+
+  it('capitalizes multi-word types correctly', () => {
+    expect(deriveSurfaceTitle(makeTree('textField'))).toBe('TextField');
+  });
+});
+
+describe('getSurfaceIcon', () => {
+  it('returns chart icon for chart type', () => {
+    expect(getSurfaceIcon(makeTree('chart'))).toBe('\u{1F4CA}');
+  });
+
+  it('returns table icon for table type', () => {
+    expect(getSurfaceIcon(makeTree('table'))).toBe('\u{1F4CB}');
+  });
+
+  it('returns form icon for form type', () => {
+    expect(getSurfaceIcon(makeTree('form'))).toBe('\u{1F4DD}');
+  });
+
+  it('returns card icon for card type', () => {
+    expect(getSurfaceIcon(makeTree('card'))).toBe('\u{1F4C4}');
+  });
+
+  it('returns metric icon for metric type', () => {
+    expect(getSurfaceIcon(makeTree('metric'))).toBe('\u{1F4CF}');
+  });
+
+  it('returns list icon for list type', () => {
+    expect(getSurfaceIcon(makeTree('list'))).toBe('\u{1F4CB}');
+  });
+
+  it('returns tabs icon for tabs type', () => {
+    expect(getSurfaceIcon(makeTree('tabs'))).toBe('\u{1F4C2}');
+  });
+
+  it('returns default icon for unknown types', () => {
+    expect(getSurfaceIcon(makeTree('text'))).toBe('\u25A0');
+  });
+
+  it('returns default icon for empty tree', () => {
+    expect(getSurfaceIcon([])).toBe('\u25A0');
+  });
+});

tests/renderer/a2ui/inlineSurfaceAssociation.test.ts

@@ -0,0 +1,45 @@
+import { describe, expect, it } from 'vitest';
+import { computeTurnIndex } from '../../../src/renderer/a2ui/surfaceAssociation';
+import type { ChatMessage } from '../../../src/main/shared/electronApi';
+
+function msg(role: ChatMessage['role'], id = `msg-${Date.now()}-${Math.random()}`): ChatMessage {
+  return { id, conversationId: 'conv-1', role, content: '', createdAt: new Date().toISOString() };
+}
+
+describe('computeTurnIndex', () => {
+  it('returns 0 for the first user message', () => {
+    const messages = [msg('user')];
+    expect(computeTurnIndex(messages, 0)).toBe(0);
+  });
+
+  it('returns 0 for an assistant message following the first user message', () => {
+    const messages = [msg('user'), msg('assistant')];
+    expect(computeTurnIndex(messages, 1)).toBe(0);
+  });
+
+  it('increments turn index for each user message', () => {
+    const messages = [msg('user'), msg('assistant'), msg('user'), msg('assistant')];
+    expect(computeTurnIndex(messages, 0)).toBe(0);
+    expect(computeTurnIndex(messages, 1)).toBe(0);
+    expect(computeTurnIndex(messages, 2)).toBe(1);
+    expect(computeTurnIndex(messages, 3)).toBe(1);
+  });
+
+  it('skips system and tool messages in turn counting', () => {
+    const messages = [msg('system'), msg('user'), msg('tool'), msg('assistant')];
+    expect(computeTurnIndex(messages, 1)).toBe(0); // user
+    expect(computeTurnIndex(messages, 3)).toBe(0); // assistant after first user
+  });
+
+  it('returns -1 when no user messages precede the index', () => {
+    const messages = [msg('system'), msg('assistant')];
+    expect(computeTurnIndex(messages, 0)).toBe(-1);
+    expect(computeTurnIndex(messages, 1)).toBe(-1);
+  });
+
+  it('handles multiple assistant messages in the same turn', () => {
+    const messages = [msg('user'), msg('assistant'), msg('assistant')];
+    expect(computeTurnIndex(messages, 1)).toBe(0);
+    expect(computeTurnIndex(messages, 2)).toBe(0);
+  });
+});