From c3673766aecdb988bb4e811376d4f1f1e18f1e0f Mon Sep 17 00:00:00 2001
From: blackhao <13851610112@163.com>
Date: Tue, 9 Dec 2025 15:05:04 -0600
Subject: some fix on trace merging

---
 frontend/src/components/Sidebar.tsx       | 240 +++++++++++++++---------
 frontend/src/components/nodes/LLMNode.tsx |  74 ++++----
 frontend/src/store/flowStore.ts           | 298 ++++++++++++++++++++++++------
 3 files changed, 432 insertions(+), 180 deletions(-)

(limited to 'frontend/src')

diff --git a/frontend/src/components/Sidebar.tsx b/frontend/src/components/Sidebar.tsx
index 5516629..06c8704 100644
--- a/frontend/src/components/Sidebar.tsx
+++ b/frontend/src/components/Sidebar.tsx
@@ -544,56 +544,118 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
     return true;
   };
 
-  // Helper: Check if specific trace path upstream has complete nodes
-  const checkTracePathComplete = (nodeId: string, traceId: string, visited: Set<string> = new Set()): boolean => {
-    if (visited.has(nodeId)) return true;
-    visited.add(nodeId);
-    
-    // Find the incoming edge that carries this trace
-    const incomingEdge = edges.find(e => 
-      e.target === nodeId && 
-      e.sourceHandle === `trace-${traceId}`
-    );
-    
-    if (!incomingEdge) return true; // Reached head of trace segment
-    
+  // Helper: find incoming edge for a given trace ID (with fallbacks)
+  const findIncomingEdgeForTrace = (nodeId: string, traceId: string): Edge | null => {
+    // 1) exact match by sourceHandle
+    let edge = edges.find(e => e.target === nodeId && e.sourceHandle === `trace-${traceId}`);
+    if (edge) return edge;
+    // 2) fallback: any incoming edge whose source has this trace in outgoingTraces
+    edge = edges.find(e => {
+      if (e.target !== nodeId) return false;
+      const src = nodes.find(n => n.id === e.source);
+      return src?.data.outgoingTraces?.some((t: Trace) => t.id === traceId);
+    });
+    return edge || null;
+  };
+
+  // Helper: get source trace IDs for a merged trace on a given node (supports propagated merged traces)
+  const getMergedSourceIds = (nodeId: string, traceId: string): string[] => {
+    const node = nodes.find(n => n.id === nodeId);
+    if (!node) return [];
+    const mergedLocal = node.data.mergedTraces?.find((m: MergedTrace) => m.id === traceId);
+    if (mergedLocal) return mergedLocal.sourceTraceIds || [];
+    const incomingMatch = node.data.traces?.find((t: Trace) => t.id === traceId);
+    if (incomingMatch?.isMerged && incomingMatch.sourceTraceIds) return incomingMatch.sourceTraceIds;
+    const outgoingMatch = node.data.outgoingTraces?.find((t: Trace) => t.id === traceId);
+    if (outgoingMatch?.isMerged && outgoingMatch.sourceTraceIds) return outgoingMatch.sourceTraceIds;
+    return [];
+  };
+
+  // Recursive: Check if specific trace path upstream has complete nodes (supports multi-level merged)
+  const checkTracePathComplete = (
+    nodeId: string,
+    traceId: string,
+    visited: Set<string> = new Set()
+  ): boolean => {
+    const visitKey = `${nodeId}-${traceId}`;
+    if (visited.has(visitKey)) return true;
+    visited.add(visitKey);
+
+    // Determine if this node is the merge owner or just receiving a propagated merged trace
+    const localMerge = nodes.find(n => n.id === nodeId)?.data.mergedTraces?.some(m => m.id === traceId);
+    const localParents = getMergedSourceIds(nodeId, traceId);
+
+    const incomingEdge = findIncomingEdgeForTrace(nodeId, traceId);
+    if (!incomingEdge) {
+      // If no incoming edge and this node owns the merge, check parents from here
+      if (localMerge && localParents.length > 0) {
+        for (const pid of localParents) {
+          if (!checkTracePathComplete(nodeId, pid, visited)) return false;
+        }
+        return true;
+      }
+      return true; // head
+    }
+
     const sourceNode = nodes.find(n => n.id === incomingEdge.source);
     if (!sourceNode || sourceNode.data.disabled) return true;
-    
-    // Check if source node is complete
-    if (!sourceNode.data.userPrompt || !sourceNode.data.response) {
-      return false; // Found incomplete node
+
+    // If merged at sourceNode (or propagated merged), recurse into each parent from the merge owner
+    const parentIds = localMerge ? localParents : getMergedSourceIds(sourceNode.id, traceId);
+    if (parentIds.length > 0) {
+      const mergeOwnerId = localMerge ? nodeId : sourceNode.id;
+      for (const pid of parentIds) {
+        if (!checkTracePathComplete(mergeOwnerId, pid, visited)) return false;
+      }
+      return true;
     }
-    
-    // Continue upstream
+
+    // Regular trace: check node content then continue upstream
+    if (!sourceNode.data.userPrompt || !sourceNode.data.response) return false;
     return checkTracePathComplete(sourceNode.id, traceId, visited);
   };
 
-  // Helper: Find the first empty node on a specific trace path
-  const findEmptyNodeOnTrace = (nodeId: string, traceId: string, visited: Set<string> = new Set()): string | null => {
-    if (visited.has(nodeId)) return null;
-    visited.add(nodeId);
-    
-    const incomingEdge = edges.find(e => 
-      e.target === nodeId && 
-      e.sourceHandle === `trace-${traceId}`
-    );
-    
-    if (!incomingEdge) return null;
-    
+  // Recursive: Find the first empty node on a specific trace path (supports multi-level merged)
+  const findEmptyNodeOnTrace = (
+    nodeId: string,
+    traceId: string,
+    visited: Set<string> = new Set()
+  ): string | null => {
+    const visitKey = `${nodeId}-${traceId}`;
+    if (visited.has(visitKey)) return null;
+    visited.add(visitKey);
+
+    // Determine if this node owns the merge or just receives propagated merged trace
+    const localMerge = nodes.find(n => n.id === nodeId)?.data.mergedTraces?.some(m => m.id === traceId);
+    const localParents = getMergedSourceIds(nodeId, traceId);
+
+    const incomingEdge = findIncomingEdgeForTrace(nodeId, traceId);
+    if (!incomingEdge) {
+      if (localMerge && localParents.length > 0) {
+        for (const pid of localParents) {
+          const upstreamEmpty = findEmptyNodeOnTrace(nodeId, pid, visited);
+          if (upstreamEmpty) return upstreamEmpty;
+        }
+      }
+      return null;
+    }
+
     const sourceNode = nodes.find(n => n.id === incomingEdge.source);
     if (!sourceNode || sourceNode.data.disabled) return null;
-    
-    // Recursively check upstream first (find the furthest empty node)
-    const upstreamEmpty = findEmptyNodeOnTrace(sourceNode.id, traceId, visited);
-    if (upstreamEmpty) return upstreamEmpty;
-    
-    // If no further upstream empty, check this node
+
+    const parentIds = localMerge ? localParents : getMergedSourceIds(sourceNode.id, traceId);
+    if (parentIds.length > 0) {
+      const mergeOwnerId = localMerge ? nodeId : sourceNode.id;
+      for (const pid of parentIds) {
+        const upstreamEmpty = findEmptyNodeOnTrace(mergeOwnerId, pid, visited);
+        if (upstreamEmpty) return upstreamEmpty;
+      }
+    }
+
     if (!sourceNode.data.userPrompt || !sourceNode.data.response) {
       return sourceNode.id;
     }
-    
-    return null;
+    return findEmptyNodeOnTrace(sourceNode.id, traceId, visited);
   };
 
   // Check if all active traces are complete (for main Run Node button)
@@ -603,23 +665,10 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
     const activeTraceIds = selectedNode.data.activeTraceIds || [];
     if (activeTraceIds.length === 0) return { complete: true };
     
-    // Check upstream nodes ONLY for active traces
+    // Check upstream nodes ONLY for active traces (supports merged trace recursion)
     for (const traceId of activeTraceIds) {
-      // Check if it's a merged trace
-      const merged = selectedNode.data.mergedTraces?.find((m: MergedTrace) => m.id === traceId);
-      
-      if (merged) {
-        // For merged trace, check all source traces
-        for (const sourceId of merged.sourceTraceIds) {
-          if (!checkTracePathComplete(selectedNode.id, sourceId)) {
-            return { complete: false, incompleteTraceId: 'upstream' };
-          }
-        }
-      } else {
-        // For regular trace
-        if (!checkTracePathComplete(selectedNode.id, traceId)) {
-          return { complete: false, incompleteTraceId: 'upstream' };
-        }
+      if (!checkTracePathComplete(selectedNode.id, traceId)) {
+        return { complete: false, incompleteTraceId: 'upstream' };
       }
     }
     
@@ -632,12 +681,11 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
       }
     }
     
-    // Check merged traces content
-    const mergedTraces = selectedNode.data.mergedTraces || [];
+    // Check merged traces content (including propagated merged traces)
     for (const traceId of activeTraceIds) {
-      const merged = mergedTraces.find((m: MergedTrace) => m.id === traceId);
-      if (merged) {
-        for (const sourceId of merged.sourceTraceIds) {
+      const sourceIds = getMergedSourceIds(selectedNode.id, traceId);
+      if (sourceIds.length > 0) {
+        for (const sourceId of sourceIds) {
           const sourceTrace = incomingTraces.find((t: Trace) => t.id === sourceId);
           if (sourceTrace && !isTraceComplete(sourceTrace)) {
             return { complete: false, incompleteTraceId: sourceId };
@@ -655,18 +703,7 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
     const activeTraceIds = selectedNode.data.activeTraceIds || [];
     
     for (const traceId of activeTraceIds) {
-      let emptyNodeId: string | null = null;
-      
-      const merged = selectedNode.data.mergedTraces?.find((m: MergedTrace) => m.id === traceId);
-      if (merged) {
-        for (const sourceId of merged.sourceTraceIds) {
-          emptyNodeId = findEmptyNodeOnTrace(selectedNode.id, sourceId);
-          if (emptyNodeId) break;
-        }
-      } else {
-        emptyNodeId = findEmptyNodeOnTrace(selectedNode.id, traceId);
-      }
-      
+      const emptyNodeId = findEmptyNodeOnTrace(selectedNode.id, traceId);
       if (emptyNodeId) {
         const emptyNode = nodes.find(n => n.id === emptyNodeId);
         if (emptyNode) {
@@ -1237,8 +1274,32 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
                         
                         <div className="flex-1">
                           <div className="flex items-center gap-2">
-                            <div className="w-2 h-2 rounded-full" style={{ backgroundColor: trace.color }}></div>
-                            <span className={`font-mono text-xs ${isDark ? 'text-gray-500' : 'text-gray-400'}`}>#{trace.id.slice(-4)}</span>
+                            {trace.isMerged ? (
+                              // Merged Trace Rendering (for propagated merged traces)
+                              <div className="flex -space-x-1 shrink-0">
+                                {(trace.mergedColors || [trace.color]).slice(0, 3).map((color, idx) => (
+                                  <div 
+                                    key={idx}
+                                    className="w-2 h-2 rounded-full border-2"
+                                    style={{ backgroundColor: color, borderColor: isDark ? '#1f2937' : '#fff' }}
+                                  />
+                                ))}
+                                {(trace.mergedColors?.length || 0) > 3 && (
+                                  <div className={`w-2 h-2 rounded-full flex items-center justify-center text-[6px] ${
+                                    isDark ? 'bg-gray-700 text-gray-400' : 'bg-gray-200 text-gray-500'
+                                  }`}>
+                                    +
+                                  </div>
+                                )}
+                              </div>
+                            ) : (
+                              // Regular Trace Rendering
+                              <div className="w-2 h-2 rounded-full" style={{ backgroundColor: trace.color }}></div>
+                            )}
+                            
+                            <span className={`font-mono text-xs ${isDark ? 'text-gray-500' : 'text-gray-400'}`}>
+                              {trace.isMerged ? 'Merged ' : ''}#{trace.id.slice(-4)}
+                            </span>
                             {!isComplete && (
                               <span className="text-[9px] text-orange-500">(incomplete)</span>
                             )}
@@ -1359,10 +1420,13 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
                             e.stopPropagation();
                             openMergedQuickChat(merged);
                           }}
+                          disabled={!isComplete}
                           className={`p-1 rounded shrink-0 ${
-                            isDark ? 'hover:bg-purple-900 text-gray-500 hover:text-purple-400' : 'hover:bg-purple-100 text-gray-400 hover:text-purple-600'
+                            isComplete 
+                              ? isDark ? 'hover:bg-purple-900 text-gray-500 hover:text-purple-400' : 'hover:bg-purple-100 text-gray-400 hover:text-purple-600'
+                              : 'text-gray-500 cursor-not-allowed opacity-50'
                           }`}
-                          title="Quick Chat with merged context"
+                          title={isComplete ? "Quick Chat with merged context" : "Trace incomplete"}
                         >
                           <MessageCircle size={12} />
                         </button>
@@ -1473,12 +1537,18 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
                   <textarea
                     value={editedResponse}
                     onChange={(e) => setEditedResponse(e.target.value)}
-                    className="w-full border border-blue-300 rounded-md p-2 text-sm min-h-[200px] font-mono focus:ring-2 focus:ring-blue-500"
+                    className={`w-full border rounded-md p-2 text-sm min-h-[200px] font-mono focus:ring-2 focus:ring-blue-500 ${
+                      isDark 
+                        ? 'bg-gray-800 border-gray-600 text-gray-200 placeholder-gray-500' 
+                        : 'bg-white border-blue-300 text-gray-900'
+                    }`}
                   />
                   <div className="flex gap-2 justify-end">
                     <button
                       onClick={handleCancelEdit}
-                      className="px-3 py-1 text-sm text-gray-600 hover:bg-gray-100 rounded flex items-center gap-1"
+                      className={`px-3 py-1 text-sm rounded flex items-center gap-1 ${
+                        isDark ? 'text-gray-400 hover:bg-gray-800' : 'text-gray-600 hover:bg-gray-100'
+                      }`}
                     >
                       <X size={14} /> Cancel
                     </button>
@@ -1493,8 +1563,8 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
               ) : (
                 <div className={`p-3 rounded-md border min-h-[150px] text-sm prose prose-sm max-w-none ${
                   isDark 
-                    ? 'bg-gray-900 border-gray-700 prose-invert' 
-                    : 'bg-gray-50 border-gray-200'
+                    ? 'bg-gray-900 border-gray-700 prose-invert text-gray-200' 
+                    : 'bg-gray-50 border-gray-200 text-gray-900'
                 }`}>
                   <ReactMarkdown>{selectedNode.data.response || (streamingNodeId === selectedNode.id ? streamBuffer : '')}</ReactMarkdown>
                 </div>
@@ -1506,22 +1576,24 @@ const Sidebar: React.FC<SidebarProps> = ({ isOpen, onToggle, onInteract }) => {
         {activeTab === 'settings' && (
           <div className="space-y-4">
             <div>
-              <label className="block text-sm font-medium text-gray-700 mb-1">Merge Strategy</label>
+              <label className={`block text-sm font-medium mb-1 ${isDark ? 'text-gray-300' : 'text-gray-700'}`}>Merge Strategy</label>
               <select 
                 value={selectedNode.data.mergeStrategy || 'smart'} 
                 onChange={(e) => handleChange('mergeStrategy', e.target.value)}
-                className="w-full border border-gray-300 rounded-md p-2 text-sm"
+                className={`w-full border rounded-md p-2 text-sm ${
+                  isDark ? 'bg-gray-700 border-gray-600 text-gray-200' : 'border-gray-300 bg-white text-gray-900'
+                }`}
               >
                 <option value="smart">Smart (Auto-merge roles)</option>
                 <option value="raw">Raw (Concatenate)</option>
               </select>
-              <p className="text-xs text-gray-500 mt-1">
+              <p className={`text-xs mt-1 ${isDark ? 'text-gray-500' : 'text-gray-500'}`}>
                 Smart merge combines consecutive messages from the same role to avoid API errors.
               </p>
             </div>
 
             <div>
-              <label className="block text-sm font-medium text-gray-700 mb-1">
+              <label className={`block text-sm font-medium mb-1 ${isDark ? 'text-gray-300' : 'text-gray-700'}`}>
                 Temperature ({selectedNode.data.temperature})
                 {[
                   'gpt-5', 'gpt-5-chat-latest', 'gpt-5-mini', 'gpt-5-nano', 
diff --git a/frontend/src/components/nodes/LLMNode.tsx b/frontend/src/components/nodes/LLMNode.tsx
index 8105070..d2e1293 100644
--- a/frontend/src/components/nodes/LLMNode.tsx
+++ b/frontend/src/components/nodes/LLMNode.tsx
@@ -217,43 +217,8 @@ const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
            })
          }
          
-         {/* Prepend input handles for merged traces - only show if merged trace has downstream */}
-         {data.mergedTraces && data.mergedTraces
-           .filter((merged: MergedTrace) => {
-             // Check if this merged trace has any outgoing edges
-             return edges.some(e => 
-               e.source === id && e.sourceHandle === `trace-${merged.id}`
-             );
-           })
-           .map((merged: MergedTrace) => {
-           const connectedEdge = edges.find(e => e.target === id && e.targetHandle === `prepend-${merged.id}`);
-           const edgeColor = connectedEdge?.style?.stroke as string;
-           
-           // Create gradient for merged trace handle
-           const colors = merged.colors.length > 0 ? merged.colors : ['#888'];
-           const gradientStops = colors.map((color, idx) => 
-             `${color} ${(idx / colors.length) * 100}%, ${color} ${((idx + 1) / colors.length) * 100}%`
-           ).join(', ');
-           const stripeGradient = `linear-gradient(45deg, ${gradientStops})`;
-           
-           return (
-             <div key={`prepend-${merged.id}`} className="relative h-4 w-4 my-1" title={`Prepend context to merged: ${merged.id}`}>
-               <Handle 
-                  type="target" 
-                  position={Position.Left} 
-                  id={`prepend-${merged.id}`}
-                  className="!w-3 !h-3 !left-[-6px] !border-2"
-                  style={{ 
-                    top: '50%', 
-                    transform: 'translateY(-50%)',
-                    background: edgeColor || stripeGradient,
-                    borderColor: isDark ? '#374151' : '#fff',
-                    borderStyle: 'dashed'
-                  }}
-               />
-             </div>
-           );
-         })}
+         {/* Prepend input handles for merged traces - REMOVED per user request */}
+         {/* Users should prepend to parent traces instead */}
       </div>
 
       {/* Dynamic Outputs (Traces) */}
@@ -272,6 +237,16 @@ const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
               return !hasOutgoingEdge;
             })
             .map((trace: Trace) => {
+              // Handle merged trace visualization
+              let backgroundStyle = trace.color;
+              if (trace.isMerged && trace.mergedColors && trace.mergedColors.length > 0) {
+                const colors = trace.mergedColors;
+                const gradientStops = colors.map((color, idx) => 
+                  `${color} ${(idx / colors.length) * 100}%, ${color} ${((idx + 1) / colors.length) * 100}%`
+                ).join(', ');
+                backgroundStyle = `linear-gradient(45deg, ${gradientStops})`;
+              }
+
               return (
                 <div key={`continue-${trace.id}`} className="relative h-4 w-4 my-1" title={`Continue trace: ${trace.id}`}>
                   <Handle 
@@ -280,7 +255,7 @@ const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
                      id={`trace-${trace.id}`}
                      className="!w-3 !h-3 !right-[-6px]"
                      style={{ 
-                       backgroundColor: trace.color,
+                       background: backgroundStyle,
                        top: '50%', 
                        transform: 'translateY(-50%)',
                        border: `2px dashed ${isDark ? '#374151' : '#fff'}`
@@ -294,8 +269,9 @@ const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
           {/* Only show traces that have actual downstream edges */}
           {data.outgoingTraces && data.outgoingTraces
             .filter(trace => {
-              // Exclude merged traces - they have their own section
-              if (trace.id.startsWith('merged-')) return false;
+              // Check if this is a locally created merged trace (should be shown in Part 2)
+              const isLocallyMerged = data.mergedTraces?.some(m => m.id === trace.id);
+              if (isLocallyMerged) return false;
               
               // Only show if there's an actual downstream edge using this trace
               const hasDownstream = edges.some(e => 
@@ -303,7 +279,18 @@ const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
               );
               return hasDownstream;
             })
-            .map((trace) => (
+            .map((trace) => {
+             // Handle merged trace visualization
+             let backgroundStyle = trace.color;
+             if (trace.isMerged && trace.mergedColors && trace.mergedColors.length > 0) {
+               const colors = trace.mergedColors;
+               const gradientStops = colors.map((color, idx) => 
+                 `${color} ${(idx / colors.length) * 100}%, ${color} ${((idx + 1) / colors.length) * 100}%`
+               ).join(', ');
+               backgroundStyle = `linear-gradient(45deg, ${gradientStops})`;
+             }
+
+             return (
              <div key={trace.id} className="relative h-4 w-4 my-1" title={`Trace: ${trace.id}`}>
                 <Handle 
                    type="source" 
@@ -311,13 +298,14 @@ const LLMNode = ({ id, data, selected }: NodeProps<NodeData>) => {
                    id={`trace-${trace.id}`}
                    className="!w-3 !h-3 !right-[-6px]"
                    style={{ 
-                     backgroundColor: trace.color,
+                     background: backgroundStyle,
                      top: '50%', 
                      transform: 'translateY(-50%)' 
                    }}
                 />
              </div>
-          ))}
+             );
+          })}
           
           {/* 2. Merged Trace Handles (with alternating color stripes) */}
           {data.mergedTraces && data.mergedTraces.map((merged: MergedTrace) => {
diff --git a/frontend/src/store/flowStore.ts b/frontend/src/store/flowStore.ts
index 49a8ece..5ed66e6 100644
--- a/frontend/src/store/flowStore.ts
+++ b/frontend/src/store/flowStore.ts
@@ -30,6 +30,10 @@ export interface Trace {
   sourceNodeId: string;
   color: string;
   messages: Message[];
+  // Optional merged trace info for visual propagation
+  isMerged?: boolean;
+  mergedColors?: string[];
+  sourceTraceIds?: string[];
 }
 
 // Merge strategy types
@@ -200,12 +204,8 @@ const useFlowStore = create<FlowState>((set, get) => ({
     
     const isOverlapping = (testX: number, testY: number) => {
       return nodes.some(node => {
-        // Use the same estimated dimensions for existing nodes too
-        // Ideally we would know their actual dimensions, but this is a safe approximation
         const nodeX = node.position.x;
         const nodeY = node.position.y;
-        
-        // Check for overlap
         return !(testX + nodeWidth + padding < nodeX || 
                  testX > nodeX + nodeWidth + padding ||
                  testY + nodeHeight + padding < nodeY || 
@@ -383,7 +383,12 @@ const useFlowStore = create<FlowState>((set, get) => ({
     }
     
     // Helper to trace back the path of a trace by following edges upstream
-    const duplicateTracePath = (traceId: string, forkAtNodeId: string): { newTraceId: string, newEdges: Edge[], firstNodeId: string } | null => {
+    const duplicateTracePath = (
+      traceId: string, 
+      forkAtNodeId: string,
+      traceOwnerNodeId?: string,
+      pendingEdges: Edge[] = []
+    ): { newTraceId: string, newEdges: Edge[], firstNodeId: string } | null => {
       // Trace back from forkAtNodeId to find the origin of this trace
       // We follow incoming edges that match the trace pattern
       
@@ -393,10 +398,10 @@ const useFlowStore = create<FlowState>((set, get) => ({
       
       // Trace backwards through incoming edges
       while (true) {
-        // Find incoming edge to current node that's part of this trace
+        // Find incoming edge to current node that carries THIS trace ID
         const incomingEdge = edges.find(e => 
           e.target === currentNodeId && 
-          e.sourceHandle?.startsWith('trace-')
+          e.sourceHandle === `trace-${traceId}`
         );
         
         if (!incomingEdge) break; // Reached the start of the trace
@@ -413,14 +418,13 @@ const useFlowStore = create<FlowState>((set, get) => ({
       const firstNode = nodes.find(n => n.id === firstNodeId);
       if (!firstNode) return null;
       
-      // Create a new trace ID for the duplicated path
-      const timestamp = Date.now();
-      const newTraceId = `fork-${firstNodeId}-${timestamp}`;
+      // Create a new trace ID for the duplicated path (guarantee uniqueness even within the same ms)
+      const uniq = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+      const newTraceId = `fork-${firstNodeId}-${uniq}`;
       const newTraceColor = getStableColor(newTraceId);
       
       // Create new edges for the entire path
       const newEdges: Edge[] = [];
-      let evolvedTraceId = newTraceId;
       
       // Track which input handles we're creating for new edges
       const newInputHandles: Map<string, number> = new Map();
@@ -432,26 +436,30 @@ const useFlowStore = create<FlowState>((set, get) => ({
         
         // Find the next available input handle for the target node
         // Count existing edges to this node + any new edges we're creating
-        const existingEdgesToTarget = edges.filter(e => e.target === toNodeId).length;
+        const existingEdgesToTarget =
+          edges.filter(e => e.target === toNodeId).length +
+          pendingEdges.filter(e => e.target === toNodeId).length;
         const newEdgesToTarget = newInputHandles.get(toNodeId) || 0;
         const nextInputIndex = existingEdgesToTarget + newEdgesToTarget;
         newInputHandles.set(toNodeId, newEdgesToTarget + 1);
         
         newEdges.push({
-          id: `edge-fork-${timestamp}-${i}`,
+          id: `edge-fork-${uniq}-${i}`,
           source: fromNodeId,
           target: toNodeId,
-          sourceHandle: `trace-${evolvedTraceId}`,
+          sourceHandle: `trace-${newTraceId}`,
           targetHandle: `input-${nextInputIndex}`,
           style: { stroke: newTraceColor, strokeWidth: 2 }
         });
-        
-        // Evolve trace ID for next edge
-        evolvedTraceId = `${evolvedTraceId}_${toNodeId}`;
       }
       
       // Find the messages up to the fork point
-      const originalTrace = sourceNode.data.outgoingTraces?.find(t => t.id === traceId);
+      const traceOwnerNode = traceOwnerNodeId
+        ? nodes.find(n => n.id === traceOwnerNodeId)
+        : sourceNode;
+      if (!traceOwnerNode) return null;
+
+      const originalTrace = traceOwnerNode.data.outgoingTraces?.find(t => t.id === traceId);
       const messagesUpToFork = originalTrace?.messages || [];
       
       // Add the new trace as a forked trace on the first node
@@ -468,6 +476,161 @@ const useFlowStore = create<FlowState>((set, get) => ({
       
       return { newTraceId, newEdges, firstNodeId };
     };
+
+    // Helper to duplicate the downstream segment of a trace from a start node to an end node
+    const duplicateDownstreamSegment = (
+      originalTraceId: string,
+      startNodeId: string,
+      endNodeId: string,
+      newTraceId: string,
+      newTraceColor: string,
+      newTraceColors: string[]
+    ): Edge[] | null => {
+      const segmentEdges: Edge[] = [];
+      let currentNodeId = startNodeId;
+      const visitedEdgeIds = new Set<string>();
+
+      while (currentNodeId !== endNodeId) {
+        const nextEdge = edges.find(
+          (e) => e.source === currentNodeId && e.sourceHandle === `trace-${originalTraceId}`
+        );
+
+        if (!nextEdge || visitedEdgeIds.has(nextEdge.id)) {
+          return null;
+        }
+
+        segmentEdges.push(nextEdge);
+        visitedEdgeIds.add(nextEdge.id);
+        currentNodeId = nextEdge.target;
+      }
+
+      const newEdges: Edge[] = [];
+      const newInputCounts: Map<string, number> = new Map();
+      const segmentTimestamp = Date.now();
+
+      segmentEdges.forEach((edge, index) => {
+        const targetNodeId = edge.target;
+        const existingEdgesToTarget = edges.filter((e) => e.target === targetNodeId).length;
+        const additionalEdges = newInputCounts.get(targetNodeId) || 0;
+        const nextInputIndex = existingEdgesToTarget + additionalEdges;
+        newInputCounts.set(targetNodeId, additionalEdges + 1);
+
+        newEdges.push({
+          id: `edge-merged-seg-${segmentTimestamp}-${index}`,
+          source: edge.source,
+          target: edge.target,
+          sourceHandle: `trace-${newTraceId}`,
+          targetHandle: `input-${nextInputIndex}`,
+          type: 'merged',
+          style: { stroke: newTraceColor, strokeWidth: 2 },
+          data: { isMerged: true, colors: newTraceColors }
+        });
+      });
+
+      return newEdges;
+    };
+
+    // Helper to duplicate a merged trace by cloning its parent traces and creating a new merged branch
+    const duplicateMergedTraceBranch = (
+      mergedTrace: Trace,
+      forkAtNodeId: string
+    ): { newTraceId: string; newEdges: Edge[]; color: string } | null => {
+      const mergeNodeId = mergedTrace.sourceNodeId;
+      const mergeNode = nodes.find((n) => n.id === mergeNodeId);
+      if (!mergeNode) return null;
+
+      const mergedDef =
+        mergeNode.data.mergedTraces?.find((m: MergedTrace) => m.id === mergedTrace.id) || null;
+      const parentTraceIds = mergedTrace.sourceTraceIds || mergedDef?.sourceTraceIds || [];
+      if (parentTraceIds.length === 0) return null;
+
+      let accumulatedEdges: Edge[] = [];
+      const newParentTraceIds: string[] = [];
+      const parentOverrides: Trace[] = [];
+
+      for (const parentId of parentTraceIds) {
+        const originalParentTrace = mergeNode.data.traces?.find((t: Trace) => t.id === parentId);
+
+        if (originalParentTrace?.isMerged && originalParentTrace.sourceTraceIds?.length) {
+          const nestedDuplicate = duplicateMergedTraceBranch(originalParentTrace, mergeNodeId);
+          if (!nestedDuplicate) {
+            return null;
+          }
+          accumulatedEdges = accumulatedEdges.concat(nestedDuplicate.newEdges);
+          newParentTraceIds.push(nestedDuplicate.newTraceId);
+
+          parentOverrides.push({
+            ...originalParentTrace,
+            id: nestedDuplicate.newTraceId,
+          });
+
+          continue;
+        }
+
+        const duplicateResult = duplicateTracePath(parentId, mergeNodeId, mergeNodeId, accumulatedEdges);
+        if (!duplicateResult) {
+          return null;
+        }
+        accumulatedEdges = accumulatedEdges.concat(duplicateResult.newEdges);
+        newParentTraceIds.push(duplicateResult.newTraceId);
+
+        if (originalParentTrace) {
+          parentOverrides.push({
+            ...originalParentTrace,
+            id: duplicateResult.newTraceId,
+          });
+        }
+      }
+
+      const strategy = mergedDef?.strategy || 'trace_order';
+      const uniqMerged = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+      const newMergedId = `merged-${mergeNodeId}-${uniqMerged}`;
+      const newColors =
+        parentOverrides.length > 0
+          ? parentOverrides.map((t) => t.color).filter((c): c is string => Boolean(c))
+          : mergedTrace.mergedColors ?? [];
+
+      const overrideTraces = parentOverrides.length > 0 ? parentOverrides : undefined;
+      const mergedMessages = get().computeMergedMessages(
+        mergeNodeId,
+        newParentTraceIds,
+        strategy,
+        overrideTraces
+      );
+
+      const newMergedDefinition: MergedTrace = {
+        id: newMergedId,
+        sourceNodeId: mergeNodeId,
+        sourceTraceIds: newParentTraceIds,
+        strategy,
+        colors: newColors.length ? newColors : [mergedTrace.color],
+        messages: mergedMessages,
+      };
+
+      const existingMerged = mergeNode.data.mergedTraces || [];
+      get().updateNodeData(mergeNodeId, {
+        mergedTraces: [...existingMerged, newMergedDefinition],
+      });
+
+      const newMergedColor = newColors[0] || mergedTrace.color || getStableColor(newMergedId);
+      const downstreamEdges = duplicateDownstreamSegment(
+        mergedTrace.id,
+        mergeNodeId,
+        forkAtNodeId,
+        newMergedId,
+        newMergedColor,
+        newColors.length ? newColors : [mergedTrace.color]
+      );
+      if (!downstreamEdges) return null;
+
+      accumulatedEdges = accumulatedEdges.concat(downstreamEdges);
+
+      return {
+        newTraceId: newMergedId,
+        newEdges: accumulatedEdges,
+        color: newMergedColor,
+      };
+    };
     
     // Helper to create a simple forked trace (for new-trace handle or first connection)
     const createSimpleForkTrace = () => {
@@ -527,30 +690,40 @@ const useFlowStore = create<FlowState>((set, get) => ({
     );
     
     if (existingEdgeFromHandle && connection.sourceHandle?.startsWith('trace-')) {
-      // This handle already has a connection - need to duplicate the entire upstream trace path
       const originalTraceId = connection.sourceHandle.replace('trace-', '');
+      const traceMeta = sourceNode.data.outgoingTraces?.find((t: Trace) => t.id === originalTraceId);
+
+      if (traceMeta?.isMerged && traceMeta.sourceTraceIds && traceMeta.sourceTraceIds.length > 0) {
+        const mergedDuplicate = duplicateMergedTraceBranch(traceMeta, connection.source!);
+        if (mergedDuplicate) {
+          set({
+            edges: [
+              ...get().edges,
+              ...mergedDuplicate.newEdges,
+              {
+              id: `edge-${connection.source}-${connection.target}-${Date.now()}`,
+              source: connection.source!,
+              target: connection.target!,
+              sourceHandle: `trace-${mergedDuplicate.newTraceId}`,
+              targetHandle: connection.targetHandle,
+              type: 'merged',
+              style: { stroke: mergedDuplicate.color, strokeWidth: 2 },
+              data: { isMerged: true, colors: traceMeta.mergedColors || [] }
+              } as Edge
+            ],
+          });
+          
+          setTimeout(() => get().propagateTraces(), 0);
+          return;
+        }
+      }
+
       const duplicateResult = duplicateTracePath(originalTraceId, connection.source!);
       
       if (duplicateResult) {
-        // Add all the duplicated edges plus the new connection
-        const { newTraceId, newEdges, firstNodeId } = duplicateResult;
+        const { newTraceId, newEdges } = duplicateResult;
         const newTraceColor = getStableColor(newTraceId);
         
-        // Calculate the evolved trace ID at the fork node
-        // The trace evolves through each node: trace-A -> trace-A_B -> trace-A_B_C
-        // We need to build the evolved ID based on the path
-        let evolvedTraceId = newTraceId;
-        
-        // Find the path from first node to fork node by looking at the new edges
-        for (const edge of newEdges) {
-          if (edge.target === connection.source) {
-            // This edge ends at our fork node, so the evolved trace ID is after this edge
-            evolvedTraceId = `${evolvedTraceId}_${edge.target}`;
-            break;
-          }
-          evolvedTraceId = `${evolvedTraceId}_${edge.target}`;
-        }
-        
         set({
           edges: [
             ...get().edges,
@@ -559,7 +732,7 @@ const useFlowStore = create<FlowState>((set, get) => ({
               id: `edge-${connection.source}-${connection.target}-${Date.now()}`,
               source: connection.source!,
               target: connection.target!,
-              sourceHandle: `trace-${evolvedTraceId}`,
+              sourceHandle: `trace-${newTraceId}`,
               targetHandle: connection.targetHandle,
               style: { stroke: newTraceColor, strokeWidth: 2 }
             } as Edge
@@ -569,7 +742,6 @@ const useFlowStore = create<FlowState>((set, get) => ({
         setTimeout(() => get().propagateTraces(), 0);
         return;
       } else {
-        // Fallback to simple fork if path duplication fails
         const newForkTrace = createSimpleForkTrace();
         
         set({
@@ -1561,24 +1733,38 @@ const useFlowStore = create<FlowState>((set, get) => ({
              const newHandleId = `trace-${matchedTrace.id}`;
              
              // Check if this is a merged trace (need gradient)
-             const isMergedTrace = matchedTrace.id.startsWith('merged-');
+             // Use the new properties on Trace object
+             const isMergedTrace = matchedTrace.isMerged || matchedTrace.id.startsWith('merged-');
+             const mergedColors = matchedTrace.mergedColors || [];
+             
+             // If colors not on trace, try to find in parent node's mergedTraces (for originator)
+             let finalColors = mergedColors;
+             if (isMergedTrace && finalColors.length === 0) {
       const parentNode = nodes.find(n => n.id === edge.source);
-             const mergedTraceData = isMergedTrace 
-               ? parentNode?.data.mergedTraces?.find((m: MergedTrace) => m.id === matchedTrace.id)
-               : null;
+                const mergedData = parentNode?.data.mergedTraces?.find((m: MergedTrace) => m.id === matchedTrace.id);
+                if (mergedData) finalColors = mergedData.colors;
+             }
              
              // Create gradient for merged traces
              let gradient: string | undefined;
-             if (mergedTraceData && mergedTraceData.colors.length > 0) {
-               const colors = mergedTraceData.colors;
-               const gradientStops = colors.map((color: string, idx: number) => 
-                 `${color} ${(idx / colors.length) * 100}%, ${color} ${((idx + 1) / colors.length) * 100}%`
+             if (finalColors.length > 0) {
+               const gradientStops = finalColors.map((color: string, idx: number) => 
+                 `${color} ${(idx / finalColors.length) * 100}%, ${color} ${((idx + 1) / finalColors.length) * 100}%`
                ).join(', ');
                gradient = `linear-gradient(90deg, ${gradientStops})`;
              }
              
              // Check if we need to update
-             if (currentEdge.sourceHandle !== newHandleId || currentEdge.style?.stroke !== matchedTrace.color) {
+             // Update if handle changed OR color changed OR merged status/colors changed
+             const currentIsMerged = currentEdge.data?.isMerged;
+             const currentColors = currentEdge.data?.colors;
+             const colorsChanged = JSON.stringify(currentColors) !== JSON.stringify(finalColors);
+             
+             if (currentEdge.sourceHandle !== newHandleId || 
+                 currentEdge.style?.stroke !== matchedTrace.color ||
+                 currentIsMerged !== isMergedTrace ||
+                 colorsChanged) {
+                 
                 updatedEdges[edgeIndex] = {
                   ...currentEdge,
                   sourceHandle: newHandleId,
@@ -1588,7 +1774,7 @@ const useFlowStore = create<FlowState>((set, get) => ({
                     ...currentEdge.data, 
                     gradient, 
                     isMerged: isMergedTrace,
-                    colors: mergedTraceData?.colors || []
+                    colors: finalColors
                   }
                 };
                 edgesChanged = true;
@@ -1749,7 +1935,10 @@ const useFlowStore = create<FlowState>((set, get) => ({
             id: merged.id,
             sourceNodeId: node.id,
             color: updatedColors[0] || getStableColor(merged.id),
-            messages: mergedMessages
+            messages: mergedMessages,
+            isMerged: true,
+            mergedColors: updatedColors,
+            sourceTraceIds: merged.sourceTraceIds
           };
           
           myOutgoingTraces.push(mergedOutgoing);
@@ -1775,19 +1964,22 @@ const useFlowStore = create<FlowState>((set, get) => ({
     });
     
     // Bulk Update Store
+    const uniqTraces = (list: Trace[]) => Array.from(new Map(list.map(t => [t.id, t])).values());
+    const uniqMerged = (list: MergedTrace[]) => Array.from(new Map(list.map(m => [m.id, m])).values());
+
     set(state => ({
       edges: updatedEdges,
       nodes: state.nodes.map(n => {
-        const traces = nodeIncomingTraces.get(n.id) || [];
-        const outTraces = nodeOutgoingTraces.get(n.id) || [];
+        const traces = uniqTraces(nodeIncomingTraces.get(n.id) || []);
+        const outTraces = uniqTraces(nodeOutgoingTraces.get(n.id) || []);
         const mergedToDelete = nodeMergedTracesToDelete.get(n.id) || [];
         const updatedMerged = nodeUpdatedMergedTraces.get(n.id);
         const cleanedForks = nodeForkedTracesToClean.get(n.id);
         
         // Filter out disconnected merged traces and update messages for remaining ones
-        let filteredMergedTraces = (n.data.mergedTraces || []).filter(
+        let filteredMergedTraces = uniqMerged((n.data.mergedTraces || []).filter(
           (m: MergedTrace) => !mergedToDelete.includes(m.id)
-        );
+        ));
         
         // Apply updated messages and colors to merged traces
         if (updatedMerged && updatedMerged.size > 0) {
-- 
cgit v1.2.3