feat: add a subagent frame

agentlite/src/agentlite/mcp.py

@@ -0,0 +1,214 @@
+"""MCP (Model Context Protocol) integration for AgentLite.
+
+This module provides integration with MCP servers, allowing agents to use
+tools from external MCP-compatible servers.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from agentlite.message import TextPart
+from agentlite.tool import CallableTool, ToolOk, ToolResult, ToolError
+
+if TYPE_CHECKING:
+    pass
+
+
+class MCPClient:
+    """Client for connecting to MCP servers.
+
+    This client allows you to connect to MCP servers and load their tools
+    into AgentLite agents.
+
+    Example:
+        >>> client = MCPClient()
+        >>> await client.connect_stdio(
+        ...     "npx", ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
+        ... )
+        >>> tools = await client.load_tools()
+        >>> agent = Agent(provider=provider, tools=tools)
+    """
+
+    def __init__(self):
+        """Initialize the MCP client."""
+        self._client: Any | None = None
+        self._connected = False
+
+    def _check_fastmcp(self) -> None:
+        """Check if fastmcp is installed."""
+        try:
+            import fastmcp  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "MCP support requires 'fastmcp' package. Install with: pip install agentlite[mcp]"
+            ) from e
+
+    async def connect_stdio(
+        self,
+        command: str,
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+    ) -> None:
+        """Connect to an MCP server via stdio.
+
+        Args:
+            command: The command to run.
+            args: Optional arguments for the command.
+            env: Optional environment variables.
+
+        Raises:
+            RuntimeError: If already connected.
+            ConnectionError: If the connection fails.
+        """
+        if self._connected:
+            raise RuntimeError("Already connected to an MCP server")
+
+        try:
+            from fastmcp import Client
+            from fastmcp.client.transports import PythonStdioTransport
+
+            transport = PythonStdioTransport(
+                command_or_script=command,
+                args=args or [],
+                env=env,
+            )
+            self._client = Client(transport)
+            self._connected = True
+        except Exception as e:
+            raise ConnectionError(f"Failed to connect to MCP server: {e}") from e
+
+    async def connect_sse(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        """Connect to an MCP server via Server-Sent Events (SSE).
+
+        Args:
+            url: The SSE endpoint URL.
+            headers: Optional headers to include in requests.
+
+        Raises:
+            RuntimeError: If already connected.
+            ConnectionError: If the connection fails.
+        """
+        if self._connected:
+            raise RuntimeError("Already connected to an MCP server")
+
+        try:
+            from fastmcp import Client
+            from fastmcp.client.transports import SSETransport
+
+            transport = SSETransport(url=url, headers=headers)
+            self._client = Client(transport)
+            self._connected = True
+        except Exception as e:
+            raise ConnectionError(f"Failed to connect to MCP server: {e}") from e
+
+    async def load_tools(self) -> list[CallableTool]:
+        """Load tools from the connected MCP server.
+
+        Returns:
+            A list of CallableTool instances wrapping the MCP tools.
+
+        Raises:
+            RuntimeError: If not connected to an MCP server.
+        """
+        if not self._connected or self._client is None:
+            raise RuntimeError("Not connected to an MCP server")
+
+        tools: list[CallableTool] = []
+
+        try:
+            async with self._client as client:
+                mcp_tools = await client.list_tools()
+
+                for mcp_tool in mcp_tools:
+                    tool = _MCPTool(
+                        client=self._client,
+                        name=mcp_tool.name,
+                        description=mcp_tool.description or "No description provided",
+                        parameters=mcp_tool.inputSchema,
+                    )
+                    tools.append(tool)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MCP tools: {e}") from e
+
+        return tools
+
+    async def close(self) -> None:
+        """Close the connection to the MCP server."""
+        if self._client is not None:
+            try:
+                await self._client.close()
+            except Exception:
+                pass
+            finally:
+                self._client = None
+                self._connected = False
+
+    async def __aenter__(self) -> MCPClient:
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        """Async context manager exit."""
+        await self.close()
+
+
+class _MCPTool(CallableTool):
+    """Wrapper for MCP tools."""
+
+    def __init__(
+        self,
+        client: Any,
+        name: str,
+        description: str,
+        parameters: dict[str, Any],
+    ):
+        """Initialize the MCP tool wrapper.
+
+        Args:
+            client: The MCP client.
+            name: The tool name.
+            description: The tool description.
+            parameters: The JSON schema for tool parameters.
+        """
+        self._client = client
+        super().__init__(
+            name=name,
+            description=description,
+            parameters=parameters,
+        )
+
+    async def __call__(self, **kwargs: Any) -> ToolResult:
+        """Execute the MCP tool.
+
+        Args:
+            **kwargs: The tool arguments.
+
+        Returns:
+            The tool result.
+        """
+        try:
+            async with self._client as client:
+                result = await client.call_tool(self.name, kwargs)
+
+                # Convert MCP result to ToolResult
+                content_parts = []
+                for content in result.content:
+                    if hasattr(content, "text"):
+                        content_parts.append(content.text)
+                    else:
+                        content_parts.append(str(content))
+
+                output = "\n".join(content_parts)
+
+                if result.isError:
+                    return ToolError(message=output or "Tool execution failed")
+
+                return ToolOk(output=output)
+        except Exception as e:
+            return ToolError(message=f"MCP tool execution failed: {e}")