"""MCP tool wrappers for Beaver's tool contract.""" from __future__ import annotations import asyncio from dataclasses import dataclass import json from typing import Any, Awaitable, Callable from beaver.tools.base import BaseTool, ToolContext, ToolResult, ToolSpec def _tool_schema(tool_def: Any) -> dict[str, Any]: schema = getattr(tool_def, "inputSchema", None) or getattr(tool_def, "input_schema", None) if isinstance(schema, dict): return schema return {"type": "object", "properties": {}} def _tool_name(tool_def: Any) -> str: return str(getattr(tool_def, "name", "") or "") def _tool_description(tool_def: Any) -> str: return str(getattr(tool_def, "description", "") or _tool_name(tool_def)) def _mcp_result_to_text(result: Any) -> str: parts: list[str] = [] for block in list(getattr(result, "content", []) or []): text = getattr(block, "text", None) parts.append(str(text if text is not None else block)) if not parts and getattr(result, "structuredContent", None) is not None: return json.dumps(getattr(result, "structuredContent"), ensure_ascii=False, indent=2) return "\n".join(parts) or "(no output)" @dataclass(slots=True) class MCPToolWrapper(BaseTool): server_id: str tool_def: Any call_tool: Callable[[str, dict[str, Any]], Awaitable[Any]] tool_timeout: int = 30 sensitive: bool = False kind: str = "online" category: str = "online" display_name: str = "" @property def original_name(self) -> str: return _tool_name(self.tool_def) @property def spec(self) -> ToolSpec: return ToolSpec( name=f"mcp_{self.server_id}_{self.original_name}", description=_tool_description(self.tool_def), input_schema=_tool_schema(self.tool_def), toolset=f"mcp-{self.server_id}", metadata={ "server_id": self.server_id, "original_tool_name": self.original_name, "kind": self.kind, "category": self.category, "display_name": self.display_name or self.server_id, "transport": "mcp", }, ) async def invoke(self, arguments: dict[str, Any], context: ToolContext) -> ToolResult: try: result = await asyncio.wait_for( self.call_tool(self.original_name, dict(arguments or {})), timeout=max(1, int(self.tool_timeout or 30)), ) return ToolResult( success=True, content=_mcp_result_to_text(result), tool_name=self.spec.name, raw_output=result, ) except Exception as exc: return ToolResult( success=False, content=f"MCP tool {self.server_id}.{self.original_name} failed: {exc}", tool_name=self.spec.name, error=str(exc), )