Fix Node Replacement Issues on IO Rename

Author: camielvs

src/components/Editor/IOEditor/InputValueEditor/FormFields/FormFields.tsx

@@ -68,14 +68,12 @@ const NameField = ({
   onBlur,
   error,
   disabled,
-  autoFocus = false,
 }: {
   inputName: string;
   onNameChange: (value: string) => void;
   onBlur?: () => void;
   error?: string | null;
   disabled?: boolean;
-  autoFocus?: boolean;
 }) => (
   <FormField label="Name" id={`input-name-${inputName}`}>
     <Input
@@ -88,7 +86,6 @@ const NameField = ({
       className={cn("text-sm", {
         "border-red-500 focus:border-red-500": !!error,
       })}
-      autoFocus={autoFocus}
     />
     <Activity mode={error ? "visible" : "hidden"}>
       <div className="text-xs text-red-500 mt-1">{error}</div>

src/components/Editor/IOEditor/InputValueEditor/InputValueEditor.test.tsx

@@ -7,9 +7,9 @@ import { InputValueEditor } from "./InputValueEditor";
 
 // Mock all the dependencies
 const mockSetComponentSpec = vi.fn();
-const mockTransferSelection = vi.fn();
 const mockNotify = vi.fn();
 const mockClearContent = vi.fn();
+const mockUpdateRefId = vi.fn();
 
 vi.mock("@/providers/ComponentSpecProvider", () => ({
   useComponentSpec: () => ({
@@ -36,9 +36,12 @@ vi.mock("@/providers/ContextPanelProvider", () => ({
   }),
 }));
 
-vi.mock("@/hooks/useNodeSelectionTransfer", () => ({
-  useNodeSelectionTransfer: () => ({
-    transferSelection: mockTransferSelection,
+vi.mock("@/hooks/useNodeManager", () => ({
+  useNodeManager: () => ({
+    updateRefId: mockUpdateRefId,
+    getNodeId: vi.fn(),
+    getHandleNodeId: vi.fn(),
+    nodeManager: {},
   }),
 }));
 
@@ -110,7 +113,7 @@ describe("InputValueEditor", () => {
 
     // Verify that the component spec was updated and selection was transferred
     expect(mockSetComponentSpec).toHaveBeenCalled();
-    expect(mockTransferSelection).toHaveBeenCalledWith("TestInput", "NewName");
+    expect(mockUpdateRefId).toHaveBeenCalledWith("TestInput", "NewName");
   });
 
   it("shows validation error when renaming to existing input name", () => {

src/components/Editor/IOEditor/InputValueEditor/InputValueEditor.tsx

@@ -10,7 +10,6 @@ import { BlockStack } from "@/components/ui/layout";
 import { Heading, Paragraph } from "@/components/ui/typography";
 import useConfirmationDialog from "@/hooks/useConfirmationDialog";
 import { useNodeManager } from "@/hooks/useNodeManager";
-import { useNodeSelectionTransfer } from "@/hooks/useNodeSelectionTransfer";
 import useToastNotification from "@/hooks/useToastNotification";
 import { useComponentSpec } from "@/providers/ComponentSpecProvider";
 import { useContextPanel } from "@/providers/ContextPanelProvider";
@@ -31,19 +30,9 @@ export const InputValueEditor = ({
   input,
   disabled = false,
 }: InputValueEditorProps) => {
-  const { getInputNodeId } = useNodeManager();
-
-  const inputNameToNodeId = useCallback(
-    (inputName: string): string => {
-      const inputId = inputNameToInputId(inputName);
-      return getInputNodeId(inputId);
-    },
-    [getInputNodeId],
-  );
+  const { updateRefId } = useNodeManager();
 
   const notify = useToastNotification();
-
-  const { transferSelection } = useNodeSelectionTransfer(inputNameToNodeId);
   const { componentSpec, setComponentSpec } = useComponentSpec();
   const { clearContent } = useContextPanel();
   const {
@@ -111,11 +100,13 @@ export const InputValueEditor = ({
         newName,
       );
 
-      transferSelection(oldName, newName);
+      const oldInputId = inputNameToInputId(oldName);
+      const newInputId = inputNameToInputId(newName);
+      updateRefId(oldInputId, newInputId);
 
       return updatedComponentSpec;
     },
-    [componentSpec, transferSelection],
+    [componentSpec, updateRefId],
   );
 
   const handleValueChange = useCallback((value: string) => {
@@ -267,7 +258,6 @@ export const InputValueEditor = ({
         onBlur={handleBlur}
         error={validationError}
         disabled={disabled}
-        autoFocus={!disabled}
       />
 
       <TextField

src/components/Editor/IOEditor/OutputNameEditor/OutputNameEditor.tsx

@@ -8,10 +8,10 @@ import { BlockStack, InlineStack } from "@/components/ui/layout";
 import { Heading, Paragraph } from "@/components/ui/typography";
 import useConfirmationDialog from "@/hooks/useConfirmationDialog";
 import { useNodeManager } from "@/hooks/useNodeManager";
-import { useNodeSelectionTransfer } from "@/hooks/useNodeSelectionTransfer";
 import { useComponentSpec } from "@/providers/ComponentSpecProvider";
 import { useContextPanel } from "@/providers/ContextPanelProvider";
 import { type OutputSpec } from "@/utils/componentSpec";
+import { outputNameToOutputId } from "@/utils/nodes/conversions";
 
 import { type OutputConnectedDetails } from "../../utils/getOutputConnectedDetails";
 import { updateOutputNameOnComponentSpec } from "../../utils/updateOutputNameOnComponentSpec";
@@ -29,17 +29,8 @@ export const OutputNameEditor = ({
   disabled,
   connectedDetails,
 }: OutputNameEditorProps) => {
-  const { getOutputNodeId } = useNodeManager();
+  const { updateRefId } = useNodeManager();
 
-  const outputNameToNodeId = useCallback(
-    (outputName: string): string => {
-      const outputId = outputNameToNodeId(outputName);
-      return getOutputNodeId(outputId);
-    },
-    [getOutputNodeId],
-  );
-
-  const { transferSelection } = useNodeSelectionTransfer(outputNameToNodeId);
   const { setComponentSpec, componentSpec } = useComponentSpec();
   const { clearContent } = useContextPanel();
   const {
@@ -65,11 +56,13 @@ export const OutputNameEditor = ({
         newName,
       );
 
-      transferSelection(oldName, newName);
+      const oldInputId = outputNameToOutputId(oldName);
+      const newInputId = outputNameToOutputId(newName);
+      updateRefId(oldInputId, newInputId);
 
       return updatedComponentSpec;
     },
-    [componentSpec, setComponentSpec],
+    [componentSpec, updateRefId],
   );
 
   const saveChanges = useCallback(() => {

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

@@ -80,13 +80,7 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
   useEffect(() => {
     if (selected) {
       if (input && isInput) {
-        setContent(
-          <InputValueEditor
-            input={input}
-            key={input.name}
-            disabled={readOnly}
-          />,
-        );
+        setContent(<InputValueEditor input={input} disabled={readOnly} />);
       }
 
       if (output && isOutput) {
@@ -98,19 +92,18 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
           <OutputNameEditor
             output={output}
             connectedDetails={outputConnectedDetails}
-            key={output.name}
             disabled={readOnly}
           />,
         );
       }
     }
+  }, [selected, readOnly, input, output, isInput, isOutput, graphSpec]);
 
+  useEffect(() => {
     return () => {
-      if (selected) {
-        clearContent();
-      }
+      clearContent();
     };
-  }, [input, output, selected, readOnly]);
+  }, []);
 
   const connectedOutput = getOutputConnectedDetails(graphSpec, spec.name);
   const outputConnectedValue = connectedOutput.outputName;