CRDT for node positions is working.

README.md

@@ -1,11 +1,11 @@
 # LynxKite 2024
 
-This is an experimental rewrite of [LynxKite](https://github.com/lynxkite/lynxkite).
-It is not compatible with the original LynxKite. The primary goals of this rewrite are:
-- Target GPU clusters instead of Hadoop clusters.
-  We use Python instead of Scala, RAPIDS instead of Apache Spark.
-- More extensible backend. Make it easy to add new LynxKite boxes.
-  Make it easy to use our frontend for other purposes, configuring and executing other pipelines.
+This is an experimental rewrite of [LynxKite](https://github.com/lynxkite/lynxkite). It is not compatible with the
+original LynxKite. The primary goals of this rewrite are:
+
+- Target GPU clusters instead of Hadoop clusters. We use Python instead of Scala, RAPIDS instead of Apache Spark.
+- More extensible backend. Make it easy to add new LynxKite boxes. Make it easy to use our frontend for other purposes,
+  configuring and executing other pipelines.
 
 Current status: **PROTOTYPE**
 
@@ -15,6 +15,7 @@ To run the backend:
 
 ```bash
 pip install -r requirements.txt
+PYTHONPATH=. pydantic2ts --module server.workspace --output ./web/src/apiTypes.ts
 uvicorn server.main:app --reload
 ```
 

requirements.txt

@@ -4,6 +4,7 @@ networkx
 numpy
 orjson
 pandas
+pydantic-to-typescript
 scipy
 uvicorn[standard]
 pycrdt

web/package-lock.json

@@ -11,6 +11,8 @@
         "@iconify-json/tabler": "^1.2.10",
         "@svgr/core": "^8.1.0",
         "@svgr/plugin-jsx": "^8.1.0",
+        "@syncedstore/core": "^0.6.0",
+        "@syncedstore/react": "^0.6.0",
         "@types/node": "^22.10.1",
         "@xyflow/react": "^12.3.5",
         "daisyui": "^4.12.20",
@@ -713,6 +715,24 @@
         "node": ">=14"
       }
     },
+    "node_modules/@reactivedata/react": {
+      "version": "0.2.2",
+      "resolved": "https://registry.npmjs.org/@reactivedata/react/-/react-0.2.2.tgz",
+      "integrity": "sha512-fJ8qoHRbicQnVcnwfHa9FO73mbHFfl590xpuGM4Mo1P2ClaDATfl9oUrrySE+tbA//M9cn8De8HTwjoNqt91sw==",
+      "license": "MIT",
+      "dependencies": {
+        "@reactivedata/reactive": "^0.2.2"
+      },
+      "peerDependencies": {
+        "react": "^16.8.0 || ^17 || ^18"
+      }
+    },
+    "node_modules/@reactivedata/reactive": {
+      "version": "0.2.2",
+      "resolved": "https://registry.npmjs.org/@reactivedata/reactive/-/reactive-0.2.2.tgz",
+      "integrity": "sha512-KnINM/Sng25QAv6sHkJO9q/XyslLegCF5jTsTSVu+AouY3uZDVf4Am99xNCqsfqFZFvnTBBDvCsHNdvTVGvPEA==",
+      "license": "MIT"
+    },
     "node_modules/@rollup/rollup-darwin-arm64": {
       "version": "4.28.1",
       "resolved": "https://registry.npmjs.org/@rollup/rollup-darwin-arm64/-/rollup-darwin-arm64-4.28.1.tgz",
@@ -1013,6 +1033,40 @@
         "@swc/counter": "^0.1.3"
       }
     },
+    "node_modules/@syncedstore/core": {
+      "version": "0.6.0",
+      "resolved": "https://registry.npmjs.org/@syncedstore/core/-/core-0.6.0.tgz",
+      "integrity": "sha512-6TtjEoYJsceYi8u1oRecXwbbLmjHaU0S7HvVfOaEdDfphZLGm/faVuA2fpazqc28F0yIFGvYzvPEBUJn9vqRNw==",
+      "license": "MIT",
+      "dependencies": {
+        "@reactivedata/reactive": "^0.2.0",
+        "@syncedstore/yjs-reactive-bindings": "^0.6.0"
+      },
+      "peerDependencies": {
+        "yjs": "^13.5.13"
+      }
+    },
+    "node_modules/@syncedstore/react": {
+      "version": "0.6.0",
+      "resolved": "https://registry.npmjs.org/@syncedstore/react/-/react-0.6.0.tgz",
+      "integrity": "sha512-pc8ycnBuH2wp7Td5nGzP9Dn2mEW9Yg1qF0yGdjJ5UdgLEwfRkq+8/hdQl+alcoSribv9St/Bi5hJckX8GVsAbQ==",
+      "license": "MIT",
+      "dependencies": {
+        "@reactivedata/react": "^0.2.1"
+      },
+      "peerDependencies": {
+        "@syncedstore/core": "*"
+      }
+    },
+    "node_modules/@syncedstore/yjs-reactive-bindings": {
+      "version": "0.6.0",
+      "resolved": "https://registry.npmjs.org/@syncedstore/yjs-reactive-bindings/-/yjs-reactive-bindings-0.6.0.tgz",
+      "integrity": "sha512-VF78h0J4iOt79YU9d6j5E6bFKu7WXYuiI2ue9ZnA+T4SNVn8viRvg0AHm3NqHzudZZUgYT3dpnbv1/ZmU7yPZQ==",
+      "license": "MIT",
+      "peerDependencies": {
+        "yjs": "^13.5.13"
+      }
+    },
     "node_modules/@types/cookie": {
       "version": "0.6.0",
       "resolved": "https://registry.npmjs.org/@types/cookie/-/cookie-0.6.0.tgz",

web/package.json

@@ -13,6 +13,8 @@
     "@iconify-json/tabler": "^1.2.10",
     "@svgr/core": "^8.1.0",
     "@svgr/plugin-jsx": "^8.1.0",
+    "@syncedstore/core": "^0.6.0",
+    "@syncedstore/react": "^0.6.0",
     "@types/node": "^22.10.1",
     "@xyflow/react": "^12.3.5",
     "daisyui": "^4.12.20",

web/src/workspace/Workspace.tsx

@@ -1,7 +1,7 @@
 // The LynxKite workspace editor.
 import { useParams } from "react-router";
 import useSWR from 'swr';
-import { useMemo, useCallback } from "react";
+import { useMemo, useCallback, useState } from "react";
 import favicon from '../assets/favicon.ico';
 import {
   ReactFlow,
@@ -9,7 +9,10 @@ import {
   MiniMap,
   MarkerType,
   useReactFlow,
+  useUpdateNodeInternals,
   ReactFlowProvider,
+  applyEdgeChanges,
+  applyNodeChanges,
   type XYPosition,
   type Node,
   type Edge,
@@ -22,8 +25,8 @@ import ArrowBack from '~icons/tabler/arrow-back.jsx';
 import Backspace from '~icons/tabler/backspace.jsx';
 // @ts-ignore
 import Atom from '~icons/tabler/atom.jsx';
-// import { syncedStore, getYjsDoc } from "@syncedstore/core";
-// import { useSyncedStore } from "@syncedstore/react";
+import { syncedStore, getYjsDoc } from "@syncedstore/core";
+import { useSyncedStore } from "@syncedstore/react";
 import { WebsocketProvider } from "y-websocket";
 import NodeWithParams from './nodes/NodeWithParams';
 // import NodeWithVisualization from './NodeWithVisualization';
@@ -37,10 +40,6 @@ import { LynxKiteState } from './LynxKiteState';
 import '@xyflow/react/dist/style.css';
 import { Workspace } from "../apiTypes.ts";
 
-
-import { useShallow } from 'zustand/react/shallow';
-import { useStore, selector, doc } from './store';
-
 export default function (props: any) {
   return (
     <ReactFlowProvider>
@@ -51,17 +50,55 @@ export default function (props: any) {
 
 
 function LynxKiteFlow() {
+  const updateNodeInternals = useUpdateNodeInternals();
   const { screenToFlowPosition } = useReactFlow();
-  const store = useStore(
-    useShallow(selector),
-  );
+  const [nodes, setNodes] = useState([] as Node[]);
+  const [edges, setEdges] = useState([] as Edge[]);
   const { path } = useParams();
 
-  // const sstore = syncedStore({ workspace: {} });
-  // const doc = getYjsDoc(sstore);
-  const wsProvider = new WebsocketProvider("ws://localhost:8000/ws/crdt", path!, doc);
+  const sstore = syncedStore({ workspace: {} });
+  const doc = getYjsDoc(sstore);
+  const wsProvider = useMemo(() => new WebsocketProvider("ws://localhost:8000/ws/crdt", path!, doc), [path]);
   wsProvider; // Just to disable the lint warning. The life cycle of this object is a mystery.
-  // const state: { workspace: Workspace } = useSyncedStore(sstore);
+  const state: { workspace: Workspace } = useSyncedStore(sstore);
+  const onNodesChange = useCallback(
+    (changes: any[]) => {
+      setNodes((nds) => applyNodeChanges(changes, nds));
+      for (const ch of changes) {
+        if (ch.type === 'position') {
+          const node = state.workspace?.nodes?.find((n) => n.id === ch.id);
+          if (node) {
+            node.position = ch.position;
+          }
+        }
+      }
+    },
+    [],
+  );
+  const onEdgesChange = useCallback(
+    (changes: any[]) => setEdges((eds) => applyEdgeChanges(changes, eds)),
+    [],
+  );
+  if (state?.workspace?.nodes && JSON.stringify(nodes) !== JSON.stringify([...state.workspace.nodes as Node[]])) {
+    const updated = Object.fromEntries(state.workspace.nodes.map((n) => [n.id, n]));
+    const oldNodes = Object.fromEntries(nodes.map((n) => [n.id, n]));
+    const updatedNodes = nodes.filter(n => updated[n.id]).map((n) => ({ ...n, ...updated[n.id] })) as Node[];
+    const newNodes = state.workspace.nodes.filter((n) => !oldNodes[n.id]);
+    const allNodes = [...updatedNodes, ...newNodes];
+    if (JSON.stringify(allNodes) !== JSON.stringify(nodes)) {
+      setNodes(allNodes as Node[]);
+    }
+  }
+  if (state?.workspace?.edges && JSON.stringify(edges) !== JSON.stringify([...state.workspace.edges as Edge[]])) {
+    const updated = Object.fromEntries(state.workspace.edges.map((e) => [e.id, e]));
+    const oldEdges = Object.fromEntries(edges.map((e) => [e.id, e]));
+    const updatedEdges = edges.filter(e => updated[e.id]).map((e) => ({ ...e, ...updated[e.id] })) as Edge[];
+    const newEdges = state.workspace.edges.filter((e) => !oldEdges[e.id]);
+    const allEdges = [...updatedEdges, ...newEdges];
+    if (JSON.stringify(allEdges) !== JSON.stringify(edges)) {
+      setEdges(allEdges as Edge[]);
+    }
+  }
 
   const fetcher = (resource: string, init?: RequestInit) => fetch(resource, init).then(res => res.json());
   const catalog = useSWR('/api/catalog', fetcher);
@@ -80,8 +117,8 @@ function LynxKiteFlow() {
         </div>
         <EnvironmentSelector
           options={Object.keys(catalog.data || {})}
-          value={store.env}
-          onChange={(env) => store.setEnv(env)}
+          value={state.workspace.env!}
+          onChange={(env) => { state.workspace.env = env; }}
         />
         <div className="tools text-secondary">
           <a href=""><Atom /></a>
@@ -90,20 +127,19 @@ function LynxKiteFlow() {
         </div>
       </div>
       <div style={{ height: "100%", width: '100vw' }}>
-        <LynxKiteState.Provider value={store}>
+        <LynxKiteState.Provider value={state.workspace}>
           <ReactFlow
-            nodes={store.nodes as Node[]}
-            edges={store.edges}
+            nodes={nodes}
+            edges={edges}
             nodeTypes={nodeTypes} fitView
-            onNodesChange={store.onNodesChange}
-            onEdgesChange={store.onEdgesChange}
+            onNodesChange={onNodesChange}
+            onEdgesChange={onEdgesChange}
             proOptions={{ hideAttribution: true }}
             maxZoom={3}
             minZoom={0.3}
             defaultEdgeOptions={{ markerEnd: { type: MarkerType.Arrow } }}
           >
             <Controls />
-            <MiniMap />
             {/* {#if nodeSearchSettings}
           <NodeSearch pos={nodeSearchSettings.pos} boxes={nodeSearchSettings.boxes} on:cancel={closeNodeSearch} on:add={addNode} />
           {/if} */}

web/src/workspace/store.ts

@@ -1,69 +0,0 @@
-// Like described in https://reactflow.dev/learn/advanced-use/state-management
-// but with https://github.com/joebobmiles/zustand-middleware-yjs/ added.
-import {
-  type Edge,
-  type Node,
-  type OnNodesChange,
-  type OnEdgesChange,
-  type OnConnect,
-} from '@xyflow/react';
-import * as apiTypes from "../apiTypes.ts";
-import * as Y from "yjs";
-import yjs from "zustand-middleware-yjs";
-export const doc = new Y.Doc();
-
-export type FlowState = {
-  env: string;
-  nodes: apiTypes.WorkspaceNode[];
-  edges: apiTypes.WorkspaceEdge[];
-  onNodesChange: OnNodesChange<Node>;
-  onEdgesChange: OnEdgesChange;
-  onConnect: OnConnect;
-  setNodes: (nodes: Node[]) => void;
-  setEdges: (edges: Edge[]) => void;
-  setEnv: (env: string) => void;
-};
-
-import { create } from 'zustand';
-import { addEdge, applyNodeChanges, applyEdgeChanges } from '@xyflow/react';
-
-export const useStore = create<FlowState>(yjs(doc, "shared", (set: any, get: any) => ({
-  env: 'LynxKite',
-  nodes: [],
-  edges: [],
-  onNodesChange: (changes: any[]) => {
-    set({
-      nodes: applyNodeChanges(changes, get().nodes),
-    });
-  },
-  onEdgesChange: (changes: any[]) => {
-    set({
-      edges: applyEdgeChanges(changes, get().edges),
-    });
-  },
-  onConnect: (connection: any) => {
-    set({
-      edges: addEdge(connection, get().edges),
-    });
-  },
-  setNodes: (nodes: Node[]) => {
-    console.log("setNodes", { nodes });
-    set({ nodes });
-  },
-  setEdges: (edges: Edge[]) => {
-    set({ edges });
-  },
-  setEnv: (env: string) => {
-    set({ env });
-  },
-})));
-
-export const selector = (state: FlowState) => ({
-  nodes: state.nodes,
-  edges: state.edges,
-  env: state.env,
-  onNodesChange: state.onNodesChange,
-  onEdgesChange: state.onEdgesChange,
-  onConnect: state.onConnect,
-  setEnv: state.setEnv,
-});