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 +++++++++++++++++++++++-------------
 1 file changed, 156 insertions(+), 84 deletions(-)

(limited to 'frontend/src/components/Sidebar.tsx')

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', 
-- 
cgit v1.2.3

