llama-stack 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

llama_stack/providers/utils/tools/mcp.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
+import hashlib
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from enum import Enum
@@ -73,6 +75,207 @@ class MCPProtol(Enum):
     SSE = 2
 
 
+class MCPSessionManager:
+    """Manages MCP session lifecycle within a request scope.
+
+    This class caches MCP sessions by (endpoint, headers_hash) to avoid redundant
+    connection establishment and tools/list calls when making multiple tool
+    invocations to the same MCP server within a single request.
+
+    Fix for GitHub issue #4452: MCP tools/list called redundantly before every
+    tool invocation.
+
+    Usage:
+        async with MCPSessionManager() as session_manager:
+            # Multiple tool calls will reuse the same session
+            result1 = await invoke_mcp_tool(..., session_manager=session_manager)
+            result2 = await invoke_mcp_tool(..., session_manager=session_manager)
+    """
+
+    def __init__(self):
+        # Cache of active sessions: key -> (session, client_context, session_context)
+        self._sessions: dict[str, tuple[ClientSession, Any, Any]] = {}
+        # Locks to prevent concurrent session creation for the same key
+        self._locks: dict[str, asyncio.Lock] = {}
+        # Global lock for managing the locks dict
+        self._global_lock = asyncio.Lock()
+
+    def _make_key(self, endpoint: str, headers: dict[str, str]) -> str:
+        """Create a cache key from endpoint and headers."""
+        # Sort headers for consistent hashing
+        headers_str = str(sorted(headers.items()))
+        headers_hash = hashlib.sha256(headers_str.encode()).hexdigest()[:16]
+        return f"{endpoint}:{headers_hash}"
+
+    async def _get_lock(self, key: str) -> asyncio.Lock:
+        """Get or create a lock for a specific cache key."""
+        async with self._global_lock:
+            if key not in self._locks:
+                self._locks[key] = asyncio.Lock()
+            return self._locks[key]
+
+    async def get_session(self, endpoint: str, headers: dict[str, str]) -> ClientSession:
+        """Get or create an MCP session for the given endpoint and headers.
+
+        Args:
+            endpoint: MCP server endpoint URL
+            headers: Headers including authorization
+
+        Returns:
+            An initialized ClientSession ready for tool calls
+        """
+        key = self._make_key(endpoint, headers)
+
+        # Check if session already exists (fast path)
+        if key in self._sessions:
+            session, _, _ = self._sessions[key]
+            return session
+
+        # Acquire lock for this specific key to prevent concurrent creation
+        lock = await self._get_lock(key)
+        async with lock:
+            # Double-check after acquiring lock
+            if key in self._sessions:
+                session, _, _ = self._sessions[key]
+                return session
+
+            # Create new session
+            session, client_ctx, session_ctx = await self._create_session(endpoint, headers)
+            self._sessions[key] = (session, client_ctx, session_ctx)
+            logger.debug(f"Created new MCP session for {endpoint} (key: {key[:32]}...)")
+            return session
+
+    async def _create_session(self, endpoint: str, headers: dict[str, str]) -> tuple[ClientSession, Any, Any]:
+        """Create a new MCP session.
+
+        Returns:
+            Tuple of (session, client_context, session_context) for lifecycle management
+        """
+        # Use the same protocol detection logic as client_wrapper
+        connection_strategies = [MCPProtol.STREAMABLE_HTTP, MCPProtol.SSE]
+        mcp_protocol = protocol_cache.get(endpoint, default=MCPProtol.UNKNOWN)
+        if mcp_protocol == MCPProtol.SSE:
+            connection_strategies = [MCPProtol.SSE, MCPProtol.STREAMABLE_HTTP]
+
+        last_exception: BaseException | None = None
+
+        for i, strategy in enumerate(connection_strategies):
+            try:
+                client = streamablehttp_client
+                if strategy == MCPProtol.SSE:
+                    client = cast(Any, sse_client)
+
+                # Enter the client context manager manually
+                client_ctx = client(endpoint, headers=headers)
+                client_streams = await client_ctx.__aenter__()
+
+                try:
+                    # Enter the session context manager manually
+                    session = ClientSession(read_stream=client_streams[0], write_stream=client_streams[1])
+                    session_ctx = session
+                    await session.__aenter__()
+
+                    try:
+                        await session.initialize()
+                        protocol_cache[endpoint] = strategy
+                        return session, client_ctx, session_ctx
+                    except BaseException:
+                        await session.__aexit__(None, None, None)
+                        raise
+                except BaseException:
+                    await client_ctx.__aexit__(None, None, None)
+                    raise
+
+            except* httpx.HTTPStatusError as eg:
+                for exc in eg.exceptions:
+                    err = cast(httpx.HTTPStatusError, exc)
+                    if err.response.status_code == 401:
+                        raise AuthenticationRequiredError(exc) from exc
+                if i == len(connection_strategies) - 1:
+                    raise
+                last_exception = eg
+            except* httpx.ConnectError as eg:
+                if i == len(connection_strategies) - 1:
+                    error_msg = f"Failed to connect to MCP server at {endpoint}: Connection refused"
+                    logger.error(f"MCP connection error: {error_msg}")
+                    raise ConnectionError(error_msg) from eg
+                else:
+                    logger.warning(
+                        f"failed to connect to MCP server at {endpoint} via {strategy.name}, "
+                        f"falling back to {connection_strategies[i + 1].name}"
+                    )
+                last_exception = eg
+            except* httpx.TimeoutException as eg:
+                if i == len(connection_strategies) - 1:
+                    error_msg = f"MCP server at {endpoint} timed out"
+                    logger.error(f"MCP timeout error: {error_msg}")
+                    raise TimeoutError(error_msg) from eg
+                else:
+                    logger.warning(
+                        f"MCP server at {endpoint} timed out via {strategy.name}, "
+                        f"falling back to {connection_strategies[i + 1].name}"
+                    )
+                last_exception = eg
+            except* httpx.RequestError as eg:
+                if i == len(connection_strategies) - 1:
+                    exc_msg = str(eg.exceptions[0]) if eg.exceptions else "Unknown error"
+                    error_msg = f"Network error connecting to MCP server at {endpoint}: {exc_msg}"
+                    logger.error(f"MCP network error: {error_msg}")
+                    raise ConnectionError(error_msg) from eg
+                else:
+                    logger.warning(
+                        f"network error connecting to MCP server at {endpoint} via {strategy.name}, "
+                        f"falling back to {connection_strategies[i + 1].name}"
+                    )
+                last_exception = eg
+            except* McpError:
+                if i < len(connection_strategies) - 1:
+                    logger.warning(
+                        f"failed to connect via {strategy.name}, falling back to {connection_strategies[i + 1].name}"
+                    )
+                else:
+                    raise
+
+        # Should not reach here, but just in case
+        if last_exception:
+            raise last_exception
+        raise RuntimeError(f"Failed to create MCP session for {endpoint}")
+
+    async def close_all(self) -> None:
+        """Close all cached sessions.
+
+        Should be called at the end of a request to clean up resources.
+
+        Note: We catch BaseException (not just Exception) because:
+        1. CancelledError is a BaseException and can occur during cleanup
+        2. anyio cancel scope errors can occur if cleanup runs in a different
+           task context than where the session was created
+        These are expected in streaming response scenarios and are handled gracefully.
+        """
+        errors = []
+        session_count = len(self._sessions)
+        for key, (session, client_ctx, _) in list(self._sessions.items()):
+            try:
+                await session.__aexit__(None, None, None)
+            except BaseException as e:
+                # Debug level since these errors are expected in streaming scenarios
+                # where cleanup runs in a different async context than session creation
+                logger.debug(f"Error closing MCP session {key}: {e}")
+                errors.append(e)
+            try:
+                await client_ctx.__aexit__(None, None, None)
+            except BaseException as e:
+                logger.debug(f"Error closing MCP client context {key}: {e}")
+                errors.append(e)
+
+        self._sessions.clear()
+        self._locks.clear()
+        logger.debug(f"Closed {session_count} MCP sessions")
+
+        if errors:
+            logger.debug(f"Encountered {len(errors)} errors while closing MCP sessions (expected in streaming)")
+
+
 @asynccontextmanager
 async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerator[ClientSession, Any]:
     # we use a ttl'd dict to cache the happy path protocol for each endpoint
@@ -151,6 +354,7 @@ async def list_mcp_tools(
     endpoint: str,
     headers: dict[str, str] | None = None,
     authorization: str | None = None,
+    session_manager: MCPSessionManager | None = None,
 ) -> ListToolDefsResponse:
     """List tools available from an MCP server.
 
@@ -158,6 +362,10 @@ async def list_mcp_tools(
         endpoint: MCP server endpoint URL
         headers: Optional base headers to include
         authorization: Optional OAuth access token (just the token, not "Bearer <token>")
+        session_manager: Optional MCPSessionManager for session reuse within a request.
+            When provided, sessions are cached and reused, avoiding redundant session
+            creation when list_mcp_tools and invoke_mcp_tool are called for the same
+            server within a request. (Fix for #4452)
 
     Returns:
         List of tool definitions from the MCP server
@@ -169,7 +377,9 @@ async def list_mcp_tools(
     final_headers = prepare_mcp_headers(headers, authorization)
 
     tools = []
-    async with client_wrapper(endpoint, final_headers) as session:
+
+    # Helper function to process session and list tools
+    async def _list_tools_from_session(session):
         tools_result = await session.list_tools()
         for tool in tools_result.tools:
             tools.append(
@@ -183,15 +393,51 @@ async def list_mcp_tools(
                     },
                 )
             )
+
+    # If a session manager is provided, use it for session reuse (fix for #4452)
+    if session_manager is not None:
+        session = await session_manager.get_session(endpoint, final_headers)
+        await _list_tools_from_session(session)
+    else:
+        # Fallback to original behavior: create a new session for this call
+        async with client_wrapper(endpoint, final_headers) as session:
+            await _list_tools_from_session(session)
+
     return ListToolDefsResponse(data=tools)
 
 
+def _parse_mcp_result(result) -> ToolInvocationResult:
+    """Parse MCP tool call result into ToolInvocationResult.
+
+    Args:
+        result: The raw MCP tool call result
+
+    Returns:
+        ToolInvocationResult with parsed content
+    """
+    content: list[InterleavedContentItem] = []
+    for item in result.content:
+        if isinstance(item, mcp_types.TextContent):
+            content.append(TextContentItem(text=item.text))
+        elif isinstance(item, mcp_types.ImageContent):
+            content.append(ImageContentItem(image=_URLOrData(data=item.data)))
+        elif isinstance(item, mcp_types.EmbeddedResource):
+            logger.warning(f"EmbeddedResource is not supported: {item}")
+        else:
+            raise ValueError(f"Unknown content type: {type(item)}")
+    return ToolInvocationResult(
+        content=content,
+        error_code=1 if result.isError else 0,
+    )
+
+
 async def invoke_mcp_tool(
     endpoint: str,
     tool_name: str,
     kwargs: dict[str, Any],
     headers: dict[str, str] | None = None,
     authorization: str | None = None,
+    session_manager: MCPSessionManager | None = None,
 ) -> ToolInvocationResult:
     """Invoke an MCP tool with the given arguments.
 
@@ -201,6 +447,9 @@ async def invoke_mcp_tool(
         kwargs: Tool invocation arguments
         headers: Optional base headers to include
         authorization: Optional OAuth access token (just the token, not "Bearer <token>")
+        session_manager: Optional MCPSessionManager for session reuse within a request.
+            When provided, sessions are cached and reused for multiple tool calls to
+            the same endpoint, avoiding redundant tools/list calls. (Fix for #4452)
 
     Returns:
         Tool invocation result with content and error information
@@ -211,20 +460,13 @@ async def invoke_mcp_tool(
     # Prepare headers with authorization handling
     final_headers = prepare_mcp_headers(headers, authorization)
 
-    async with client_wrapper(endpoint, final_headers) as session:
+    # If a session manager is provided, use it for session reuse (fix for #4452)
+    if session_manager is not None:
+        session = await session_manager.get_session(endpoint, final_headers)
         result = await session.call_tool(tool_name, kwargs)
+        return _parse_mcp_result(result)
 
-        content: list[InterleavedContentItem] = []
-        for item in result.content:
-            if isinstance(item, mcp_types.TextContent):
-                content.append(TextContentItem(text=item.text))
-            elif isinstance(item, mcp_types.ImageContent):
-                content.append(ImageContentItem(image=_URLOrData(data=item.data)))
-            elif isinstance(item, mcp_types.EmbeddedResource):
-                logger.warning(f"EmbeddedResource is not supported: {item}")
-            else:
-                raise ValueError(f"Unknown content type: {type(item)}")
-        return ToolInvocationResult(
-            content=content,
-            error_code=1 if result.isError else 0,
-        )
+    # Fallback to original behavior: create a new session for each call
+    async with client_wrapper(endpoint, final_headers) as session:
+        result = await session.call_tool(tool_name, kwargs)
+        return _parse_mcp_result(result)

{llama_stack-0.4.2.dist-info → llama_stack-0.4.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llama_stack
-Version: 0.4.2
+Version: 0.4.4
 Summary: Llama Stack
 Author-email: Meta Llama <llama-oss@meta.com>
 License: MIT
@@ -46,7 +46,7 @@ Requires-Dist: psycopg2-binary
 Requires-Dist: tornado>=6.5.3
 Requires-Dist: urllib3>=2.6.3
 Provides-Extra: client
-Requires-Dist: llama-stack-client==0.4.2; extra == "client"
+Requires-Dist: llama-stack-client==0.4.4; extra == "client"
 Dynamic: license-file
 
 # Llama Stack