diff --git a/src/backend/base/langflow/api/v1/endpoints.py b/src/backend/base/langflow/api/v1/endpoints.py
index c12c2f470..1fcaf1b36 100644
--- a/src/backend/base/langflow/api/v1/endpoints.py
+++ b/src/backend/base/langflow/api/v1/endpoints.py
@@ -615,6 +615,9 @@ async def custom_component_update(
             component,
             user_id=user.id,
         )
+
+        component_node["tool_mode"] = code_request.tool_mode
+
         if hasattr(cc_instance, "set_attributes"):
             template = code_request.get_template()
             params = {}
diff --git a/src/backend/base/langflow/api/v1/schemas.py b/src/backend/base/langflow/api/v1/schemas.py
index 87b137bfd..20ab9c668 100644
--- a/src/backend/base/langflow/api/v1/schemas.py
+++ b/src/backend/base/langflow/api/v1/schemas.py
@@ -204,6 +204,7 @@ class UpdateCustomComponentRequest(CustomComponentRequest):
     field: str
     field_value: str | int | float | bool | dict | list | None = None
     template: dict
+    tool_mode: bool = False
 
     def get_template(self):
         return dotdict(self.template)
diff --git a/src/backend/base/langflow/custom/custom_component/component.py b/src/backend/base/langflow/custom/custom_component/component.py
index 96a4cd3b0..d76e47adc 100644
--- a/src/backend/base/langflow/custom/custom_component/component.py
+++ b/src/backend/base/langflow/custom/custom_component/component.py
@@ -396,10 +396,10 @@ class Component(CustomComponent):
 
     def run_and_validate_update_outputs(self, frontend_node: dict, field_name: str, field_value: Any):
         frontend_node = self.update_outputs(frontend_node, field_name, field_value)
-        if field_name == "tool_mode":
-            # Replace all outputs with the tool_output value if tool_mode is True
-            # else replace it with the original outputs
-            frontend_node["outputs"] = [self._build_tool_output()] if field_value else frontend_node["outputs"]
+        if field_name == "tool_mode" or frontend_node.get("tool_mode"):
+            is_tool_mode = field_value or frontend_node.get("tool_mode")
+            frontend_node["outputs"] = [self._build_tool_output()] if is_tool_mode else frontend_node["outputs"]
+
         return self._validate_frontend_node(frontend_node)
 
     def _validate_frontend_node(self, frontend_node: dict):
diff --git a/src/frontend/src/CustomNodes/GenericNode/components/NodeInputField/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/NodeInputField/index.tsx
index ad5cc66f5..c2fc549be 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/NodeInputField/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/NodeInputField/index.tsx
@@ -49,6 +49,7 @@ export default function NodeInputField({
     node: data.node!,
     nodeId: data.id,
     parameterId: name,
+    tool_mode: data.node!.tool_mode ?? false,
   });
   const setFilterEdge = useFlowStore((state) => state.setFilterEdge);
   const { handleNodeClass } = useHandleNodeClass(data.id);
diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index e6a8ad6a8..b2388ebfe 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -17,7 +17,10 @@ import { useShortcutsStore } from "../../stores/shortcuts";
 import { useTypesStore } from "../../stores/typesStore";
 import { OutputFieldType, VertexBuildTypeAPI } from "../../types/api";
 import { NodeDataType } from "../../types/flow";
-import { scapedJSONStringfy } from "../../utils/reactflowUtils";
+import {
+  checkHasToolMode,
+  scapedJSONStringfy,
+} from "../../utils/reactflowUtils";
 import { classNames, cn } from "../../utils/utils";
 import { getNodeInputColors } from "../helpers/get-node-input-colors";
 import { getNodeInputColorsName } from "../helpers/get-node-input-colors-name";
@@ -332,9 +335,7 @@ export default function GenericNode({
     return null;
   };
 
-  const hasToolMode =
-    data.node?.template &&
-    Object.values(data.node.template).some((field) => field.tool_mode);
+  const hasToolMode = checkHasToolMode(data.node?.template ?? {});
 
   return (
     <div className={cn(isOutdated && !isUserEdited ? "relative -mt-10" : "")}>
diff --git a/src/frontend/src/CustomNodes/helpers/mutate-template.ts b/src/frontend/src/CustomNodes/helpers/mutate-template.ts
index 2967d885a..b1bd6e0c4 100644
--- a/src/frontend/src/CustomNodes/helpers/mutate-template.ts
+++ b/src/frontend/src/CustomNodes/helpers/mutate-template.ts
@@ -19,6 +19,7 @@ export const mutateTemplate = debounce(
     >,
     setErrorData,
     parameterName?: string,
+    callback?: () => void,
   ) => {
     try {
       const newNode = cloneDeep(node);
@@ -31,6 +32,7 @@ export const mutateTemplate = debounce(
         newNode.outputs = newTemplate.outputs;
       }
       setNodeClass(newNode);
+      callback?.();
     } catch (e) {
       const error = e as ResponseErrorDetailAPI;
       setErrorData({
diff --git a/src/frontend/src/CustomNodes/hooks/use-handle-new-value.tsx b/src/frontend/src/CustomNodes/hooks/use-handle-new-value.tsx
index 9ecb1dc53..0244ef209 100644
--- a/src/frontend/src/CustomNodes/hooks/use-handle-new-value.tsx
+++ b/src/frontend/src/CustomNodes/hooks/use-handle-new-value.tsx
@@ -41,6 +41,7 @@ const useHandleOnNewValue = ({
     parameterId: name,
     nodeId: nodeId,
     node: node,
+    tool_mode: node.tool_mode ?? false,
   });
 
   const handleOnNewValue: handleOnNewValueType = async (changes, options?) => {
diff --git a/src/frontend/src/components/core/parameterRenderComponent/components/inputComponent/components/popover/index.tsx b/src/frontend/src/components/core/parameterRenderComponent/components/inputComponent/components/popover/index.tsx
index 8147ce84c..f13a46127 100644
--- a/src/frontend/src/components/core/parameterRenderComponent/components/inputComponent/components/popover/index.tsx
+++ b/src/frontend/src/components/core/parameterRenderComponent/components/inputComponent/components/popover/index.tsx
@@ -62,6 +62,8 @@ const CustomInputPopover = ({
     }
   };
 
+  console.log(placeholder, value, id);
+
   return (
     <Popover modal open={showOptions} onOpenChange={setShowOptions}>
       <PopoverAnchor>
diff --git a/src/frontend/src/components/core/parameterRenderComponent/components/inputGlobalComponent/index.tsx b/src/frontend/src/components/core/parameterRenderComponent/components/inputGlobalComponent/index.tsx
index 8decc5aaa..fd9075311 100644
--- a/src/frontend/src/components/core/parameterRenderComponent/components/inputGlobalComponent/index.tsx
+++ b/src/frontend/src/components/core/parameterRenderComponent/components/inputGlobalComponent/index.tsx
@@ -69,6 +69,7 @@ export default function InputGlobalComponent({
       });
     }
   }
+
   return (
     <InputComponent
       nodeStyle
diff --git a/src/frontend/src/components/core/parameterRenderComponent/components/refreshParameterComponent/index.tsx b/src/frontend/src/components/core/parameterRenderComponent/components/refreshParameterComponent/index.tsx
index ed11de425..c548ab1e3 100644
--- a/src/frontend/src/components/core/parameterRenderComponent/components/refreshParameterComponent/index.tsx
+++ b/src/frontend/src/components/core/parameterRenderComponent/components/refreshParameterComponent/index.tsx
@@ -30,6 +30,7 @@ export function RefreshParameterComponent({
     parameterId: name,
     nodeId: nodeId,
     node: nodeClass,
+    tool_mode: nodeClass.tool_mode ?? false,
   });
 
   const setErrorData = useAlertStore((state) => state.setErrorData);
diff --git a/src/frontend/src/controllers/API/queries/nodes/use-post-template-value.ts b/src/frontend/src/controllers/API/queries/nodes/use-post-template-value.ts
index 40d751078..98b518f67 100644
--- a/src/frontend/src/controllers/API/queries/nodes/use-post-template-value.ts
+++ b/src/frontend/src/controllers/API/queries/nodes/use-post-template-value.ts
@@ -10,12 +10,14 @@ import { UseRequestProcessor } from "../../services/request-processor";
 
 interface IPostTemplateValue {
   value: any;
+  tool_mode?: boolean;
 }
 
 interface IPostTemplateValueParams {
   node: APIClassType;
   nodeId: string;
   parameterId: string;
+  tool_mode: boolean;
 }
 
 export const usePostTemplateValue: useMutationFunctionType<
@@ -23,7 +25,7 @@ export const usePostTemplateValue: useMutationFunctionType<
   IPostTemplateValue,
   APIClassType,
   ResponseErrorDetailAPI
-> = ({ parameterId, nodeId, node }, options?) => {
+> = ({ parameterId, nodeId, node, tool_mode }, options?) => {
   const { mutate } = UseRequestProcessor();
 
   const postTemplateValueFn = async (
@@ -32,7 +34,6 @@ export const usePostTemplateValue: useMutationFunctionType<
     const template = node.template;
 
     if (!template) return;
-
     const response = await api.post<APIClassType>(
       getURL("CUSTOM_COMPONENT", { update: "update" }),
       {
@@ -40,6 +41,7 @@ export const usePostTemplateValue: useMutationFunctionType<
         template: template,
         field: parameterId,
         field_value: payload.value,
+        tool_mode: tool_mode,
       },
     );
 
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 3cd3f3303..7796e21fd 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -5,6 +5,7 @@ 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";
@@ -32,6 +33,7 @@ import { useStoreStore } from "../../../../stores/storeStore";
 import { nodeToolbarPropsType } from "../../../../types/components";
 import { FlowType } from "../../../../types/flow";
 import {
+  checkHasToolMode,
   createFlowComponent,
   downloadNode,
   expandGroupNode,
@@ -71,15 +73,20 @@ export default function NodeToolbarComponent({
   const [openModal, setOpenModal] = useState(false);
   const isGroup = data.node?.flow ? true : false;
   const frozen = data.node?.frozen ?? false;
+  const currentFlow = useFlowStore((state) => state.currentFlow);
 
   const addFlow = useAddFlow();
 
+  const { mutate: patchUpdateFlow } = usePatchUpdateFlow();
+
   const isMinimal = countHandlesFn(data) <= 1 && numberOfOutputHandles <= 1;
   function activateToolMode() {
     const newValue = !toolMode;
     setToolMode(newValue);
 
     updateToolMode(data.id, newValue);
+    data.node!.tool_mode = newValue;
+
     mutateTemplate(
       newValue,
       data.node!,
@@ -87,7 +94,24 @@ export default function NodeToolbarComponent({
       postToolModeValue,
       setErrorData,
       "tool_mode",
+      () => {
+        const node = currentFlow?.data?.nodes.find(
+          (node) => node.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!,
+          data: currentFlow?.data!,
+          description: currentFlow?.description!,
+          folder_id: currentFlow?.folder_id!,
+          endpoint_name: currentFlow?.endpoint_name!,
+        });
+      },
     );
+
     updateNodeInternals(data.id);
   }
   function minimize() {
@@ -145,9 +169,7 @@ export default function NodeToolbarComponent({
   }
   // Check if any of the data.node.template fields have tool_mode as True
   // if so we can show the tool mode button
-  const hasToolMode =
-    data.node?.template &&
-    Object.values(data.node.template).some((field) => field.tool_mode);
+  const hasToolMode = checkHasToolMode(data.node?.template ?? {});
 
   function openDocs() {
     if (data.node?.documentation) {
@@ -370,6 +392,7 @@ export default function NodeToolbarComponent({
     node: data.node!,
     nodeId: data.id,
     parameterId: "tool_mode",
+    tool_mode: data.node!.tool_mode ?? false,
   });
 
   return (
diff --git a/src/frontend/src/utils/reactflowUtils.ts b/src/frontend/src/utils/reactflowUtils.ts
index fc902c36b..7386e59ff 100644
--- a/src/frontend/src/utils/reactflowUtils.ts
+++ b/src/frontend/src/utils/reactflowUtils.ts
@@ -1736,3 +1736,7 @@ export function someFlowTemplateFields(
     });
   });
 }
+
+export function checkHasToolMode(template: APITemplateType) {
+  return template && Object.values(template).some((field) => field.tool_mode);
+}