right click logicHEADmain

3 files changed, 187 insertions, 41 deletions

diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index 1eaafec..8c52751 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -1,16 +1,19 @@
-import { useCallback, useRef } from 'react';
+import { useCallback, useRef, useState } from 'react';
 import ReactFlow, { 
   Background, 
   Controls, 
   MiniMap, 
   ReactFlowProvider,
   Panel,
-  useReactFlow
+  useReactFlow,
+  type Node,
+  type Edge
 } from 'reactflow';
 import 'reactflow/dist/style.css';
 import useFlowStore from './store/flowStore';
 import LLMNode from './components/nodes/LLMNode';
 import Sidebar from './components/Sidebar';
+import { ContextMenu } from './components/ContextMenu';
 import { Plus } from 'lucide-react';
 
 const nodeTypes = {
@@ -25,59 +28,44 @@ function Flow() {
     onEdgesChange, 
     onConnect, 
     addNode,
+    deleteEdge,
+    deleteNode,
+    deleteBranch,
     setSelectedNode 
   } = useFlowStore();
 
   const reactFlowWrapper = useRef<HTMLDivElement>(null);
   const { project } = useReactFlow();
+  const [menu, setMenu] = useState<{ x: number; y: number; type: 'pane' | 'node' | 'edge'; id?: string } | null>(null);
 
-  const handleAddNode = () => {
-    const id = `node_${Date.now()}`;
-    addNode({
-      id,
-      type: 'llmNode',
-      position: { x: Math.random() * 400, y: Math.random() * 400 },
-      data: { 
-        label: 'New Question', 
-        model: 'gpt-4o', 
-        temperature: 0.7,
-        systemPrompt: '',
-        userPrompt: '',
-        mergeStrategy: 'smart',
-        messages: [],
-        traces: [],
-        outgoingTraces: [],
-        forkedTraces: [],
-        response: '',
-        status: 'idle',
-        inputs: 1
-      },
-    });
+  const onPaneClick = () => {
+    setSelectedNode(null);
+    setMenu(null);
   };
 
-  const onNodeClick = (_: any, node: any) => {
-    setSelectedNode(node.id);
+  const handlePaneContextMenu = (event: React.MouseEvent) => {
+     event.preventDefault();
+     setMenu({ x: event.clientX, y: event.clientY, type: 'pane' });
   };
 
-  const onPaneClick = () => {
-    setSelectedNode(null);
+  const handleNodeContextMenu = (event: React.MouseEvent, node: Node) => {
+     event.preventDefault();
+     setMenu({ x: event.clientX, y: event.clientY, type: 'node', id: node.id });
   };
 
-  const onPaneContextMenu = (event: React.MouseEvent) => {
+  const handleEdgeContextMenu = (event: React.MouseEvent, edge: Edge) => {
      event.preventDefault();
-     const bounds = reactFlowWrapper.current?.getBoundingClientRect();
-     if (!bounds) return;
-     
-     const position = project({
-       x: event.clientX - bounds.left,
-       y: event.clientY - bounds.top
-     });
+     setMenu({ x: event.clientX, y: event.clientY, type: 'edge', id: edge.id });
+  };
 
-     const id = `node_${Date.now()}`;
-     addNode({
+  const handleAddNode = (position?: { x: number, y: number }) => {
+    const id = `node_${Date.now()}`;
+    const pos = position || { x: Math.random() * 400, y: Math.random() * 400 };
+    
+    addNode({
       id,
       type: 'llmNode',
-      position,
+      position: pos,
       data: { 
         label: 'New Question', 
         model: 'gpt-4o', 
@@ -94,6 +82,11 @@ function Flow() {
         inputs: 1
       },
     });
+    setMenu(null);
+  };
+
+  const onNodeClick = (_: any, node: Node) => {
+    setSelectedNode(node.id);
   };
 
   return (
@@ -108,7 +101,9 @@ function Flow() {
           nodeTypes={nodeTypes}
           onNodeClick={onNodeClick}
           onPaneClick={onPaneClick}
-          onPaneContextMenu={onPaneContextMenu} // Use Right Click to add node for now
+          onPaneContextMenu={handlePaneContextMenu}
+          onNodeContextMenu={handleNodeContextMenu}
+          onEdgeContextMenu={handleEdgeContextMenu}
           fitView
         >
           <Background color="#aaa" gap={16} />
@@ -116,13 +111,54 @@ function Flow() {
           <MiniMap />
           <Panel position="top-left">
             <button 
-              onClick={handleAddNode}
+              onClick={() => handleAddNode()}
               className="bg-white px-4 py-2 rounded-md shadow-md font-medium text-gray-700 hover:bg-gray-50 flex items-center gap-2"
             >
               <Plus size={16} /> Add Block
             </button>
           </Panel>
         </ReactFlow>
+        
+        {menu && (
+          <ContextMenu 
+            x={menu.x} 
+            y={menu.y} 
+            onClose={() => setMenu(null)}
+            items={
+              menu.type === 'pane' ? [
+                { 
+                  label: 'Add Node', 
+                  onClick: () => {
+                    const bounds = reactFlowWrapper.current?.getBoundingClientRect();
+                    if (bounds) {
+                      const position = project({
+                        x: menu.x - bounds.left,
+                        y: menu.y - bounds.top
+                      });
+                      handleAddNode(position);
+                    }
+                  } 
+                }
+              ] : menu.type === 'node' ? [
+                { 
+                  label: 'Delete Node (Cascade)', 
+                  danger: true,
+                  onClick: () => menu.id && deleteBranch(menu.id)
+                }
+              ] : [
+                { 
+                  label: 'Disconnect', 
+                  onClick: () => menu.id && deleteEdge(menu.id) 
+                },
+                { 
+                  label: 'Delete Branch', 
+                  danger: true,
+                  onClick: () => menu.id && deleteBranch(undefined, menu.id) 
+                }
+              ]
+            }
+          />
+        )}
       </div>
       <Sidebar />
     </div>
diff --git a/frontend/src/components/ContextMenu.tsx b/frontend/src/components/ContextMenu.tsx
new file mode 100644
index 0000000..459641b
--- /dev/null
+++ b/frontend/src/components/ContextMenu.tsx
@@ -0,0 +1,32 @@
+import React from 'react';
+
+interface ContextMenuProps {
+  x: number;
+  y: number;
+  items: { label: string; onClick: () => void; danger?: boolean }[];
+  onClose: () => void;
+}
+
+export const ContextMenu: React.FC<ContextMenuProps> = ({ x, y, items, onClose }) => {
+  return (
+    <div 
+      className="fixed z-50 bg-white border border-gray-200 shadow-lg rounded-md py-1 min-w-[150px]"
+      style={{ top: y, left: x }}
+      onClick={(e) => e.stopPropagation()} // Prevent click through
+    >
+      {items.map((item, idx) => (
+        <button
+          key={idx}
+          className={`w-full text-left px-4 py-2 text-sm hover:bg-gray-100 ${item.danger ? 'text-red-600 hover:bg-red-50' : 'text-gray-700'}`}
+          onClick={() => {
+            item.onClick();
+            onClose();
+          }}
+        >
+          {item.label}
+        </button>
+      ))}
+    </div>
+  );
+};
+
diff --git a/frontend/src/store/flowStore.ts b/frontend/src/store/flowStore.ts
index 6083a36..d2114aa 100644
--- a/frontend/src/store/flowStore.ts
+++ b/frontend/src/store/flowStore.ts
@@ -68,6 +68,11 @@ interface FlowState {
   
   getActiveContext: (nodeId: string) => Message[]; 
   
+  // Actions
+  deleteEdge: (edgeId: string) => void;
+  deleteNode: (nodeId: string) => void;
+  deleteBranch: (startNodeId?: string, startEdgeId?: string) => void;
+
   propagateTraces: () => void;
 }
 
@@ -196,6 +201,79 @@ const useFlowStore = create<FlowState>((set, get) => ({
     return contextMessages;
   },
 
+  deleteEdge: (edgeId: string) => {
+    set({
+      edges: get().edges.filter(e => e.id !== edgeId)
+    });
+    get().propagateTraces();
+  },
+
+  deleteNode: (nodeId: string) => {
+    set({
+      nodes: get().nodes.filter(n => n.id !== nodeId),
+      edges: get().edges.filter(e => e.source !== nodeId && e.target !== nodeId)
+    });
+    get().propagateTraces();
+  },
+
+  deleteBranch: (startNodeId?: string, startEdgeId?: string) => {
+    const { edges, nodes } = get();
+    // We ONLY delete edges, NOT nodes.
+    const edgesToDelete = new Set<string>();
+
+    // Helper to traverse downstream EDGES based on Trace Dependency
+    const traverse = (currentEdge: Edge) => {
+      if (edgesToDelete.has(currentEdge.id)) return;
+      edgesToDelete.add(currentEdge.id);
+      
+      const targetNodeId = currentEdge.target;
+      // Identify the trace ID carried by this edge
+      const traceId = currentEdge.sourceHandle?.replace('trace-', '');
+      if (!traceId) return; 
+      
+      // Look for outgoing edges from the target node that carry the EVOLUTION of this trace.
+      // Our logic generates next trace ID as: `${traceId}_${targetNodeId}`
+      const expectedNextTraceId = `${traceId}_${targetNodeId}`;
+      
+      const outgoing = edges.filter(e => e.source === targetNodeId);
+      outgoing.forEach(nextEdge => {
+        // If the outgoing edge carries the evolved trace, delete it too
+        if (nextEdge.sourceHandle === `trace-${expectedNextTraceId}`) {
+          traverse(nextEdge);
+        }
+      });
+    };
+
+    if (startNodeId) {
+      // If deleting a node, we delete ALL outgoing edges recursively.
+      // Because all traces passing through this node are broken.
+      // But we can't use `traverse` directly because we don't have a single start edge.
+      // We just start traverse on ALL outgoing edges of this node.
+      const initialOutgoing = edges.filter(e => e.source === startNodeId);
+      initialOutgoing.forEach(e => traverse(e));
+      
+      // Also delete incoming to this node
+      const incomingToNode = edges.filter(e => e.target === startNodeId);
+      incomingToNode.forEach(e => edgesToDelete.add(e.id));
+      
+      set({
+        nodes: nodes.filter(n => n.id !== startNodeId),
+        edges: edges.filter(e => !edgesToDelete.has(e.id))
+      });
+    } else if (startEdgeId) {
+      const startEdge = edges.find(e => e.id === startEdgeId);
+      if (startEdge) {
+        traverse(startEdge);
+      }
+      
+      set({
+        edges: edges.filter(e => !edgesToDelete.has(e.id))
+      });
+    }
+
+    get().propagateTraces();
+  },
+
   propagateTraces: () => {
     const { nodes, edges } = get();