diff --git a/packages/ui/src/components/Visualization/Custom/ContextMenu/ItemAddStep.tsx b/packages/ui/src/components/Visualization/Custom/ContextMenu/ItemAddStep.tsx
index 16b84d67c..08acbcea9 100644
--- a/packages/ui/src/components/Visualization/Custom/ContextMenu/ItemAddStep.tsx
+++ b/packages/ui/src/components/Visualization/Custom/ContextMenu/ItemAddStep.tsx
@@ -1,53 +1,19 @@
 import { ContextMenuItem } from '@patternfly/react-topology';
-import { FunctionComponent, PropsWithChildren, useCallback, useContext } from 'react';
+import { FunctionComponent, PropsWithChildren } from 'react';
 import { IDataTestID } from '../../../../models';
-import {
-  AddStepMode,
-  IVisualizationNode,
-  IVisualizationNodeData,
-} from '../../../../models/visualization/base-visual-entity';
-import { CatalogModalContext } from '../../../../providers/catalog-modal.provider';
-import { EntitiesContext } from '../../../../providers/entities.provider';
-import { EntitiesContextResult } from '../../../../hooks';
+import { AddStepMode, IVisualizationNode } from '../../../../models/visualization/base-visual-entity';
+import { useAddStep } from '../hooks/add-step.hook';
 
 interface ItemAddStepProps extends PropsWithChildren<IDataTestID> {
   mode: AddStepMode.PrependStep | AddStepMode.AppendStep;
   vizNode: IVisualizationNode;
 }
 
-export const addNode = async (
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  catalogModalContext: any,
-  entitiesContext: EntitiesContextResult | null,
-  vizNode: IVisualizationNode<IVisualizationNodeData>,
-  mode: AddStepMode = AddStepMode.AppendStep,
-) => {
-  if (!vizNode || !entitiesContext) return;
-
-  /** Get compatible nodes and the location where can be introduced */
-  const compatibleNodes = entitiesContext.camelResource.getCompatibleComponents(mode, vizNode.data);
-
-  /** Open Catalog modal, filtering the compatible nodes */
-  const definedComponent = await catalogModalContext?.getNewComponent(compatibleNodes);
-  if (!definedComponent) return;
-
-  /** Add new node to the entities */
-  vizNode.addBaseEntityStep(definedComponent, mode);
-
-  /** Update entity */
-  entitiesContext.updateEntitiesFromCamelResource();
-};
-
 export const ItemAddStep: FunctionComponent<ItemAddStepProps> = (props) => {
-  const entitiesContext = useContext(EntitiesContext);
-  const catalogModalContext = useContext(CatalogModalContext);
-
-  const onAddNode = useCallback(async () => {
-    addNode(catalogModalContext, entitiesContext, props.vizNode, props.mode);
-  }, [catalogModalContext, entitiesContext, props.mode, props.vizNode]);
+  const { onAddStep } = useAddStep(props.vizNode, props.mode);
 
   return (
-    <ContextMenuItem onClick={onAddNode} data-testid={props['data-testid']}>
+    <ContextMenuItem onClick={onAddStep} data-testid={props['data-testid']}>
       {props.children}
     </ContextMenuItem>
   );
diff --git a/packages/ui/src/components/Visualization/Custom/Edge/EdgeEnd.tsx b/packages/ui/src/components/Visualization/Custom/Edge/EdgeEnd.tsx
index 7a39b22e3..a1a2c1a63 100644
--- a/packages/ui/src/components/Visualization/Custom/Edge/EdgeEnd.tsx
+++ b/packages/ui/src/components/Visualization/Custom/Edge/EdgeEnd.tsx
@@ -8,11 +8,10 @@ import {
   isEdge,
   observer,
 } from '@patternfly/react-topology';
-import { FunctionComponent, useCallback, useContext } from 'react';
+import { FunctionComponent } from 'react';
 import { IVisualizationNode } from '../../../../models';
-import { CatalogModalContext, EntitiesContext } from '../../../../providers';
-import { addNode } from '../ContextMenu/ItemAddStep';
 import { LayoutType } from '../../Canvas';
+import { useAddStep } from '../hooks/add-step.hook';
 
 type DefaultEdgeProps = Parameters<typeof DefaultEdge>[0];
 interface EdgeEndProps extends DefaultEdgeProps {
@@ -24,15 +23,11 @@ export const EdgeEndWithButton: FunctionComponent<EdgeEndProps> = observer(({ el
   if (!isEdge(element)) {
     throw new Error('EdgeEndWithButton must be used only on Edge elements');
   }
-  const entitiesContext = useContext(EntitiesContext);
-  const catalogModalContext = useContext(CatalogModalContext);
   const vizNode: IVisualizationNode = element.getSource().getData()?.vizNode;
   const isHorizontal = element.getGraph().getLayout() === LayoutType.DagreHorizontal;
   const endPoint = element.getEndPoint();
 
-  const onAdd = useCallback(() => {
-    addNode(catalogModalContext, entitiesContext, vizNode);
-  }, [catalogModalContext, entitiesContext, vizNode]);
+  const { onAddStep } = useAddStep(vizNode);
 
   let x = endPoint.x;
   let y = endPoint.y;
@@ -50,7 +45,7 @@ export const EdgeEndWithButton: FunctionComponent<EdgeEndProps> = observer(({ el
       {...rest}
     >
       <g data-testid={`custom-edge__${element?.getId()}`}>
-        <Decorator showBackground radius={14} x={x} y={y} icon={<PlusIcon />} onClick={onAdd} />
+        <Decorator showBackground radius={14} x={x} y={y} icon={<PlusIcon />} onClick={onAddStep} />
       </g>
     </DefaultEdge>
   );
diff --git a/packages/ui/src/components/Visualization/Custom/hooks/add-step.hook.tsx b/packages/ui/src/components/Visualization/Custom/hooks/add-step.hook.tsx
new file mode 100644
index 000000000..3cacb62e0
--- /dev/null
+++ b/packages/ui/src/components/Visualization/Custom/hooks/add-step.hook.tsx
@@ -0,0 +1,38 @@
+import { useCallback, useContext, useMemo } from 'react';
+import { AddStepMode, IVisualizationNode } from '../../../../models/visualization/base-visual-entity';
+import { CatalogModalContext } from '../../../../providers';
+import { EntitiesContext } from '../../../../providers/entities.provider';
+
+export const useAddStep = (
+  vizNode: IVisualizationNode,
+  mode: AddStepMode.PrependStep | AddStepMode.AppendStep = AddStepMode.AppendStep,
+) => {
+  const entitiesContext = useContext(EntitiesContext);
+  const catalogModalContext = useContext(CatalogModalContext);
+
+  const onAddStep = useCallback(async () => {
+    if (!vizNode || !entitiesContext) return;
+
+    /** Get compatible nodes and the location where can be introduced */
+    const compatibleNodes = entitiesContext.camelResource.getCompatibleComponents(mode, vizNode.data);
+
+    /** Open Catalog modal, filtering the compatible nodes */
+    const definedComponent = await catalogModalContext?.getNewComponent(compatibleNodes);
+    if (!definedComponent) return;
+
+    /** Add new node to the entities */
+    vizNode.addBaseEntityStep(definedComponent, mode);
+
+    /** Update entity */
+    entitiesContext.updateEntitiesFromCamelResource();
+  }, [catalogModalContext, entitiesContext, mode, vizNode]);
+
+  const value = useMemo(
+    () => ({
+      onAddStep,
+    }),
+    [onAddStep],
+  );
+
+  return value;
+};