1  # Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  #
  3  # Licensed under the Apache License, Version 2.0 (the "License");
  4  # you may not use this file except in compliance with the License.
  5  # You may obtain a copy of the License at
  6  #
  7  #     http://www.apache.org/licenses/LICENSE-2.0
  8  #
  9  # Unless required by applicable law or agreed to in writing, software
 10  # distributed under the License is distributed on an "AS IS" BASIS,
 11  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  # See the License for the specific language governing permissions and
 13  # limitations under the License.
 14  #
 15  # Requirement: Any integration or derivative work must explicitly attribute
 16  # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  # documentation or user interface, as detailed in the NOTICE file.
 18  
 19  from typing import Any, Literal
 20  
 21  from tools.registry import register_tool
 22  from utils.mcp_tools import MCPTools
 23  from utils.tool_context import ToolContext
 24  
 25  
 26  async def _get_mcp_manager(
 27      url: str | None, transport: str, context: ToolContext | None
 28  ) -> MCPTools:
 29      """Helper to get MCPTools manager, either from context or creating a new one."""
 30      if url:
 31          return MCPTools(url=url, transport=transport)  # type: ignore
 32  
 33      if context and context.tool_dispatcher and context.tool_dispatcher.mcp_server_url:
 34          manager = await context.tool_dispatcher._ensure_mcp_manager()
 35          if manager:
 36              return manager
 37  
 38      raise ValueError("MCP server URL is required (either via parameter or global configuration)")
 39  
 40  
 41  @register_tool(sandbox_execution=False)
 42  async def call_mcp_tool(
 43      tool_name: str,
 44      url: str | None = None,
 45      transport: Literal["sse", "streamable-http"] = "sse",
 46      context: ToolContext | None = None,
 47      **kwargs,
 48  ) -> dict[str, Any]:
 49      """Call a remote MCP server tool."""
 50      try:
 51          manager = await _get_mcp_manager(url, transport, context)
 52          ret = await manager.call_remote_tool(tool_name, **kwargs)
 53          return {"tool_name": tool_name, "tool_result": ret}
 54      except Exception as e:
 55          return {"error": str(e)}
 56  
 57  
 58  @register_tool(sandbox_execution=False)
 59  async def list_mcp_tools(
 60      url: str | None = None,
 61      transport: Literal["sse", "streamable-http"] = "sse",
 62      context: ToolContext | None = None,
 63  ) -> dict[str, Any]:
 64      """List tools available on a remote MCP server."""
 65      try:
 66          manager = await _get_mcp_manager(url, transport, context)
 67          ret = await manager.describe_mcp_tools()
 68          return {"tools": ret}
 69      except Exception as e:
 70          return {"error": str(e)}
 71  
 72  
 73  @register_tool(sandbox_execution=False)
 74  async def list_mcp_prompts(
 75      url: str | None = None,
 76      transport: Literal["sse", "streamable-http"] = "sse",
 77      context: ToolContext | None = None,
 78  ) -> dict[str, Any]:
 79      """List prompts available on a remote MCP server."""
 80      try:
 81          manager = await _get_mcp_manager(url, transport, context)
 82          ret = await manager.list_remote_prompts()
 83          return {"prompts": ret}
 84      except Exception as e:
 85          return {"error": str(e)}
 86  
 87  
 88  @register_tool(sandbox_execution=False)
 89  async def list_mcp_resources(
 90      url: str | None = None,
 91      transport: Literal["sse", "streamable-http"] = "sse",
 92      context: ToolContext | None = None,
 93  ) -> dict[str, Any]:
 94      """List resources available on a remote MCP server."""
 95      try:
 96          manager = await _get_mcp_manager(url, transport, context)
 97          ret = await manager.list_remote_resources()
 98          return {"resources": ret}
 99      except Exception as e:
100          return {"error": str(e)}
101  
102  
103  # Alias for backward compatibility
104  # mcp_tool = call_mcp_tool