hud-python 0.4.24__tar.gz → 0.4.25__tar.gz

{hud_python-0.4.24 → hud_python-0.4.25}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hud-python
-Version: 0.4.24
+Version: 0.4.25
 Summary: SDK for the HUD platform.
 Project-URL: Homepage, https://github.com/hud-evals/hud-python
 Project-URL: Bug Tracker, https://github.com/hud-evals/hud-python/issues

{hud_python-0.4.24 → hud_python-0.4.25}/hud/clients/mcp_use.py

@@ -15,7 +15,7 @@ from hud.types import MCPToolCall, MCPToolResult
 from hud.version import __version__ as hud_version
 
 from .base import BaseHUDClient
-from .utils.retry import retry_with_backoff
+from .utils.mcp_use_retry import patch_all_sessions
 
 logger = logging.getLogger(__name__)
 
@@ -64,6 +64,10 @@ class MCPUseHUDClient(BaseHUDClient):
             self._sessions = await self._client.create_all_sessions()
             logger.info("Created %d MCP sessions", len(self._sessions))
 
+            # Patch all sessions with retry logic
+            patch_all_sessions(self._sessions)
+            logger.debug("Applied retry logic to all MCP sessions")
+
             # Configure validation for all sessions based on client setting
             try:
                 for session in self._sessions.values():
@@ -128,11 +132,8 @@ class MCPUseHUDClient(BaseHUDClient):
                     logger.warning("Client session not initialized for %s", server_name)
                     continue
 
-                # List tools with retry logic for HTTP errors
-                tools_result = await retry_with_backoff(
-                    session.connector.client_session.list_tools,
-                    operation_name=f"list_tools_{server_name}",
-                )
+                # List tools (retry logic is handled at transport level)
+                tools_result = await session.connector.client_session.list_tools()
 
                 logger.info(
                     "Discovered %d tools from '%s': %s",
@@ -206,12 +207,10 @@ class MCPUseHUDClient(BaseHUDClient):
         if session.connector.client_session is None:
             raise ValueError(f"Client session not initialized for {server_name}")
 
-        # Call tool with retry logic for HTTP errors (502, 503, 504)
-        result = await retry_with_backoff(
-            session.connector.client_session.call_tool,
+        # Call tool (retry logic is handled at transport level)
+        result = await session.connector.client_session.call_tool(
             name=original_tool.name,  # Use original tool name, not prefixed
             arguments=tool_call.arguments or {},
-            operation_name=f"call_tool_{original_tool.name}",
         )
 
         if self.verbose:
@@ -239,10 +238,8 @@ class MCPUseHUDClient(BaseHUDClient):
                     continue
                 # Prefer standard method name if available
                 if hasattr(session.connector.client_session, "list_resources"):
-                    resources = await retry_with_backoff(
-                        session.connector.client_session.list_resources,
-                        operation_name=f"list_resources_{server_name}",
-                    )
+                    # List resources (retry logic is handled at transport level)
+                    resources = await session.connector.client_session.list_resources()
                 else:
                     # If the client doesn't support resource listing, skip
                     continue
@@ -272,11 +269,8 @@ class MCPUseHUDClient(BaseHUDClient):
                 resource_uri = AnyUrl(uri) if isinstance(uri, str) else uri
                 # Prefer read_resource; fall back to list_resources if needed
                 if hasattr(session.connector.client_session, "read_resource"):
-                    result = await retry_with_backoff(
-                        session.connector.client_session.read_resource,
-                        resource_uri,
-                        operation_name=f"read_resource_{server_name}",
-                    )
+                    # Read resource (retry logic is handled at transport level)
+                    result = await session.connector.client_session.read_resource(resource_uri)
                 else:
                     # Fallback path for older clients: not supported in strict typing
                     raise AttributeError("read_resource not available")

hud_python-0.4.25/hud/clients/tests/test_mcp_use_retry.py

@@ -0,0 +1,378 @@
+"""Tests for MCP-use client retry functionality."""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+import requests
+from mcp import types
+
+from hud.clients.mcp_use import MCPUseHUDClient
+from hud.clients.utils.mcp_use_retry import (
+    create_async_retry_wrapper,
+    create_retry_session,
+    patch_all_sessions,
+    patch_mcp_session_http_client,
+)
+from hud.types import MCPToolCall
+
+
+class TestRetrySession:
+    """Test the retry session creation."""
+
+    def test_create_retry_session(self):
+        """Test that retry session is configured correctly."""
+        session = create_retry_session(
+            max_retries=5,
+            retry_status_codes=(500, 502, 503, 504),
+            retry_delay=0.5,
+            backoff_factor=2.0,
+        )
+
+        # Check that session has adapters mounted
+        assert "http://" in session.adapters
+        assert "https://" in session.adapters
+
+        # Check adapter configuration
+        adapter = session.adapters["http://"]
+        assert adapter.max_retries.total == 5
+        assert 500 in adapter.max_retries.status_forcelist
+        assert 502 in adapter.max_retries.status_forcelist
+        assert adapter.max_retries.backoff_factor == 2.0
+
+    def test_retry_session_default_values(self):
+        """Test retry session with default values."""
+        session = create_retry_session()
+
+        adapter = session.adapters["https://"]
+        assert adapter.max_retries.total == 3
+        assert 502 in adapter.max_retries.status_forcelist
+        assert 503 in adapter.max_retries.status_forcelist
+        assert 504 in adapter.max_retries.status_forcelist
+
+
+class TestAsyncRetryWrapper:
+    """Test the async retry wrapper functionality."""
+
+    @pytest.mark.asyncio
+    async def test_retry_on_error_status_codes(self):
+        """Test that async wrapper retries on specific status codes."""
+        call_count = 0
+
+        async def mock_func(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+
+            # First two calls fail, third succeeds
+            if call_count < 3:
+                result = Mock()
+                result.status_code = 503  # Service unavailable
+                return result
+
+            result = Mock()
+            result.status_code = 200
+            return result
+
+        wrapped = create_async_retry_wrapper(
+            mock_func,
+            max_retries=3,
+            retry_status_codes=(503,),
+            retry_delay=0.01,  # Short delay for testing
+        )
+
+        result = await wrapped()
+        assert call_count == 3
+        assert result.status_code == 200
+
+    @pytest.mark.asyncio
+    async def test_retry_on_exception(self):
+        """Test that async wrapper retries on exceptions with status codes."""
+        call_count = 0
+
+        async def mock_func(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+
+            if call_count < 3:
+                raise Exception("HTTP 503 Service Unavailable")
+
+            return Mock(status_code=200)
+
+        wrapped = create_async_retry_wrapper(
+            mock_func,
+            max_retries=3,
+            retry_status_codes=(503,),
+            retry_delay=0.01,
+        )
+
+        result = await wrapped()
+        assert call_count == 3
+        assert result.status_code == 200
+
+    @pytest.mark.asyncio
+    async def test_no_retry_on_success(self):
+        """Test that successful calls don't trigger retries."""
+        call_count = 0
+
+        async def mock_func(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            return Mock(status_code=200)
+
+        wrapped = create_async_retry_wrapper(mock_func)
+
+        result = await wrapped()
+        assert call_count == 1
+        assert result.status_code == 200
+
+    @pytest.mark.asyncio
+    async def test_max_retries_exceeded(self):
+        """Test that retries stop after max attempts."""
+        call_count = 0
+
+        async def mock_func(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            raise Exception("HTTP 503 Service Unavailable")
+
+        wrapped = create_async_retry_wrapper(
+            mock_func,
+            max_retries=2,
+            retry_status_codes=(503,),
+            retry_delay=0.01,
+        )
+
+        with pytest.raises(Exception) as exc_info:
+            await wrapped()
+
+        assert "503" in str(exc_info.value)
+        assert call_count == 3  # Initial + 2 retries
+
+
+class TestSessionPatching:
+    """Test the session patching functionality."""
+
+    def test_patch_sync_session(self):
+        """Test patching a synchronous session."""
+        # Create mock session with connector
+        mock_session = Mock()
+        mock_session.connector = Mock()
+        mock_session.connector._connection_manager = Mock()
+        mock_session.connector._connection_manager._session = requests.Session()
+
+        # Patch the session
+        patch_mcp_session_http_client(mock_session)
+
+        # Verify the session was replaced with retry-enabled one
+        patched_session = mock_session.connector._connection_manager._session
+        assert "http://" in patched_session.adapters
+        assert "https://" in patched_session.adapters
+
+        # Check that it has retry configuration
+        adapter = patched_session.adapters["http://"]
+        assert hasattr(adapter, "max_retries")
+
+    @pytest.mark.asyncio
+    async def test_patch_async_session(self):
+        """Test patching an async session."""
+        # Create mock async session
+        mock_session = Mock()
+        mock_session.connector = Mock()
+        mock_session.connector.client_session = Mock()
+
+        async def mock_send_request(*args, **kwargs):
+            return Mock(status_code=200)
+
+        mock_session.connector.client_session._send_request = mock_send_request
+
+        # Patch the session
+        patch_mcp_session_http_client(mock_session)
+
+        # Verify _send_request was wrapped
+        wrapped_func = mock_session.connector.client_session._send_request
+        assert wrapped_func != mock_send_request  # Function was replaced
+
+        # Test that wrapped function still works
+        result = await wrapped_func()
+        assert result.status_code == 200
+
+    def test_patch_all_sessions(self):
+        """Test patching multiple sessions."""
+        # Create mock sessions
+        session1 = Mock()
+        session1.connector = Mock()
+        session1.connector._connection_manager = Mock()
+        session1.connector._connection_manager.session = requests.Session()
+
+        session2 = Mock()
+        session2.connector = Mock()
+        session2.connector.client_session = Mock()
+        session2.connector.client_session._send_request = AsyncMock()
+
+        sessions = {"server1": session1, "server2": session2}
+
+        # Patch all sessions
+        patch_all_sessions(sessions)
+
+        # Verify both were patched
+        assert "http://" in session1.connector._connection_manager.session.adapters
+        assert session2.connector.client_session._send_request != AsyncMock
+
+
+class TestMCPUseClientRetry:
+    """Test retry functionality integrated into MCPUseHUDClient."""
+
+    @pytest.mark.asyncio
+    async def test_client_applies_retry_on_connect(self):
+        """Test that MCPUseHUDClient applies retry logic during connection."""
+        config = {"test_server": {"url": "http://localhost:8080"}}
+        client = MCPUseHUDClient(config)
+
+        # Mock the MCPUseClient and session creation
+        with patch("hud.clients.mcp_use.MCPUseClient") as MockMCPUseClient:
+            mock_client = Mock()
+            MockMCPUseClient.from_dict.return_value = mock_client
+
+            # Create mock session
+            mock_session = Mock()
+            mock_session.connector = Mock()
+            mock_session.connector.client_session = Mock()
+            mock_session.connector.client_session._send_request = AsyncMock()
+            mock_session.connector.client_session.list_tools = AsyncMock(
+                return_value=Mock(tools=[])
+            )
+
+            mock_client.create_all_sessions = AsyncMock(return_value={"test_server": mock_session})
+
+            # Initialize client (which applies retry logic)
+            await client.initialize()
+
+            # Verify session was created and patched
+            assert len(client._sessions) == 1
+            assert "test_server" in client._sessions
+
+    @pytest.mark.asyncio
+    async def test_tool_call_with_retry(self):
+        """Test that tool calls work with retry logic."""
+        config = {"test_server": {"url": "http://localhost:8080"}}
+        client = MCPUseHUDClient(config)
+
+        with patch("hud.clients.mcp_use.MCPUseClient") as MockMCPUseClient:
+            mock_client = Mock()
+            MockMCPUseClient.from_dict.return_value = mock_client
+
+            # Create mock session
+            mock_session = Mock()
+            mock_session.connector = Mock()
+            mock_session.connector.client_session = Mock()
+
+            # Mock tool listing
+            test_tool = types.Tool(
+                name="test_tool",
+                description="Test tool",
+                inputSchema={"type": "object"},
+            )
+            mock_session.connector.client_session.list_tools = AsyncMock(
+                return_value=Mock(tools=[test_tool])
+            )
+
+            # Mock tool call with simulated retry
+            call_count = 0
+
+            async def mock_call_tool(name, arguments):
+                nonlocal call_count
+                call_count += 1
+
+                # First call fails, second succeeds
+                if call_count == 1:
+                    raise Exception("HTTP 503 Service Unavailable")
+
+                return Mock(
+                    content=[types.TextContent(type="text", text="Success")],
+                    isError=False,
+                    structuredContent=None,
+                )
+
+            mock_session.connector.client_session.call_tool = mock_call_tool
+            mock_session.connector.client_session._send_request = AsyncMock()
+
+            mock_client.create_all_sessions = AsyncMock(return_value={"test_server": mock_session})
+
+            # Initialize and call tool
+            await client.initialize()
+
+            # Wrap call_tool with retry for this test
+            original_call = mock_session.connector.client_session.call_tool
+            mock_session.connector.client_session.call_tool = create_async_retry_wrapper(
+                original_call,
+                max_retries=2,
+                retry_status_codes=(503,),
+                retry_delay=0.01,
+            )
+
+            result = await client.call_tool(MCPToolCall(name="test_tool", arguments={}))
+
+            # Verify retry worked
+            assert call_count == 2  # Failed once, then succeeded
+            assert not result.isError
+            assert result.content[0].text == "Success"
+
+    @pytest.mark.asyncio
+    async def test_resource_read_with_retry(self):
+        """Test that resource reading works with retry logic."""
+        config = {"test_server": {"url": "http://localhost:8080"}}
+        client = MCPUseHUDClient(config)
+
+        with patch("hud.clients.mcp_use.MCPUseClient") as MockMCPUseClient:
+            mock_client = Mock()
+            MockMCPUseClient.from_dict.return_value = mock_client
+
+            # Create mock session
+            mock_session = Mock()
+            mock_session.connector = Mock()
+            mock_session.connector.client_session = Mock()
+            mock_session.connector.client_session.list_tools = AsyncMock(
+                return_value=Mock(tools=[])
+            )
+
+            # Mock resource read with simulated retry
+            call_count = 0
+
+            async def mock_read_resource(uri):
+                nonlocal call_count
+                call_count += 1
+
+                # First call fails, second succeeds
+                if call_count == 1:
+                    raise Exception("HTTP 502 Bad Gateway")
+
+                return Mock(contents=[Mock(text='{"status": "ok"}')])
+
+            mock_session.connector.client_session.read_resource = mock_read_resource
+            mock_session.connector.client_session._send_request = AsyncMock()
+
+            mock_client.create_all_sessions = AsyncMock(return_value={"test_server": mock_session})
+
+            # Initialize
+            await client.initialize()
+
+            # Wrap read_resource with retry for this test
+            original_read = mock_session.connector.client_session.read_resource
+            mock_session.connector.client_session.read_resource = create_async_retry_wrapper(
+                original_read,
+                max_retries=2,
+                retry_status_codes=(502,),
+                retry_delay=0.01,
+            )
+
+            result = await client.read_resource("test://resource")
+
+            # Verify retry worked
+            assert call_count == 2  # Failed once, then succeeded
+            assert result is not None
+            assert result.contents[0].text == '{"status": "ok"}'
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

hud_python-0.4.25/hud/clients/utils/mcp_use_retry.py

@@ -0,0 +1,201 @@
+"""Retry wrapper for MCP-use HTTP transport.
+
+This module provides a transport-level retry mechanism for MCP-use,
+similar to the approach used in FastMCP.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Any, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def create_retry_session(
+    max_retries: int = 3,
+    retry_status_codes: tuple[int, ...] = (502, 503, 504),
+    retry_delay: float = 1.0,
+    backoff_factor: float = 2.0,
+) -> requests.Session:
+    """
+    Create a requests session with retry logic.
+
+    Args:
+        max_retries: Maximum number of retry attempts
+        retry_status_codes: HTTP status codes to retry
+        retry_delay: Initial delay between retries in seconds
+        backoff_factor: Multiplier for exponential backoff
+
+    Returns:
+        Configured requests.Session with retry logic
+    """
+    session = requests.Session()
+
+    # Configure retry strategy
+    retry = Retry(
+        total=max_retries,
+        backoff_factor=backoff_factor,
+        status_forcelist=list(retry_status_codes),
+        # Allow retries on all methods
+        allowed_methods=["HEAD", "GET", "OPTIONS", "POST", "PUT", "DELETE", "PATCH"],
+        # Respect Retry-After header if present
+        respect_retry_after_header=True,
+    )
+
+    # Create adapter with retry strategy
+    adapter = HTTPAdapter(max_retries=retry)
+
+    # Mount adapter for both HTTP and HTTPS
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+
+    logger.debug(
+        "Created retry session with max_retries=%d, status_codes=%s, backoff_factor=%.1f",
+        max_retries,
+        retry_status_codes,
+        backoff_factor,
+    )
+
+    return session
+
+
+def patch_mcp_session_http_client(session: Any) -> None:
+    """
+    Patch an MCP-use session to use HTTP retry logic.
+
+    This function attempts to replace the HTTP client used by an MCP session
+    with one that has retry logic enabled.
+
+    Args:
+        session: MCP-use session to patch
+    """
+    try:
+        # Check if session has a connector with an HTTP client
+        if hasattr(session, "connector"):
+            connector = session.connector
+
+            # For HTTP connectors, patch the underlying HTTP client
+            if hasattr(connector, "_connection_manager"):
+                manager = connector._connection_manager
+
+                # If it's using requests, replace the session
+                if hasattr(manager, "_session") or hasattr(manager, "session"):
+                    retry_session = create_retry_session()
+
+                    # Try different attribute names
+                    if hasattr(manager, "_session"):
+                        manager._session = retry_session
+                        logger.debug("Patched connection manager's _session with retry logic")
+                    elif hasattr(manager, "session"):
+                        manager.session = retry_session
+                        logger.debug("Patched connection manager's session with retry logic")
+
+            # Also check for client_session (async variant)
+            if hasattr(connector, "client_session") and connector.client_session:
+                client = connector.client_session
+
+                # Wrap the async HTTP methods with retry logic
+                if hasattr(client, "_send_request"):
+                    original_send = client._send_request
+                    client._send_request = create_async_retry_wrapper(original_send)
+                    logger.debug("Wrapped client_session._send_request with retry logic")
+
+    except Exception as e:
+        logger.warning("Could not patch MCP session with retry logic: %s", e)
+
+
+def create_async_retry_wrapper(
+    func: Callable[..., Any],
+    max_retries: int = 3,
+    retry_status_codes: tuple[int, ...] = (502, 503, 504),
+    retry_delay: float = 1.0,
+    backoff_factor: float = 2.0,
+) -> Callable[..., Any]:
+    """
+    Create an async wrapper that adds retry logic to a function.
+
+    Args:
+        func: The async function to wrap
+        max_retries: Maximum number of retry attempts
+        retry_status_codes: HTTP status codes to retry
+        retry_delay: Initial delay between retries
+        backoff_factor: Multiplier for exponential backoff
+
+    Returns:
+        Wrapped function with retry logic
+    """
+
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        last_exception = None
+        delay = retry_delay
+
+        for attempt in range(max_retries + 1):
+            try:
+                result = await func(*args, **kwargs)
+
+                # Check if result has a status code that should trigger retry
+                if (
+                    hasattr(result, "status_code")
+                    and result.status_code in retry_status_codes
+                    and attempt < max_retries
+                ):
+                    logger.warning(
+                        "HTTP %d error (attempt %d/%d), retrying in %.1fs",
+                        result.status_code,
+                        attempt + 1,
+                        max_retries + 1,
+                        delay,
+                    )
+                    await asyncio.sleep(delay)
+                    delay *= backoff_factor
+                    continue
+
+                return result
+
+            except Exception as e:
+                # Check if it's an HTTP error that should be retried
+                error_str = str(e)
+                should_retry = any(str(code) in error_str for code in retry_status_codes)
+
+                if should_retry and attempt < max_retries:
+                    logger.warning(
+                        "Error '%s' (attempt %d/%d), retrying in %.1fs",
+                        e,
+                        attempt + 1,
+                        max_retries + 1,
+                        delay,
+                    )
+                    await asyncio.sleep(delay)
+                    delay *= backoff_factor
+                    last_exception = e
+                else:
+                    raise
+
+        # If we exhausted retries, raise the last exception
+        if last_exception:
+            raise last_exception
+
+    return wrapper
+
+
+def patch_all_sessions(sessions: dict[str, Any]) -> None:
+    """
+    Apply retry logic to all MCP sessions.
+
+    Args:
+        sessions: Dictionary of session name to session object
+    """
+    for name, session in sessions.items():
+        logger.debug("Patching session '%s' with retry logic", name)
+        patch_mcp_session_http_client(session)