Stricter Typing on Nodes

Author: camielvs

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

@@ -306,7 +306,7 @@ const FlowCanvas = ({
   const nodeData = useMemo(
     () => ({
       connectable: !readOnly && !!nodesConnectable,
-      readOnly,
+      readOnly: !!readOnly,
       nodeCallbacks,
     }),
     [readOnly, nodesConnectable, nodeCallbacks],

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

@@ -3,9 +3,10 @@ import { memo, useMemo } from "react";
 
 import { cn } from "@/lib/utils";
 import { TaskNodeProvider } from "@/providers/TaskNodeProvider";
-import type { TaskNodeData } from "@/types/taskNode";
+import type { TaskNodeData } from "@/types/nodes";
 import type { ComponentReference } from "@/utils/componentSpec";
 import { generateTaskSpec } from "@/utils/nodes/generateTaskSpec";
+import { createEmptyTaskCallbacks } from "@/utils/nodes/taskCallbackUtils";
 
 import { TaskNodeCard } from "../TaskNode/TaskNodeCard";
 
@@ -74,5 +75,9 @@ const generateGhostTaskNodeData = (
     taskSpec,
     taskId: ghostTaskId,
     isGhost: true,
+    readOnly: true,
+    highlighted: false,
+    connectable: false,
+    callbacks: createEmptyTaskCallbacks(),
   };
 };

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

@@ -2,7 +2,7 @@ import { type NodeProps, useReactFlow } from "@xyflow/react";
 import { memo, type PropsWithChildren, useMemo } from "react";
 
 import { cn } from "@/lib/utils";
-import type { HintNodeData } from "@/types/hintNode";
+import type { HintNodeData } from "@/types/nodes";
 
 const HintNode = ({ data }: NodeProps) => {
   const { getZoom } = useReactFlow();

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

@@ -10,16 +10,12 @@ import { Paragraph } from "@/components/ui/typography";
 import { cn } from "@/lib/utils";
 import { useComponentSpec } from "@/providers/ComponentSpecProvider";
 import { useContextPanel } from "@/providers/ContextPanelProvider";
+import type { IONodeData } from "@/types/nodes";
+import type { InputSpec, TypeSpecType } from "@/utils/componentSpec";
 
 interface IONodeProps {
   type: "input" | "output";
-  data: {
-    label: string;
-    value?: string;
-    default?: string;
-    type?: string;
-    readOnly?: boolean;
-  };
+  data: IONodeData;
   selected: boolean;
   deletable: boolean;
 }
@@ -28,10 +24,12 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
   const { graphSpec, componentSpec } = useComponentSpec();
   const { setContent, clearContent } = useContextPanel();
 
+  const { spec, readOnly } = data;
+
   const isInput = type === "input";
   const isOutput = type === "output";
 
-  const readOnly = !!data.readOnly;
+  const inputSpec = isInput ? (spec as InputSpec) : null;
 
   const handleType = isInput ? "source" : "target";
   const handlePosition = isInput ? Position.Right : Position.Left;
@@ -44,13 +42,13 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
   const borderColor = selected ? selectedColor : defaultColor;
 
   const input = useMemo(
-    () => componentSpec.inputs?.find((input) => input.name === data.label),
-    [componentSpec.inputs, data.label],
+    () => componentSpec.inputs?.find((input) => input.name === spec.name),
+    [componentSpec.inputs, spec.name],
   );
 
   const output = useMemo(
-    () => componentSpec.outputs?.find((output) => output.name === data.label),
-    [componentSpec.outputs, data.label],
+    () => componentSpec.outputs?.find((output) => output.name === spec.name),
+    [componentSpec.outputs, spec.name],
   );
 
   useEffect(() => {
@@ -88,7 +86,7 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
     };
   }, [input, output, selected, readOnly]);
 
-  const connectedOutput = getOutputConnectedDetails(graphSpec, data.label);
+  const connectedOutput = getOutputConnectedDetails(graphSpec, spec.name);
   const outputConnectedValue = connectedOutput.outputName;
   const outputConnectedType = connectedOutput.outputType;
   const outputConnectedTaskId = connectedOutput.taskId;
@@ -98,30 +96,28 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
 
   const handleClassName = isInput ? "translate-x-1.5" : "-translate-x-1.5";
 
-  const hasDataValue = !!data.value;
-  const hasDataDefault = !!data.default;
-
-  const inputValue = hasDataValue
-    ? data.value
-    : hasDataDefault
-      ? data.default
+  const inputValue = inputSpec?.value
+    ? inputSpec?.value
+    : inputSpec?.default
+      ? inputSpec?.default
       : null;
-
   const outputValue = outputConnectedValue ?? null;
-
   const value = isInput ? inputValue : outputValue;
 
+  const inputType = getTypeDisplayString(inputSpec?.type);
+  const outputType = getTypeDisplayString(outputConnectedType);
+  const typeValue = isInput ? inputType : outputType;
+
   return (
     <Card className={cn("border-2 max-w-[300px] p-0", borderColor)}>
       <CardHeader className="px-2 py-2.5">
-        <CardTitle className="break-words">{data.label}</CardTitle>
+        <CardTitle className="break-words">{spec.name}</CardTitle>
       </CardHeader>
       <CardContent className="p-2 max-w-[250px]">
         <BlockStack gap="2">
           {/* type */}
           <Paragraph size="xs" font="mono" className="truncate text-slate-700">
-            <span className="font-bold">Type:</span>{" "}
-            {outputConnectedType ?? data.type ?? "Any"}
+            <span className="font-bold">Type:</span> {typeValue}
           </Paragraph>
 
           {!!outputConnectedTaskId && (
@@ -161,3 +157,13 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
 };
 
 export default memo(IONode);
+
+const getTypeDisplayString = (type: TypeSpecType | undefined): string => {
+  if (!type) return "Any";
+
+  if (typeof type === "string") {
+    return type;
+  }
+
+  return JSON.stringify(type);
+};

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

@@ -4,7 +4,7 @@ import { memo, useMemo } from "react";
 import type { ContainerExecutionStatus } from "@/api/types.gen";
 import { useComponentSpec } from "@/providers/ComponentSpecProvider";
 import { TaskNodeProvider } from "@/providers/TaskNodeProvider";
-import type { TaskNodeData } from "@/types/taskNode";
+import type { TaskNodeData } from "@/types/nodes";
 
 import { StatusIndicator } from "./StatusIndicator";
 import { TaskNodeCard } from "./TaskNodeCard";

src/components/shared/ReactFlow/FlowCanvas/utils/addTask.ts

@@ -1,12 +1,13 @@
 import type { XYPosition } from "@xyflow/react";
 
-import type { TaskType } from "@/types/taskNode";
-import type {
-  ComponentSpec,
-  GraphSpec,
-  InputSpec,
-  OutputSpec,
-  TaskSpec,
+import type { TaskType } from "@/types/nodes";
+import {
+  type ComponentSpec,
+  type GraphSpec,
+  type InputSpec,
+  isGraphImplementation,
+  type OutputSpec,
+  type TaskSpec,
 } from "@/utils/componentSpec";
 import { deepClone } from "@/utils/deepClone";
 import {
@@ -23,7 +24,7 @@ const addTask = (
 ): ComponentSpec => {
   const newComponentSpec = deepClone(componentSpec);
 
-  if (!("graph" in newComponentSpec.implementation)) {
+  if (!isGraphImplementation(newComponentSpec.implementation)) {
     console.error("Implementation does not contain a graph.");
     return newComponentSpec;
   }

src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.test.ts

@@ -1,7 +1,7 @@
 import type { Node } from "@xyflow/react";
 import { describe, expect, it, vi } from "vitest";
 
-import type { NodeCallbacks, TaskNodeData } from "@/types/taskNode";
+import type { TaskNodeData } from "@/types/nodes";
 import type {
   ComponentSpec,
   InputSpec,
@@ -85,14 +85,6 @@ const createMockTaskNodeCallbacks = () => ({
   onUpgrade: vi.fn(),
 });
 
-const createMockNodeCallbacks = (): NodeCallbacks => ({
-  setArguments: vi.fn(),
-  setAnnotations: vi.fn(),
-  onDelete: vi.fn(),
-  onDuplicate: vi.fn(),
-  onUpgrade: vi.fn(),
-});
-
 const createMockTaskNode = (
   taskId: string,
   taskSpec: TaskSpec,
@@ -106,8 +98,10 @@ const createMockTaskNode = (
     taskId,
     label: "Test Task",
     highlighted: false,
+    readOnly: false,
+    isGhost: false,
+    connectable: true,
     callbacks: createMockTaskNodeCallbacks(),
-    nodeCallbacks: createMockNodeCallbacks(),
   },
   selected: false,
   dragging: false,

src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.ts

@@ -1,6 +1,6 @@
 import { type Node, type XYPosition } from "@xyflow/react";
 
-import type { TaskNodeData } from "@/types/taskNode";
+import type { IONodeData, NodeData, TaskNodeData } from "@/types/nodes";
 import {
   type ComponentSpec,
   type GraphInputArgument,
@@ -23,6 +23,7 @@ import {
   taskIdToNodeId,
 } from "@/utils/nodes/nodeIdUtils";
 import { setPositionInAnnotations } from "@/utils/nodes/setPositionInAnnotations";
+import { convertTaskCallbacksToNodeCallbacks } from "@/utils/nodes/taskCallbackUtils";
 import {
   getUniqueInputName,
   getUniqueOutputName,
@@ -277,17 +278,19 @@ export const duplicateNodes = (
         return null;
       }
 
-      const originalNodeData = originalNode.data as TaskNodeData;
-
       if (originalNode.type === "task") {
         const newTaskId = nodeIdToTaskId(newNodeId);
 
         const newTaskSpec = updatedGraphSpec.tasks[newTaskId];
 
-        const newNode = createTaskNode(
-          [newTaskId, newTaskSpec],
-          originalNodeData,
-        );
+        const taskData = originalNode.data as TaskNodeData;
+        const nodeData: NodeData = {
+          readOnly: taskData.readOnly,
+          connectable: taskData.connectable,
+          callbacks: convertTaskCallbacksToNodeCallbacks(taskData.callbacks),
+        };
+
+        const newNode = createTaskNode([newTaskId, newTaskSpec], nodeData);
 
         newNode.id = newNodeId;
         newNode.selected = false;
@@ -313,7 +316,12 @@ export const duplicateNodes = (
           return null;
         }
 
-        const newNode = createInputNode(newInputSpec, originalNodeData);
+        const inputData = originalNode.data as IONodeData;
+        const nodeData: NodeData = {
+          readOnly: inputData.readOnly,
+        };
+
+        const newNode = createInputNode(newInputSpec, nodeData);
 
         newNode.id = newNodeId;
         newNode.selected = false;
@@ -339,7 +347,12 @@ export const duplicateNodes = (
           return null;
         }
 
-        const newNode = createOutputNode(newOutputSpec, originalNodeData);
+        const outputData = originalNode.data as IONodeData;
+        const nodeData: NodeData = {
+          readOnly: outputData.readOnly,
+        };
+
+        const newNode = createOutputNode(newOutputSpec, nodeData);
 
         newNode.id = newNodeId;
 

src/components/shared/ReactFlow/FlowCanvas/utils/getTaskFromEvent.ts

@@ -1,6 +1,6 @@
 import type { DragEvent } from "react";
 
-import type { TaskType } from "@/types/taskNode";
+import type { TaskType } from "@/types/nodes";
 import type { TaskSpec } from "@/utils/componentSpec";
 
 export const getTaskFromEvent = (event: DragEvent) => {

src/hooks/useHintNode.ts

@@ -4,7 +4,7 @@ import { useMemo } from "react";
 
 import { useBetaFlagValue } from "@/components/shared/Settings/useBetaFlags";
 import { useComponentLibrary } from "@/providers/ComponentLibraryProvider";
-import type { HintNodeData } from "@/types/hintNode";
+import type { HintNodeData } from "@/types/nodes";
 
 const HINT_NODE_ID = "hint-node";