agentpool 2.1.9__py3-none-any.whl → 2.2.3__py3-none-any.whl

agentpool_toolsets/mcp_discovery/toolset.py

@@ -0,0 +1,454 @@
+# mypy: disable-error-code="import-not-found,unused-ignore"
+"""MCP Discovery Toolset - dynamic MCP server exploration and tool execution.
+
+This toolset provides a stable interface for agents to discover and use MCP servers
+on-demand without needing to preload all tools upfront. This preserves prompt cache
+stability while enabling access to the entire MCP ecosystem.
+
+The toolset exposes three main capabilities:
+1. search_mcp_servers - Semantic search over 1000+ servers (uses pre-built index)
+2. list_mcp_tools - Get tools from a specific server (connects on-demand)
+3. call_mcp_tool - Execute a tool on any server (reuses connections)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from pydantic import HttpUrl
+
+from agentpool.log import get_logger
+from agentpool.mcp_server.client import MCPClient
+from agentpool.mcp_server.registries.official_registry_client import (
+    MCPRegistryClient,
+    MCPRegistryError,
+)
+from agentpool.resource_providers import ResourceProvider
+
+
+if TYPE_CHECKING:
+    from agentpool.agents.context import AgentContext
+    from agentpool.mcp_server.registries.official_registry_client import RegistryServer
+    from agentpool.tools.base import Tool
+    from agentpool_config.mcp_server import MCPServerConfig
+
+
+logger = get_logger(__name__)
+
+# Path to the pre-built semantic search index (parquet file loaded into LanceDB at runtime)
+PARQUET_PATH = Path(__file__).parent / "data" / "mcp_servers.parquet"
+
+
+class MCPDiscoveryToolset(ResourceProvider):
+    """Toolset for dynamic MCP server discovery and tool execution.
+
+    This toolset allows agents to:
+    - Search the MCP registry for servers by keyword
+    - List tools available on a specific server
+    - Call tools on any server without preloading
+
+    Connections are managed lazily - servers are only connected when needed,
+    and connections are kept alive for the session duration.
+    """
+
+    def __init__(
+        self,
+        name: str = "mcp_discovery",
+        registry_url: str = "https://registry.modelcontextprotocol.io",
+        allowed_servers: list[str] | None = None,
+        blocked_servers: list[str] | None = None,
+    ) -> None:
+        """Initialize the MCP Discovery toolset.
+
+        Args:
+            name: Name for this toolset provider
+            registry_url: Base URL for the MCP registry API
+            allowed_servers: If set, only these server names can be used
+            blocked_servers: Server names that cannot be used
+        """
+        super().__init__(name=name)
+        self._registry_url = registry_url
+        self._registry: MCPRegistryClient | None = None
+        self._connections: dict[str, MCPClient] = {}
+        self._server_cache: dict[str, RegistryServer] = {}
+        self._tools_cache: dict[str, list[dict[str, Any]]] = {}
+        self._allowed_servers = set(allowed_servers) if allowed_servers else None
+        self._blocked_servers = set(blocked_servers) if blocked_servers else set()
+        self._tools: list[Tool] | None = None
+        # Lazy-loaded semantic search components
+        self._db: Any = None
+        self._table: Any = None
+        self._embed_model: Any = None
+        self._tmpdir: str | None = None
+
+    def _get_registry(self) -> MCPRegistryClient:
+        """Get or create the registry client."""
+        if self._registry is None:
+            self._registry = MCPRegistryClient(base_url=self._registry_url)
+        return self._registry
+
+    def _get_search_index(self) -> Any:
+        """Get or create the LanceDB search index from parquet file."""
+        if self._table is not None:
+            return self._table
+
+        import tempfile
+
+        import lancedb  # type: ignore[import-untyped]
+        import pyarrow as pa  # type: ignore[import-untyped]
+        import pyarrow.parquet as pq  # type: ignore[import-untyped]
+
+        if not PARQUET_PATH.exists():
+            msg = f"MCP registry index not found at {PARQUET_PATH}. Run build_mcp_registry_index.py"
+            raise FileNotFoundError(msg)
+
+        # Load parquet
+        arrow_table = pq.read_table(PARQUET_PATH)
+
+        # Convert vector column to fixed-size list for LanceDB vector search
+        # LanceDB requires fixed-size vectors, not variable-length lists
+        vectors = arrow_table.column("vector").to_pylist()
+        if vectors:
+            vec_dim = len(vectors[0])
+            fixed_vectors = pa.FixedSizeListArray.from_arrays(
+                pa.array([v for vec in vectors for v in vec], type=pa.float32()),
+                list_size=vec_dim,
+            )
+            # Replace the vector column
+            col_idx = arrow_table.schema.get_field_index("vector")
+            arrow_table = arrow_table.set_column(col_idx, "vector", fixed_vectors)
+
+        # Use a temp directory for LanceDB (it needs a path but we're loading from parquet)
+        if self._db is None:
+            self._tmpdir = tempfile.mkdtemp(prefix="mcp_discovery_")
+            self._db = lancedb.connect(self._tmpdir)
+
+        self._table = self._db.create_table("servers", arrow_table, mode="overwrite")
+        return self._table
+
+    def _get_embed_model(self) -> Any:
+        """Get or create the FastEmbed model for query embedding."""
+        if self._embed_model is not None:
+            return self._embed_model
+
+        from fastembed import TextEmbedding
+
+        self._embed_model = TextEmbedding("BAAI/bge-small-en-v1.5")
+        return self._embed_model
+
+    def _is_server_allowed(self, server_name: str) -> bool:
+        """Check if a server is allowed to be used."""
+        if server_name in self._blocked_servers:
+            return False
+        if self._allowed_servers is not None:
+            return server_name in self._allowed_servers
+        return True
+
+    async def _get_server_config(self, server_name: str) -> MCPServerConfig:
+        """Get connection config for a server from the registry."""
+        from agentpool_config.mcp_server import (
+            SSEMCPServerConfig,
+            StreamableHTTPMCPServerConfig,
+        )
+
+        # Check cache first
+        server: RegistryServer
+        if server_name in self._server_cache:
+            server = self._server_cache[server_name]
+        else:
+            # Use list_servers and find by name - more reliable than get_server
+            registry = self._get_registry()
+            all_servers = await registry.list_servers()
+            found_server: RegistryServer | None = None
+            for s in all_servers:
+                if s.name == server_name:
+                    found_server = s
+                    break
+            if found_server is None:
+                msg = f"Server {server_name!r} not found in registry"
+                raise MCPRegistryError(msg)
+            server = found_server
+            self._server_cache[server_name] = server
+
+        # Find a usable remote endpoint
+        for remote in server.remotes:
+            if remote.type == "sse":
+                return SSEMCPServerConfig(url=HttpUrl(remote.url))
+            if remote.type in ("streamable-http", "http"):
+                return StreamableHTTPMCPServerConfig(url=HttpUrl(remote.url))
+
+        msg = f"No supported remote transport for server {server_name!r}"
+        raise MCPRegistryError(msg)
+
+    async def _get_connection(self, server_name: str) -> MCPClient:
+        """Get or create a connection to a server."""
+        if server_name in self._connections:
+            client = self._connections[server_name]
+            if client.connected:
+                return client
+            # Connection dropped, remove and reconnect
+            del self._connections[server_name]
+
+        config = await self._get_server_config(server_name)
+        client = MCPClient(config=config)
+        await client.__aenter__()
+        self._connections[server_name] = client
+        logger.info("Connected to MCP server", server=server_name)
+        return client
+
+    async def _close_connections(self) -> None:
+        """Close all active connections."""
+        for name, client in list(self._connections.items()):
+            try:
+                await client.__aexit__(None, None, None)
+                logger.debug("Closed connection", server=name)
+            except Exception as e:  # noqa: BLE001
+                logger.warning("Error closing connection", server=name, error=e)
+        self._connections.clear()
+
+    async def get_tools(self) -> list[Tool]:
+        """Get the discovery tools."""
+        if self._tools is not None:
+            return self._tools
+
+        self._tools = [
+            self.create_tool(
+                self.search_mcp_servers,
+                category="search",
+                read_only=True,
+                idempotent=True,
+            ),
+            self.create_tool(
+                self.list_mcp_tools,
+                category="search",
+                read_only=True,
+                idempotent=True,
+            ),
+            self.create_tool(
+                self.call_mcp_tool,
+                category="execute",
+                open_world=True,
+            ),
+        ]
+        return self._tools
+
+    async def search_mcp_servers(  # noqa: D417
+        self,
+        agent_ctx: AgentContext,
+        query: str,
+        max_results: int = 10,
+    ) -> str:
+        """Search the MCP registry for servers matching a query.
+
+        Uses semantic search over 1000+ indexed MCP servers. The search understands
+        meaning, not just keywords - e.g., "web scraping" finds crawlers too.
+
+        Args:
+            query: Search term (e.g., "github issues", "database sql", "file system")
+            max_results: Maximum number of results to return
+
+        Returns:
+            List of matching servers with names and descriptions
+        """
+        await agent_ctx.events.tool_call_start(
+            title=f"Searching MCP servers: {query}",
+            kind="search",
+        )
+
+        try:
+            # Get embedding for query
+            model = self._get_embed_model()
+            query_embedding = next(iter(model.embed([query]))).tolist()
+
+            # Search the index
+            table = self._get_search_index()
+            results = table.search(query_embedding).limit(max_results * 2).to_arrow()
+
+            if len(results) == 0:
+                return f"No MCP servers found matching '{query}'"
+
+            # Format results, filtering by allowed/blocked
+            lines = [f"Found MCP servers matching '{query}':\n"]
+            count = 0
+            for i in range(len(results)):
+                name = results["name"][i].as_py()
+
+                # Filter by allowed/blocked
+                if not self._is_server_allowed(name):
+                    continue
+
+                desc = results["description"][i].as_py()
+                version = results["version"][i].as_py()
+                has_remote = results["has_remote"][i].as_py()
+                remote_types = results["remote_types"][i].as_py()
+
+                lines.append(f"**{name}** (v{version})")
+                lines.append(f"  {desc}")
+                if has_remote and remote_types:
+                    lines.append(f"  Transports: {remote_types}")
+                lines.append("")
+
+                count += 1
+                if count >= max_results:
+                    break
+
+            if count == 0:
+                return f"No MCP servers found matching '{query}'"
+
+            lines[0] = f"Found {count} MCP servers matching '{query}':\n"
+            return "\n".join(lines)
+
+        except FileNotFoundError as e:
+            return f"Error: {e}"
+        except Exception as e:
+            logger.exception("Error searching MCP servers")
+            return f"Error searching MCP servers: {e}"
+
+    async def list_mcp_tools(  # noqa: D417
+        self,
+        agent_ctx: AgentContext,
+        server_name: str,
+    ) -> str:
+        """List all tools available on a specific MCP server.
+
+        This connects to the server (if not already connected) and retrieves
+        the list of available tools with their descriptions and parameters.
+
+        Args:
+            server_name: Name of the MCP server (e.g., "com.github/github")
+
+        Returns:
+            List of tools with names, descriptions, and parameter schemas
+        """
+        await agent_ctx.events.tool_call_start(
+            title=f"Listing tools from: {server_name}",
+            kind="search",
+        )
+
+        if not self._is_server_allowed(server_name):
+            return f"Error: Server '{server_name}' is not allowed"
+
+        try:
+            # Check tools cache first
+            if server_name in self._tools_cache:
+                tools_data = self._tools_cache[server_name]
+            else:
+                client = await self._get_connection(server_name)
+                mcp_tools = await client.list_tools()
+
+                # Convert to serializable format and cache
+                tools_data = []
+                for tool in mcp_tools:
+                    tool_info: dict[str, Any] = {
+                        "name": tool.name,
+                        "description": tool.description or "No description",
+                    }
+                    # Include parameter info
+                    if tool.inputSchema:
+                        props = tool.inputSchema.get("properties", {})
+                        required = set(tool.inputSchema.get("required", []))
+                        params = []
+                        for pname, pschema in props.items():
+                            param_str = pname
+                            if pname in required:
+                                param_str += " (required)"
+                            if "type" in pschema:
+                                param_str += f": {pschema['type']}"
+                            if "description" in pschema:
+                                param_str += f" - {pschema['description']}"
+                            params.append(param_str)
+                        if params:
+                            tool_info["parameters"] = params
+                    tools_data.append(tool_info)
+
+                self._tools_cache[server_name] = tools_data
+
+            if not tools_data:
+                return f"No tools found on server '{server_name}'"
+
+            # Format output
+            lines = [f"Tools available on **{server_name}** ({len(tools_data)} tools):\n"]
+            for tool_info in tools_data:
+                lines.append(f"### {tool_info['name']}")
+                lines.append(f"{tool_info['description']}")
+                if "parameters" in tool_info:
+                    lines.append("Parameters:")
+                    lines.extend(f"  - {param}" for param in tool_info["parameters"])
+                lines.append("")
+
+            return "\n".join(lines)
+
+        except MCPRegistryError as e:
+            return f"Error: {e}"
+        except Exception as e:
+            logger.exception("Error listing MCP tools", server=server_name)
+            return f"Error listing tools from '{server_name}': {e}"
+
+    async def call_mcp_tool(  # noqa: D417
+        self,
+        agent_ctx: AgentContext,
+        server_name: str,
+        tool_name: str,
+        arguments: dict[str, Any] | None = None,
+    ) -> str | Any:
+        """Call a tool on an MCP server.
+
+        Use this to execute a specific tool on an MCP server. The server
+        connection is reused if already established.
+
+        Args:
+            server_name: Name of the MCP server (e.g., "com.github/github")
+            tool_name: Name of the tool to call
+            arguments: Arguments to pass to the tool
+
+        Returns:
+            The result from the tool execution
+        """
+        await agent_ctx.events.tool_call_start(
+            title=f"Calling {tool_name} on {server_name}",
+            kind="execute",
+        )
+
+        if not self._is_server_allowed(server_name):
+            return f"Error: Server '{server_name}' is not allowed"
+
+        try:
+            client = await self._get_connection(server_name)
+
+            # Extract text content from result
+            from mcp.types import TextContent
+
+            result = await client._client.call_tool(tool_name, arguments or {})
+
+            # Process result content
+            text_parts = [
+                content.text for content in result.content if isinstance(content, TextContent)
+            ]
+
+            if text_parts:
+                return "\n".join(text_parts)
+
+            # Return raw result if no text content
+            return str(result)
+
+        except MCPRegistryError as e:
+            return f"Error: {e}"
+        except Exception as e:
+            logger.exception("Error calling MCP tool", server=server_name, tool=tool_name)
+            return f"Error calling '{tool_name}' on '{server_name}': {e}"
+
+    async def cleanup(self) -> None:
+        """Clean up resources."""
+        await self._close_connections()
+        if self._registry:
+            await self._registry.close()
+            self._registry = None
+        # Clean up temp directory used for LanceDB
+        if self._tmpdir:
+            import shutil
+
+            shutil.rmtree(self._tmpdir, ignore_errors=True)
+            self._tmpdir = None
+        self._db = None
+        self._table = None

agentpool_toolsets/mcp_run_toolset.py

@@ -3,21 +3,37 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Literal, Self, cast
 
 from schemez import OpenAIFunctionDefinition
 
+from agentpool.log import get_logger
 from agentpool.resource_providers import ResourceProvider
 
 
 if TYPE_CHECKING:
+    from types import TracebackType
+
+    from mcp import ClientSession
     from mcp.types import CallToolResult
 
     from agentpool.tools.base import Tool
 
+logger = get_logger(__name__)
+
 
 class McpRunTools(ResourceProvider):
-    """Provider for MCP.run tools."""
+    """Provider for MCP.run tools.
+
+    Maintains a persistent SSE connection to MCP.run to receive tool change
+    notifications. Use as an async context manager to ensure proper cleanup.
+
+    Example:
+        async with McpRunTools("default") as provider:
+            tools = await provider.get_tools()
+    """
+
+    kind: Literal["mcp_run"] = "mcp_run"
 
     def __init__(self, entity_id: str, session_id: str | None = None) -> None:
         from mcp_run import Client, ClientConfig  # type: ignore[import-untyped]
@@ -27,6 +43,51 @@ class McpRunTools(ResourceProvider):
         config = ClientConfig()
         self.client = Client(session_id=id_, config=config)
         self._tools: list[Tool] | None = None
+        self._session: ClientSession | None = None
+        self._mcp_client_ctx: Any = None  # Context manager for persistent connection
+
+    async def __aenter__(self) -> Self:
+        """Start persistent SSE connection."""
+        # Create MCPClient from mcp_run
+        mcp_client = self.client.mcp_sse()
+        self._mcp_client_ctx = mcp_client.connect()
+        self._session = await self._mcp_client_ctx.__aenter__()
+
+        # Set up notification handler for tool changes
+        # The MCP ClientSession dispatches notifications via _received_notification
+        # We monkey-patch it to intercept ToolListChangedNotification
+        original_handler = self._session._received_notification
+
+        async def notification_handler(notification: Any) -> None:
+            from mcp.types import ToolListChangedNotification
+
+            if isinstance(notification.root, ToolListChangedNotification):
+                logger.info("MCP.run tool list changed notification received")
+                await self._on_tools_changed()
+            # Call original handler
+            await original_handler(notification)
+
+        self._session._received_notification = notification_handler  # type: ignore[method-assign]
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Close persistent SSE connection."""
+        if self._mcp_client_ctx:
+            await self._mcp_client_ctx.__aexit__(exc_type, exc_val, exc_tb)
+            self._mcp_client_ctx = None
+            self._session = None
+
+    async def _on_tools_changed(self) -> None:
+        """Handle tool list change notification."""
+        logger.info("MCP.run tools changed, refreshing cache")
+        self._tools = None
+        await self.tools_changed.emit(self.create_change_event("tools"))
 
     async def get_tools(self) -> list[Tool]:
         """Get tools from MCP.run."""
@@ -36,10 +97,20 @@ class McpRunTools(ResourceProvider):
 
         self._tools = []
         for name, tool in self.client.tools.items():
-
-            async def run(tool_name: str = name, **input_dict: Any) -> CallToolResult:
-                async with self.client.mcp_sse().connect() as session:
-                    return await session.call_tool(tool_name, arguments=input_dict)  # type: ignore[no-any-return]
+            # Capture session for use in tool calls
+            session = self._session
+
+            async def run(
+                tool_name: str = name,
+                _session: ClientSession | None = session,
+                **input_dict: Any,
+            ) -> CallToolResult:
+                if _session is not None:
+                    # Use persistent session if available
+                    return await _session.call_tool(tool_name, arguments=input_dict)
+                # Fallback to creating a new connection (when not using context manager)
+                async with self.client.mcp_sse().connect() as new_session:
+                    return await new_session.call_tool(tool_name, arguments=input_dict)  # type: ignore[no-any-return]
 
             run.__name__ = name
             wrapped_tool = self.create_tool(
@@ -49,6 +120,13 @@ class McpRunTools(ResourceProvider):
 
         return self._tools
 
+    async def refresh_tools(self) -> None:
+        """Manually refresh tools from MCP.run and emit change event."""
+        logger.info("Manually refreshing MCP.run tools")
+        self._tools = None
+        await self.get_tools()
+        await self.tools_changed.emit(self.create_change_event("tools"))
+
 
 if __name__ == "__main__":
     import anyio