From 0478f36e368b055aa88cfc51f958e57e4fa7750c Mon Sep 17 00:00:00 2001 From: Yongteng Lei Date: Fri, 27 Jun 2025 10:23:34 +0800 Subject: [PATCH] Feat: allow users to choose which MCP tools are enabled (#8519) ### What problem does this PR solve? Allow users to choose which MCP tools are enabled. ### Type of change - [x] New Feature (non-breaking change which adds functionality) --- api/apps/mcp_server_app.py | 48 ++++++++++++++++++++++++++++++++++---- 1 file changed, 43 insertions(+), 5 deletions(-) diff --git a/api/apps/mcp_server_app.py b/api/apps/mcp_server_app.py index 3d29debdb..34bd9dfb2 100644 --- a/api/apps/mcp_server_app.py +++ b/api/apps/mcp_server_app.py @@ -85,7 +85,7 @@ def create() -> Response: @manager.route("/update", methods=["POST"]) # noqa: F821 @login_required -@validate_request("id") +@validate_request("mcp_id") def update() -> Response: req = request.get_json() @@ -96,7 +96,7 @@ def update() -> Response: if server_name and len(server_name.encode("utf-8")) > 255: return get_data_error_result(message=f"Invaild MCP name or length is {len(server_name)} which is large than 255.") - mcp_id = req.get("id", "") + mcp_id = req.get("mcp_id", "") e, mcp_server = MCPServerService.get_by_id(mcp_id) if not e or mcp_server.tenant_id != current_user.id: return get_data_error_result(message=f"Cannot find MCP server {mcp_id} for user {current_user.id}") @@ -106,8 +106,10 @@ def update() -> Response: try: req["tenant_id"] = current_user.id + req.pop("mcp_id", None) + req["id"] = mcp_id - if not MCPServerService.filter_update([MCPServer.id == req["id"], MCPServer.tenant_id == req["tenant_id"]], req): + if not MCPServerService.filter_update([MCPServer.id == mcp_id, MCPServer.tenant_id == current_user.id], req): return get_data_error_result(message="Failed to updated MCP server.") e, updated_mcp = MCPServerService.get_by_id(req["id"]) @@ -240,11 +242,23 @@ def list_tools() -> Response: if e and mcp_server.tenant_id == current_user.id: server_key = mcp_server.id + cached_tools = mcp_server.variables.get("tools", {}) + tool_call_session = MCPToolCallSession(mcp_server, mcp_server.variables) tool_call_sessions.append(tool_call_session) - tools = tool_call_session.get_tools(timeout) - results[server_key] = [tool.model_dump() for tool in tools] + try: + tools = tool_call_session.get_tools(timeout) + except Exception: + tools = [] + + results[server_key] = [] + for tool in tools: + tool_dict = tool.model_dump() + cached_tool = cached_tools.get(tool_dict["name"]) + + tool_dict["enabled"] = cached_tool.get("enabled") if cached_tool and "enabled" in cached_tool else True + results[server_key].append(tool_dict) # PERF: blocking call to close sessions — consider moving to background thread or task queue close_multiple_mcp_toolcall_sessions(tool_call_sessions) @@ -284,3 +298,27 @@ def test_tool() -> Response: return get_json_result(data=result) except Exception as e: return server_error_response(e) + + +@manager.route("/cache_tools", methods=["POST"]) # noqa: F821 +@login_required +@validate_request("mcp_id", "tools") +def cache_tool() -> Response: + req = request.get_json() + mcp_id = req.get("mcp_id", "") + if not mcp_id: + return get_data_error_result(message="No MCP server ID provided.") + tools = req.get("tools", []) + + e, mcp_server = MCPServerService.get_by_id(mcp_id) + if not e or mcp_server.tenant_id != current_user.id: + return get_data_error_result(message=f"Cannot find MCP server {mcp_id} for user {current_user.id}") + + variables = mcp_server.variables + tools = {tool["name"]: tool for tool in tools if isinstance(tool, dict) and "name" in tool} + variables["tools"] = tools + + if not MCPServerService.filter_update([MCPServer.id == mcp_id, MCPServer.tenant_id == current_user.id], {"variables": variables}): + return get_data_error_result(message="Failed to updated MCP server.") + + return get_json_result(data=tools)