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
This commit is contained in:
Lucas Oliveira 2025-06-23 11:15:22 -03:00 committed by GitHub
commit 4a09655f2f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 282 additions and 83 deletions

View file

@ -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

View file

@ -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)

View file

@ -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):

View file

@ -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<string, any>): 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<any[]>([]);
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 (
<div className="flex w-full flex-col gap-2">
{options && (
<>
{options.length > 0 ? (
<Button
variant="primary"
size="xs"
role="combobox"
onClick={handleOpenListSelectionDialog}
className="dropdown-component-outline input-edit-node w-full py-2"
data-testid="mcp-server-dropdown"
disabled={disabled}
{options == null || options.length > 0 || showSaveButton ? (
<div className="flex w-full gap-2">
<Button
variant={!showSaveButton ? "primary" : "secondary"}
size="xs"
role="combobox"
onClick={
!showSaveButton
? handleOpenListSelectionDialog
: handleRemoveButtonClick
}
className={cn(
!showSaveButton
? "dropdown-component-outline input-edit-node"
: "",
"w-full py-2",
)}
data-testid="mcp-server-dropdown"
disabled={disabled || !options}
>
<div
className={cn(
"flex w-full items-center justify-start text-sm font-normal",
)}
>
<div
className={cn(
"flex w-full items-center justify-start text-sm font-normal",
)}
>
<span className="truncate">
{selectedItem[0]?.name
<span className="truncate">
{!options
? "Loading servers..."
: selectedItem[0]?.name
? selectedItem[0]?.name
: "Select a server..."}
</span>
<ForwardedIconComponent
name="ChevronsUpDown"
className="ml-auto h-5 w-5 text-muted-foreground"
/>
</div>
</Button>
) : (
<Button size="sm" onClick={handleAddButtonClick}>
<span>Add MCP Server</span>
</span>
<ForwardedIconComponent
name={!showSaveButton ? "ChevronsUpDown" : "X"}
className="ml-auto h-5 w-5 text-muted-foreground"
/>
</div>
</Button>
{showSaveButton && (
<Button
variant="primary"
size="iconMd"
className="px-2.5"
onClick={handleSaveButtonClick}
data-testid="save-mcp-server-button"
>
<ForwardedIconComponent
name="Save"
className="h-5 w-5 text-muted-foreground"
/>
</Button>
)}
</div>
) : (
<Button size="sm" onClick={handleAddButtonClick}>
<span>Add MCP Server</span>
</Button>
)}
{options && (
<>
<ListSelectionComponent
open={open}
onClose={handleCloseListSelectionDialog}
@ -105,7 +198,7 @@ export default function McpComponent({
options={options}
limit={1}
id={id}
value={value}
value={name}
editNode={editNode}
headerSearchPlaceholder="Search MCP Servers..."
handleOnNewValue={handleOnNewValue}
@ -113,13 +206,13 @@ export default function McpComponent({
addButtonText="Add MCP Server"
onAddButtonClick={handleAddButtonClick}
/>
<AddMcpServerModal
open={addOpen}
setOpen={setAddOpen}
onSuccess={handleSuccess}
/>
</>
)}
<AddMcpServerModal
open={addOpen}
setOpen={setAddOpen}
onSuccess={handleSuccess}
/>
</div>
);
}

View file

@ -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<MCPServerInfoType>;
type getMCPServersResponse = Array<MCPServerInfoType>;
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<getMCPServersResponse>(
`${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<getMCPServersResponse>(
`${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;
};

View file

@ -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");

View file

@ -102,8 +102,9 @@ export default function MCPServersPage() {
<div className="flex items-center gap-2">
<span className="text-sm font-medium">{server.name}</span>
<span className="text-mmd text-muted-foreground">
{server.toolsCount} action
{server.toolsCount === 1 ? "" : "s"}
{server.toolsCount === null
? "Loading..."
: `${server.toolsCount} action${server.toolsCount === 1 ? "" : "s"}`}
</span>
</div>
<DropdownMenu>

View file

@ -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 = {

View file

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