fix(mcp): stabilize Tool Mode activation (#14188)

* fix(mcp): stabilize tool mode activation

* fix(frontend): prevent stale tool mode refresh

Author: erichare

src/backend/tests/unit/components/models_and_agents/test_mcp_component_dynamic.py

@@ -22,6 +22,13 @@
 from lfx.components.models_and_agents.mcp_component import MCPToolsComponent
 
 
+def test_tool_mode_capability_is_available_before_server_discovery() -> None:
+    """The toolbar must not wait for an MCP connection to discover Tool Mode."""
+    tool_placeholder = next(input_ for input_ in MCPToolsComponent.inputs if input_.name == "tool_placeholder")
+
+    assert tool_placeholder.tool_mode is True
+
+
 def _make_tool(name: str) -> MagicMock:
     tool = MagicMock()
     tool.name = name
@@ -324,6 +331,14 @@ def _build_config(**overrides):
         config.update(overrides)
         return config
 
+    @pytest.mark.asyncio
+    async def test_clearing_server_keeps_tool_mode_capability(self) -> None:
+        component = MCPToolsComponent()
+
+        build_config = await component.update_build_config(self._build_config(), {}, "mcp_server")
+
+        assert build_config["tool_placeholder"]["tool_mode"] is True
+
     @pytest.mark.asyncio
     async def test_refresh_bypasses_existing_options(self) -> None:
         component = MCPToolsComponent()

src/frontend/src/CustomNodes/helpers/__tests__/mutate-template.test.ts

@@ -0,0 +1,92 @@
+import type { APIClassType } from "@/types/api";
+import { mutateTemplate } from "../mutate-template";
+
+describe("mutateTemplate", () => {
+  afterEach(() => {
+    jest.useRealTimers();
+  });
+
+  it("sends Tool Mode changes immediately", async () => {
+    const node = {
+      template: {
+        code: { value: "component source" },
+      },
+      outputs: [],
+      tool_mode: false,
+    } as unknown as APIClassType;
+    const updatedNode = {
+      template: node.template,
+      outputs: [],
+      last_updated: "2026-07-21T16:30:00.000Z",
+    } as unknown as APIClassType;
+    const mutateAsync = jest.fn().mockResolvedValue(updatedNode);
+    const setNodeClass = jest.fn();
+
+    await mutateTemplate(
+      true,
+      "mcp-tools-node",
+      node,
+      setNodeClass,
+      { mutateAsync } as never,
+      jest.fn(),
+      "tool_mode",
+      jest.fn(),
+      true,
+    );
+
+    expect(mutateAsync).toHaveBeenCalledWith(
+      expect.objectContaining({ value: true, tool_mode: true }),
+    );
+    expect(setNodeClass).toHaveBeenCalledWith(
+      expect.objectContaining({ tool_mode: true }),
+    );
+  });
+
+  it("cancels a stale Toolset metadata refresh when Tool Mode changes", async () => {
+    jest.useFakeTimers();
+    const node = {
+      template: {
+        code: { value: "component source" },
+        tools_metadata: { value: [{ name: "fetch_content" }] },
+      },
+      outputs: [{ name: "component_as_tool" }],
+      tool_mode: true,
+    } as unknown as APIClassType;
+    const metadataMutateAsync = jest.fn();
+    const toolModeMutateAsync = jest.fn().mockResolvedValue({
+      template: node.template,
+      outputs: [],
+      last_updated: "2026-07-21T16:31:00.000Z",
+    });
+
+    await mutateTemplate(
+      node.template.tools_metadata.value,
+      "url-node",
+      node,
+      jest.fn(),
+      { mutateAsync: metadataMutateAsync } as never,
+      jest.fn(),
+      "tools_metadata",
+      jest.fn(),
+      true,
+    );
+
+    await mutateTemplate(
+      false,
+      "url-node",
+      node,
+      jest.fn(),
+      { mutateAsync: toolModeMutateAsync } as never,
+      jest.fn(),
+      "tool_mode",
+      jest.fn(),
+      false,
+    );
+    await jest.runAllTimersAsync();
+
+    expect(toolModeMutateAsync).toHaveBeenCalledWith(
+      expect.objectContaining({ value: false, tool_mode: false }),
+    );
+    expect(metadataMutateAsync).not.toHaveBeenCalled();
+  });
+});

src/frontend/src/CustomNodes/helpers/mutate-template.ts

@@ -1,8 +1,8 @@
 import type { UseMutationResult } from "@tanstack/react-query";
 import { cloneDeep, debounce } from "lodash";
 import { SAVE_DEBOUNCE_TIME } from "@/constants/constants";
-import i18n from "../../i18n";
 import type { APIClassType, ResponseErrorDetailAPI } from "@/types/api";
+import i18n from "../../i18n";
 import { updateHiddenOutputs } from "./update-hidden-outputs";
 
 // Map to store debounced functions for each node ID + parameter combination
@@ -16,6 +16,7 @@ export const mutateTemplate = async (
   postTemplateValue: UseMutationResult<
     APIClassType | undefined,
     ResponseErrorDetailAPI,
+    // biome-ignore lint/suspicious/noExplicitAny: legacy mutation payload
     any
   >,
   setErrorData,
@@ -38,6 +39,7 @@ export const mutateTemplate = async (
           postTemplateValue: UseMutationResult<
             APIClassType | undefined,
             ResponseErrorDetailAPI,
+            // biome-ignore lint/suspicious/noExplicitAny: legacy mutation payload
             any
           >,
           setErrorData,
@@ -89,7 +91,17 @@ export const mutateTemplate = async (
     );
   }
 
-  debouncedFunctions.get(debounceKey)?.(
+  // Enabling Tool Mode mounts the tools_metadata field, which queues its own
+  // debounced refresh. If the user turns Tool Mode off before that refresh
+  // runs, the queued request still carries tool_mode=true and can restore the
+  // Toolset output after the off response. The explicit toggle supersedes that
+  // pending metadata refresh.
+  if (parameterName === "tool_mode") {
+    debouncedFunctions.get(`${nodeId}-tools_metadata`)?.cancel();
+  }
+
+  const debouncedFunction = debouncedFunctions.get(debounceKey);
+  debouncedFunction?.(
     newValue,
     node,
     setNodeClass,
@@ -100,4 +112,11 @@ export const mutateTemplate = async (
     toolMode,
     isRefresh,
   );
+
+  // Tool Mode is a discrete toggle, so delaying it like a text input leaves
+  // the node in its previous output shape and gives slower refresh responses
+  // a chance to repaint the toggle with stale state.
+  if (parameterName === "tool_mode") {
+    await debouncedFunction?.flush();
+  }
 };

src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/__tests__/config-transition.test.tsx

@@ -1,9 +1,12 @@
-import { act, render, screen } from "@testing-library/react";
+import { act, fireEvent, render, screen } from "@testing-library/react";
 import { useUtilityStore } from "@/stores/utilityStore";
 import NodeToolbarComponent from "../index";
 
 const mockFreezeAllVertices = jest.fn();
 const mockAddFlow = jest.fn();
+const mockCheckHasToolMode = jest.fn(() => false);
+const mockMutateTemplate = jest.fn();
+const mockPostToolModeValue = { isPending: false };
 const mockSetNoticeData = jest.fn();
 const mockSetErrorData = jest.fn();
 const mockSetSuccessData = jest.fn();
@@ -12,6 +15,10 @@ jest.mock("@xyflow/react", () => ({
   useUpdateNodeInternals: () => jest.fn(),
 }));
 
+jest.mock("@/CustomNodes/helpers/mutate-template", () => ({
+  mutateTemplate: (...args: unknown[]) => mockMutateTemplate(...args),
+}));
+
 jest.mock("@/CustomNodes/hooks/use-handle-new-value", () => ({
   __esModule: true,
   default: () => ({
@@ -59,7 +66,7 @@ jest.mock("@/components/ui/button", () => ({
 }));
 
 jest.mock("@/controllers/API/queries/nodes/use-post-template-value", () => ({
-  usePostTemplateValue: () => ({}),
+  usePostTemplateValue: () => mockPostToolModeValue,
 }));
 
 jest.mock("@/controllers/API/queries/vertex", () => ({
@@ -160,7 +167,7 @@ jest.mock("../../../../../components/ui/select-custom", () => ({
 }));
 
 jest.mock("../../../../../utils/reactflowUtils", () => ({
-  checkHasToolMode: jest.fn(() => false),
+  checkHasToolMode: () => mockCheckHasToolMode(),
   createFlowComponent: jest.fn((data) => ({
     name: data.id,
   })),
@@ -217,8 +224,16 @@ const getProps = () => ({
     node: {
       display_name: "Prompt",
       description: "Prompt node",
+      documentation: "",
       template: {
-        code: { value: "print('hello')" },
+        code: {
+          value: "print('hello')",
+          type: "code",
+          required: true,
+          list: false,
+          show: true,
+          readonly: false,
+        },
       },
       outputs: [],
       frozen: false,
@@ -237,6 +252,8 @@ const getProps = () => ({
 describe("NodeToolbarComponent config transitions", () => {
   beforeEach(() => {
     jest.clearAllMocks();
+    mockCheckHasToolMode.mockReturnValue(false);
+    mockPostToolModeValue.isPending = false;
     act(() => {
       useUtilityStore.setState({ allowCustomComponents: false });
     });
@@ -263,4 +280,40 @@ describe("NodeToolbarComponent config transitions", () => {
 
     expect(screen.queryByTestId("code-button-modal")).not.toBeInTheDocument();
   });
+
+  it("keeps an optimistic Tool Mode toggle on while its update is pending", () => {
+    mockCheckHasToolMode.mockReturnValue(true);
+    const props = {
+      ...getProps(),
+      data: {
+        ...getProps().data,
+        node: {
+          ...getProps().data.node,
+          tool_mode: false,
+        },
+      },
+    };
+    const { rerender } = render(<NodeToolbarComponent {...props} />);
+    const toolModeButton = screen.getByTestId("tool-mode-button");
+
+    fireEvent.click(screen.getByText("Tool Mode"));
+    expect(toolModeButton).toHaveClass("text-primary");
+
+    mockPostToolModeValue.isPending = true;
+    rerender(
+      <NodeToolbarComponent
+        {...props}
+        data={{
+          ...props.data,
+          node: {
+            ...props.data.node,
+            outputs: [],
+            tool_mode: false,
+          },
+        }}
+      />,
+    );
+
+    expect(screen.getByTestId("tool-mode-button")).toHaveClass("text-primary");
+  });
 });

src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx

@@ -139,6 +139,12 @@ const NodeToolbarComponent = memo(
     );
 
     useEffect(() => {
+      // Keep the optimistic toggle state while the server rebuilds the node.
+      // Other in-flight field refreshes can update the same node first with
+      // its previous tool_mode value; syncing that transient value makes the
+      // toggle appear to turn itself off.
+      if (postToolModeValue.isPending) return;
+
       if (data.node?.tool_mode !== undefined) {
         setToolMode(
           data.node?.tool_mode ||
@@ -148,7 +154,7 @@ const NodeToolbarComponent = memo(
             false,
         );
       }
-    }, [data.node?.tool_mode, data.node?.outputs]);
+    }, [data.node?.tool_mode, data.node?.outputs, postToolModeValue.isPending]);
 
     const { handleNodeClass: handleNodeClassHook } = useHandleNodeClass(
       data.id,

src/lfx/src/lfx/_assets/component_index.json

@@ -316,7 +316,10 @@ def map_outputs(self) -> None:
             info="Placeholder for the tool",
             value="",
             show=False,
-            tool_mode=False,
+            # Tool Mode is a capability of the component, not of the selected
+            # server. Advertising it up front keeps the toolbar from waiting
+            # for the MCP server's tool-discovery request to finish.
+            tool_mode=True,
         ),
     ]
 
@@ -663,7 +666,7 @@ async def update_build_config(self, build_config: dict, field_value: str, field_
                     build_config["tool"]["options"] = []
                     build_config["tool"]["value"] = ""
                     build_config["tool"]["placeholder"] = ""
-                    build_config["tool_placeholder"]["tool_mode"] = False
+                    build_config["tool_placeholder"]["tool_mode"] = True
                     self.remove_non_default_keys(build_config)
                     return build_config