From 4a09655f2f41acfe2bc610f3b3bcfe00a8817788 Mon Sep 17 00:00:00 2001 From: Lucas Oliveira <62335616+lucaseduoli@users.noreply.github.com> Date: Mon, 23 Jun 2025 11:15:22 -0300 Subject: [PATCH] fix: implemented cached values and temporary MCP servers on MCP component (#8628) * Added actionCount to fetch only servers without actionCount * Updated queries and uses to use servers without action data first, and then to fetch them * removed comment * updated constants * Added loading dropdown * Make options persist * Implemented new value format for McpComponent and implemented saving and removing temp Mcp Server if config is existent * Changed value type * Implemented cache and saving the server config * Fixed mcp server test * fix backend formatting * fixed lint * Added await * Fixed save button not appearing when no servers are available * added condition to only show save button when options is not null --- src/backend/base/langflow/api/v2/mcp.py | 10 +- .../langflow/components/data/mcp_component.py | 70 +++++-- src/backend/base/langflow/inputs/inputs.py | 2 +- .../components/mcpComponent/index.tsx | 189 +++++++++++++----- .../API/queries/mcp/use-get-mcp-servers.ts | 53 ++++- .../src/modals/addMcpServerModal/index.tsx | 7 +- .../pages/MCPServersPage/index.tsx | 5 +- src/frontend/src/types/mcp/index.ts | 7 +- .../extended/features/mcp-server.spec.ts | 22 ++ 9 files changed, 282 insertions(+), 83 deletions(-) diff --git a/src/backend/base/langflow/api/v2/mcp.py b/src/backend/base/langflow/api/v2/mcp.py index cd743ef88..de650a0f3 100644 --- a/src/backend/base/langflow/api/v2/mcp.py +++ b/src/backend/base/langflow/api/v2/mcp.py @@ -101,27 +101,29 @@ async def get_servers( session: DbSession, storage_service=Depends(get_storage_service), settings_service=Depends(get_settings_service), + action_count: bool | None = None, ): """Get the list of available servers.""" import asyncio server_list = await get_server_list(current_user, session, storage_service, settings_service) + if not action_count: + # Return only the server names, with mode and toolsCount as None + return [{"name": server_name, "mode": None, "toolsCount": None} for server_name in server_list["mcpServers"]] + # Check all of the tool counts for each server concurrently async def check_server(server_name: str) -> dict: - server_info = {"name": server_name, "mode": "", "toolsCount": 0} + server_info: dict[str, str | int | None] = {"name": server_name, "mode": None, "toolsCount": None} try: mode, tool_list, _ = await update_tools( server_name=server_name, server_config=server_list["mcpServers"][server_name], ) - - # Get the server configuration server_info["mode"] = mode.lower() server_info["toolsCount"] = len(tool_list) except Exception as e: # noqa: BLE001 logger.exception(f"Error checking server {server_name}: {e}") - return server_info # Run all server checks concurrently diff --git a/src/backend/base/langflow/components/data/mcp_component.py b/src/backend/base/langflow/components/data/mcp_component.py index 41a29aec9..685709143 100644 --- a/src/backend/base/langflow/components/data/mcp_component.py +++ b/src/backend/base/langflow/components/data/mcp_component.py @@ -8,13 +8,14 @@ from langflow.base.mcp.util import ( create_input_schema_from_json_schema, update_tools, ) -from langflow.custom.custom_component.component import Component +from langflow.custom.custom_component.component_with_cache import ComponentWithCache from langflow.inputs.inputs import InputTypes from langflow.io import DropdownInput, McpInput, MessageTextInput, Output # Import McpInput from langflow.io from langflow.io.schema import flatten_schema, schema_to_langflow_inputs from langflow.logging import logger from langflow.schema.dataframe import DataFrame from langflow.services.auth.utils import create_user_longterm_token +from langflow.services.cache.utils import CacheMiss # Import get_server from the backend API from langflow.services.database.models.user.crud import get_user_by_id @@ -59,7 +60,7 @@ def maybe_unflatten_dict(flat: dict[str, Any]) -> dict[str, Any]: return nested -class MCPToolsComponent(Component): +class MCPToolsComponent(ComponentWithCache): schema_inputs: list = [] stdio_client: MCPStdioClient = MCPStdioClient() sse_client: MCPSseClient = MCPSseClient() @@ -136,17 +137,35 @@ class MCPToolsComponent(Component): else: return schema_inputs - async def update_tool_list(self): - server_name = getattr(self, "mcp_server", None) + async def update_tool_list(self, mcp_server_value=None): + # Accepts mcp_server_value as dict {name, config} or uses self.mcp_server + mcp_server = mcp_server_value if mcp_server_value is not None else getattr(self, "mcp_server", None) + server_name = None + server_config_from_value = None + if isinstance(mcp_server, dict): + server_name = mcp_server.get("name") + server_config_from_value = mcp_server.get("config") + else: + server_name = mcp_server if not server_name: self.tools = [] - return [] + return [], {"name": server_name, "config": server_config_from_value} + + # Use shared cache if available + cached = self._shared_component_cache.get(server_name) + if not isinstance(cached, CacheMiss): + self.tools = cached["tools"] + self.tool_names = cached["tool_names"] + self._tool_cache = cached["tool_cache"] + server_config_from_value = cached["config"] + return self.tools, {"name": server_name, "config": server_config_from_value} try: async for db in get_session(): user_id, _ = await create_user_longterm_token(db) current_user = await get_user_by_id(db, user_id) + # Try to get server config from DB/API server_config = await get_server( server_name, current_user, @@ -155,9 +174,13 @@ class MCPToolsComponent(Component): settings_service=get_settings_service(), ) + # If get_server returns empty but we have a config, use it + if not server_config and server_config_from_value: + server_config = server_config_from_value + if not server_config: self.tools = [] - return [] + return [], {"name": server_name, "config": server_config} _, tool_list, tool_cache = await update_tools( server_name=server_name, @@ -168,7 +191,18 @@ class MCPToolsComponent(Component): self.tool_names = [tool.name for tool in tool_list if hasattr(tool, "name")] self._tool_cache = tool_cache - return tool_list + self.tools = tool_list + # Cache the result using shared cache + self._shared_component_cache.set( + server_name, + { + "tools": tool_list, + "tool_names": self.tool_names, + "tool_cache": tool_cache, + "config": server_config, + }, + ) + return tool_list, {"name": server_name, "config": server_config} except Exception as e: msg = f"Error updating tool list: {e!s}" logger.exception(msg) @@ -181,7 +215,7 @@ class MCPToolsComponent(Component): try: if len(self.tools) == 0: try: - self.tools = await self.update_tool_list() + self.tools, build_config["mcp_server"]["value"] = await self.update_tool_list() except ValueError: build_config["tool"]["options"] = [] build_config["tool"]["value"] = "" @@ -208,7 +242,9 @@ class MCPToolsComponent(Component): return build_config elif field_name == "mcp_server": try: - self.tools = await self.update_tool_list() + # field_value is now a dict {name, config} + mcp_server_value = field_value + self.tools, build_config["mcp_server"]["value"] = await self.update_tool_list(mcp_server_value) except ValueError: if not build_config["tools_metadata"]["show"]: build_config["tool"]["show"] = True @@ -231,7 +267,7 @@ class MCPToolsComponent(Component): build_config["tool"]["value"] = "" elif field_name == "tool_mode": try: - self.tools = await self.update_tool_list() + self.tools, build_config["mcp_server"]["value"] = await self.update_tool_list() except ValueError: if not build_config["tools_metadata"]["show"]: build_config["tool"]["show"] = True @@ -294,7 +330,7 @@ class MCPToolsComponent(Component): async def _update_tool_config(self, build_config: dict, tool_name: str) -> None: """Update tool configuration with proper error handling.""" if not self.tools: - self.tools = await self.update_tool_list() + self.tools, build_config["mcp_server"]["value"] = await self.update_tool_list() if not tool_name: return @@ -361,7 +397,7 @@ class MCPToolsComponent(Component): async def build_output(self) -> DataFrame: """Build output with improved error handling and validation.""" try: - self.tools = await self.update_tool_list() + self.tools, _ = await self.update_tool_list() if self.tool != "": exec_tool = self._tool_cache[self.tool] tool_args = self.get_inputs_for_all_tools(self.tools)[self.tool] @@ -388,11 +424,5 @@ class MCPToolsComponent(Component): async def _get_tools(self): """Get cached tools or update if necessary.""" - # if not self.tools: - if not self.mcp_server: - msg = "MCP Server is not set" - self.tools = [] - self.tool_names = [] - logger.exception(msg) - - return await self.update_tool_list() + mcp_server = getattr(self, "mcp_server", None) + return await self.update_tool_list(mcp_server) diff --git a/src/backend/base/langflow/inputs/inputs.py b/src/backend/base/langflow/inputs/inputs.py index 0837addd7..72e1f92c6 100644 --- a/src/backend/base/langflow/inputs/inputs.py +++ b/src/backend/base/langflow/inputs/inputs.py @@ -633,7 +633,7 @@ class McpInput(BaseInputMixin, MetadataTraceMixin): """ field_type: SerializableFieldTypes = FieldTypes.MCP - value: str = Field(default="") + value: dict[str, Any] = Field(default_factory=dict) class LinkInput(BaseInputMixin, LinkMixin): diff --git a/src/frontend/src/components/core/parameterRenderComponent/components/mcpComponent/index.tsx b/src/frontend/src/components/core/parameterRenderComponent/components/mcpComponent/index.tsx index f3c85d78f..44645ecc7 100644 --- a/src/frontend/src/components/core/parameterRenderComponent/components/mcpComponent/index.tsx +++ b/src/frontend/src/components/core/parameterRenderComponent/components/mcpComponent/index.tsx @@ -1,5 +1,7 @@ +import { useAddMCPServer } from "@/controllers/API/queries/mcp/use-add-mcp-server"; import { useGetMCPServers } from "@/controllers/API/queries/mcp/use-get-mcp-servers"; import AddMcpServerModal from "@/modals/addMcpServerModal"; +import useAlertStore from "@/stores/alertStore"; import { useEffect, useMemo, useRef, useState } from "react"; import ListSelectionComponent from "../../../../../CustomNodes/GenericNode/components/ListSelectionComponent"; import { cn } from "../../../../../utils/utils"; @@ -15,37 +17,67 @@ export default function McpComponent({ id = "", }: InputProps): JSX.Element { const { data: mcpServers } = useGetMCPServers(); + const { mutate: addMcpServer } = useAddMCPServer(); + const setErrorData = useAlertStore((state) => state.setErrorData); const options = useMemo( () => mcpServers?.map((server) => ({ name: server.name, - description: !server.toolsCount - ? "No actions found" - : `${server.toolsCount} action${server.toolsCount === 1 ? "" : "s"}`, + description: + server.toolsCount === null + ? "Loading..." + : !server.toolsCount + ? "No actions found" + : `${server.toolsCount} action${server.toolsCount === 1 ? "" : "s"}`, })), [mcpServers], ); const [open, setOpen] = useState(false); const [addOpen, setAddOpen] = useState(false); const [selectedItem, setSelectedItem] = useState([]); + const { name, config } = useMemo( + () => value ?? { name: "", config: {} }, + [value], + ); // Initialize selected item from value on mount or value/options change + const selectedOption = useMemo( + () => + name + ? (options?.find((option) => option.name === name) ?? { name: null }) + : null, + [name, options], + ); + useEffect(() => { - const selectedOption = value - ? options?.find((option) => option.name === value) + if (!options) return; + const selectedOption = name + ? options?.find((option) => option.name === name) : null; - setSelectedItem( - selectedOption ? [{ name: selectedOption.name }] : [{ name: "" }], - ); - if (value !== selectedOption?.name) { - handleOnNewValue({ value: "" }, { skipSnapshot: true }); + + if ( + name !== selectedOption?.name && + Object.keys(config ?? {}).length === 0 + ) { + setSelectedItem( + selectedOption ? [{ name: selectedOption.name }] : [{ name: "" }], + ); + handleOnNewValue( + { value: { name: "", config: {} } }, + { skipSnapshot: true }, + ); + return; } - }, [value, options]); + setSelectedItem([{ name }]); + }, [name, options]); // Handle selection from dialog const handleSelection = (item: any) => { setSelectedItem([{ name: item.name }]); - handleOnNewValue({ value: item.name }, { skipSnapshot: true }); + handleOnNewValue( + { value: { name: item.name, config: {} } }, + { skipSnapshot: true }, + ); setOpen(false); }; @@ -53,49 +85,110 @@ export default function McpComponent({ setAddOpen(true); }; - const handleOpenListSelectionDialog = () => setOpen(true); + const handleSaveButtonClick = () => { + addMcpServer( + { + name, + ...(config ?? {}), + }, + { + onSuccess: () => { + handleSuccess(name); + }, + onError: (error) => { + setErrorData({ + title: "Error adding MCP server", + list: [error.message], + }); + }, + }, + ); + }; + + const handleRemoveButtonClick = () => { + handleOnNewValue({ value: { name: "", config: {} } }); + }; + + const handleOpenListSelectionDialog = () => { + setOpen(true); + }; const handleCloseListSelectionDialog = () => setOpen(false); const handleSuccess = (server: string) => { - handleOnNewValue({ value: server }); + handleOnNewValue({ value: { name: server, config: {} } }); setOpen(false); }; + const showSaveButton = useMemo(() => { + return ( + !selectedOption?.name && + Object.keys(config ?? {}).length > 0 && + options !== null + ); + }, [selectedOption, config]); + return (
- {options && ( - <> - {options.length > 0 ? ( - - ) : ( -
+ + {showSaveButton && ( + )} + + ) : ( + + )} + {options && ( + <> + )} - ); } diff --git a/src/frontend/src/controllers/API/queries/mcp/use-get-mcp-servers.ts b/src/frontend/src/controllers/API/queries/mcp/use-get-mcp-servers.ts index ba27f67de..4647838cb 100644 --- a/src/frontend/src/controllers/API/queries/mcp/use-get-mcp-servers.ts +++ b/src/frontend/src/controllers/API/queries/mcp/use-get-mcp-servers.ts @@ -1,21 +1,52 @@ import { useQueryFunctionType } from "@/types/api"; import { MCPServerInfoType } from "@/types/mcp"; +import { useEffect } from "react"; import { api } from "../../api"; import { getURL } from "../../helpers/constants"; import { UseRequestProcessor } from "../../services/request-processor"; +// This type is now updated to allow nulls for mode/toolsCount +// type getMCPServersResponse = Array; + type getMCPServersResponse = Array; export const useGetMCPServers: useQueryFunctionType< undefined, getMCPServersResponse > = (options) => { - const { query } = UseRequestProcessor(); + const { query, queryClient } = UseRequestProcessor(); + // First fetch: action_count=false (fast) const responseFn = async () => { try { const { data } = await api.get( - `${getURL("MCP_SERVERS", undefined, true)}`, + `${getURL("MCP_SERVERS", undefined, true)}?action_count=false`, + ); + // Merge with cached data to preserve non-null mode/toolsCount + const cachedData = queryClient.getQueryData(["useGetMCPServers"]) as + | getMCPServersResponse + | undefined; + if (cachedData && Array.isArray(cachedData)) { + const merged = data.map((server) => { + const cached = cachedData.find((s) => s.name === server.name); + return cached && (cached.toolsCount !== null || cached.mode !== null) + ? { ...server, toolsCount: cached.toolsCount, mode: cached.mode } + : server; + }); + return merged; + } + return data; + } catch (error) { + console.error(error); + return []; + } + }; + + // Second fetch: action_count=true (slow, updates mode/toolsCount) + const fetchWithCounts = async () => { + try { + const { data } = await api.get( + `${getURL("MCP_SERVERS", undefined, true)}?action_count=true`, ); return data; } catch (error) { @@ -28,5 +59,23 @@ export const useGetMCPServers: useQueryFunctionType< ...options, }); + useEffect(() => { + if (queryResult.data && queryResult.data.length > 0) { + fetchWithCounts().then((countsData) => { + if (!countsData || countsData.length === 0) return; + // Merge by name + queryClient.setQueryData( + ["useGetMCPServers"], + (oldData: getMCPServersResponse = []) => { + return oldData.map((server) => { + const updated = countsData.find((s) => s.name === server.name); + return updated ? { ...server, ...updated } : server; + }); + }, + ); + }); + } + }, [queryResult.data]); + return queryResult; }; diff --git a/src/frontend/src/modals/addMcpServerModal/index.tsx b/src/frontend/src/modals/addMcpServerModal/index.tsx index 14b5af2cb..82607c9c7 100644 --- a/src/frontend/src/modals/addMcpServerModal/index.tsx +++ b/src/frontend/src/modals/addMcpServerModal/index.tsx @@ -12,6 +12,7 @@ import { TabsTrigger, } from "@/components/ui/tabs-button"; import { Textarea } from "@/components/ui/textarea"; +import { MAX_MCP_SERVER_NAME_LENGTH } from "@/constants/constants"; import { useAddMCPServer } from "@/controllers/API/queries/mcp/use-add-mcp-server"; import { usePatchMCPServer } from "@/controllers/API/queries/mcp/use-patch-mcp-server"; import { CustomLink } from "@/customization/components/custom-link"; @@ -134,7 +135,7 @@ export default function AddMcpServerModal({ "snake_case", "no_blank", "lowercase", - ]).slice(0, 30); + ]).slice(0, MAX_MCP_SERVER_NAME_LENGTH); try { await modifyMCPServer({ name, @@ -168,7 +169,7 @@ export default function AddMcpServerModal({ "snake_case", "no_blank", "lowercase", - ]).slice(0, 30); + ]).slice(0, MAX_MCP_SERVER_NAME_LENGTH); try { await modifyMCPServer({ name, @@ -202,7 +203,7 @@ export default function AddMcpServerModal({ "snake_case", "no_blank", "lowercase", - ]).slice(0, 30), + ]).slice(0, MAX_MCP_SERVER_NAME_LENGTH), })); } catch (e: any) { setError(e.message || "Invalid input"); diff --git a/src/frontend/src/pages/SettingsPage/pages/MCPServersPage/index.tsx b/src/frontend/src/pages/SettingsPage/pages/MCPServersPage/index.tsx index fcc33d550..beb9d249c 100644 --- a/src/frontend/src/pages/SettingsPage/pages/MCPServersPage/index.tsx +++ b/src/frontend/src/pages/SettingsPage/pages/MCPServersPage/index.tsx @@ -102,8 +102,9 @@ export default function MCPServersPage() {
{server.name} - {server.toolsCount} action - {server.toolsCount === 1 ? "" : "s"} + {server.toolsCount === null + ? "Loading..." + : `${server.toolsCount} action${server.toolsCount === 1 ? "" : "s"}`}
diff --git a/src/frontend/src/types/mcp/index.ts b/src/frontend/src/types/mcp/index.ts index 78bd78bc2..029bda550 100644 --- a/src/frontend/src/types/mcp/index.ts +++ b/src/frontend/src/types/mcp/index.ts @@ -9,10 +9,11 @@ export type MCPSettingsType = { }; export type MCPServerInfoType = { - id: string; + id?: string; name: string; - description: string; - toolsCount: number; + description?: string; + mode: string | null; + toolsCount: number | null; }; export type MCPServerType = { diff --git a/src/frontend/tests/extended/features/mcp-server.spec.ts b/src/frontend/tests/extended/features/mcp-server.spec.ts index b7eb5c97f..689474870 100644 --- a/src/frontend/tests/extended/features/mcp-server.spec.ts +++ b/src/frontend/tests/extended/features/mcp-server.spec.ts @@ -182,5 +182,27 @@ test( await expect(page.getByText("test_server")).not.toBeVisible({ timeout: 3000, }); + + await awaitBootstrapTest(page, { skipModal: true }); + await page.getByText("Untitled document").first().click(); + + await page.waitForTimeout(1000); + + await page.waitForSelector('[data-testid="save-mcp-server-button"]', { + timeout: 10000, + }); + + await page.getByTestId("save-mcp-server-button").click({ timeout: 10000 }); + + await page.waitForTimeout(1000); + + await expect(page.getByTestId("save-mcp-server-button")).toBeHidden({ + timeout: 10000, + }); + + await page.getByTestId("mcp-server-dropdown").click({ timeout: 10000 }); + await expect(page.getByText("test_server")).toHaveCount(2, { + timeout: 10000, + }); }, );