1 files changed, 77 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>