go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/projects/nodes/static/_nextjs/src/store/store.ts

/**
 * Copyright (c) 2024 - Present. Will Charczuk. All rights reserved.
 * Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.
 */
import { createWithEqualityFn } from 'zustand/traditional';
import {
  Connection,
  Edge,
  EdgeChange,
  Node,
  NodeChange,
  addEdge,
  applyNodeChanges,
  applyEdgeChanges,
  Position,
  MarkerType,
  getConnectedEdges,
  NodeRemoveChange,
  XYPosition,
} from 'reactflow';
import * as api from '../api/nodes';
import * as graphApi from '../api/graphs';
import * as workflowApi from '../api/workflow';
import { NodeData } from './nodeData';
import { ApiEffect } from './apiEffect';
import { nodeTypes, valueArrayTypes } from '../refdata/nodeTypes';

export type StoreState = {
  isLoading: boolean;
  graph: graphApi.Graph | null;
  graphs: graphApi.Graph[];
  nodes: Node<NodeData>[];
  watchedNodes: api.Node[];
  edges: Edge[];

  onApiEffect: (changes: ApiEffect[]) => Promise<any[]>;

  onNodesChange: (changes: NodeChange[]) => void;
  onEdgesChange: (changes: EdgeChange[]) => void;
  onConnect: (connection: Connection) => Promise<void>;
  onDisconnect: (edge: Edge) => Promise<void>;

  applyNodeDataChangeToStore: (id: string, change: (n0d: NodeData) => NodeData) => void;

  onRefresh: (graphId: string) => Promise<void>;
  onStabilize: () => Promise<workflowApi.WorkflowRunStatus>;
};

const useStateStore = createWithEqualityFn<StoreState>((set, get) => {
  // NOTE(wc): a word on change management!
  // react will only propagate a change to a state object if the _reference_ changes, that is
  // you have to create an entirely new object with the desired changes for the change to
  // propagate through react elements.
  //
  // e.g.: if you have a state reference `foo`:
  //  `foo.bar = "baz"` will not propagate!, you have to do `foo = {...foo, bar: "baz"}`

  // applyNodeDataChangeToStore will apply a change to a given node's data and update the reference for that node's data.
  const applyNodeDataChangeToStore = (id: string, change: (n0d: NodeData) => NodeData) => {
    set({
      nodes: get().nodes.map(n => {
        if (n.id === id) {
          const n0d = Object.assign({}, n.data)
          const n1d = change(n0d)
          n.data = n1d
        }
        return n
      })
    })
  }

  const filterWatchedNodes = (nodes: api.Node[]): api.Node[] => {
    const disalowedValueTypes = new Set([...valueArrayTypes, 'table']);
    const allowedNodeTypes = new Set(['var', 'observer']);
    const flowNodeIsDisallowed = (n: api.Node): boolean => {
      const node_type = n.metadata.node_type;
      const input_type = n.metadata.input_type || '';
      const output_type = n.metadata.output_type || '';
      const nodeWatched = n.metadata.watched === undefined ? false : n.metadata.watched;
      const nodeTypeIsAllowed = allowedNodeTypes.has(node_type);
      const nodeValueTypeIsAllowed = ((node_type === 'observer' && !disalowedValueTypes.has(input_type)) || (node_type === 'var' && !disalowedValueTypes.has(output_type)))
      return nodeTypeIsAllowed && nodeValueTypeIsAllowed && nodeWatched;
    }
    return nodes.filter(flowNodeIsDisallowed).sort((a, b) => a.label.localeCompare(b.label))
  }

  const graphId = () => {
    return get().graph?.id || '';
  }

  // onApiEffect will make changes to both the underlying data store via the API _and_ apply those changes to
  // the client memory mapped store. you should use this directly (instead of the reactflow methods) where you can!
  const onApiEffect = async (changes: ApiEffect[]) => {
    return Promise.all(changes.map(async c => {
      switch (c.type) {
        case 'add-node': {
          const nodeId = await api.postNode(graphId(), c.node);
          const node: api.Node = { ...Object.assign({}, c.node), id: nodeId }
          const { onApiEffect, graph } = get()
          const nodeType = nodeTypes[c.node.metadata.node_type];
          if (nodeType.canSetValue) {
            if (c.value !== undefined) {
              await api.putNodeValue(graphId(), nodeId, c.value)
            }
          }
          set({
            nodes: [...get().nodes, flowNodeFromApiNode(node, { onApiEffect, graph })],
          })
          return nodeId
        }
        case 'duplicate-node': {
          const nodeType = nodeTypes[c.node.metadata.node_type];
          let value = null;
          if (nodeType.canSetValue) {
            value = await api.getNodeValue(graphId(), c.node.id)
          }
          const [nodeId] = await onApiEffect([{
            type: 'add-node',
            node: {
              ...c.node,
              label: c.node.label + '-copy',
              metadata: {
                ...c.node.metadata,
                position_x: c.node.metadata.position_x + 350,
                position_y: c.node.metadata.position_y,
              }
            },
            value: value
          }])
          return nodeId
        }
        case 'set-stale':
          return api.postNodeStale(graphId(), c.node_id)
        case 'unlink-nodes':
          await api.deleteEdge(graphId(), c)
          const edgeId = formatEdgeId(c)
          set({
            edges: get().edges.filter((e) => e.id !== edgeId)
          })
          return
        case 'link-nodes':
          await api.postEdge(graphId(), c)
          set({
            edges: addEdge(flowEdgeFromApiEdge(c), get().edges),
          });
          return
        case 'remove-node':
          await api.deleteNode(graphId(), c.node_id)
          const affectedNode = get().nodes.find(fn => fn.id === c.node_id)
          if (!affectedNode) {
            throw new Error(`remove-node; affected node not found: ${c.node_id}`)
          }
          const connectedEdges = getConnectedEdges([affectedNode], get().edges);
          set({
            nodes: get().nodes.filter(n => n.id !== c.node_id),
            edges: get().edges.filter(e => !connectedEdges.includes(e)),
            watchedNodes: [...get().watchedNodes.filter(n => n.id !== c.node_id)],
          })
          return
        case 'update-label':
          await api.patchNode(graphId(), c.node_id, { label: c.label });
          applyNodeDataChangeToStore(c.node_id, (n0d) => {
            n0d.node.label = c.label
            return n0d
          })
          return
        case 'update-expression':
          await api.patchNode(graphId(), c.node_id, { expression: c.expression });
          applyNodeDataChangeToStore(c.node_id, (n0d) => {
            n0d.node.metadata.expression = c.expression
            return n0d
          })
          return
        case 'update-value': {
          await api.putNodeValue(graphId(), c.node_id, c.value);
          applyNodeDataChangeToStore(c.node_id, (n0d) => {
            n0d.refreshedAt = new Date()
            return n0d
          })
          return
        }
        case 'update-position':
          await api.patchNode(graphId(), c.node_id, { position_x: c.position.x, position_y: c.position.y })
          applyNodeDataChangeToStore(c.node_id, (n0d) => {
            n0d.node.metadata.position_x = c.position.x;
            n0d.node.metadata.position_y = c.position.y;
            return n0d
          })
          return
        case 'update-size':
          await api.patchNode(graphId(), c.node_id, { display_height: c.height, display_width: c.width })
          applyNodeDataChangeToStore(c.node_id, (n0d) => {
            n0d.node.metadata.height = c.height;
            n0d.node.metadata.width = c.width;
            return n0d
          })
          return
        case 'update-collapsed-all':
          await api.patchNodes(graphId(), { collapsed: c.collapsed });
          set({
            nodes: get().nodes.map(n => {
              const n0d = Object.assign({}, n.data)
              n0d.node.metadata.collapsed = c.collapsed
              n.data = n0d
              return n
            })
          })
          return
        case 'update-collapsed':
          await api.patchNode(graphId(), c.node_id, { collapsed: c.collapsed });
          applyNodeDataChangeToStore(c.node_id, (n0d) => {
            n0d.node.metadata.collapsed = c.collapsed
            return n0d
          })
          return
        case 'update-watched':
          await api.patchNode(graphId(), c.node_id, { watched: c.watched });
          applyNodeDataChangeToStore(c.node_id, (n0d) => {
            n0d.node.metadata.watched = c.watched
            return n0d
          })
          if (c.watched) {
            const newWatched = get().nodes.find(n => n.id === c.node_id)?.data.node
            if (newWatched === undefined) {
              return
            }
            set({
              watchedNodes: [...get().watchedNodes, newWatched],
            })
          } else {
            set({
              watchedNodes: [...get().watchedNodes.filter(n => n.id !== c.node_id)],
            })
          }

          return
        case 'update-watched-collapsed':
          await api.patchNode(graphId(), c.node_id, { watched_collapsed: c.watched_collapsed });
          applyNodeDataChangeToStore(c.node_id, (n0d) => {
            n0d.node.metadata.watched_collapsed = c.watched_collapsed
            return n0d
          })
          return
        case 'update-graph-viewport': {
          await graphApi.patchGraph(c.graph_id, {
            "viewport_x": c.viewport.x,
            "viewport_y": c.viewport.y,
            "viewport_zoom": c.viewport.zoom,
          });
          const { graph, graphs } = get();
          if (graph !== null && graph.id === c.graph_id) {
            set({
              graph: {
                ...graph,
                metadata: {
                  ...graph?.metadata,
                  viewport_x: c.viewport.x,
                  viewport_y: c.viewport.y,
                  viewport_zoom: c.viewport.zoom,
                },
              },
              graphs: [...graphs.map(g => {
                if (g.id === c.graph_id) {
                  const newObj = Object.assign({}, g)
                  newObj.metadata.viewport_x = c.viewport.x
                  newObj.metadata.viewport_y = c.viewport.y
                  newObj.metadata.viewport_zoom = c.viewport.zoom
                  return newObj
                }
                return g
              })],
            })
          }
          return
        }
        case 'update-graph-label': {
          await graphApi.patchGraph(c.graph_id, {
            "label": c.label,
          });
          const { graph, graphs } = get();
          if (graph !== null && graph.id === c.graph_id) {
            set({
              graph: {
                ...graph,
                label: c.label,
              },
              graphs: [...graphs.map(g => {
                if (g.id === c.graph_id) {
                  const newObj = Object.assign({}, g)
                  newObj.label = c.label
                  return newObj
                }
                return g
              })],
            })
          }
          return
        }
        case 'delete-graph': {
          await graphApi.deleteGraph(c.graph_id);
          const [graph, graphs] = await Promise.all([
            graphApi.getGraphActive(),
            graphApi.getGraphs(),
          ]);
          set({
            graph: graph,
            graphs: graphs,
          })
        }
      }
    }))
  }

  const onNodesChange = async (changes: NodeChange[]) => {
    set({
      nodes: applyNodeChanges(changes, get().nodes),
      edges: changes.filter(c => c.type == 'remove').reduce((acc, change) => {
        const removedNode = change as NodeRemoveChange
        const affectedNode = get().nodes.find(fn => fn.id === removedNode.id)
        if (affectedNode) {
          const connectedEdges = getConnectedEdges([affectedNode], get().edges);
          return acc.filter((edge) => !connectedEdges.includes(edge));
        }
        return acc;
      }, get().edges)
    });
  }

  const onEdgesChange = (changes: EdgeChange[]) => {
    set({
      edges: applyEdgeChanges(changes, get().edges),
    });
  }

  const onConnect = async (connection: Connection) => {
    const targetNode = get().nodes.find(n => n.id === connection.target);
    if (!targetNode) {
      return
    }
    const sourceNode = get().nodes.find(n => n.id === connection.source);
    if (!sourceNode) {
      return
    }

    const targetNodeType = nodeTypes[targetNode.data.node.metadata.node_type];
    const targetNodeHandle = targetNodeType.inputs.find(i => i.id === connection.targetHandle)
    if (!targetNodeHandle) {
      return
    }

    if (targetNodeHandle.multiplicity === 'single') {
      // we have to _unlink_ any existing inputs here first!
      // find the existing edge.
      const existingEdge = get().edges.find(e => e.target === connection.target && e.targetHandle === connection.targetHandle)
      if (existingEdge !== undefined) {
        await onApiEffect([{
          type: 'unlink-nodes',
          parent_id: existingEdge.source,
          child_id: existingEdge.target,
          child_input_name: existingEdge.targetHandle === undefined ? null : existingEdge.targetHandle,
        }])
      }
    }
    await onApiEffect([{
      type: 'link-nodes',
      parent_id: connection.source || '',
      child_id: connection.target || '',
      child_input_name: connection.targetHandle === undefined ? null : connection.targetHandle,
    }])
  }

  const onDisconnect = async (edge: Edge) => {
    await api.deleteEdge(graphId(), { child_id: edge.target || '', parent_id: edge.source || '', child_input_name: edge.targetHandle === undefined ? null : edge.targetHandle })
    set({
      edges: get().edges.filter((e) => e.id !== edge.id)
    })
  }

  const onRefresh = async (graphId: string) => {
    set({ isLoading: true })
    const graph = await graphApi.getGraph(graphId);
    if (graph === null) {
      set({
        isLoading: false,
        graph: graph,
      })
      return
    }
    const [graphs, apiNodes, apiEdges] = await Promise.all([
      graphApi.getGraphs(),
      api.getNodes(graph.id),
      api.getEdges(graph.id),
    ])
    const flowNodes = apiNodes.map(apiNode => flowNodeFromApiNode(apiNode, { onApiEffect, graph }));
    const watchedNodes = filterWatchedNodes(apiNodes);
    set({
      isLoading: false,
      graph: graph,
      graphs: graphs,
      nodes: flowNodes,
      watchedNodes: watchedNodes,
      edges: flowEdgesFromApiEdges(apiEdges),
    })
  }

  const onStabilize = async () => {
    const workflowRun = await workflowApi.stabilize(graphId(), true);
    const workflowRunResult = await workflowApi.pollUntilFinished(workflowRun.workflow_id, workflowRun.run_id)

    // get the graph again to fetch the upated stabilization number
    const graph = await graphApi.getGraph(graphId());
    set({
      nodes: [...get().nodes.map(n => {
        const dataCopy = Object.assign({}, n.data)
        dataCopy.refreshedAt = new Date()
        dataCopy.stabilizationNum = graph.stabilization_num
        n.data = dataCopy
        return n
      })],
      graph,
    })
    if (workflowRunResult.failure !== null) {
      throw new Error(workflowRunResult.failure?.cause.message)
    }
    return workflowRunResult
  }

  return {
    isLoading: true,
    graph: graphApi.emptyGraph,
    graphs: [],
    nodes: [],
    edges: [],
    watchedNodes: [],
    onApiEffect,
    onNodesChange,
    onEdgesChange,
    onConnect,
    onDisconnect,
    onRefresh,
    onStabilize,
    applyNodeDataChangeToStore,
  }
}, Object.is);

const nodeDefaults = {
  sourcePosition: Position.Right,
  targetPosition: Position.Left,
}

export function apiNodeFromFlowNodeData(data: NodeData, pos: XYPosition): api.Node {
  return {
    ...data.node,
    metadata: {
      ...data.node?.metadata,
      position_x: pos.x,
      position_y: pos.y,
    }
  }
}

export const flowNodeFromApiNode = (node: api.Node, state: Pick<StoreState, 'onApiEffect' | 'graph'>): Node<NodeData> => {
  const resizable = node.metadata.node_type === "observer" && node.metadata.input_type === "svg"
  return {
    ...nodeDefaults,
    id: node.id,
    position: {
      x: node.metadata.position_x,
      y: node.metadata.position_y,
    },
    parentNode: node.metadata.parent_node_id,
    type: 'nodeCard',
    resizing: resizable,
    dragHandle: '.bp5-section-header',
    style: {
      ...(resizable ? { height: node.metadata.height || 0, width: node.metadata.width } : {})
    },
    data: {
      node: node,
      refreshedAt: new Date(),
      stabilizationNum: state.graph?.stabilization_num,
      getValue: async () => {
        return api.getNodeValue(state.graph?.id || '', node.id)
      },
      onCollapse: async (collapsed: boolean) => {
        await state.onApiEffect([{
          type: 'update-collapsed',
          node_id: node.id,
          collapsed: collapsed,
        }])
      },
    },
  }
}

export const defaultEdgeOptions = {
  type: 'smoothstep',
  animated: true,
  markerEnd: {
    type: MarkerType.ArrowClosed,
  },
}

export const flowEdgesFromApiEdges = (edges: api.Edge[]): Edge[] => {
  const output: Edge[] = []
  for (const e of edges) {
    output.push(flowEdgeFromApiEdge(e))
  }
  return output
}

export const formatEdgeId = (edge: api.Edge): string => {
  return `${edge.parent_id}->${edge.child_id}.${edge.child_input_name}`
}

export const flowEdgeFromApiEdge = (edge: api.Edge): Edge => {
  return {
    id: formatEdgeId(edge),
    source: edge.parent_id,
    target: edge.child_id,
    targetHandle: edge.child_input_name,
    className: `edge`,
    ...defaultEdgeOptions,
  }
}

export default useStateStore;