agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/tools/mcp/multi_mcp.py

@@ -1,9 +1,12 @@
+import inspect
+import time
+import warnings
 import weakref
 from contextlib import AsyncExitStack
 from dataclasses import asdict
 from datetime import timedelta
 from types import TracebackType
-from typing import List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 
 from agno.tools import Toolkit
 from agno.tools.function import Function
@@ -11,6 +14,11 @@ from agno.tools.mcp.params import SSEClientParams, StreamableHTTPClientParams
 from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.utils.mcp import get_entrypoint_for_tool, prepare_command
 
+if TYPE_CHECKING:
+    from agno.agent import Agent
+    from agno.run import RunContext
+    from agno.team.team import Team
+
 try:
     from mcp import ClientSession, StdioServerParameters
     from mcp.client.sse import sse_client
@@ -47,6 +55,7 @@ class MultiMCPTools(Toolkit):
         exclude_tools: Optional[list[str]] = None,
         refresh_connection: bool = False,
         allow_partial_failure: bool = False,
+        header_provider: Optional[Callable[..., dict[str, Any]]] = None,
         **kwargs,
     ):
         """
@@ -64,7 +73,14 @@ class MultiMCPTools(Toolkit):
             exclude_tools: Optional list of tool names to exclude (if None, excludes none).
             allow_partial_failure: If True, allows toolkit to initialize even if some MCP servers fail to connect. If False, any failure will raise an exception.
             refresh_connection: If True, the connection and tools will be refreshed on each run
+            header_provider: Header provider function for all servers. Takes RunContext and returns dict of HTTP headers.
         """
+        warnings.warn(
+            "The MultiMCPTools class is deprecated and will be removed in a future version. Please use multiple MCPTools instances instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
         super().__init__(name="MultiMCPTools", **kwargs)
 
         if urls_transports is not None:
@@ -86,6 +102,16 @@ class MultiMCPTools(Toolkit):
         self.exclude_tools = exclude_tools
         self.refresh_connection = refresh_connection
 
+        self.header_provider = header_provider
+
+        # Validate header_provider signature
+        if header_provider:
+            try:
+                # Just verify we can inspect the signature - no parameter requirements
+                inspect.signature(header_provider)
+            except Exception as e:
+                log_warning(f"Could not validate header_provider signature: {e}")
+
         if server_params_list is None and commands is None and urls is None:
             raise ValueError("Either server_params_list or commands or urls must be provided")
 
@@ -130,6 +156,14 @@ class MultiMCPTools(Toolkit):
         self._connection_task = None
         self._successful_connections = 0
         self._sessions: list[ClientSession] = []
+        self._session_to_server_idx: Dict[int, int] = {}  # Maps session list index to server params index
+
+        # Session management for per-agent-run sessions with dynamic headers
+        # For MultiMCP, we track sessions per (run_id, server_idx) since we have multiple servers
+        # Maps (run_id, server_idx) to (session, timestamp) for TTL-based cleanup
+        self._run_sessions: Dict[Tuple[str, int], Tuple[ClientSession, float]] = {}
+        self._run_session_contexts: Dict[Tuple[str, int], Any] = {}  # Maps (run_id, server_idx) to context managers
+        self._session_ttl_seconds: float = 300.0  # 5 minutes default TTL
 
         self.allow_partial_failure = allow_partial_failure
 
@@ -153,6 +187,205 @@ class MultiMCPTools(Toolkit):
         except (RuntimeError, BaseException):
             return False
 
+    def _call_header_provider(
+        self,
+        run_context: Optional["RunContext"] = None,
+        agent: Optional["Agent"] = None,
+        team: Optional["Team"] = None,
+    ) -> dict[str, Any]:
+        """Call the header_provider with run_context, agent, and/or team based on its signature.
+
+        Args:
+            run_context: The RunContext for the current agent run
+            agent: The Agent instance (if running within an agent)
+            team: The Team instance (if running within a team)
+
+        Returns:
+            dict[str, Any]: The headers returned by the header_provider
+        """
+        header_provider = getattr(self, "header_provider", None)
+        if header_provider is None:
+            return {}
+
+        try:
+            sig = inspect.signature(header_provider)
+            param_names = set(sig.parameters.keys())
+
+            # Build kwargs based on what the function accepts
+            call_kwargs: dict[str, Any] = {}
+
+            if "run_context" in param_names:
+                call_kwargs["run_context"] = run_context
+            if "agent" in param_names:
+                call_kwargs["agent"] = agent
+            if "team" in param_names:
+                call_kwargs["team"] = team
+
+            # Check if function accepts **kwargs (VAR_KEYWORD)
+            has_var_keyword = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values())
+
+            if has_var_keyword:
+                # Pass all available context to **kwargs
+                call_kwargs = {"run_context": run_context, "agent": agent, "team": team}
+                return header_provider(**call_kwargs)
+            elif call_kwargs:
+                return header_provider(**call_kwargs)
+            else:
+                # Function takes no recognized parameters - check for positional
+                positional_params = [
+                    p
+                    for p in sig.parameters.values()
+                    if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+                ]
+                if positional_params:
+                    # Legacy support: pass run_context as first positional arg
+                    return header_provider(run_context)
+                else:
+                    # Function takes no parameters
+                    return header_provider()
+        except Exception as e:
+            log_warning(f"Error calling header_provider: {e}")
+            return {}
+
+    async def _cleanup_stale_sessions(self) -> None:
+        """Clean up sessions older than TTL to prevent memory leaks."""
+        if not self._run_sessions:
+            return
+
+        now = time.time()
+        stale_keys = [
+            cache_key
+            for cache_key, (_, created_at) in self._run_sessions.items()
+            if now - created_at > self._session_ttl_seconds
+        ]
+
+        for run_id, server_idx in stale_keys:
+            log_debug(f"Cleaning up stale session for run_id={run_id}, server_idx={server_idx}")
+            await self.cleanup_run_session(run_id, server_idx)
+
+    async def get_session_for_run(
+        self,
+        run_context: Optional["RunContext"] = None,
+        server_idx: int = 0,
+        agent: Optional["Agent"] = None,
+        team: Optional["Team"] = None,
+    ) -> ClientSession:
+        """
+        Get or create a session for the given run_context and server index.
+
+        If header_provider is configured and run_context is provided, this creates
+        a new session with dynamic headers for this specific agent run and server.
+
+        Args:
+            run_context: The RunContext containing user_id, metadata, etc.
+            server_idx: Index of the server in self._sessions list
+            agent: The Agent instance (if running within an agent)
+            team: The Team instance (if running within a team)
+
+        Returns:
+            ClientSession: Either the default session or a per-run session with dynamic headers
+        """
+        # If no header_provider or no run_context, use the default session
+        if not self.header_provider or not run_context:
+            # Return the default session for this server
+            if server_idx < len(self._sessions):
+                return self._sessions[server_idx]
+            raise ValueError(f"Server index {server_idx} out of range")
+
+        # Lazy cleanup of stale sessions
+        await self._cleanup_stale_sessions()
+
+        # Check if we already have a session for this (run_id, server_idx)
+        run_id = run_context.run_id
+        cache_key = (run_id, server_idx)
+        if cache_key in self._run_sessions:
+            session, _ = self._run_sessions[cache_key]
+            return session
+
+        # Create a new session with dynamic headers for this run and server
+        log_debug(f"Creating new session for run_id={run_id}, server_idx={server_idx} with dynamic headers")
+
+        # Generate dynamic headers from the provider
+        dynamic_headers = self._call_header_provider(run_context=run_context, agent=agent, team=team)
+
+        # Get the server params for this server index
+        if server_idx >= len(self.server_params_list):
+            raise ValueError(f"Server index {server_idx} out of range")
+
+        server_params = self.server_params_list[server_idx]
+
+        # Create new session with merged headers based on transport type
+        if isinstance(server_params, SSEClientParams):
+            params_dict = asdict(server_params)
+            existing_headers = params_dict.get("headers") or {}
+            params_dict["headers"] = {**existing_headers, **dynamic_headers}
+
+            context = sse_client(**params_dict)  # type: ignore
+            client_timeout = min(self.timeout_seconds, params_dict.get("timeout", self.timeout_seconds))
+
+        elif isinstance(server_params, StreamableHTTPClientParams):
+            params_dict = asdict(server_params)
+            existing_headers = params_dict.get("headers") or {}
+            params_dict["headers"] = {**existing_headers, **dynamic_headers}
+
+            context = streamablehttp_client(**params_dict)  # type: ignore
+            params_timeout = params_dict.get("timeout", self.timeout_seconds)
+            if isinstance(params_timeout, timedelta):
+                params_timeout = int(params_timeout.total_seconds())
+            client_timeout = min(self.timeout_seconds, params_timeout)
+        else:
+            # stdio doesn't support headers, fall back to default session
+            log_warning(
+                f"Cannot use dynamic headers with stdio transport for server {server_idx}, using default session"
+            )
+            if server_idx < len(self._sessions):
+                return self._sessions[server_idx]
+            raise ValueError(f"Server index {server_idx} out of range")
+
+        # Enter the context and create session
+        session_params = await context.__aenter__()  # type: ignore
+        read, write = session_params[0:2]
+
+        session_context = ClientSession(read, write, read_timeout_seconds=timedelta(seconds=client_timeout))  # type: ignore
+        session = await session_context.__aenter__()  # type: ignore
+
+        # Initialize the session
+        await session.initialize()
+
+        # Store the session with timestamp and context for cleanup
+        self._run_sessions[cache_key] = (session, time.time())
+        self._run_session_contexts[cache_key] = (context, session_context)
+
+        return session
+
+    async def cleanup_run_session(self, run_id: str, server_idx: int) -> None:
+        """Clean up a per-run session."""
+        cache_key = (run_id, server_idx)
+        if cache_key not in self._run_sessions:
+            return
+
+        try:
+            context, session_context = self._run_session_contexts[cache_key]
+
+            # Exit session context - silently ignore errors
+            try:
+                await session_context.__aexit__(None, None, None)
+            except (RuntimeError, Exception):
+                pass  # Silently ignore
+
+            # Exit transport context - silently ignore errors
+            try:
+                await context.__aexit__(None, None, None)
+            except (RuntimeError, Exception):
+                pass  # Silently ignore
+
+        except Exception:
+            pass  # Silently ignore all cleanup errors
+        finally:
+            # Remove from cache
+            self._run_sessions.pop(cache_key, None)
+            self._run_session_contexts.pop(cache_key, None)
+
     async def connect(self, force: bool = False):
         """Initialize a MultiMCPTools instance and connect to the MCP servers"""
 
@@ -214,7 +447,7 @@ class MultiMCPTools(Toolkit):
 
         server_connection_errors = []
 
-        for server_params in self.server_params_list:
+        for server_idx, server_params in enumerate(self.server_params_list):
             try:
                 # Handle stdio connections
                 if isinstance(server_params, StdioServerParameters):
@@ -223,7 +456,7 @@ class MultiMCPTools(Toolkit):
                     session = await self._async_exit_stack.enter_async_context(
                         ClientSession(read, write, read_timeout_seconds=timedelta(seconds=self.timeout_seconds))
                     )
-                    await self.initialize(session)
+                    await self.initialize(session, server_idx)
                     self._successful_connections += 1
 
                 # Handle SSE connections
@@ -233,7 +466,7 @@ class MultiMCPTools(Toolkit):
                     )
                     read, write = client_connection
                     session = await self._async_exit_stack.enter_async_context(ClientSession(read, write))
-                    await self.initialize(session)
+                    await self.initialize(session, server_idx)
                     self._successful_connections += 1
 
                 # Handle Streamable HTTP connections
@@ -243,7 +476,7 @@ class MultiMCPTools(Toolkit):
                     )
                     read, write = client_connection[0:2]
                     session = await self._async_exit_stack.enter_async_context(ClientSession(read, write))
-                    await self.initialize(session)
+                    await self.initialize(session, server_idx)
                     self._successful_connections += 1
 
             except Exception as e:
@@ -268,13 +501,26 @@ class MultiMCPTools(Toolkit):
         if not self._initialized:
             return
 
-        try:
-            await self._async_exit_stack.aclose()
-            self._sessions = []
-            self._successful_connections = 0
+        import warnings
 
-        except (RuntimeError, BaseException) as e:
-            log_error(f"Failed to close MCP connections: {e}")
+        # Suppress async generator cleanup warnings
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning, message=".*async_generator.*")
+            warnings.filterwarnings("ignore", message=".*cancel scope.*")
+
+            try:
+                # Clean up all per-run sessions first
+                cache_keys = list(self._run_sessions.keys())
+                for run_id, server_idx in cache_keys:
+                    await self.cleanup_run_session(run_id, server_idx)
+
+                # Clean up main sessions
+                await self._async_exit_stack.aclose()
+                self._sessions = []
+                self._successful_connections = 0
+
+            except (RuntimeError, BaseException):
+                pass  # Silently ignore all cleanup errors
 
         self._initialized = False
 
@@ -298,7 +544,7 @@ class MultiMCPTools(Toolkit):
         self._successful_connections = 0
 
     async def build_tools(self) -> None:
-        for session in self._sessions:
+        for session_list_idx, session in enumerate(self._sessions):
             # Get the list of tools from the MCP server
             available_tools = await session.list_tools()
 
@@ -314,7 +560,12 @@ class MultiMCPTools(Toolkit):
             for tool in filtered_tools:
                 try:
                     # Get an entrypoint for the tool
-                    entrypoint = get_entrypoint_for_tool(tool, session)
+                    entrypoint = get_entrypoint_for_tool(
+                        tool=tool,
+                        session=session,
+                        mcp_tools_instance=self,  # Pass self to enable dynamic headers
+                        server_idx=session_list_idx,  # Pass session list index for session lookup
+                    )
 
                     # Create a Function for the tool
                     f = Function(
@@ -333,14 +584,18 @@ class MultiMCPTools(Toolkit):
                     log_error(f"Failed to register tool {tool.name}: {e}")
                     raise
 
-    async def initialize(self, session: ClientSession) -> None:
+    async def initialize(self, session: ClientSession, server_idx: int = 0) -> None:
         """Initialize the MCP toolkit by getting available tools from the MCP server"""
 
         try:
             # Initialize the session if not already initialized
             await session.initialize()
 
+            # Track which server index this session belongs to
+            session_list_idx = len(self._sessions)
             self._sessions.append(session)
+            self._session_to_server_idx[session_list_idx] = server_idx
+
             self._initialized = True
         except Exception as e:
             log_error(f"Failed to get MCP tools: {e}")

agno/tools/mem0.py

@@ -2,6 +2,7 @@ import json
 from os import getenv
 from typing import Any, Dict, List, Optional, Union
 
+from agno.run import RunContext
 from agno.tools import Toolkit
 from agno.utils.log import log_debug, log_error, log_warning
 
@@ -68,13 +69,13 @@ class Mem0Tools(Toolkit):
     def _get_user_id(
         self,
         method_name: str,
-        session_state: Dict[str, Any],
+        run_context: RunContext,
     ) -> str:
         """Resolve the user ID"""
         resolved_user_id = self.user_id
         if not resolved_user_id:
             try:
-                resolved_user_id = session_state.get("current_user_id")
+                resolved_user_id = run_context.user_id
             except Exception:
                 pass
         if not resolved_user_id:
@@ -85,7 +86,7 @@ class Mem0Tools(Toolkit):
 
     def add_memory(
         self,
-        session_state,
+        run_context: RunContext,
         content: Union[str, Dict[str, str]],
     ) -> str:
         """Add facts to the user's memory.
@@ -98,7 +99,7 @@ class Mem0Tools(Toolkit):
             str: JSON-encoded Mem0 response or an error message.
         """
 
-        resolved_user_id = self._get_user_id("add_memory", session_state=session_state)
+        resolved_user_id = self._get_user_id("add_memory", run_context=run_context)
         if isinstance(resolved_user_id, str) and resolved_user_id.startswith("Error in add_memory:"):
             return resolved_user_id
         try:
@@ -121,12 +122,12 @@ class Mem0Tools(Toolkit):
 
     def search_memory(
         self,
-        session_state: Dict[str, Any],
+        run_context: RunContext,
         query: str,
     ) -> str:
         """Semantic search for *query* across the user's stored memories."""
 
-        resolved_user_id = self._get_user_id("search_memory", session_state=session_state)
+        resolved_user_id = self._get_user_id("search_memory", run_context=run_context)
         if isinstance(resolved_user_id, str) and resolved_user_id.startswith("Error in search_memory:"):
             return resolved_user_id
         try:
@@ -151,10 +152,10 @@ class Mem0Tools(Toolkit):
             log_error(f"Error searching memory: {e}")
             return f"Error searching memory: {e}"
 
-    def get_all_memories(self, session_state: Dict[str, Any]) -> str:
+    def get_all_memories(self, run_context: RunContext) -> str:
         """Return **all** memories for the current user as a JSON string."""
 
-        resolved_user_id = self._get_user_id("get_all_memories", session_state=session_state)
+        resolved_user_id = self._get_user_id("get_all_memories", run_context=run_context)
         if isinstance(resolved_user_id, str) and resolved_user_id.startswith("Error in get_all_memories:"):
             return resolved_user_id
         try:
@@ -177,10 +178,10 @@ class Mem0Tools(Toolkit):
             log_error(f"Error getting all memories: {e}")
             return f"Error getting all memories: {e}"
 
-    def delete_all_memories(self, session_state: Dict[str, Any]) -> str:
+    def delete_all_memories(self, run_context: RunContext) -> str:
         """Delete *all* memories associated with the current user"""
 
-        resolved_user_id = self._get_user_id("delete_all_memories", session_state=session_state)
+        resolved_user_id = self._get_user_id("delete_all_memories", run_context=run_context)
         if isinstance(resolved_user_id, str) and resolved_user_id.startswith("Error in delete_all_memories:"):
             error_msg = resolved_user_id
             log_error(error_msg)