diff --git a/frontend/src/components/visual-editor/v2/Diagrams.tsx b/frontend/src/components/visual-editor/v2/Diagrams.tsx
index 0e9aac32..b85f4ac0 100644
--- a/frontend/src/components/visual-editor/v2/Diagrams.tsx
+++ b/frontend/src/components/visual-editor/v2/Diagrams.tsx
@@ -48,7 +48,6 @@ import { IBlock } from "@/types/block.types";
 import { ICategory } from "@/types/category.types";
 import { BlockPorts } from "@/types/visual-editor.types";
 
-
 import BlockDialog from "../BlockDialog";
 import { ZOOM_LEVEL } from "../constants";
 import { useVisualEditor } from "../hooks/useVisualEditor";
@@ -141,6 +140,7 @@ const Diagrams = () => {
 
       if (id) {
         setSelectedCategoryId?.(id);
+        setSelectedBlockId(undefined); // Reset selected block when switching categories, resetting edit & remove buttons
       }
     }
   };