feat(mcp): add experimental agent managed connection support

Author: dbschmigelski

src/strands/_async.py

@@ -0,0 +1,31 @@
+"""Private async execution utilities."""
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from typing import Awaitable, Callable, TypeVar
+
+T = TypeVar("T")
+
+
+def run_async(async_func: Callable[[], Awaitable[T]]) -> T:
+    """Run an async function in a separate thread to avoid event loop conflicts.
+
+    This utility handles the common pattern of running async code from sync contexts
+    by using ThreadPoolExecutor to isolate the async execution.
+
+    Args:
+        async_func: A callable that returns an awaitable
+
+    Returns:
+        The result of the async function
+    """
+
+    async def execute_async() -> T:
+        return await async_func()
+
+    def execute() -> T:
+        return asyncio.run(execute_async())
+
+    with ThreadPoolExecutor() as executor:
+        future = executor.submit(execute)
+        return future.result()

src/strands/agent/agent.py

@@ -9,11 +9,9 @@
 2. Method-style for direct tool access: `agent.tool.tool_name(param1="value")`
 """
 
-import asyncio
 import json
 import logging
 import random
-from concurrent.futures import ThreadPoolExecutor
 from typing import (
     Any,
     AsyncGenerator,
@@ -31,7 +29,9 @@
 from pydantic import BaseModel
 
 from .. import _identifier
+from .._async import run_async
 from ..event_loop.event_loop import event_loop_cycle
+from ..experimental.tools import ToolProvider
 from ..handlers.callback_handler import PrintingCallbackHandler, null_callback_handler
 from ..hooks import (
     AfterInvocationEvent,
@@ -160,12 +160,7 @@ async def acall() -> ToolResult:
 
                     return tool_results[0]
 
-                def tcall() -> ToolResult:
-                    return asyncio.run(acall())
-
-                with ThreadPoolExecutor() as executor:
-                    future = executor.submit(tcall)
-                    tool_result = future.result()
+                tool_result = run_async(acall)
 
                 if record_direct_tool_call is not None:
                     should_record_direct_tool_call = record_direct_tool_call
@@ -208,7 +203,7 @@ def __init__(
         self,
         model: Union[Model, str, None] = None,
         messages: Optional[Messages] = None,
-        tools: Optional[list[Union[str, dict[str, str], Any]]] = None,
+        tools: Optional[list[Union[str, dict[str, str], ToolProvider, Any]]] = None,
         system_prompt: Optional[str] = None,
         callback_handler: Optional[
             Union[Callable[..., Any], _DefaultCallbackHandlerSentinel]
@@ -240,7 +235,8 @@ def __init__(
                 - File paths (e.g., "/path/to/tool.py")
                 - Imported Python modules (e.g., from strands_tools import current_time)
                 - Dictionaries with name/path keys (e.g., {"name": "tool_name", "path": "/path/to/tool.py"})
-                - Functions decorated with `@strands.tool` decorator.
+                - Functions decorated with `@strands.tool` decorator
+                - ToolProvider instances for managed tool collections
 
                 If provided, only these tools will be available. If None, all tools will be available.
             system_prompt: System prompt to guide model behavior.
@@ -333,6 +329,9 @@ def __init__(
         else:
             self.state = AgentState()
 
+        # Track cleanup state
+        self._cleanup_called = False
+
         self.tool_caller = Agent.ToolCaller(self)
 
         self.hooks = HookRegistry()
@@ -399,13 +398,7 @@ def __call__(self, prompt: AgentInput = None, **kwargs: Any) -> AgentResult:
                 - metrics: Performance metrics from the event loop
                 - state: The final state of the event loop
         """
-
-        def execute() -> AgentResult:
-            return asyncio.run(self.invoke_async(prompt, **kwargs))
-
-        with ThreadPoolExecutor() as executor:
-            future = executor.submit(execute)
-            return future.result()
+        return run_async(lambda: self.invoke_async(prompt, **kwargs))
 
     async def invoke_async(self, prompt: AgentInput = None, **kwargs: Any) -> AgentResult:
         """Process a natural language prompt through the agent's event loop.
@@ -459,13 +452,7 @@ def structured_output(self, output_model: Type[T], prompt: AgentInput = None) ->
         Raises:
             ValueError: If no conversation history or prompt is provided.
         """
-
-        def execute() -> T:
-            return asyncio.run(self.structured_output_async(output_model, prompt))
-
-        with ThreadPoolExecutor() as executor:
-            future = executor.submit(execute)
-            return future.result()
+        return run_async(lambda: self.structured_output_async(output_model, prompt))
 
     async def structured_output_async(self, output_model: Type[T], prompt: AgentInput = None) -> T:
         """This method allows you to get structured output from the agent.
@@ -524,6 +511,69 @@ async def structured_output_async(self, output_model: Type[T], prompt: AgentInpu
             finally:
                 self.hooks.invoke_callbacks(AfterInvocationEvent(agent=self))
 
+    def cleanup(self) -> None:
+        """Clean up resources used by the agent.
+
+        This method cleans up all tool providers that require explicit cleanup,
+        such as MCP clients. It should be called when the agent is no longer needed
+        to ensure proper resource cleanup.
+
+        Note: This method uses a "belt and braces" approach with automatic cleanup
+        through __del__ as a fallback, but explicit cleanup is recommended.
+        """
+        run_async(self.cleanup_async)
+
+    async def cleanup_async(self) -> None:
+        """Asynchronously clean up resources used by the agent.
+
+        This method cleans up all tool providers that require explicit cleanup,
+        such as MCP clients. It should be called when the agent is no longer needed
+        to ensure proper resource cleanup.
+
+        Note: This method uses a "belt and braces" approach with automatic cleanup
+        through __del__ as a fallback, but explicit cleanup is recommended.
+        """
+        if self._cleanup_called:
+            return
+
+        logger.debug("agent_id=<%s> | cleaning up agent resources", self.agent_id)
+
+        for provider in self.tool_registry.tool_providers:
+            try:
+                await provider.cleanup()
+                logger.debug(
+                    "agent_id=<%s>, provider=<%s> | cleaned up tool provider", self.agent_id, type(provider).__name__
+                )
+            except Exception as e:
+                logger.warning(
+                    "agent_id=<%s>, provider=<%s>, error=<%s> | failed to cleanup tool provider",
+                    self.agent_id,
+                    type(provider).__name__,
+                    e,
+                )
+
+        self._cleanup_called = True
+        logger.debug("agent_id=<%s> | agent cleanup complete", self.agent_id)
+
+    def __del__(self) -> None:
+        """Automatic cleanup when agent is garbage collected.
+
+        This serves as a fallback cleanup mechanism, but explicit cleanup() is preferred.
+        """
+        try:
+            if self._cleanup_called or not self.tool_registry.tool_providers:
+                return
+
+            logger.warning(
+                "agent_id=<%s> | Agent cleanup called via __del__. "
+                "Consider calling agent.cleanup() explicitly for better resource management.",
+                self.agent_id,
+            )
+            self.cleanup()
+        except Exception as e:
+            # Log exceptions during garbage collection cleanup for debugging
+            logger.debug("agent_id=<%s>, error=<%s> | exception during __del__ cleanup", self.agent_id, e)
+
     async def stream_async(
         self,
         prompt: AgentInput = None,

src/strands/experimental/__init__.py

@@ -2,3 +2,7 @@
 
 This module implements experimental features that are subject to change in future revisions without notice.
 """
+
+from . import tools
+
+__all__ = ["tools"]

src/strands/experimental/tools/__init__.py

@@ -0,0 +1,5 @@
+"""Experimental tools package."""
+
+from .tool_provider import ToolProvider
+
+__all__ = ["ToolProvider"]

src/strands/experimental/tools/mcp/__init__.py

@@ -0,0 +1,5 @@
+"""Experimental MCP Tool Provider."""
+
+from .mcp_tool_provider import MCPToolProvider, ToolFilters
+
+__all__ = ["MCPToolProvider", "ToolFilters"]

src/strands/experimental/tools/mcp/mcp_tool_provider.py

@@ -0,0 +1,180 @@
+"""MCP Tool Provider implementation."""
+
+import logging
+from typing import Callable, Optional, Pattern, Sequence, Union
+
+from typing_extensions import TypedDict
+
+from ....tools.mcp.mcp_agent_tool import MCPAgentTool
+from ....tools.mcp.mcp_client import MCPClient
+from ....types.exceptions import ToolProviderException
+from ....types.tools import AgentTool
+from ..tool_provider import ToolProvider
+
+logger = logging.getLogger(__name__)
+
+_ToolFilterCallback = Callable[[AgentTool], bool]
+_ToolFilterPattern = Union[str, Pattern[str], _ToolFilterCallback]
+
+
+class ToolFilters(TypedDict, total=False):
+    """Filters for controlling which MCP tools are loaded and available.
+
+    Tools are filtered in this order:
+    1. If 'allowed' is specified, only tools matching these patterns are included
+    2. Tools matching 'rejected' patterns are then excluded
+    3. If the result exceeds 'max_tools', it's truncated
+    """
+
+    allowed: list[_ToolFilterPattern]
+    rejected: list[_ToolFilterPattern]
+    max_tools: int
+
+
+class MCPToolProvider(ToolProvider):
+    """Tool provider for MCP clients with managed lifecycle."""
+
+    def __init__(
+        self, *, client: MCPClient, tool_filters: Optional[ToolFilters] = None, disambiguator: Optional[str] = None
+    ) -> None:
+        """Initialize with an MCP client.
+
+        Args:
+            client: The MCP client to manage.
+            tool_filters: Optional filters to apply to tools.
+            disambiguator: Optional prefix for tool names.
+        """
+        logger.debug(
+            "tool_filters=<%s>, disambiguator=<%s> | initializing MCPToolProvider", tool_filters, disambiguator
+        )
+        self._client = client
+        self._tool_filters = tool_filters
+        self._disambiguator = disambiguator
+        self._tools: Optional[list[MCPAgentTool]] = None  # None = not loaded yet, [] = loaded but empty
+        self._started = False
+
+    async def load_tools(self) -> Sequence[AgentTool]:
+        """Load and return tools from the MCP client.
+
+        Returns:
+            List of tools from the MCP server.
+        """
+        logger.debug("started=<%s>, cached_tools=<%s> | loading tools", self._started, self._tools is not None)
+
+        if not self._started:
+            try:
+                logger.debug("starting MCP client")
+                self._client.start()
+                self._started = True
+                logger.debug("MCP client started successfully")
+            except Exception as e:
+                logger.error("error=<%s> | failed to start MCP client", e)
+                raise ToolProviderException(f"Failed to start MCP client: {e}") from e
+
+        if self._tools is None:
+            logger.debug("loading tools from MCP server")
+            self._tools = []
+            pagination_token = None
+            page_count = 0
+
+            # Determine max_tools limit for early termination
+            max_tools_limit = None
+            if self._tool_filters and "max_tools" in self._tool_filters:
+                max_tools_limit = self._tool_filters["max_tools"]
+                logger.debug("max_tools_limit=<%d> | will stop when reached", max_tools_limit)
+
+            while True:
+                logger.debug("page=<%d>, token=<%s> | fetching tools page", page_count, pagination_token)
+                paginated_tools = self._client.list_tools_sync(pagination_token)
+
+                # Process each tool as we get it
+                for tool in paginated_tools:
+                    # Apply filters
+                    if self._should_include_tool(tool):
+                        # Apply disambiguation if needed
+                        processed_tool = self._apply_disambiguation(tool)
+                        self._tools.append(processed_tool)
+
+                        # Check if we've reached max_tools limit
+                        if max_tools_limit is not None and len(self._tools) >= max_tools_limit:
+                            logger.debug("max_tools_reached=<%d> | stopping pagination early", len(self._tools))
+                            return self._tools
+
+                logger.debug(
+                    "page=<%d>, page_tools=<%d>, total_filtered=<%d> | processed page",
+                    page_count,
+                    len(paginated_tools),
+                    len(self._tools),
+                )
+
+                pagination_token = paginated_tools.pagination_token
+                page_count += 1
+
+                if pagination_token is None:
+                    break
+
+            logger.debug("final_tools=<%d> | loading complete", len(self._tools))
+
+        return self._tools
+
+    def _should_include_tool(self, tool: MCPAgentTool) -> bool:
+        """Check if a tool should be included based on allowed/rejected filters."""
+        if not self._tool_filters:
+            return True
+
+        # Apply allowed filter
+        if "allowed" in self._tool_filters:
+            if not self._matches_patterns(tool, self._tool_filters["allowed"]):
+                return False
+
+        # Apply rejected filter
+        if "rejected" in self._tool_filters:
+            if self._matches_patterns(tool, self._tool_filters["rejected"]):
+                return False
+
+        return True
+
+    def _apply_disambiguation(self, tool: MCPAgentTool) -> MCPAgentTool:
+        """Apply disambiguation to a single tool if needed."""
+        if not self._disambiguator:
+            return tool
+
+        # Create new tool with disambiguated agent name but preserve original MCP name
+        old_name = tool.tool_name
+        new_agent_name = f"{self._disambiguator}_{tool.mcp_tool.name}"
+        new_tool = MCPAgentTool(tool.mcp_tool, tool.mcp_client, agent_facing_tool_name=new_agent_name)
+        logger.debug("tool_rename=<%s->%s> | renamed tool", old_name, new_agent_name)
+        return new_tool
+
+    def _matches_patterns(self, tool: MCPAgentTool, patterns: list[_ToolFilterPattern]) -> bool:
+        """Check if tool matches any of the given patterns."""
+        for pattern in patterns:
+            if callable(pattern):
+                if pattern(tool):
+                    return True
+            elif hasattr(pattern, "match") and hasattr(pattern, "pattern"):
+                if pattern.match(tool.tool_name):
+                    return True
+            elif isinstance(pattern, str):
+                if pattern == tool.tool_name:
+                    return True
+        return False
+
+    async def cleanup(self) -> None:
+        """Clean up the MCP client connection."""
+        if not self._started:
+            return
+
+        logger.debug("cleaning up MCP client")
+        try:
+            logger.debug("stopping MCP client")
+            self._client.stop(None, None, None)
+            logger.debug("MCP client stopped successfully")
+        except Exception as e:
+            logger.error("error=<%s> | failed to cleanup MCP client", e)
+            raise ToolProviderException(f"Failed to cleanup MCP client: {e}") from e
+
+        # Only reset state if cleanup succeeded
+        self._started = False
+        self._tools = None
+        logger.debug("MCP client cleanup complete")