some fix on trace merging

2 files changed, 187 insertions, 127 deletions

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) => {