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

Code before:
                    )
                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.
        """
        errors = []
        for key, (session, client_ctx, _) in list(self._sessions.items()):
            try:
                await session.__aexit__(None, None, None)
            except Exception as e:
                logger.warning(f"Error closing MCP session {key}: {e}")
                errors.append(e)
            try:
                await client_ctx.__aexit__(None, None, None)
            except Exception as e:
                logger.warning(f"Error closing MCP client context {key}: {e}")
                errors.append(e)

        self._sessions.clear()
        self._locks.clear()
        logger.debug(f"Closed {len(self._sessions)} MCP sessions")

        if errors:
            logger.warning(f"Encountered {len(errors)} errors while closing MCP sessions")

    async def __aenter__(self) -> "MCPSessionManager":
        """Enter the async context manager."""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Exit the async context manager and close all sessions."""
        await self.close_all()


@asynccontextmanager
async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerator[ClientSession, Any]:

Review comment (security):
Not sure about this one, so feel free to close, but would it be possible to re-construct authorization headers from the session key hashes? If so, we probably should not log them, or at least not log them in their entirety.

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