diff --git a/src/backend/base/langflow/components/helpers/id_generator.py b/src/backend/base/langflow/components/helpers/id_generator.py
index fd9393841..ebfbcb2e6 100644
--- a/src/backend/base/langflow/components/helpers/id_generator.py
+++ b/src/backend/base/langflow/components/helpers/id_generator.py
@@ -21,6 +21,7 @@ class IDGeneratorComponent(Component):
             display_name="Value",
             info="The generated unique ID.",
             refresh_button=True,
+            tool_mode=True,
         ),
     ]
 
diff --git a/src/frontend/tests/extended/regression/general-bugs-component-as-tool-shortcut.spec.ts b/src/frontend/tests/extended/regression/general-bugs-component-as-tool-shortcut.spec.ts
index bc6312aec..1ec1e5ccf 100644
--- a/src/frontend/tests/extended/regression/general-bugs-component-as-tool-shortcut.spec.ts
+++ b/src/frontend/tests/extended/regression/general-bugs-component-as-tool-shortcut.spec.ts
@@ -32,7 +32,12 @@ test(
     await page.getByTestId("title-ID Generator").click();
     await page.keyboard.press("ControlOrMeta+Shift+m");
 
-    expect(await page.getByText("Toolset", { exact: true }).count()).toBe(0);
+    await page.waitForSelector('text="Toolset"', {
+      timeout: 3000,
+    });
+    expect(
+      await page.getByText("Toolset", { exact: true }).count(),
+    ).toBeGreaterThan(0);
 
     await page.getByTestId("title-ID Generator").click();
 
@@ -43,29 +48,26 @@ test(
     await page.getByTestId("code-button-modal").click();
 
     let code = await extractAndCleanCode(page);
-    code = code!.replace(
-      "refresh_button=True,",
-      "refresh_button=True,\n tool_mode=True,",
-    );
+    let updatedCode = code!.replace("tool_mode=True", "tool_mode=False");
+
+    expect(updatedCode).not.toBe(code);
 
     await page.locator("textarea").last().press(`ControlOrMeta+a`);
     await page.keyboard.press("Backspace");
-    await page.locator("textarea").last().fill(code);
+    await page.locator("textarea").last().fill(updatedCode);
+    const customComponentPromise = page.waitForResponse("**/custom_component");
     await page.locator('//*[@id="checkAndSaveBtn"]').click();
+    const customComponentResponse = await customComponentPromise;
+    // check if the response is 200
+    expect(customComponentResponse?.status()).toBe(200);
 
-    await page.waitForSelector('[data-testid="fit_view"]', {
+    await page.waitForSelector('[data-testid="title-ID Generator"]', {
       timeout: 3000,
     });
 
     await page.getByTestId("title-ID Generator").click();
     await page.keyboard.press("ControlOrMeta+Shift+m");
 
-    await page.waitForSelector('text="Toolset"', {
-      timeout: 3000,
-    });
-
-    expect(
-      await page.getByText("Toolset", { exact: true }).count(),
-    ).toBeGreaterThan(0);
+    expect(await page.getByText("Toolset", { exact: true }).count()).toBe(0);
   },
 );