agentpool 2.1.9__py3-none-any.whl

agentpool/mcp_server/manager.py

@@ -0,0 +1,232 @@
+"""MCP server management for AgentPool."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+from contextlib import AsyncExitStack
+from typing import TYPE_CHECKING, Any, Self, cast
+
+import anyio
+from pydantic_ai import BinaryContent, BinaryImage, UsageLimits
+
+from agentpool.log import get_logger
+from agentpool.resource_providers import AggregatingResourceProvider, ResourceProvider
+from agentpool.resource_providers.mcp_provider import MCPResourceProvider
+from agentpool_config.mcp_server import BaseMCPServerConfig
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from types import TracebackType
+
+    from mcp import types
+    from mcp.shared.context import RequestContext
+    from mcp.types import SamplingMessage
+    from pydantic_ai import UserContent
+
+    from agentpool_config.mcp_server import MCPServerConfig
+
+
+logger = get_logger(__name__)
+
+
+class MCPManager:
+    """Manages MCP server connections and distributes resource providers."""
+
+    def __init__(
+        self,
+        name: str = "mcp",
+        owner: str | None = None,
+        servers: Sequence[MCPServerConfig | str] | None = None,
+        accessible_roots: list[str] | None = None,
+    ) -> None:
+        self.name = name
+        self.owner = owner
+        self.servers: list[MCPServerConfig] = []
+        for server in servers or []:
+            self.add_server_config(server)
+        self.providers: list[MCPResourceProvider] = []
+        self.aggregating_provider = AggregatingResourceProvider(
+            providers=cast(list[ResourceProvider], self.providers),
+            name=f"{name}_aggregated",
+        )
+        self.exit_stack = AsyncExitStack()
+        self._accessible_roots = accessible_roots
+
+    def add_server_config(self, cfg: MCPServerConfig | str) -> None:
+        """Add a new MCP server to the manager."""
+        resolved = BaseMCPServerConfig.from_string(cfg) if isinstance(cfg, str) else cfg
+        self.servers.append(resolved)
+
+    def __repr__(self) -> str:
+        return f"MCPManager(name={self.name!r}, servers={len(self.servers)})"
+
+    async def __aenter__(self) -> Self:
+        try:
+            if tasks := [self._setup_server(server) for server in self.servers]:
+                await asyncio.gather(*tasks)
+        except Exception as e:
+            await self.__aexit__(type(e), e, e.__traceback__)
+            msg = "Failed to initialize MCP manager"
+            raise RuntimeError(msg) from e
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.cleanup()
+
+    async def _sampling_callback(
+        self,
+        messages: list[SamplingMessage],
+        params: types.CreateMessageRequestParams,
+        context: RequestContext[Any, Any, Any],
+    ) -> str:
+        """Handle MCP sampling by creating a new agent with specified preferences."""
+        from mcp import types
+
+        from agentpool.agents import Agent
+
+        # Convert messages to prompts for the agent
+        prompts: list[UserContent] = []
+        for mcp_msg in messages:
+            match mcp_msg.content:
+                case types.TextContent(text=text):
+                    prompts.append(text)
+                case types.ImageContent(data=data, mimeType=mime_type):
+                    binary_data = base64.b64decode(data)
+                    prompts.append(BinaryImage(data=binary_data, media_type=mime_type))
+                case types.AudioContent(data=data, mimeType=mime_type):
+                    binary_data = base64.b64decode(data)
+                    prompts.append(BinaryContent(data=binary_data, media_type=mime_type))
+
+        # Extract model from preferences
+        model = None
+        if (prefs := params.modelPreferences) and prefs.hints and prefs.hints[0].name:
+            model = prefs.hints[0].name
+        # Create usage limits from sampling parameters
+        limits = UsageLimits(output_tokens_limit=params.maxTokens, request_limit=1)
+        # TODO: Apply temperature from params.temperature
+        sys_prompt = params.systemPrompt or ""
+        agent = Agent(name="sampling-agent", model=model, system_prompt=sys_prompt, session=False)
+        try:
+            async with agent:
+                result = await agent.run(*prompts, store_history=False, usage_limits=limits)
+                return result.content
+
+        except Exception as e:
+            logger.exception("Sampling failed")
+            return f"Sampling failed: {e!s}"
+
+    async def _setup_server(self, config: MCPServerConfig) -> None:
+        """Set up a single MCP server resource provider."""
+        if not config.enabled:
+            return
+
+        provider = MCPResourceProvider(
+            server=config,
+            name=f"{self.name}_{config.client_id}",
+            owner=self.owner,
+            source="pool" if self.owner == "pool" else "node",
+            sampling_callback=self._sampling_callback,
+            accessible_roots=self._accessible_roots,
+        )
+        provider = await self.exit_stack.enter_async_context(provider)
+        self.providers.append(provider)
+
+    def get_mcp_providers(self) -> list[MCPResourceProvider]:
+        """Get all MCP resource providers managed by this manager."""
+        return list(self.providers)
+
+    def get_aggregating_provider(self) -> AggregatingResourceProvider:
+        """Get the aggregating provider that contains all MCP providers."""
+        return self.aggregating_provider
+
+    async def setup_server_runtime(self, config: MCPServerConfig) -> MCPResourceProvider:
+        """Set up a single MCP server at runtime while manager is running.
+
+        Returns:
+            The newly created and initialized MCPResourceProvider
+        """
+        if not config.enabled:
+            msg = f"Server config {config.client_id} is disabled"
+            raise ValueError(msg)
+
+        # Add the config first
+        self.add_server_config(config)
+        provider = MCPResourceProvider(
+            server=config,
+            name=f"{self.name}_{config.client_id}",
+            owner=self.owner,
+            source="pool" if self.owner == "pool" else "node",
+            sampling_callback=self._sampling_callback,
+            accessible_roots=self._accessible_roots,
+        )
+        provider = await self.exit_stack.enter_async_context(provider)
+        self.providers.append(provider)
+        # Note: AggregatingResourceProvider automatically sees the new provider
+        # since it references self.providers list
+
+        return provider
+
+    async def cleanup(self) -> None:
+        """Clean up all MCP connections and providers."""
+        try:
+            try:
+                # Clean up exit stack (which includes MCP providers)
+                await self.exit_stack.aclose()
+            except RuntimeError as e:
+                if "different task" in str(e):
+                    # Handle task context mismatch
+                    current_task = asyncio.current_task()
+                    if current_task:
+                        loop = asyncio.get_running_loop()
+                        await loop.create_task(self.exit_stack.aclose())
+                else:
+                    raise
+
+            self.providers.clear()
+
+        except Exception as e:
+            msg = "Error during MCP manager cleanup"
+            logger.exception(msg, exc_info=e)
+            raise RuntimeError(msg) from e
+
+    @property
+    def active_servers(self) -> list[str]:
+        """Get IDs of active servers."""
+        return [provider.server.client_id for provider in self.providers]
+
+
+if __name__ == "__main__":
+    from agentpool_config.mcp_server import StdioMCPServerConfig
+
+    cfg = StdioMCPServerConfig(
+        command="uv",
+        args=["run", "/home/phil65/dev/oss/agentpool/tests/mcp_server/server.py"],
+    )
+
+    async def main() -> None:
+        manager = MCPManager(servers=[cfg])
+        async with manager:
+            providers = manager.get_mcp_providers()
+            print(f"Found {len(providers)} providers")
+            provider = providers[0]
+            prompts = await provider.get_prompts()
+            print(f"Found prompts: {prompts}")
+            # Test static prompt (no arguments)
+            static_prompt = next(p for p in prompts if p.name == "static_prompt")
+            print(f"\n--- Testing static prompt: {static_prompt} ---")
+            components = await static_prompt.get_components()
+            assert components, "No prompt components found"
+            print(f"Found {len(components)} prompt components:")
+            for i, component in enumerate(components):
+                comp_type = type(component).__name__
+                print(f"  {i + 1}. {comp_type}: {component.content}")
+
+    anyio.run(main)

agentpool/mcp_server/message_handler.py

@@ -0,0 +1,164 @@
+"""FastMCP message handler for agentpool."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from mcp.shared.session import RequestResponder
+    import mcp.types
+
+    from agentpool.mcp_server import MCPClient
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class MCPMessageHandler:
+    """Custom message handler that bridges FastMCP to agentpool notifications."""
+
+    client: MCPClient
+    """The MCP client instance."""
+    tool_change_callback: Callable[[], Awaitable[None]] | None = None
+    """Tool change callback."""
+    prompt_change_callback: Callable[[], Awaitable[None]] | None = None
+    """Prompt change callback."""
+    resource_change_callback: Callable[[], Awaitable[None]] | None = None
+    """Resource change callback."""
+
+    async def __call__(
+        self,
+        message: RequestResponder[mcp.types.ServerRequest, mcp.types.ClientResult]
+        | mcp.types.ServerNotification
+        | Exception,
+    ) -> None:
+        """Handle FastMCP messages by dispatching to appropriate handlers."""
+        from mcp.shared.session import RequestResponder
+        import mcp.types
+
+        await self.on_message(message)
+        match message:
+            # requests
+            case RequestResponder() as responder:
+                await self.on_request(responder)
+                # Handle specific requests
+                root = responder.request.root
+                match root:
+                    case mcp.types.PingRequest():
+                        await self.on_ping(root)
+                    case mcp.types.ListRootsRequest():
+                        await self.on_list_roots(root)
+                    case mcp.types.CreateMessageRequest():
+                        await self.on_create_message(root)
+
+            case mcp.types.ServerNotification() as notification:
+                await self.on_notification(notification)
+                root = notification.root
+                match root:
+                    case mcp.types.CancelledNotification():
+                        await self.on_cancelled(root)
+                    case mcp.types.ProgressNotification():
+                        await self.on_progress(root)
+                    case mcp.types.LoggingMessageNotification():
+                        await self.on_logging_message(root)
+                    case mcp.types.ToolListChangedNotification():
+                        await self.on_tool_list_changed(root)
+                    case mcp.types.ResourceListChangedNotification():
+                        await self.on_resource_list_changed(root)
+                    case mcp.types.PromptListChangedNotification():
+                        await self.on_prompt_list_changed(root)
+                    case mcp.types.ResourceUpdatedNotification():
+                        await self.on_resource_updated(root)
+                    case mcp.types.ElicitCompleteNotification():
+                        await self.on_elicit_complete(root)
+
+            case Exception():
+                await self.on_exception(message)
+
+    async def on_message(
+        self,
+        message: RequestResponder[mcp.types.ServerRequest, mcp.types.ClientResult]
+        | mcp.types.ServerNotification
+        | Exception,
+    ) -> None:
+        """Handle generic messages."""
+
+    async def on_request(
+        self, message: RequestResponder[mcp.types.ServerRequest, mcp.types.ClientResult]
+    ) -> None:
+        """Handle requests."""
+
+    async def on_notification(self, message: mcp.types.ServerNotification) -> None:
+        """Handle server notifications."""
+
+    async def on_tool_list_changed(self, message: mcp.types.ToolListChangedNotification) -> None:
+        """Handle tool list changes."""
+        logger.info("MCP tool list changed", message=message)
+        # Call the tool change callback if provided
+        if self.tool_change_callback:
+            await self.tool_change_callback()
+
+    async def on_resource_list_changed(
+        self, message: mcp.types.ResourceListChangedNotification
+    ) -> None:
+        """Handle resource list changes."""
+        logger.info("MCP resource list changed", message=message)
+        # Call the resource change callback if provided
+        if self.resource_change_callback:
+            await self.resource_change_callback()
+
+    async def on_resource_updated(self, message: mcp.types.ResourceUpdatedNotification) -> None:
+        """Handle resource updates."""
+        # ResourceUpdatedNotification has uri directly, not in params
+        logger.info("MCP resource updated", uri=getattr(message, "uri", "unknown"))
+
+    async def on_progress(self, message: mcp.types.ProgressNotification) -> None:
+        """Handle progress notifications with proper context."""
+        # Note: Progress notifications from MCP servers are now handled per-tool-call
+        # with the contextual progress handler, so global notifications are ignored
+
+    async def on_prompt_list_changed(
+        self, message: mcp.types.PromptListChangedNotification
+    ) -> None:
+        """Handle prompt list changes."""
+        logger.info("MCP prompt list changed", message=message)
+        # Call the prompt change callback if provided
+        if self.prompt_change_callback:
+            await self.prompt_change_callback()
+
+    async def on_cancelled(self, message: mcp.types.CancelledNotification) -> None:
+        """Handle cancelled operations."""
+        logger.info("MCP operation cancelled", message=message)
+
+    async def on_logging_message(self, message: mcp.types.LoggingMessageNotification) -> None:
+        """Handle server log messages."""
+        # This is handled by _log_handler, but keep for completeness
+
+    async def on_exception(self, message: Exception) -> None:
+        """Handle exceptions."""
+        logger.error("MCP client exception", error=message)
+
+    async def on_ping(self, message: mcp.types.PingRequest) -> None:
+        """Handle ping requests."""
+
+    async def on_list_roots(self, message: mcp.types.ListRootsRequest) -> None:
+        """Handle list roots requests."""
+
+    async def on_create_message(self, message: mcp.types.CreateMessageRequest) -> None:
+        """Handle create message requests."""
+
+    async def on_elicit_complete(self, message: mcp.types.ElicitCompleteNotification) -> None:
+        """Handle elicitation completion notifications.
+
+        Sent by servers when a URL mode elicitation completes out-of-band.
+        """
+        logger.info(
+            "MCP elicitation completed",
+            elicitation_id=message.params.elicitationId,
+        )

agentpool/mcp_server/registries/__init__.py

@@ -0,0 +1 @@
+"""MCP registry clients."""