diff --git a/src/frontend/src/components/core/parameterRenderComponent/components/toggleShadComponent/index.tsx b/src/frontend/src/components/core/parameterRenderComponent/components/toggleShadComponent/index.tsx
index be2f8e6c7..e17c33613 100644
--- a/src/frontend/src/components/core/parameterRenderComponent/components/toggleShadComponent/index.tsx
+++ b/src/frontend/src/components/core/parameterRenderComponent/components/toggleShadComponent/index.tsx
@@ -36,21 +36,23 @@ export default function ToggleShadComponent({
   }
 
   return (
-    <Switch
-      id={id}
-      data-testid={id}
-      style={{
-        transform: `scaleX(${scaleX}) scaleY(${scaleY})`,
-      }}
-      disabled={disabled}
-      className=""
-      checked={value}
-      onCheckedChange={(isEnabled: boolean) => {
-        const data = showToogle
-          ? { advanced: !isEnabled }
-          : { value: isEnabled };
-        handleOnNewValue(data);
-      }}
-    ></Switch>
+    <div onClick={(e) => e.stopPropagation()}>
+      <Switch
+        id={id}
+        data-testid={id}
+        style={{
+          transform: `scaleX(${scaleX}) scaleY(${scaleY})`,
+        }}
+        disabled={disabled}
+        className=""
+        checked={value}
+        onCheckedChange={(isEnabled: boolean) => {
+          const data = showToogle
+            ? { advanced: !isEnabled }
+            : { value: isEnabled };
+          handleOnNewValue(data);
+        }}
+      />
+    </div>
   );
 }
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index e5be9a2a7..e8dbe9148 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -3,12 +3,12 @@ import { mutateTemplate } from "@/CustomNodes/helpers/mutate-template";
 import useHandleOnNewValue from "@/CustomNodes/hooks/use-handle-new-value";
 import useHandleNodeClass from "@/CustomNodes/hooks/use-handle-node-class";
 import ShadTooltip from "@/components/common/shadTooltipComponent";
+import ToggleShadComponent from "@/components/core/parameterRenderComponent/components/toggleShadComponent";
 import { Button } from "@/components/ui/button";
 import { usePatchUpdateFlow } from "@/controllers/API/queries/flows/use-patch-update-flow";
 import { usePostTemplateValue } from "@/controllers/API/queries/nodes/use-post-template-value";
 import { usePostRetrieveVertexOrder } from "@/controllers/API/queries/vertex";
 import useAddFlow from "@/hooks/flows/use-add-flow";
-import CodeAreaModal from "@/modals/codeAreaModal";
 import { APIClassType } from "@/types/api";
 import _, { cloneDeep } from "lodash";
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
@@ -20,9 +20,6 @@ import {
   SelectItem,
   SelectTrigger,
 } from "../../../../components/ui/select-custom";
-import ConfirmationModal from "../../../../modals/confirmationModal";
-import EditNodeModal from "../../../../modals/editNodeModal";
-import ShareModal from "../../../../modals/shareModal";
 import useAlertStore from "../../../../stores/alertStore";
 import { useDarkStore } from "../../../../stores/darkStore";
 import useFlowStore from "../../../../stores/flowStore";
@@ -42,6 +39,7 @@ import { cn, getNodeLength, openInNewTab } from "../../../../utils/utils";
 import { ToolbarButton } from "./components/toolbar-button";
 import ToolbarModals from "./components/toolbar-modals";
 import useShortcuts from "./hooks/use-shortcuts";
+import ShortcutDisplay from "./shortcutDisplay";
 import ToolbarSelectItem from "./toolbarSelectItem";
 
 const NodeToolbarComponent = memo(
@@ -77,6 +75,7 @@ const NodeToolbarComponent = memo(
     const [openModal, setOpenModal] = useState(false);
     const frozen = data.node?.frozen ?? false;
     const currentFlow = useFlowStore((state) => state.currentFlow);
+    const updateNodeInternals = useUpdateNodeInternals();
 
     const paste = useFlowStore((state) => state.paste);
     const nodes = useFlowStore((state) => state.nodes);
@@ -96,6 +95,28 @@ const NodeToolbarComponent = memo(
     });
     const updateToolMode = useFlowStore((state) => state.updateToolMode);
 
+    const flowDataNodes = useMemo(
+      () => currentFlow?.data?.nodes,
+      [currentFlow],
+    );
+
+    const node = useMemo(
+      () => flowDataNodes?.find((n) => n.id === data.id),
+      [flowDataNodes, data.id],
+    );
+
+    const index = useMemo(
+      () => flowDataNodes?.indexOf(node!)!,
+      [flowDataNodes, node],
+    );
+
+    const postToolModeValue = usePostTemplateValue({
+      node: data.node!,
+      nodeId: data.id,
+      parameterId: "tool_mode",
+      tool_mode: data.node!.tool_mode ?? false,
+    });
+
     const isSaved = flows?.some((flow) =>
       Object.values(flow).includes(data.node?.display_name!),
     );
@@ -142,11 +163,22 @@ const NodeToolbarComponent = memo(
       return hasComponentAsTool ?? false;
     });
 
-    const handleActivateToolMode = useCallback(() => {
-      const newValue = !toolMode;
+    const { handleNodeClass: handleNodeClassHook } = useHandleNodeClass(
+      data.id,
+    );
+
+    const handleNodeClass = (newNodeClass: APIClassType, type: string) => {
+      handleNodeClassHook(newNodeClass, type);
+    };
+
+    const handleActivateToolMode = () => {
+      const newValue = !flowDataNodes![index]!.data.node.tool_mode;
+
       updateToolMode(data.id, newValue);
       data.node!.tool_mode = newValue;
 
+      setToolMode(newValue);
+
       mutateTemplate(
         newValue,
         data.node!,
@@ -155,10 +187,7 @@ const NodeToolbarComponent = memo(
         setErrorData,
         "tool_mode",
         () => {
-          const node = currentFlow?.data?.nodes.find((n) => n.id === data.id);
-          const index = currentFlow?.data?.nodes.indexOf(node!)!;
           currentFlow!.data!.nodes[index]!.data.node.tool_mode = newValue;
-
           patchUpdateFlow({
             id: currentFlow?.id!,
             name: currentFlow?.name!,
@@ -171,7 +200,8 @@ const NodeToolbarComponent = memo(
       );
 
       updateNodeInternals(data.id);
-    }, [toolMode, data, currentFlow]);
+      return newValue;
+    };
 
     const handleMinimize = useCallback(() => {
       if (isMinimal || !showNode) {
@@ -288,7 +318,6 @@ const NodeToolbarComponent = memo(
         onCloseAdvancedModal!(false);
       }
     }, [showModalAdvanced]);
-    const updateNodeInternals = useUpdateNodeInternals();
 
     const setLastCopiedSelection = useFlowStore(
       (state) => state.setLastCopiedSelection,
@@ -400,14 +429,6 @@ const NodeToolbarComponent = memo(
       handleOnNewValueHook({ value });
     };
 
-    const { handleNodeClass: handleNodeClassHook } = useHandleNodeClass(
-      data.id,
-    );
-
-    const handleNodeClass = (newNodeClass: APIClassType, type: string) => {
-      handleNodeClassHook(newNodeClass, type);
-    };
-
     const selectTriggerRef = useRef(null);
 
     const handleButtonClick = () => {
@@ -418,13 +439,6 @@ const NodeToolbarComponent = memo(
       setOpenShowMoreOptions && setOpenShowMoreOptions(open);
     };
 
-    const postToolModeValue = usePostTemplateValue({
-      node: data.node!,
-      nodeId: data.id,
-      parameterId: "tool_mode",
-      tool_mode: data.node!.tool_mode ?? false,
-    });
-
     const renderToolbarButtons = useMemo(
       () => (
         <>
@@ -468,21 +482,52 @@ const NodeToolbarComponent = memo(
             />
           )}
           {hasToolMode && (
-            <ToolbarButton
-              icon="Hammer"
-              label="Tool Mode"
-              onClick={() => {
-                takeSnapshot();
-                handleSelectChange("toolMode");
-              }}
-              shortcut={shortcuts.find((s) =>
-                s.name.toLowerCase().startsWith("tool mode"),
-              )}
-              className={cn(
-                "node-toolbar-buttons h-[2rem]",
-                toolMode && "text-primary",
-              )}
-            />
+            <ShadTooltip
+              content={
+                <ShortcutDisplay
+                  {...shortcuts.find(
+                    ({ name }) => name.toLowerCase() === "tool mode",
+                  )!}
+                />
+              }
+              side="top"
+            >
+              <Button
+                className={cn(
+                  "node-toolbar-buttons h-[2rem]",
+                  toolMode && "text-primary",
+                )}
+                variant="ghost"
+                onClick={(event) => {
+                  event.preventDefault();
+                  takeSnapshot();
+                  handleSelectChange("toolMode");
+                }}
+                size="node-toolbar"
+                data-testid="tool-mode-button"
+              >
+                <IconComponent
+                  name="Hammer"
+                  className={cn(
+                    "h-4 w-4 transition-all",
+                    toolMode ? "text-primary" : "",
+                  )}
+                />
+                <span className="text-[13px] font-medium">Tool Mode</span>
+                <ToggleShadComponent
+                  value={toolMode}
+                  editNode={false}
+                  handleOnNewValue={() => {
+                    takeSnapshot();
+                    handleSelectChange("toolMode");
+                  }}
+                  disabled={false}
+                  size="medium"
+                  showToogle={false}
+                  id="tool-mode-toggle"
+                />
+              </Button>
+            </ShadTooltip>
           )}
         </>
       ),
diff --git a/src/frontend/tests/extended/features/tool-mode.spec.ts b/src/frontend/tests/extended/features/tool-mode.spec.ts
new file mode 100644
index 000000000..b71cdce91
--- /dev/null
+++ b/src/frontend/tests/extended/features/tool-mode.spec.ts
@@ -0,0 +1,132 @@
+import { expect, test } from "@playwright/test";
+import { awaitBootstrapTest } from "../../utils/await-bootstrap-test";
+
+test(
+  "User should be able to use components as tool",
+  { tag: ["@release"] },
+  async ({ page }) => {
+    await awaitBootstrapTest(page);
+    await page.getByTestId("blank-flow").click();
+    await page.waitForSelector('[data-testid="disclosure-vector stores"]', {
+      timeout: 3000,
+      state: "visible",
+    });
+
+    await page.getByTestId("disclosure-vector stores").click();
+    await page.waitForSelector('[data-testid="vectorstoresAstra DB"]', {
+      timeout: 3000,
+      state: "visible",
+    });
+    await page
+      .getByTestId("vectorstoresAstra DB")
+      .hover()
+      .then(async () => {
+        await page.getByTestId("add-component-button-astra-db").click();
+      });
+
+    await page.getByTestId("generic-node-title-arrangement").click();
+
+    await page.keyboard.press("ControlOrMeta+Shift+m");
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "visible",
+    });
+
+    expect(await page.getByText("toolset").count()).toBeGreaterThan(0);
+
+    await page.keyboard.press("ControlOrMeta+Shift+m");
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "hidden",
+    });
+
+    expect(await page.getByText("toolset").count()).toBe(0);
+
+    await page.getByTestId("tool-mode-button").click();
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "visible",
+    });
+
+    expect(await page.getByText("toolset").count()).toBeGreaterThan(0);
+
+    await page.getByTestId("tool-mode-button").click();
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "hidden",
+    });
+
+    expect(await page.getByText("toolset").count()).toBe(0);
+
+    await page.getByTestId("tool-mode-button").click();
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "visible",
+    });
+
+    expect(await page.getByText("toolset").count()).toBeGreaterThan(0);
+
+    await page.getByTestId("tool-mode-button").click();
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "hidden",
+    });
+
+    expect(await page.getByText("toolset").count()).toBe(0);
+
+    await page.getByTestId("tool-mode-button").click();
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "visible",
+    });
+
+    expect(await page.getByText("toolset").count()).toBeGreaterThan(0);
+
+    await page.getByTestId("tool-mode-button").click();
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "hidden",
+    });
+
+    expect(await page.getByText("toolset").count()).toBe(0);
+
+    await page.getByTestId("tool-mode-button").click();
+
+    await page.waitForSelector("text=toolset", {
+      timeout: 3000,
+      state: "visible",
+    });
+
+    await page.getByTestId("disclosure-vector stores").click();
+
+    await page.getByTestId("disclosure-agents").click();
+
+    await page.waitForSelector('[data-testid="agentsAgent"]', {
+      timeout: 3000,
+      state: "visible",
+    });
+    await page
+      .getByTestId("agentsAgent")
+      .hover()
+      .then(async () => {
+        await page.getByTestId("add-component-button-agent").click();
+      });
+
+    await page
+      .getByTestId("handle-astradb-shownode-toolset-right")
+      .first()
+      .click();
+
+    await page.getByTestId("handle-agent-shownode-tools-left").first().click();
+
+    expect(await page.locator(".react-flow__edge").count()).toBeGreaterThan(0);
+  },
+);