Move IO Node Handles into Node Manager

Author: camielvs

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

@@ -27,7 +27,7 @@ interface IONodeProps {
 }
 
 const IONode = ({ type, data, selected = false }: IONodeProps) => {
-  const { getInputNodeId, getOutputNodeId } = useNodeManager();
+  const { getNodeId, getHandleNodeId } = useNodeManager();
   const { graphSpec, componentSpec } = useComponentSpec();
   const { setContent, clearContent } = useContextPanel();
 
@@ -58,21 +58,24 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
     [componentSpec.outputs, spec.name],
   );
 
-  const nodeId = isInput
-    ? getInputNodeId(inputNameToInputId(spec.name))
-    : getOutputNodeId(outputNameToOutputId(spec.name));
+  const inputId = inputNameToInputId(spec.name);
+  const outputId = outputNameToOutputId(spec.name);
+  const id = isInput ? inputId : outputId;
 
-  const nodeHandleId = `${nodeId}_handle`;
+  const nodeId = getNodeId(id, type);
+
+  const handleNodeType = isInput ? "outputHandle" : "inputHandle";
+  const nodeHandleId = getHandleNodeId(id, "handle", handleNodeType);
 
   const handleHandleClick = useCallback(() => {
     if (ENABLE_DEBUG_MODE) {
-      console.log(`${isInput ? "Input" : "Output"} Node Handle clicked:`, {
-        name: isInput ? input?.name : output?.name,
+      console.log(`${type} Node Handle clicked:`, {
+        name: spec.name,
         nodeId,
         handleId: nodeHandleId,
       });
     }
-  }, [isInput, input, output, nodeId, nodeHandleId]);
+  }, [spec, nodeId, nodeHandleId]);
 
   useEffect(() => {
     if (selected) {

src/hooks/useNodeManager.ts

@@ -7,6 +7,24 @@ export const useNodeManager = () => {
   const { nodeManager } = useComponentSpec();
 
   // Get node ID utilities
+  const getNodeId = useCallback(
+    (name: string, type: "input" | "output" | "task"): string => {
+      return nodeManager.getNodeId(name, type);
+    },
+    [nodeManager],
+  );
+
+  const getHandleNodeId = useCallback(
+    (
+      taskId: string,
+      handleName: string,
+      handleType: "inputHandle" | "outputHandle",
+    ): string => {
+      return nodeManager.getHandleNodeId(taskId, handleName, handleType);
+    },
+    [nodeManager],
+  );
+
   const getTaskNodeId = useCallback(
     (taskId: string): string => {
       return nodeManager.getNodeId(taskId, "task");
@@ -91,6 +109,8 @@ export const useNodeManager = () => {
   return useMemo(
     () => ({
       // Core operations
+      getNodeId,
+      getHandleNodeId,
       getTaskId,
       updateTaskId,
       removeNode,
@@ -109,6 +129,8 @@ export const useNodeManager = () => {
       nodeManager,
     }),
     [
+      getNodeId,
+      getHandleNodeId,
       getTaskId,
       updateTaskId,
       removeNode,

src/nodeManager.ts

@@ -4,6 +4,10 @@ import {
   type ComponentSpec,
   isGraphImplementation,
 } from "./utils/componentSpec";
+import {
+  inputNameToInputId,
+  outputNameToOutputId,
+} from "./utils/nodes/conversions";
 
 export type NodeType =
   | "task"
@@ -248,12 +252,27 @@ export class NodeManager {
     }
 
     // Graph-level inputs and outputs
-    componentSpec.inputs?.forEach((input) =>
-      currentTasks.get("input")?.add(input.name),
-    );
-    componentSpec.outputs?.forEach((output) =>
-      currentTasks.get("output")?.add(output.name),
-    );
+    componentSpec.inputs?.forEach((input) => {
+      const inputId = inputNameToInputId(input.name);
+      currentTasks.get("input")?.add(input.name);
+
+      // Add to currentTaskHandles like a regular task
+      if (!currentTaskHandles.has(inputId)) {
+        currentTaskHandles.set(inputId, new Set());
+      }
+      currentTaskHandles.get(inputId)!.add("outputHandle:output");
+    });
+
+    componentSpec.outputs?.forEach((output) => {
+      const outputId = outputNameToOutputId(output.name);
+      currentTasks.get("output")?.add(output.name);
+
+      // Add to currentTaskHandles like a regular task
+      if (!currentTaskHandles.has(outputId)) {
+        currentTaskHandles.set(outputId, new Set());
+      }
+      currentTaskHandles.get(outputId)!.add("inputHandle:input");
+    });
 
     // Remove mappings for deleted tasks by type
     for (const [nodeType, taskMap] of this.taskToNodeMap) {