Source code for aidputils.agents.tools.mcp_tool
from typing import Any
import json
import asyncio
from aidputils.agents.toolkit.metrics_util import MetricUtility
from aidputils.agents.tools.base_tool import BaseTool
from aidputils.agents.tools.mcp.mcp_service import get_mcp_client, list_objects_paginated_sync
from aidputils.agents.auth.util.agent_util import _run_in_event_loop
def _mcp_error(message: str, code: int, data: dict | None = None) -> dict:
"""
Build an MCP/JSON-RPC-style error envelope.
code: HTTP status code (e.g., 400, 401, 500)
message: human-readable error message.
data: optional additional details.
"""
return {"error": {"code": code, "message": message, "data": data or {}}}
def _extract_status_code(ex: Exception) -> tuple[int | None, str | None]:
"""
Extract HTTP status code and detailed error message from an exception if available.
Returns (status_code, detailed_message) tuple.
"""
# Check direct attributes
if hasattr(ex, 'status_code'):
return ex.status_code, str(ex)
if hasattr(ex, 'code') and isinstance(ex.code, int):
return ex.code, str(ex)
if hasattr(ex, 'status') and isinstance(ex.status, int):
return ex.status, str(ex)
# Check for response attribute with status (HTTPStatusError pattern)
if hasattr(ex, 'response') and hasattr(ex.response, 'status_code'):
return ex.response.status_code, str(ex)
# Handle ExceptionGroups (asyncio TaskGroups) - needed for MCP library
# ExceptionGroups have an 'exceptions' attribute containing sub-exceptions
if hasattr(ex, 'exceptions') and isinstance(ex.exceptions, (list, tuple)):
for sub_exc in ex.exceptions:
status_code, message = _extract_status_code(sub_exc)
if status_code:
return status_code, message
return None, None
def _run_with_timeout(coro_factory, timeout_secs: float):
"""
Run a coroutine produced by coro_factory with a timeout.
This helper accepts a zero-argument callable that returns an awaitable (coro_factory)
so that the coroutine is only created when we are ready to schedule it. The coroutine
is wrapped with asyncio.wait_for(..., timeout_secs). To support environments that may
already have an active event loop (e.g., notebooks, frameworks) as well as plain
synchronous contexts, the actual execution is delegated to _run_in_event_loop which
safely bridges both cases.
Raises:
asyncio.TimeoutError: if the coroutine does not complete within timeout_secs.
"""
async def _runner():
return await asyncio.wait_for(coro_factory(), timeout=float(timeout_secs))
return _run_in_event_loop(_runner())
def _message_to_dict(msg: Any) -> dict:
"""
Convert a LangChain-style message (HumanMessage/AIMessage/SystemMessage) into a
minimal JSON-serializable dict: {"type": <str>, "content": <any>}.
If the object lacks a 'type' attribute, a lowercase class name is used as a fallback.
"""
msg_type = getattr(msg, "type", None) or msg.__class__.__name__.lower()
content = getattr(msg, "content", None)
return {"type": msg_type, "content": content}
def _json_safe_envelope(operation: str, data: Any) -> Any:
"""
Normalize MCP responses into JSON-serializable shapes.
Behavior by operation:
- list_tools / list_resources:
Expects a dict with an "items" list. Each item is normalized to include:
name, description, args_schema (aka inputSchema), out_schema (aka outputSchema).
Items may be dicts or objects; both are handled.
- get_prompt:
Expects a list of message objects; each is converted to a dict via _message_to_dict.
For unknown types:
- Lists are converted element-wise, falling back to str(obj) for non-serializable items.
- Primitives and already-serializable objects are returned as-is.
"""
# For list_tools/list_resources, data is expected to be a dict with "items"
if isinstance(data, dict):
items = data.get("items")
if isinstance(items, list):
if operation == "list_tools":
normalized = []
for tool in items:
# Accept dicts or objects; prefer dict keys when present
if isinstance(tool, dict):
name = tool.get("name")
description = tool.get("description")
args_schema = tool.get("args_schema") or tool.get("inputSchema")
out_schema = tool.get("out_schema") or tool.get("outputSchema")
else:
name = getattr(tool, "name", None)
description = getattr(tool, "description", None)
args_schema = getattr(tool, "inputSchema", None)
out_schema = getattr(tool, "outputSchema", None)
normalized.append({
"name": name,
"description": description,
"args_schema": args_schema,
"out_schema": out_schema
})
data = {**data, "items": normalized}
return data
# For get_prompt, data can be a list of Messages
if isinstance(data, list):
if operation == "get_prompt":
return [_message_to_dict(m) for m in data]
# Fallback: stringify non-serializable items
return [m if isinstance(m, (dict, list, str, int, float, bool, type(None))) else str(m) for m in data]
# Primitive or unknown types: return as-is (json.dumps will handle primitives)
return data
def _serialize_call_tool_result(result: Any) -> dict:
"""
Normalize a tool invocation result to {"output": <str>}.
Extraction order:
1) structuredContent/structured_content.result if present; JSON-encode non-primitives.
2) First content chunk's text/content if available.
3) Fallback to json.dumps(result), else str(result).
"""
output = None
try:
structured = getattr(result, "structuredContent", None) or getattr(result, "structured_content", None)
if isinstance(structured, dict):
res_val = structured.get("result")
if res_val is not None:
output = res_val if isinstance(res_val, (str, int, float, bool, type(None))) else json.dumps(res_val)
else:
output = json.dumps(structured)
except Exception:
pass
if output is None:
content = getattr(result, "content", None)
if isinstance(content, list) and len(content) > 0:
first = content[0]
text = getattr(first, "text", None) or getattr(first, "content", None)
if text is not None:
output = str(text)
if output is None:
try:
output = json.dumps(result)
except Exception:
output = str(result)
return {"output": output}
[docs]
@BaseTool.register
class MCPTool(BaseTool):
"""
Tool for interacting with an MCP server.
Supported operations via runtime_params["operation"]:
- "list_tools": Return available tools (normalized, JSON-serializable).
- "list_resources": Return available resources (normalized).
- "get_prompt": Return a list of messages for a named prompt (messages are normalized).
- "test_connection": Establish a simple MCP session to verify connectivity/auth; optional timeout_secs.
- "test_tool": Invoke a specific tool; requires 'tool_name' and optional 'arguments'.
Configuration (conf):
- endpoint (str): Base URL for the MCP server.
- auth (optional): Auth configuration passed through to the client.
- transport (str): Transport name. Defaults to "streamable_http".
Runtime params (runtime_params):
- server_name (str): The MCP server identifier.
- timeout_secs (float|int, optional): Per-call timeout.
- tool_name (str, for test_tool): Tool to execute.
- arguments (dict|any, for test_tool): Arguments; if wrapped as {"values": {...}} it is unwrapped.
- custom_headers (dict[str, str], optional): Additional HTTP headers merged on top of auth-derived headers when constructing the client.
All return values are shaped to be JSON-serializable for logging/metrics/clients.
Metrics counters (success/failure) are declared but incrementing is handled by BaseTool hooks if configured.
"""
metricsUtil = MetricUtility("MCP_TOOL")
success_counter = metricsUtil.create_counter("MCP_TOOL_SUCCESS_COUNTER")
failure_counter = metricsUtil.create_counter("MCP_TOOL_FAILURE_COUNTER")
DEFAULT_TIMEOUT = 30
@classmethod
def _invoke_tool(cls, conf, runtime_params, **context_vars):
"""
Execute an MCP operation based on runtime_params['operation'].
Parameters:
- conf: dict containing endpoint (required), auth (optional), transport (default 'streamable_http')
- runtime_params: dict with keys:
- server_name: MCP server identifier
- operation: one of {'list_tools','list_resources','get_prompt','test_connection','test_tool'}
- timeout_secs: optional per-call timeout
- tool_name, arguments: when operation == 'test_tool'
- context_vars: ignored by this tool; kept for BaseTool API compatibility.
"""
# Validate input params and configuration
if not isinstance(runtime_params, dict):
return _mcp_error("Invalid params: runtime_params must be a dict", 400)
operation = runtime_params.get("operation")
allowed_ops = {"list_tools", "list_resources", "get_prompt", "test_connection", "test_tool"}
if not isinstance(operation, str) or operation not in allowed_ops:
return _mcp_error(f"Unsupported or missing operation: {operation!r}", 400, {"allowed": sorted(allowed_ops)})
server_name = runtime_params.get("server_name")
if not isinstance(server_name, str) or not server_name.strip():
return _mcp_error("Invalid params: 'server_name' must be a non-empty string", 400)
if not isinstance(conf, dict):
return _mcp_error("Invalid params: 'conf' must be a dict", 400)
endpoint = conf.get("endpoint")
if not isinstance(endpoint, str) or not endpoint.strip():
return _mcp_error("Invalid params: conf['endpoint'] must be a non-empty string", 400)
timeout = runtime_params.get("timeout_secs")
transport = conf.get("transport", "streamable_http")
auth = conf.get("auth")
custom_headers = runtime_params.get("custom_headers")
client = get_mcp_client(server_name, endpoint, auth, transport, custom_headers=custom_headers)
match operation:
case "list_tools" | "list_resources" | "get_prompt":
try:
raw = list_objects_paginated_sync(client, runtime_params)
# Ensure the returned payload is JSON-serializable and matches expected shapes
return _json_safe_envelope(operation, raw)
except asyncio.TimeoutError:
return _mcp_error(f"Timeout while performing operation '{operation}'", 408, {"operation": operation})
except Exception as ex:
status_code, detailed_message = _extract_status_code(ex)
if status_code == 401:
message = detailed_message or "Unauthorized: Authentication required"
return _mcp_error(message, 401, {"operation": operation})
elif status_code == 403:
message = detailed_message or "Forbidden: Access denied"
return _mcp_error(message, 403, {"operation": operation})
else:
return _mcp_error(f"Server error during '{operation}': {str(ex)}", 500, {"operation": operation})
case "test_connection":
# Validate connectivity/auth by establishing a simple session.
try:
timeout = float(runtime_params.get("timeout_secs") or cls.DEFAULT_TIMEOUT)
except Exception:
return _mcp_error("Invalid params: 'timeout_secs' must be a positive number", 400)
try:
_run_with_timeout(
lambda: client.test_connection(server_name),
timeout
)
return {"output": "MCP connection successful"}
except asyncio.TimeoutError:
return _mcp_error("MCP connection failed: timed out", 408, {"operation": "test_connection"})
except Exception as ex:
status_code, detailed_message = _extract_status_code(ex)
if status_code == 401:
message = detailed_message or "Unauthorized: Authentication required"
return _mcp_error(message, 401, {"operation": "test_connection"})
elif status_code == 403:
message = detailed_message or "Forbidden: Access denied"
return _mcp_error(message, 403, {"operation": "test_connection"})
else:
return _mcp_error(f"MCP connection failed: {str(ex)}", 500, {"operation": "test_connection"})
case "test_tool":
# Invoke a specific tool exposed by the MCP server using arguments from runtime_params
tool_name = runtime_params.get("tool_name")
if not isinstance(tool_name, str) or not tool_name.strip():
return _mcp_error("Invalid params: 'tool_name' is required for operation 'test_tool'", 400)
args = runtime_params.get("arguments")
# If provided as {"values": {...}}, unwrap
if isinstance(args, dict) and "values" in args and isinstance(args["values"], dict):
args = args["values"]
try:
timeout = float(runtime_params.get("timeout_secs") or cls.DEFAULT_TIMEOUT)
except Exception:
return _mcp_error("Invalid params: 'timeout_secs' must be a positive number", 400)
try:
result = _run_with_timeout(
lambda: client.call_tool(server_name, tool_name, args),
timeout
)
return _serialize_call_tool_result(result)
except asyncio.TimeoutError:
return _mcp_error(f"MCP tool invocation failed for '{tool_name}': timed out", 408, {"operation": "test_tool", "tool_name": tool_name})
except Exception as ex:
status_code, detailed_message = _extract_status_code(ex)
if status_code == 401:
message = detailed_message or "Unauthorized: Authentication required"
return _mcp_error(message, 401, {"operation": "test_tool", "tool_name": tool_name})
elif status_code == 403:
message = detailed_message or "Forbidden: Access denied"
return _mcp_error(message, 403, {"operation": "test_tool", "tool_name": tool_name})
else:
return _mcp_error(f"MCP tool invocation failed for '{tool_name}': {str(ex)}", 500, {"operation": "test_tool", "tool_name": tool_name})
return None
@classmethod
def _format_output_as_mcp(cls, response):
"""
Format the tool response for downstream consumption.
MCPTool returns JSON-serializable structures already; this method is a passthrough.
"""
return response