some fix on trace merging

1 files changed, 31 insertions, 43 deletions

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