Browse Source

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)
tags/v0.20.0
Yongteng Lei 4 months ago
parent
commit
0478f36e36
No account linked to committer's email address
1 changed files with 43 additions and 5 deletions
  1. 43
    5
      api/apps/mcp_server_app.py

+ 43
- 5
api/apps/mcp_server_app.py View File

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

Loading…
Cancel
Save