agentpool 2.1.9__py3-none-any.whl

agentpool/mcp_server/client.py

@@ -0,0 +1,429 @@
+"""FastMCP-based client implementation for AgentPool.
+
+This module provides a client for communicating with MCP servers using FastMCP.
+It includes support for contextual progress handlers that extend FastMCP's
+standard progress callbacks with tool execution context (tool name, call ID, and input).
+
+The key innovation is the signature injection system that allows MCP tools to work
+seamlessly with PydanticAI's RunContext while providing rich progress information.
+
+Elicitation is handled via a forwarding callback pattern: a stable callback is
+registered at connection time, but it delegates to a mutable handler that can
+be swapped per tool call (allowing AgentContext.handle_elicitation to be used).
+"""
+
+from __future__ import annotations
+
+import contextlib
+from importlib.metadata import version
+import logging
+from typing import TYPE_CHECKING, Any, Self, assert_never
+
+import anyio
+from pydantic_ai import RunContext, ToolReturn
+from schemez import FunctionSchema
+
+from agentpool.agents.context import AgentContext
+from agentpool.log import get_logger
+from agentpool.mcp_server.constants import MCP_TO_LOGGING
+from agentpool.mcp_server.helpers import extract_text_content, mcp_tool_to_fn_schema
+from agentpool.mcp_server.message_handler import MCPMessageHandler
+from agentpool.tools.base import Tool
+from agentpool.utils.signatures import create_modified_signature
+from agentpool_config.mcp_server import (
+    SSEMCPServerConfig,
+    StdioMCPServerConfig,
+    StreamableHTTPMCPServerConfig,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable, Sequence
+
+    import fastmcp
+    from fastmcp.client import ClientTransport
+    from fastmcp.client.elicitation import ElicitationHandler
+    from fastmcp.client.logging import LogMessage
+    from fastmcp.client.messages import MessageHandler, MessageHandlerT
+    from fastmcp.client.sampling import SamplingHandler
+    from mcp.shared.context import RequestContext
+    from mcp.types import (
+        BlobResourceContents,
+        ContentBlock,
+        ElicitRequestParams,
+        GetPromptResult,
+        Icon,
+        Implementation,
+        Prompt as MCPPrompt,
+        Resource as MCPResource,
+        TextResourceContents,
+        Tool as MCPTool,
+    )
+    from pydantic_ai import BinaryContent
+    from upathtools.filesystems import MCPFileSystem, MCPToolsFileSystem
+
+    from agentpool_config.mcp_server import MCPServerConfig
+
+
+logger = get_logger(__name__)
+
+
+class MCPClient:
+    """FastMCP-based client for communicating with MCP servers."""
+
+    def __init__(
+        self,
+        config: MCPServerConfig,
+        sampling_callback: SamplingHandler[Any, Any] | None = None,
+        message_handler: MessageHandlerT | MessageHandler | None = None,
+        accessible_roots: list[str] | None = None,
+        tool_change_callback: Callable[[], Awaitable[None]] | None = None,
+        prompt_change_callback: Callable[[], Awaitable[None]] | None = None,
+        resource_change_callback: Callable[[], Awaitable[None]] | None = None,
+        client_name: str | None = None,
+        client_title: str | None = None,
+        client_website_url: str | None = None,
+        client_icon_path: str | None = None,
+    ) -> None:
+        # Mutable handler swapped per call_tool for dynamic elicitation
+        self._current_elicitation_handler: ElicitationHandler | None = None
+        self.config = config
+        self._sampling_callback = sampling_callback
+        # Store message handler or mark for lazy creation
+        self._message_handler = message_handler
+        self._accessible_roots = accessible_roots or []
+        self._tool_change_callback = tool_change_callback
+        self._prompt_change_callback = prompt_change_callback
+        self._resource_change_callback = resource_change_callback
+        self._client_name = client_name
+        self._client_title = client_title
+        self._client_website_url = client_website_url
+        self._client_icon_path = client_icon_path
+        self._client = self._get_client(self.config)
+
+    @property
+    def connected(self) -> bool:
+        """Check if client is connected by examining session state."""
+        return self._client.is_connected()
+
+    async def __aenter__(self) -> Self:
+        """Enter context manager."""
+        try:
+            # First attempt with configured auth
+            await self._client.__aenter__()  # type: ignore[no-untyped-call]
+
+        except Exception as first_error:
+            # OAuth fallback for HTTP/SSE if not already using OAuth
+            if not isinstance(self.config, StdioMCPServerConfig) and not self.config.auth.oauth:
+                try:
+                    with contextlib.suppress(Exception):
+                        await self._client.__aexit__(None, None, None)  # type: ignore[no-untyped-call]
+                    self._client = self._get_client(self.config, force_oauth=True)
+                    await self._client.__aenter__()  # type: ignore[no-untyped-call]
+                    logger.info("Connected with OAuth fallback")
+                except Exception:  # noqa: BLE001
+                    raise first_error from None
+            else:
+                raise
+
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        """Exit context manager and cleanup."""
+        try:
+            await self._client.__aexit__(None, None, None)  # type: ignore[no-untyped-call]
+        except Exception as e:  # noqa: BLE001
+            logger.warning("Error during FastMCP client cleanup", error=e)
+
+    def get_resource_fs(self) -> MCPFileSystem:
+        """Get a filesystem for accessing MCP resources."""
+        from upathtools.filesystems import MCPFileSystem
+
+        return MCPFileSystem(client=self._client)
+
+    def get_tools_fs(self) -> MCPToolsFileSystem:
+        """Get a filesystem for accessing MCP tools as code."""
+        from upathtools.filesystems import MCPToolsFileSystem
+
+        return MCPToolsFileSystem(client=self._client)
+
+    async def _log_handler(self, message: LogMessage) -> None:
+        """Handle server log messages."""
+        level = MCP_TO_LOGGING.get(message.level, logging.INFO)
+        logger.log(level, "MCP Server: ", data=message.data)
+
+    async def _forwarding_elicitation_callback[T](
+        self,
+        message: str,
+        response_type: type[T],
+        params: ElicitRequestParams,
+        context: RequestContext[Any, Any],
+    ) -> T | dict[str, Any] | Any:
+        """Forwarding callback that delegates to current handler.
+
+        This callback is registered once at connection time, but delegates to
+        _current_elicitation_handler which can be swapped per tool call.
+        """
+        from fastmcp.client.elicitation import ElicitResult
+
+        # Try current handler first (set per call_tool)
+        if self._current_elicitation_handler:
+            return await self._current_elicitation_handler(message, response_type, params, context)
+        # No handler available - decline by default
+        return ElicitResult(action="decline")
+
+    def _get_client(
+        self, config: MCPServerConfig, force_oauth: bool = False
+    ) -> fastmcp.Client[Any]:
+        """Create FastMCP client based on config."""
+        import fastmcp
+        from fastmcp.client import SSETransport, StreamableHttpTransport
+        from fastmcp.client.transports import StdioTransport
+        from mcp.types import Icon, Implementation
+
+        transport: ClientTransport
+        # Create transport based on config type
+        match config:
+            case StdioMCPServerConfig(command=command, args=args):
+                env = config.get_env_vars()
+                transport = StdioTransport(command=command, args=args, env=env)
+                oauth = False
+                if force_oauth:
+                    msg = "OAuth is not supported for StdioMCPServerConfig"
+                    raise ValueError(msg)
+
+            case SSEMCPServerConfig(url=url, headers=headers, auth=auth):
+                transport = SSETransport(url=url, headers=headers)
+                oauth = auth.oauth
+
+            case StreamableHTTPMCPServerConfig(url=url, headers=headers, auth=auth):
+                transport = StreamableHttpTransport(url=url, headers=headers)
+                oauth = auth.oauth
+            case _ as unreachable:
+                assert_never(unreachable)
+
+        # Create message handler if needed
+        msg_handler = self._message_handler or MCPMessageHandler(
+            self,
+            self._tool_change_callback,
+            self._prompt_change_callback,
+            self._resource_change_callback,
+        )
+
+        # Build client_info if client_name is provided
+        client_info: Implementation | None = None
+        if self._client_name:
+            icons: list[Icon] | None = None
+            if self._client_icon_path:
+                icons = [Icon(src=self._client_icon_path)]
+            client_info = Implementation(
+                name=self._client_name,
+                version=version("agentpool"),
+                title=self._client_title,
+                websiteUrl=self._client_website_url,
+                icons=icons,
+            )
+
+        return fastmcp.Client(
+            transport,
+            log_handler=self._log_handler,
+            roots=self._accessible_roots,
+            timeout=config.timeout,
+            elicitation_handler=self._forwarding_elicitation_callback,
+            sampling_handler=self._sampling_callback,
+            message_handler=msg_handler,
+            auth="oauth" if (force_oauth or oauth) else None,
+            client_info=client_info,
+        )
+
+    async def list_tools(self) -> list[MCPTool]:
+        """Get available tools directly from the server."""
+        if not self.connected:
+            msg = "Not connected to MCP server"
+            raise RuntimeError(msg)
+
+        try:
+            tools = await self._client.list_tools()
+            logger.debug("Listed tools from MCP server", num_tools=len(tools))
+        except Exception as e:  # noqa: BLE001
+            logger.warning("Failed to list tools", error=e)
+            return []
+        else:
+            return tools
+
+    async def list_prompts(self) -> list[MCPPrompt]:
+        """Get available prompts from the server."""
+        if not self.connected:
+            msg = "Not connected to MCP server"
+            raise RuntimeError(msg)
+
+        try:
+            return await self._client.list_prompts()
+        except Exception as e:  # noqa: BLE001
+            logger.debug("Failed to list prompts", error=e)
+            return []
+
+    async def list_resources(self) -> list[MCPResource]:
+        """Get available resources from the server."""
+        if not self.connected:
+            msg = "Not connected to MCP server"
+            raise RuntimeError(msg)
+
+        try:
+            return await self._client.list_resources()
+        except Exception as e:
+            msg = f"Failed to list resources: {e}"
+            raise RuntimeError(msg) from e
+
+    async def get_prompt(
+        self, name: str, arguments: dict[str, str] | None = None
+    ) -> GetPromptResult:
+        """Get a specific prompt's content."""
+        if not self.connected:
+            msg = "Not connected to MCP server"
+            raise RuntimeError(msg)
+
+        try:
+            return await self._client.get_prompt_mcp(name, arguments)
+        except Exception as e:
+            msg = f"Failed to get prompt {name!r}: {e}"
+            raise RuntimeError(msg) from e
+
+    def convert_tool(self, tool: MCPTool) -> Tool:
+        """Create a properly typed callable from MCP tool schema."""
+
+        async def tool_callable(
+            ctx: RunContext, agent_ctx: AgentContext[Any], **kwargs: Any
+        ) -> str | Any | ToolReturn:
+            """Dynamically generated MCP tool wrapper."""
+            # Filter out None values for optional params
+            schema_props = tool.inputSchema.get("properties", {})
+            required_props = set(tool.inputSchema.get("required", []))
+            filtered_kwargs = {
+                k: v
+                for k, v in kwargs.items()
+                if k in required_props or (k in schema_props and v is not None)
+            }
+            return await self.call_tool(tool.name, ctx, filtered_kwargs, agent_ctx)
+
+        # Set proper signature and annotations with both RunContext and AgentContext
+        schema = mcp_tool_to_fn_schema(tool)
+        fn_schema = FunctionSchema.from_dict(schema)
+        sig = fn_schema.to_python_signature()
+
+        tool_callable.__signature__ = create_modified_signature(  # type: ignore[attr-defined]
+            sig, inject={"ctx": RunContext, "agent_ctx": AgentContext}
+        )
+        annotations = fn_schema.get_annotations()
+        annotations["ctx"] = RunContext
+        annotations["agent_ctx"] = AgentContext
+        # Update return annotation to support multiple types
+        annotations["return"] = str | Any | ToolReturn  # type: ignore
+        tool_callable.__annotations__ = annotations
+        tool_callable.__name__ = tool.name
+        tool_callable.__doc__ = tool.description or "No description provided."
+        return Tool.from_callable(tool_callable, source="mcp")
+
+    async def call_tool(
+        self,
+        name: str,
+        run_context: RunContext,
+        arguments: dict[str, Any] | None = None,
+        agent_ctx: AgentContext[Any] | None = None,
+    ) -> ToolReturn | str | Any:
+        """Call an MCP tool with full PydanticAI return type support."""
+        if not self.connected:
+            msg = "Not connected to MCP server"
+            raise RuntimeError(msg)
+
+        # Create progress handler that bridges to AgentContext if available
+        progress_handler = None
+        if agent_ctx:
+
+            async def fastmcp_progress_handler(
+                progress: float,
+                total: float | None,
+                message: str | None,
+            ) -> None:
+                await agent_ctx.report_progress(progress, total, message or "")
+
+            progress_handler = fastmcp_progress_handler
+
+        # Set up per-call elicitation handler from AgentContext
+        if agent_ctx:
+
+            async def elicitation_handler[T](
+                message: str,
+                response_type: type[T] | None,
+                params: ElicitRequestParams,
+                context: RequestContext[Any, Any, Any],
+            ) -> T | dict[str, Any] | Any:
+                from fastmcp.client.elicitation import ElicitResult
+                from mcp.types import ElicitResult as MCPElicitResult, ErrorData
+
+                result = await agent_ctx.handle_elicitation(params)
+                match result:
+                    case MCPElicitResult(action="accept", content=content):
+                        return content
+                    case MCPElicitResult(action="cancel"):
+                        return ElicitResult(action="cancel")
+                    case MCPElicitResult(action="decline"):
+                        return ElicitResult(action="decline")
+                    case ErrorData():
+                        return ElicitResult(action="decline")
+                    case _:
+                        return ElicitResult(action="decline")
+
+            self._current_elicitation_handler = elicitation_handler
+
+        try:
+            result = await self._client.call_tool(
+                name, arguments, progress_handler=progress_handler
+            )
+            content = await self._from_mcp_content(result.content)
+            # Decision logic for return type
+            match (result.data is not None, bool(content)):
+                case (True, True):  # Both structured data and rich content -> ToolReturn
+                    return ToolReturn(return_value=result.data, content=content)
+                case (True, False):  # Only structured data -> return directly
+                    return result.data
+                case (False, True):  # Only content -> ToolReturn with content
+                    msg = "Tool executed successfully"
+                    return ToolReturn(return_value=msg, content=content)
+                case (False, False):  # Fallback to text extraction
+                    return extract_text_content(result.content)
+                case _:  # Handle unexpected cases
+                    msg = f"Unexpected MCP content: {result.content}"
+                    raise ValueError(msg)  # noqa: TRY301
+        except Exception as e:
+            msg = f"MCP tool call failed: {e}"
+            raise RuntimeError(msg) from e
+        finally:
+            # Clear per-call handler
+            self._current_elicitation_handler = None
+
+    async def _from_mcp_content(
+        self,
+        mcp_content: Sequence[ContentBlock | TextResourceContents | BlobResourceContents],
+    ) -> list[str | BinaryContent]:
+        """Convert MCP content blocks to PydanticAI content types."""
+        from agentpool.mcp_server.conversions import from_mcp_content
+
+        return await from_mcp_content(mcp_content)
+
+
+if __name__ == "__main__":
+    path = "/home/phil65/dev/oss/agentpool/tests/mcp_server/server.py"
+    # path = Path(__file__).parent / "test_mcp_server.py"
+    config = StdioMCPServerConfig(
+        command="uv",
+        args=["run", str(path)],
+    )
+
+    async def main() -> None:
+        async with MCPClient(config=config) as mcp_client:
+            # Create MCP filesystem
+            fs = mcp_client.get_resource_fs()
+            print(await fs._ls(""))
+
+    anyio.run(main)

agentpool/mcp_server/constants.py

@@ -0,0 +1,32 @@
+"""MCP related constants."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    import mcp
+
+
+MCP_TO_LOGGING: dict[mcp.LoggingLevel, int] = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "notice": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+    "alert": logging.CRITICAL,
+    "emergency": logging.CRITICAL,
+}
+
+
+# Map Python logging levels to MCP logging levels
+LOGGING_TO_MCP: dict[int, mcp.LoggingLevel] = {
+    logging.DEBUG: "debug",
+    logging.INFO: "info",
+    logging.WARNING: "warning",
+    logging.ERROR: "error",
+    logging.CRITICAL: "critical",
+}

agentpool/mcp_server/conversions.py

@@ -0,0 +1,172 @@
+"""Conversions between internal and MCP types."""
+
+from __future__ import annotations
+
+import base64
+from typing import TYPE_CHECKING, Any, assert_never
+
+from mcp.types import (
+    AudioContent,
+    BlobResourceContents,
+    EmbeddedResource,
+    ImageContent,
+    PromptMessage,
+    ResourceLink,
+    TextContent,
+    TextResourceContents,
+)
+from pydantic_ai import (
+    AudioUrl,
+    BinaryContent,
+    BinaryImage,
+    DocumentUrl,
+    FileUrl,
+    ImageUrl,
+    SystemPromptPart,
+    TextPart,
+    UserPromptPart,
+    VideoUrl,
+)
+
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from fastmcp import Client
+    from mcp.types import ContentBlock
+    from pydantic_ai import ModelRequestPart, ModelResponsePart
+
+
+logger = get_logger(__name__)
+
+
+def to_mcp_messages(
+    part: ModelRequestPart | ModelResponsePart,
+) -> list[PromptMessage]:
+    """Convert internal PromptMessage to MCP PromptMessage."""
+    messages = []
+    match part:
+        case UserPromptPart(content=str() as c):
+            content = TextContent(type="text", text=c)
+            messages.append(PromptMessage(role="user", content=content))
+        case UserPromptPart(content=content_items):
+            for item in content_items:
+                match item:
+                    case BinaryContent():
+                        if item.is_audio:
+                            encoded = base64.b64encode(item.data).decode("utf-8")
+                            audio = AudioContent(
+                                type="audio", data=encoded, mimeType=item.media_type
+                            )
+                            messages.append(PromptMessage(role="user", content=audio))
+                        elif item.is_image:
+                            encoded = base64.b64encode(item.data).decode("utf-8")
+                            image = ImageContent(
+                                type="image", data=encoded, mimeType=item.media_type
+                            )
+                        messages.append(PromptMessage(role="user", content=image))
+                    case FileUrl(url=url):
+                        content = TextContent(type="text", text=url)
+                        messages.append(PromptMessage(role="user", content=content))
+
+        case SystemPromptPart(content=msg):
+            messages.append(PromptMessage(role="user", content=TextContent(type="text", text=msg)))
+        case TextPart(content=msg):
+            messages.append(
+                PromptMessage(role="assistant", content=TextContent(type="text", text=msg))
+            )
+    return messages
+
+
+def _url_from_mime_type(uri: str, mime_type: str | None) -> FileUrl:
+    """Convert URI to appropriate pydantic-ai URL type based on MIME type."""
+    if not mime_type:
+        return DocumentUrl(url=uri)
+
+    if mime_type.startswith("image/"):
+        return ImageUrl(url=uri)
+    if mime_type.startswith("audio/"):
+        return AudioUrl(url=uri)
+    if mime_type.startswith("video/"):
+        return VideoUrl(url=uri)
+    return DocumentUrl(url=uri)
+
+
+async def from_mcp_content(
+    mcp_content: Sequence[ContentBlock | TextResourceContents | BlobResourceContents],
+    client: Client[Any] | None = None,
+) -> list[str | BinaryContent]:
+    """Convert MCP content blocks to PydanticAI content types.
+
+    If a FastMCP client is given, this function will try to resolve the ResourceLinks.
+
+    """
+    contents: list[Any] = []
+
+    for block in mcp_content:
+        match block:
+            case TextContent(text=text):
+                contents.append(text)
+            case TextResourceContents(text=text):
+                contents.append(text)
+            case ImageContent(data=data, mimeType=mime_type):
+                decoded_data = base64.b64decode(data)
+                img = BinaryImage(data=decoded_data, media_type=mime_type)
+                contents.append(img)
+            case AudioContent(data=data, mimeType=mime_type):
+                decoded_data = base64.b64decode(data)
+                content = BinaryContent(data=decoded_data, media_type=mime_type)
+                contents.append(content)
+            case BlobResourceContents(blob=blob):
+                decoded_data = base64.b64decode(blob)
+                mime = "application/octet-stream"
+                content = BinaryContent(data=decoded_data, media_type=mime)
+                contents.append(content)
+            case ResourceLink(uri=uri, mimeType=mime_type):
+                if client:
+                    try:
+                        res = await client.read_resource(uri)
+                        nested = await from_mcp_content(res, client)
+                        contents.extend(nested)
+                        continue
+                    except Exception:  # noqa: BLE001
+                        # Fallback to URL if reading fails
+                        logger.warning("Failed to read resource", uri=uri)
+                # Convert to appropriate URL type based on MIME type
+                contents.append(_url_from_mime_type(str(uri), mime_type))
+            # mypy doesnt understand exhaustivness check for "nested typing", so we nest match-case
+            case EmbeddedResource(resource=resource):
+                match resource:
+                    case TextResourceContents(text=text):
+                        contents.append(text)
+                    case BlobResourceContents() as blob_resource:
+                        contents.append(f"[Binary data: {blob_resource.mimeType}]")
+                    case _ as unreachable:
+                        assert_never(unreachable)  # ty: ignore
+            case _ as unreachable:
+                assert_never(unreachable)
+    return contents
+
+
+def content_block_as_text(content: ContentBlock) -> str:
+    match content:
+        case TextContent(text=text):
+            return text
+        case EmbeddedResource(resource=resource):
+            match resource:
+                case TextResourceContents() as text_contents:
+                    return text_contents.text
+                case BlobResourceContents() as blob_contents:
+                    return f"[Resource: {blob_contents.uri}]"
+                case _ as unreachable:
+                    assert_never(unreachable)  # ty: ignore
+        case ResourceLink(uri=uri, description=desc):
+            return f"[Resource Link: {uri}] - {desc}" if desc else f"[Resource Link: {uri}]"
+        case ImageContent(mimeType=mime_type):
+            return f"[Image: {mime_type}]"
+        case AudioContent(mimeType=mime_type):
+            return f"[Audio: {mime_type}]"
+        case _ as unreachable:
+            assert_never(unreachable)

agentpool/mcp_server/helpers.py

@@ -0,0 +1,47 @@
+"""Helper functions for MCP server client operations.
+
+This module contains stateless utility functions that support MCP tool conversion
+and content handling for PydanticAI integration.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from mcp.types import ContentBlock, Tool as MCPTool
+
+
+logger = get_logger(__name__)
+
+
+def mcp_tool_to_fn_schema(tool: MCPTool) -> dict[str, Any]:
+    """Convert MCP tool to OpenAI function schema format."""
+    return {
+        "name": tool.name,
+        "description": tool.description or "",
+        "parameters": tool.inputSchema or {"type": "object", "properties": {}},
+    }
+
+
+def extract_text_content(mcp_content: list[ContentBlock]) -> str:
+    """Extract text content from MCP content blocks.
+
+    Args:
+        mcp_content: List of MCP content blocks
+
+    Returns:
+        First available text content or fallback string
+    """
+    from mcp.types import TextContent
+
+    for block in mcp_content:
+        match block:
+            case TextContent(text=text):
+                return text
+
+    # Fallback: stringify the content
+    return str(mcp_content[0]) if mcp_content else "Tool executed successfully"