Refactor TaskNode Callback Generation

Author: camielvs

src/components/shared/ReactFlow/FlowCanvas/FlowCanvas.tsx

@@ -340,13 +340,7 @@ const FlowCanvas = ({
 
       setComponentSpec(updatedComponentSpec);
     },
-    [
-      reactFlowInstance,
-      componentSpec,
-      nodeData,
-      setComponentSpec,
-      updateOrAddNodes,
-    ],
+    [reactFlowInstance, componentSpec, setComponentSpec, updateOrAddNodes],
   );
 
   useEffect(() => {

src/types/taskNode.ts

@@ -6,26 +6,31 @@ import type {
 
 import type { Annotations } from "./annotations";
 
+export type TaskType = "task" | "input" | "output";
+
+export interface TaskNodeCreationData extends Record<string, unknown> {
+  readOnly?: boolean;
+  connectable?: boolean;
+  nodeCallbacks?: NodeCallbacks;
+}
+
 export interface TaskNodeData extends Record<string, unknown> {
   taskSpec?: TaskSpec;
   taskId?: string;
   readOnly?: boolean;
   isGhost?: boolean;
   connectable?: boolean;
   highlighted?: boolean;
-  callbacks?: TaskNodeCallbacks;
-  nodeCallbacks?: NodeCallbacks;
+  callbacks?: TaskCallbacks;
 }
 
 export type NodeAndTaskId = {
   taskId: string;
   nodeId: string;
 };
 
-export type TaskType = "task" | "input" | "output";
-
 /* Note: Optional callbacks will cause TypeScript to break when applying the callbacks to the Nodes. */
-interface TaskNodeCallbacks {
+export interface TaskCallbacks {
   setArguments: (args: Record<string, ArgumentType>) => void;
   setAnnotations: (annotations: Annotations) => void;
   onDelete: () => void;
@@ -34,13 +39,14 @@ interface TaskNodeCallbacks {
 }
 
 // Dynamic Node Callback types - every callback has a version with the node & task id added to it as an input parameter
-export type CallbackWithIds<K extends keyof TaskNodeCallbacks> =
-  TaskNodeCallbacks[K] extends (...args: infer A) => infer R
-    ? (ids: NodeAndTaskId, ...args: A) => R
-    : never;
+type CallbackWithIds<K extends keyof TaskCallbacks> = TaskCallbacks[K] extends (
+  ...args: infer A
+) => infer R
+  ? (ids: NodeAndTaskId, ...args: A) => R
+  : never;
 
 export type NodeCallbacks = {
-  [K in keyof TaskNodeCallbacks]: CallbackWithIds<K>;
+  [K in keyof TaskCallbacks]: CallbackWithIds<K>;
 };
 
 export type TaskNodeDimensions = { w: number; h: number | undefined };

src/utils/nodes/createInputNode.ts

@@ -1,12 +1,15 @@
 import { type Node } from "@xyflow/react";
 
-import type { TaskNodeData } from "@/types/taskNode";
+import type { TaskNodeCreationData } from "@/types/taskNode";
 
 import type { InputSpec } from "../componentSpec";
 import { extractPositionFromAnnotations } from "./extractPositionFromAnnotations";
 import { inputNameToNodeId } from "./nodeIdUtils";
 
-export const createInputNode = (input: InputSpec, nodeData: TaskNodeData) => {
+export const createInputNode = (
+  input: InputSpec,
+  nodeData: TaskNodeCreationData,
+) => {
   const { name, annotations, ...rest } = input;
 
   const position = extractPositionFromAnnotations(annotations);

src/utils/nodes/createNodesFromComponentSpec.ts

@@ -1,6 +1,6 @@
 import { type Node } from "@xyflow/react";
 
-import type { TaskNodeData } from "@/types/taskNode";
+import type { TaskNodeCreationData } from "@/types/taskNode";
 import {
   type ComponentSpec,
   type GraphSpec,
@@ -13,7 +13,7 @@ import { createTaskNode } from "./createTaskNode";
 
 const createNodesFromComponentSpec = (
   componentSpec: ComponentSpec,
-  nodeData: TaskNodeData,
+  nodeData: TaskNodeCreationData,
 ): Node[] => {
   if (!isGraphImplementation(componentSpec.implementation)) {
     return [];
@@ -27,15 +27,18 @@ const createNodesFromComponentSpec = (
   return [...taskNodes, ...inputNodes, ...outputNodes];
 };
 
-const createTaskNodes = (graphSpec: GraphSpec, nodeData: TaskNodeData) => {
+const createTaskNodes = (
+  graphSpec: GraphSpec,
+  nodeData: TaskNodeCreationData,
+) => {
   return Object.entries(graphSpec.tasks).map((task) =>
     createTaskNode(task, nodeData),
   );
 };
 
 const createInputNodes = (
   componentSpec: ComponentSpec,
-  nodeData: TaskNodeData,
+  nodeData: TaskNodeCreationData,
 ) => {
   return (componentSpec.inputs ?? []).map((inputSpec) =>
     createInputNode(inputSpec, nodeData),
@@ -44,7 +47,7 @@ const createInputNodes = (
 
 const createOutputNodes = (
   componentSpec: ComponentSpec,
-  nodeData: TaskNodeData,
+  nodeData: TaskNodeCreationData,
 ) => {
   return (componentSpec.outputs ?? []).map((outputSpec) =>
     createOutputNode(outputSpec, nodeData),

src/utils/nodes/createOutputNode.ts

@@ -1,14 +1,14 @@
 import { type Node } from "@xyflow/react";
 
-import type { TaskNodeData } from "@/types/taskNode";
+import type { TaskNodeCreationData } from "@/types/taskNode";
 
 import type { OutputSpec } from "../componentSpec";
 import { extractPositionFromAnnotations } from "./extractPositionFromAnnotations";
 import { outputNameToNodeId } from "./nodeIdUtils";
 
 export const createOutputNode = (
   output: OutputSpec,
-  nodeData: TaskNodeData,
+  nodeData: TaskNodeCreationData,
 ) => {
   const { name, annotations, ...rest } = output;
 

src/utils/nodes/createTaskNode.ts

@@ -1,34 +1,37 @@
 import { type Node } from "@xyflow/react";
 
-import type { TaskNodeData } from "@/types/taskNode";
+import type { TaskNodeCreationData, TaskNodeData } from "@/types/taskNode";
 
 import type { TaskSpec } from "../componentSpec";
 import { extractPositionFromAnnotations } from "./extractPositionFromAnnotations";
-import { generateDynamicNodeCallbacks } from "./generateDynamicNodeCallbacks";
+import { generateTaskCallbacks } from "./generateTaskCallbacks";
 import { taskIdToNodeId } from "./nodeIdUtils";
 
 export const createTaskNode = (
   task: [`${string}`, TaskSpec],
-  nodeData: TaskNodeData,
+  nodeData: TaskNodeCreationData,
 ) => {
   const [taskId, taskSpec] = task;
+  const { nodeCallbacks, ...data } = nodeData;
 
   const position = extractPositionFromAnnotations(taskSpec.annotations);
   const nodeId = taskIdToNodeId(taskId);
 
   // Inject the taskId and nodeId into the callbacks
-  const nodeCallbacks = nodeData.nodeCallbacks;
-  const dynamicCallbacks = generateDynamicNodeCallbacks(nodeId, nodeCallbacks);
+  const ids = { taskId, nodeId };
+  const taskCallbacks = generateTaskCallbacks(ids, nodeCallbacks);
+
+  const taskNodeData: TaskNodeData = {
+    ...data,
+    taskSpec,
+    taskId,
+    highlighted: false,
+    callbacks: taskCallbacks,
+  };
 
   return {
     id: nodeId,
-    data: {
-      ...nodeData,
-      taskSpec,
-      taskId,
-      highlighted: false,
-      callbacks: dynamicCallbacks, // Use these callbacks internally within the node
-    },
+    data: taskNodeData,
     position: position,
     type: "task",
   } as Node;

src/utils/nodes/generateDynamicNodeCallbacks.ts

@@ -0,0 +1,25 @@
+import type {
+  NodeAndTaskId,
+  NodeCallbacks,
+  TaskCallbacks,
+} from "@/types/taskNode";
+
+// Map NodeCallbacks to TaskCallbacks by injecting nodeId and taskId
+export const generateTaskCallbacks = (
+  ids: NodeAndTaskId,
+  nodeCallbacks?: NodeCallbacks,
+): TaskCallbacks => {
+  if (!nodeCallbacks) {
+    return {} as TaskCallbacks;
+  }
+
+  return {
+    setArguments: (args) => nodeCallbacks.setArguments?.(ids, args),
+    setAnnotations: (annotations) =>
+      nodeCallbacks.setAnnotations?.(ids, annotations),
+    onDelete: () => nodeCallbacks.onDelete?.(ids),
+    onDuplicate: (selected) => nodeCallbacks.onDuplicate?.(ids, selected),
+    onUpgrade: (newComponentRef) =>
+      nodeCallbacks.onUpgrade?.(ids, newComponentRef),
+  };
+};