agno 2.3.21__py3-none-any.whl → 2.3.23__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/toolkit.py

@@ -1,4 +1,5 @@
 from collections import OrderedDict
+from inspect import iscoroutinefunction
 from typing import Any, Callable, Dict, List, Optional, Sequence, Union
 
 from agno.tools.function import Function
@@ -14,6 +15,7 @@ class Toolkit:
         self,
         name: str = "toolkit",
         tools: Sequence[Union[Callable[..., Any], Function]] = [],
+        async_tools: Optional[Sequence[tuple[Callable[..., Any], str]]] = None,
         instructions: Optional[str] = None,
         add_instructions: bool = False,
         include_tools: Optional[list[str]] = None,
@@ -32,6 +34,9 @@ class Toolkit:
         Args:
             name: A descriptive name for the toolkit
             tools: List of tools to include in the toolkit (can be callables or Function objects from @tool decorator)
+            async_tools: List of (async_callable, tool_name) tuples for async variants.
+                        Used when async methods have different names than sync methods.
+                        Example: [(self.anavigate_to, "navigate_to"), (self.ascreenshot, "screenshot")]
             instructions: Instructions for the toolkit
             add_instructions: Whether to add instructions to the toolkit
             include_tools: List of tool names to include in the toolkit
@@ -47,7 +52,11 @@ class Toolkit:
         """
         self.name: str = name
         self.tools: Sequence[Union[Callable[..., Any], Function]] = tools
+        self._async_tools: Sequence[tuple[Callable[..., Any], str]] = async_tools or []
+        # Functions dict - used by agent.run() and agent.print_response()
         self.functions: Dict[str, Function] = OrderedDict()
+        # Async functions dict - used by agent.arun() and agent.aprint_response()
+        self.async_functions: Dict[str, Function] = OrderedDict()
         self.instructions: Optional[str] = instructions
         self.add_instructions: bool = add_instructions
 
@@ -71,8 +80,11 @@ class Toolkit:
         self.cache_dir: Optional[str] = cache_dir
 
         # Automatically register all methods if auto_register is True
-        if auto_register and self.tools:
-            self._register_tools()
+        if auto_register:
+            if self.tools:
+                self._register_tools()
+            if self._async_tools:
+                self._register_async_tools()
 
     def _get_tool_name(self, tool: Union[Callable[..., Any], Function]) -> str:
         """Get the name of a tool, whether it's a Function or callable."""
@@ -125,14 +137,25 @@ class Toolkit:
                 log_warning(f"Show result tool(s) not present in the toolkit: {', '.join(missing_show_result)}")
 
     def _register_tools(self) -> None:
-        """Register all tools."""
+        """Register all sync tools."""
         for tool in self.tools:
             self.register(tool)
 
+    def _register_async_tools(self) -> None:
+        """Register all async tools with their mapped names.
+
+        Async detection is automatic via iscoroutinefunction.
+        """
+        for async_func, tool_name in self._async_tools:
+            self.register(async_func, name=tool_name)
+
     def register(self, function: Union[Callable[..., Any], Function], name: Optional[str] = None) -> None:
         """Register a function with the toolkit.
 
         This method supports both regular callables and Function objects (from @tool decorator).
+        Automatically detects if the function is async (using iscoroutinefunction) and registers
+        it to the appropriate dict (functions for sync, async_functions for async).
+
         When a Function object is passed (e.g., from a @tool decorated method), it will:
         1. Extract the configuration from the Function object
         2. Look for a bound method with the same name on `self`
@@ -140,17 +163,18 @@ class Toolkit:
 
         Args:
             function: The callable or Function object to register
-            name: Optional custom name for the function
-
-        Returns:
-            The registered function
+            name: Optional custom name for the function (useful for aliasing)
         """
         try:
             # Handle Function objects (from @tool decorator)
             if isinstance(function, Function):
-                return self._register_decorated_tool(function, name)
+                # Auto-detect if this is an async function
+                is_async = function.entrypoint is not None and iscoroutinefunction(function.entrypoint)
+                return self._register_decorated_tool(function, name, is_async=is_async)
+
+            # Handle regular callables - auto-detect async
+            is_async = iscoroutinefunction(function)
 
-            # Handle regular callables
             tool_name = name or function.__name__
             if self.include_tools is not None and tool_name not in self.include_tools:
                 return
@@ -168,14 +192,19 @@ class Toolkit:
                 stop_after_tool_call=tool_name in self.stop_after_tool_call_tools,
                 show_result=tool_name in self.show_result_tools or tool_name in self.stop_after_tool_call_tools,
             )
-            self.functions[f.name] = f
-            log_debug(f"Function: {f.name} registered with {self.name}")
+
+            if is_async:
+                self.async_functions[f.name] = f
+                log_debug(f"Async function: {f.name} registered with {self.name}")
+            else:
+                self.functions[f.name] = f
+                log_debug(f"Function: {f.name} registered with {self.name}")
         except Exception as e:
             func_name = self._get_tool_name(function)
             logger.warning(f"Failed to create Function for: {func_name}")
             raise e
 
-    def _register_decorated_tool(self, function: Function, name: Optional[str] = None) -> None:
+    def _register_decorated_tool(self, function: Function, name: Optional[str] = None, is_async: bool = False) -> None:
         """Register a Function object from @tool decorator, binding it to self.
 
         When @tool decorator is used on a class method, it creates a Function with an unbound
@@ -185,6 +214,7 @@ class Toolkit:
         Args:
             function: The Function object from @tool decorator
             name: Optional custom name override
+            is_async: If True, register to async_functions dict instead of functions
         """
         import inspect
 
@@ -207,14 +237,24 @@ class Toolkit:
 
         if params and params[0] == "self":
             # Create a bound method by wrapping the function to include self
-            def make_bound_method(func, instance):
-                def bound(*args, **kwargs):
-                    return func(instance, *args, **kwargs)
+            if is_async:
+
+                def make_bound_method(func, instance):
+                    async def bound(*args, **kwargs):
+                        return await func(instance, *args, **kwargs)
+
+                    bound.__name__ = getattr(func, "__name__", tool_name)
+                    bound.__doc__ = getattr(func, "__doc__", None)
+                    return bound
+            else:
+
+                def make_bound_method(func, instance):
+                    def bound(*args, **kwargs):
+                        return func(instance, *args, **kwargs)
 
-                # Preserve function metadata for debugging
-                bound.__name__ = getattr(func, "__name__", tool_name)
-                bound.__doc__ = getattr(func, "__doc__", None)
-                return bound
+                    bound.__name__ = getattr(func, "__name__", tool_name)
+                    bound.__doc__ = getattr(func, "__doc__", None)
+                    return bound
 
             bound_method = make_bound_method(original_func, self)
         else:
@@ -251,8 +291,36 @@ class Toolkit:
             cache_dir=function.cache_dir if function.cache_dir else self.cache_dir,
             cache_ttl=function.cache_ttl if function.cache_ttl != 3600 else self.cache_ttl,
         )
-        self.functions[f.name] = f
-        log_debug(f"Function: {f.name} registered with {self.name} (from @tool decorator)")
+
+        if is_async:
+            self.async_functions[f.name] = f
+            log_debug(f"Async function: {f.name} registered with {self.name} (from @tool decorator)")
+        else:
+            self.functions[f.name] = f
+            log_debug(f"Function: {f.name} registered with {self.name} (from @tool decorator)")
+
+    def get_functions(self) -> Dict[str, Function]:
+        """Get sync functions dict.
+
+        Returns:
+            Dict of function name to Function for sync execution
+        """
+        return self.functions
+
+    def get_async_functions(self) -> Dict[str, Function]:
+        """Get functions dict optimized for async execution.
+
+        Returns a merged dict where async_functions take precedence over functions.
+        This allows async-optimized implementations to be automatically used in async contexts,
+        while falling back to sync implementations for tools without async variants.
+
+        Returns:
+            Dict of function name to Function, with async variants preferred
+        """
+        # Merge: start with sync functions, override with async variants
+        merged = OrderedDict(self.functions)
+        merged.update(self.async_functions)
+        return merged
 
     @property
     def requires_connect(self) -> bool:

agno/utils/mcp.py

@@ -1,5 +1,6 @@
 import json
 from functools import partial
+from typing import TYPE_CHECKING, Optional, Union
 from uuid import uuid4
 
 from agno.utils.log import log_debug, log_exception
@@ -15,28 +16,68 @@ except (ImportError, ModuleNotFoundError):
 from agno.media import Image
 from agno.tools.function import ToolResult
 
+if TYPE_CHECKING:
+    from agno.agent import Agent
+    from agno.run import RunContext
+    from agno.team.team import Team
+    from agno.tools.mcp.mcp import MCPTools
+    from agno.tools.mcp.multi_mcp import MultiMCPTools
 
-def get_entrypoint_for_tool(tool: MCPTool, session: ClientSession):
+
+def get_entrypoint_for_tool(
+    tool: MCPTool,
+    session: ClientSession,
+    mcp_tools_instance: Optional[Union["MCPTools", "MultiMCPTools"]] = None,
+    server_idx: int = 0,
+):
     """
     Return an entrypoint for an MCP tool.
 
     Args:
         tool: The MCP tool to create an entrypoint for
-        session: The session to use
+        session: The MCP ClientSession to use
+        mcp_tools_instance: Optional MCPTools or MultiMCPTools instance
+        server_idx: Index of the server (for MultiMCPTools)
 
     Returns:
         Callable: The entrypoint function for the tool
     """
 
-    async def call_tool(tool_name: str, **kwargs) -> ToolResult:
+    async def call_tool(
+        tool_name: str,
+        run_context: Optional["RunContext"] = None,
+        agent: Optional["Agent"] = None,
+        team: Optional["Team"] = None,
+        **kwargs,
+    ) -> ToolResult:
+        # Execute the MCP tool call
         try:
-            await session.send_ping()
-        except Exception as e:
-            log_exception(e)
+            # Get the appropriate session for this run
+            # If mcp_tools_instance has header_provider and run_context is provided,
+            # this will create/reuse a session with dynamic headers
+            if mcp_tools_instance and hasattr(mcp_tools_instance, "get_session_for_run"):
+                # Import here to avoid circular imports
+                from agno.tools.mcp.multi_mcp import MultiMCPTools
+
+                # For MultiMCPTools, pass server_idx; for MCPTools, only pass run_context
+                if isinstance(mcp_tools_instance, MultiMCPTools):
+                    active_session = await mcp_tools_instance.get_session_for_run(
+                        run_context=run_context, server_idx=server_idx, agent=agent, team=team
+                    )
+                else:
+                    active_session = await mcp_tools_instance.get_session_for_run(
+                        run_context=run_context, agent=agent, team=team
+                    )
+            else:
+                active_session = session
+
+            try:
+                await active_session.send_ping()
+            except Exception as e:
+                log_exception(e)
 
-        try:
             log_debug(f"Calling MCP Tool '{tool_name}' with args: {kwargs}")
-            result: CallToolResult = await session.call_tool(tool_name, kwargs)  # type: ignore
+            result: CallToolResult = await active_session.call_tool(tool_name, kwargs)  # type: ignore
 
             # Return an error if the tool call failed
             if result.isError: