@kiberon-labs/behave-graph-flow 1.0.0 → 3.0.0

package/src/components/Flow.tsx

@@ -1,37 +1,121 @@
-import type { GraphJSON, IRegistry } from '@kiberon-labs/behave-graph';
-import React from 'react';
-import { Background, BackgroundVariant, ReactFlow } from 'reactflow';
+import React, {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState
+} from 'react';
+import {
+  Background,
+  BackgroundVariant,
+  type Edge as ReactFlowEdge,
+  MiniMap,
+  ReactFlow,
+  type Node as ReactFlowNode,
+  type ReactFlowInstance
+} from 'reactflow';
 
 import { useBehaveGraphFlow } from '../hooks/useBehaveGraphFlow.js';
 import { useFlowHandlers } from '../hooks/useFlowHandlers.js';
-import { useGraphRunner } from '../hooks/useGraphRunner.js';
-import { useNodeSpecJson } from '../hooks/useNodeSpecJson.js';
-import CustomControls from './Controls.js';
-import { type Examples } from './modals/LoadModal.js';
-import { NodePicker } from './NodePicker.js';
-
-type FlowProps = {
-  initialGraph: GraphJSON;
-  registry: IRegistry;
-  examples: Examples;
+import { useWasdPan } from '../hooks/useWasdPan.js';
+import { useGraph } from '@/system/provider.js';
+import { useStore } from 'zustand';
+import {
+  NodeContextMenu,
+  type INodeContextMenuProps
+} from './contextMenus/node.js';
+import {
+  EdgeContextMenu,
+  type IEdgeContextMenuProps
+} from './contextMenus/edge.js';
+import {
+  SelectionContextMenu,
+  type ISelectionContextMenuProps
+} from './contextMenus/selection.js';
+import { registerDefaults } from '@/generators/registerDefaults';
+import { FloatingToolbar } from './FloatingToolbar';
+import { layerId } from '@/annotations';
+
+const REACTFLOW_NODE = 'react-flow__node';
+
+const getAnnotatedLayerId = (node: ReactFlowNode): string | undefined => {
+  if (!node.data || typeof node.data !== 'object') return undefined;
+  const record = node.data as { annotations?: Record<string, unknown> };
+  const value = record.annotations?.[layerId];
+  return typeof value === 'string' ? value : undefined;
 };
 
-export const Flow: React.FC<FlowProps> = ({
-  initialGraph: graph,
-  registry,
-  examples
-}) => {
-  const specJson = useNodeSpecJson(registry);
+const isNodeVisibleInLayers = (
+  node: ReactFlowNode,
+  params: {
+    layers: Record<string, { visible: boolean }>;
+    nodeLayers: Record<string, string>;
+    defaultLayerId: string;
+  }
+): boolean => {
+  const mappedLayerId =
+    params.nodeLayers[node.id] ??
+    getAnnotatedLayerId(node) ??
+    params.defaultLayerId;
+  const layer =
+    params.layers[mappedLayerId] ?? params.layers[params.defaultLayerId];
+  return layer?.visible !== false;
+};
 
-  const {
-    nodes,
-    edges,
-    onNodesChange,
-    onEdgesChange,
-    graphJson,
-    setGraphJson,
-    nodeTypes
-  } = useBehaveGraphFlow({
+export const Flow: React.FC = () => {
+  const session = useGraph();
+  const editor = session.editor;
+  const getGraphJson = useStore(session.flowStore, (x) => x.getGraph);
+  const specJson = useStore(editor.specStore, (x) => x.specs);
+  const showGrid = useStore(editor.systemSettings, (x) => x.showGrid);
+  const showMinimap = useStore(editor.systemSettings, (x) => x.showMinimap);
+  const snapToGrid = useStore(editor.systemSettings, (x) => x.snapGrid);
+  const gridSize = useStore(editor.systemSettings, (x) => x.gridSize);
+  const edgeTypes = useStore(session.flowStore, (x) => x.edgeTypes);
+  const nodeTypes = useStore(session.flowStore, (x) => x.nodeTypes);
+  const layers = useStore(session.layerStore, (x) => x.layers);
+  const nodeLayers = useStore(session.layerStore, (x) => x.nodeLayers);
+  const defaultLayerId = useStore(session.layerStore, (x) => x.defaultLayerId);
+
+  const ref = useRef<HTMLDivElement>(null);
+  const setRef = useStore(session.refStore, (x) => x.setRef);
+
+  // When this canvas is interacted with, mark its graph as focused so the shared
+  // side-panels (Variables, Traces, ...) bind to it. Lets two graphs be open at
+  // once without a singular "active graph" , focus simply follows interaction.
+  const focusThisGraph = useCallback(() => {
+    const active = editor.activeGraph.getState();
+    if (active.activeGraphId !== session.id) {
+      active.setActiveGraph(session.id);
+    }
+  }, [editor, session.id]);
+
+  const getReactFlowInstance = useCallback(
+    () => session.refStore.getState().getRef('reactflow'),
+    [session.refStore]
+  );
+
+  useWasdPan({ getReactFlowInstance });
+
+  // Set reactflow ref
+  const setReactflowRef = React.useCallback(
+    (reactFlowInstance: ReactFlowInstance) => {
+      if (reactFlowInstance) {
+        setRef('reactflow', reactFlowInstance);
+      }
+    },
+    [setRef]
+  );
+
+  const graph = useMemo(() => getGraphJson(), []);
+
+  // Editor-level built-in content. Idempotent per editor, so opening multiple
+  // graph tabs (each a Flow canvas) registers it once rather than per mount.
+  useEffect(() => {
+    registerDefaults(editor);
+  }, [editor]);
+
+  const { nodes, edges, onNodesChange, onEdgesChange } = useBehaveGraphFlow({
     initialGraphJson: graph,
     specJson
   });
@@ -41,11 +125,7 @@ export const Flow: React.FC<FlowProps> = ({
     handleStartConnect,
     handleStopConnect,
     handlePaneClick,
-    handlePaneContextMenu,
-    nodePickerVisibility,
-    handleAddNode,
-    closeNodePicker,
-    nodePickFilters
+    handlePaneContextMenu
   } = useFlowHandlers({
     nodes,
     onEdgesChange,
@@ -53,49 +133,181 @@ export const Flow: React.FC<FlowProps> = ({
     specJSON: specJson
   });
 
-  const { togglePlay, playing } = useGraphRunner({
-    graphJson,
-    registry
-  });
+  const [menu, setMenu] = useState<INodeContextMenuProps | null>(null);
+  const [edgeMenu, setEdgeMenu] = useState<IEdgeContextMenuProps | null>(null);
+  const [selectionMenu, setSelectionMenu] =
+    useState<ISelectionContextMenuProps | null>(null);
+  const onNodeContextMenu = useCallback(
+    (event: React.MouseEvent<Element, MouseEvent>) => {
+      // Prevent native context menu from showing
+      event.preventDefault();
+      setEdgeMenu(null);
+      setSelectionMenu(null);
+      const offset = ref.current?.getBoundingClientRect();
+      //Keep ascending till we find the .react-flow__node
+      let target = event.target as HTMLElement | null;
+      let nodeID = null;
+      while (target && !target.classList.contains(REACTFLOW_NODE)) {
+        target = target.parentElement;
+      }
+      if (target) {
+        nodeID = target.getAttribute('data-id');
+      }
+      setMenu({
+        //We should be safe here as reactflow only triggers this on nodes
+        nodeID: nodeID!,
+        top: event.clientY - (offset?.top ?? 0),
+        left: event.clientX - (offset?.left ?? 0)
+      });
+    },
+    [setMenu, setEdgeMenu]
+  );
+
+  const onEdgeContextMenu = useCallback(
+    (event: React.MouseEvent, edge: any) => {
+      event.preventDefault();
+      event.stopPropagation();
+      setMenu(null);
+      setSelectionMenu(null);
+      const offset = ref.current?.getBoundingClientRect();
+      setEdgeMenu({
+        edgeID: edge.id,
+        sourceID: edge.source,
+        targetID: edge.target,
+        top: event.clientY - (offset?.top ?? 0),
+        left: event.clientX - (offset?.left ?? 0)
+      });
+    },
+    [setEdgeMenu, setMenu]
+  );
+
+  const onSelectionContextMenu = useCallback(
+    (event: React.MouseEvent) => {
+      event.preventDefault();
+      event.stopPropagation();
+      setMenu(null);
+      setEdgeMenu(null);
+      const offset = ref.current?.getBoundingClientRect();
+      setSelectionMenu({
+        top: event.clientY - (offset?.top ?? 0),
+        left: event.clientX - (offset?.left ?? 0)
+      });
+    },
+    [setMenu, setEdgeMenu, setSelectionMenu]
+  );
+
+  const snapGrid = useMemo(() => {
+    return [gridSize, gridSize] as [number, number];
+  }, [gridSize]);
+
+  const onPaneClick = useCallback(() => {
+    setMenu(null);
+    setEdgeMenu(null);
+    setSelectionMenu(null);
+    handlePaneClick();
+  }, [setMenu, setEdgeMenu, setSelectionMenu, handlePaneClick]);
+
+  const renderedNodes = useMemo(() => {
+    return nodes.map((node) => {
+      const mappedLayerId =
+        nodeLayers[node.id] ?? getAnnotatedLayerId(node) ?? defaultLayerId;
+      const layer = layers[mappedLayerId] ?? layers[defaultLayerId];
+      if (!layer) return node;
+
+      return {
+        ...node,
+        hidden: !layer.visible,
+        style: {
+          ...(node.style ?? {}),
+          opacity: layer.opacity
+        }
+      };
+    });
+  }, [defaultLayerId, layers, nodeLayers, nodes]);
+
+  const renderedEdges = useMemo(() => {
+    const nodesById = new Map<string, ReactFlowNode>();
+    nodes.forEach((node) => {
+      nodesById.set(node.id, node);
+    });
+
+    return edges.map((edge: ReactFlowEdge) => {
+      const sourceNode = nodesById.get(edge.source);
+      const targetNode = nodesById.get(edge.target);
+
+      // If either endpoint is missing just leave the edge as-is
+      if (!sourceNode || !targetNode) return edge;
+
+      const sourceVisible = isNodeVisibleInLayers(sourceNode, {
+        layers,
+        nodeLayers,
+        defaultLayerId
+      });
+      const targetVisible = isNodeVisibleInLayers(targetNode, {
+        layers,
+        nodeLayers,
+        defaultLayerId
+      });
+
+      const shouldHide = !sourceVisible || !targetVisible;
+
+      // Return as hidden when an endpoint node is in a hidden layer.
+      // Using `hidden` rather than filtering so ReactFlow still tracks the
+      // connection and handles continue to render as connected.
+      if (shouldHide === edge.hidden) return edge;
+      return { ...edge, hidden: shouldHide };
+    });
+  }, [defaultLayerId, edges, layers, nodeLayers, nodes]);
 
   return (
-    <ReactFlow
-      nodeTypes={nodeTypes}
-      nodes={nodes}
-      edges={edges}
-      onNodesChange={onNodesChange}
-      onEdgesChange={onEdgesChange}
-      onConnect={onConnect}
-      // @ts-ignore
-      onConnectStart={handleStartConnect}
-      // @ts-ignore
-      onConnectEnd={handleStopConnect}
-      fitView
-      fitViewOptions={{ maxZoom: 1 }}
-      onPaneClick={handlePaneClick}
-      onPaneContextMenu={handlePaneContextMenu}
-    >
-      <CustomControls
-        playing={playing}
-        togglePlay={togglePlay}
-        setBehaviorGraph={setGraphJson}
-        examples={examples}
-        specJson={specJson}
-      />
-      <Background
-        variant={BackgroundVariant.Lines}
-        color="#2a2b2d"
-        style={{ backgroundColor: '#1E1F22' }}
-      />
-      {nodePickerVisibility && (
-        <NodePicker
-          position={nodePickerVisibility}
-          filters={nodePickFilters}
-          onPickNode={handleAddNode}
-          onClose={closeNodePicker}
-          specJSON={specJson}
-        />
-      )}
-    </ReactFlow>
+    <>
+      <ReactFlow
+        id={session.id}
+        style={{ flex: 1 }}
+        ref={ref}
+        onPointerDownCapture={focusThisGraph}
+        onInit={setReactflowRef}
+        nodeTypes={nodeTypes}
+        edgeTypes={edgeTypes}
+        elevateEdgesOnSelect={true}
+        nodes={renderedNodes}
+        edges={renderedEdges}
+        onSelectionContextMenu={onSelectionContextMenu}
+        onNodeContextMenu={onNodeContextMenu}
+        // @ts-ignore
+        onEdgeContextMenu={onEdgeContextMenu}
+        maxZoom={Infinity}
+        minZoom={-Infinity}
+        onNodesChange={onNodesChange}
+        onEdgesChange={onEdgesChange}
+        onConnect={onConnect}
+        // @ts-ignore
+        onConnectStart={handleStartConnect}
+        // @ts-ignore
+        onConnectEnd={handleStopConnect}
+        snapToGrid={snapToGrid}
+        snapGrid={snapGrid}
+        fitView
+        //TODO. Reconsier this prop for performance
+        onlyRenderVisibleElements={true}
+        onPaneClick={onPaneClick}
+        onPaneContextMenu={handlePaneContextMenu}
+        proOptions={{ hideAttribution: true }}
+      >
+        {showGrid && (
+          <Background
+            variant={BackgroundVariant.Lines}
+            gap={90}
+            color="#373737"
+            style={{ backgroundColor: 'var(--colors-bgCanvas)' }}
+          />
+        )}
+        {showMinimap && <MiniMap />}
+        {menu && <NodeContextMenu {...menu} />}
+        {edgeMenu && <EdgeContextMenu {...edgeMenu} />}
+        {selectionMenu && <SelectionContextMenu {...selectionMenu} />}
+        <FloatingToolbar />
+      </ReactFlow>
+    </>
   );
 };

package/src/components/contextMenus/DynamicContextMenu.tsx

@@ -0,0 +1,85 @@
+import type { CSSProperties } from 'react';
+import { useStore } from 'zustand';
+import { VscodeContextMenu } from '@vscode-elements/react-elements';
+import { useGraph } from '@/system';
+import type { CommandContext } from '@/store/commands';
+import type { ContextMenuTarget } from '@/store/contextMenu';
+
+type MenuData =
+  | { separator: true }
+  | { label: string; value: string; keybinding?: string };
+
+/**
+ * Renders a context menu for a target from the editor's contextMenu registry,
+ * dispatching selections through the command registry. The three concrete menus
+ * (node/edge/selection) are thin wrappers that supply the target + context.
+ */
+export const DynamicContextMenu = ({
+  target,
+  context,
+  style
+}: {
+  target: ContextMenuTarget;
+  context: Omit<CommandContext, 'editor' | 'session'>;
+  style?: CSSProperties;
+}) => {
+  const session = useGraph();
+  const editor = session.editor;
+
+  // Subscribe so the menu re-renders when items/commands/keymap change.
+  useStore(editor.contextMenuStore, (s) => s.items);
+  useStore(editor.commandStore, (s) => s.commands);
+  useStore(editor.hotKeyStore, (s) => s.keymap);
+
+  const hotKeys = editor.hotKeyStore.getState();
+
+  const ctx: CommandContext = { editor, session, ...context };
+
+  const items = editor.contextMenuStore
+    .getState()
+    .getItems(target)
+    .filter((i) => !i.when || i.when(ctx));
+
+  const data: MenuData[] = [];
+  let lastGroup: unknown;
+  for (const item of items) {
+    if (data.length > 0 && item.group !== lastGroup) {
+      data.push({ separator: true });
+    }
+    lastGroup = item.group;
+    // Auto-detect the shortcut from the command's live keymap binding; fall
+    // back to a static hint only when the command has no bound key.
+    const derived = item.commandId
+      ? hotKeys.getCommandKeybinding(item.commandId)
+      : undefined;
+    data.push({
+      label: typeof item.label === 'function' ? item.label(ctx) : item.label,
+      value: item.id,
+      keybinding: derived ?? item.keybinding
+    });
+  }
+
+  if (data.length === 0) return null;
+
+  const onSelect = (e: any) => {
+    const id = e?.detail?.value as string | undefined;
+    const item = items.find((i) => i.id === id);
+    if (!item) return;
+    if (item.onSelect) {
+      item.onSelect(ctx);
+      return;
+    }
+    if (item.commandId) {
+      void editor.commandStore.getState().run(item.commandId, ctx);
+    }
+  };
+
+  return (
+    <VscodeContextMenu
+      show
+      onVscContextMenuSelect={onSelect}
+      style={{ zIndex: 2000, position: 'absolute', ...style }}
+      data={data}
+    />
+  );
+};