fix: Fix tool mode switch state persistence in NodeToolbarComponent (#5316)

✨ (toggleShadComponent/index.tsx): Refactor ToggleShadComponent to wrap Switch component in a div to prevent event propagation
📝 (nodeToolbarComponent/index.tsx): Add ShortcutDisplay component to display keyboard shortcuts for tool mode button
📝 (nodeToolbarComponent/index.tsx): Remove unused imports (CodeAreaModal, ConfirmationModal, EditNodeModal, ShareModal) from NodeToolbarComponent
📝 (nodeToolbarComponent/index.tsx): Remove duplicated declaration of updateNodeInternals function in NodeToolbarComponent
📝 (nodeToolbarComponent/index.tsx): Refactor handleActivateToolMode function to update tool mode value and node data more efficiently
📝 (nodeToolbarComponent/index.tsx): Refactor handleNodeClass function to handle node class changes more efficiently
📝 (nodeToolbarComponent/index.tsx): Refactor postToolModeValue function to post tool mode value to API more efficiently
📝 (nodeToolbarComponent/index.tsx): Refactor renderToolbarButtons to include ShortcutDisplay component for tool mode button and improve button functionality

✨ (tool-mode.spec.ts): Add test for user interaction with components as tools in the application to ensure proper functionality and user experience.

src/frontend/src/components/core/parameterRenderComponent/components/toggleShadComponent/index.tsx

@@ -36,21 +36,23 @@ export default function ToggleShadComponent({
   }
 
   return (
-    <Switch
+    <div onClick={(e) => e.stopPropagation()}>
-      id={id}
+      <Switch
-      data-testid={id}
+        id={id}
-      style={{
+        data-testid={id}
-        transform: `scaleX(${scaleX}) scaleY(${scaleY})`,
+        style={{
-      }}
+          transform: `scaleX(${scaleX}) scaleY(${scaleY})`,
-      disabled={disabled}
+        }}
-      className=""
+        disabled={disabled}
-      checked={value}
+        className=""
-      onCheckedChange={(isEnabled: boolean) => {
+        checked={value}
-        const data = showToogle
+        onCheckedChange={(isEnabled: boolean) => {
-          ? { advanced: !isEnabled }
+          const data = showToogle
-          : { value: isEnabled };
+            ? { advanced: !isEnabled }
-        handleOnNewValue(data);
+            : { value: isEnabled };
-      }}
+          handleOnNewValue(data);
-    ></Switch>
+        }}
+      />
+    </div>
   );
 }

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 { handleNodeClass: handleNodeClassHook } = useHandleNodeClass(
-      const newValue = !toolMode;
+      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
+            <ShadTooltip
-              icon="Hammer"
+              content={
-              label="Tool Mode"
+                <ShortcutDisplay
-              onClick={() => {
+                  {...shortcuts.find(
-                takeSnapshot();
+                    ({ name }) => name.toLowerCase() === "tool mode",
-                handleSelectChange("toolMode");
+                  )!}
-              }}
+                />
-              shortcut={shortcuts.find((s) =>
+              }
-                s.name.toLowerCase().startsWith("tool mode"),
+              side="top"
-              )}
+            >
-              className={cn(
+              <Button
-                "node-toolbar-buttons h-[2rem]",
+                className={cn(
-                toolMode && "text-primary",
+                  "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>
           )}
         </>
       ),

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);
+  },
+);