Repository: meta-llama/llama-stack
File: src/llama_stack/providers/utils/tools/mcp.py (Python)

Code before:
    """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)

Review comment (security):
Is it possible for two users to create a session with the same endpoint and headers? What if those users shared credentials for the MCP server, but were distinct users of llama stack? Would we need to additionally include some way to identify the user in this hash?

Also, how does this behave over multiple conversation turns? Is each turn treated as its own request, or does this require session affinity?

Code after:
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)