hud-python 0.4.23__py3-none-any.whl → 0.4.25__py3-none-any.whl

hud/agents/base.py

@@ -207,6 +207,7 @@ class MCPAgent(ABC):
             else:
                 raise TypeError(f"prompt_or_task must be str or Task, got {type(prompt_or_task)}")
         except Exception as e:
+            # Always return a Trace object for any exception
             if self._is_connection_error(e):
                 # Return error trace for connection failures
                 return Trace(
@@ -215,7 +216,15 @@ class MCPAgent(ABC):
                     content=self._get_connection_error_message(e),
                     isError=True,
                 )
-            raise
+            else:
+                # Return error trace for any other exception
+                return Trace(
+                    reward=0.0,
+                    done=True,
+                    content=f"Task failed with error: {e}",
+                    isError=True,
+                    info={"error": str(e)},
+                )
         finally:
             # Cleanup auto-created resources
             await self._cleanup()
@@ -262,34 +271,53 @@ class MCPAgent(ABC):
             prompt_result = Trace(reward=0.0, done=True, content=str(e), isError=True)
             prompt_result.populate_from_context()
 
-        # Always evaluate if we have a prompt result and evaluate tool
-        if prompt_result is not None and task.evaluate_tool is not None:
+        # Always evaluate if we have evaluate tool, regardless of errors
+        if task.evaluate_tool is not None:
             try:
                 self.console.progress_log(f"Evaluating tool phase: {task.evaluate_tool}")
                 results = await self.call_tools(task.evaluate_tool)
 
                 if any(result.isError for result in results):
-                    raise RuntimeError(f"{results}")
-
-                # Extract reward and content from evaluation
-                if results:
-                    reward = find_reward(results[0])
-                    eval_content = find_content(results[0])
-
-                    # Update the prompt result with evaluation reward
-                    prompt_result.reward = reward
-
-                    # Update the prompt result with evaluation content (if available)
-                    if eval_content:
-                        # Prompt result may already have final response content, so we append to it
-                        if prompt_result.content:
-                            prompt_result.content += "\n\n" + eval_content
+                    self.console.warning_log(f"Evaluate tool returned error: {results}")
+                    # Still extract what we can from the error response
+                    if prompt_result is None:
+                        prompt_result = Trace(
+                            reward=0.0,
+                            done=True,
+                            content="Task failed before evaluation",
+                            isError=True,
+                        )
+                    prompt_result.reward = 0.0  # Default to 0 on error
+                else:
+                    # Extract reward and content from evaluation
+                    if results:
+                        reward = find_reward(results[0])
+                        eval_content = find_content(results[0])
+
+                        # Update the prompt result with evaluation reward
+                        if prompt_result is None:
+                            prompt_result = Trace(
+                                reward=reward, done=True, content=eval_content or "", isError=False
+                            )
                         else:
-                            prompt_result.content = eval_content
+                            prompt_result.reward = reward
+
+                            # Update the prompt result with evaluation content (if available)
+                            if eval_content:
+                                # Prompt result may already have final response content,
+                                # so we append to it
+                                if prompt_result.content:
+                                    prompt_result.content += "\n\n" + eval_content
+                                else:
+                                    prompt_result.content = eval_content
 
             except Exception as e:
                 self.console.error_log(f"Evaluation phase failed: {e}")
-                # Continue with the prompt result even if evaluation failed
+                # Ensure we have a result even if evaluation failed
+                if prompt_result is None:
+                    prompt_result = Trace(
+                        reward=0.0, done=True, content=f"Evaluation failed: {e}", isError=True
+                    )
 
         return (
             prompt_result

hud/agents/claude.py

@@ -196,7 +196,11 @@ class ClaudeAgent(MCPAgent):
                 response = await self.anthropic_client.beta.messages.create(**create_kwargs)
                 break
             except BadRequestError as e:
-                if e.message.startswith("prompt is too long"):
+                if (
+                    "prompt is too long" in str(e)
+                    or "request_too_large" in str(e)
+                    or e.status_code == 413
+                ):
                     logger.warning("Prompt too long, truncating message history")
                     # Keep first message and last 20 messages
                     if len(current_messages) > 21:

hud/clients/mcp_use.py

@@ -15,6 +15,7 @@ from hud.types import MCPToolCall, MCPToolResult
 from hud.version import __version__ as hud_version
 
 from .base import BaseHUDClient
+from .utils.mcp_use_retry import patch_all_sessions
 
 logger = logging.getLogger(__name__)
 
@@ -63,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():
@@ -127,7 +132,7 @@ class MCPUseHUDClient(BaseHUDClient):
                     logger.warning("Client session not initialized for %s", server_name)
                     continue
 
-                # List tools
+                # List tools (retry logic is handled at transport level)
                 tools_result = await session.connector.client_session.list_tools()
 
                 logger.info(
@@ -202,6 +207,7 @@ class MCPUseHUDClient(BaseHUDClient):
         if session.connector.client_session is None:
             raise ValueError(f"Client session not initialized for {server_name}")
 
+        # 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 {},
@@ -232,6 +238,7 @@ class MCPUseHUDClient(BaseHUDClient):
                     continue
                 # Prefer standard method name if available
                 if hasattr(session.connector.client_session, "list_resources"):
+                    # 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
@@ -262,6 +269,7 @@ 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"):
+                    # 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

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/clients/utils/__init__.py

@@ -1 +1,26 @@
 """HUD MCP client utilities."""
+
+from __future__ import annotations
+
+from .retry import (
+    DEFAULT_BACKOFF_FACTOR,
+    DEFAULT_MAX_RETRIES,
+    DEFAULT_RETRY_DELAY,
+    DEFAULT_RETRY_STATUS_CODES,
+    is_retryable_error,
+    retry_with_backoff,
+    with_retry,
+)
+from .retry_transport import RetryTransport, create_retry_httpx_client
+
+__all__ = [
+    "DEFAULT_BACKOFF_FACTOR",
+    "DEFAULT_MAX_RETRIES",
+    "DEFAULT_RETRY_DELAY",
+    "DEFAULT_RETRY_STATUS_CODES",
+    "RetryTransport",
+    "create_retry_httpx_client",
+    "is_retryable_error",
+    "retry_with_backoff",
+    "with_retry",
+]

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)

hud/clients/utils/retry.py

@@ -0,0 +1,186 @@
+"""Shared retry utilities for MCP client operations."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from functools import wraps
+from typing import TYPE_CHECKING, Any, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+from httpx import HTTPStatusError
+from mcp.shared.exceptions import McpError
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+# Default retry configuration matching requests.py
+DEFAULT_MAX_RETRIES = 4
+DEFAULT_RETRY_DELAY = 2.0
+DEFAULT_RETRY_STATUS_CODES = {502, 503, 504}
+DEFAULT_BACKOFF_FACTOR = 2.0
+
+
+def is_retryable_error(error: Exception, retry_status_codes: set[int]) -> bool:
+    """
+    Check if an error is retryable based on status codes.
+
+    Args:
+        error: The exception to check
+        retry_status_codes: Set of HTTP status codes to retry on
+
+    Returns:
+        True if the error is retryable, False otherwise
+    """
+    # Check for HTTP status errors with retryable status codes
+    if isinstance(error, HTTPStatusError):
+        return error.response.status_code in retry_status_codes
+
+    # Check for MCP errors that might wrap HTTP errors
+    if isinstance(error, McpError):
+        error_msg = str(error).lower()
+        # Check for common gateway error patterns in the message
+        for code in retry_status_codes:
+            if str(code) in error_msg:
+                return True
+        # Check for gateway error keywords
+        if any(
+            keyword in error_msg
+            for keyword in ["bad gateway", "service unavailable", "gateway timeout"]
+        ):
+            return True
+
+    # Check for generic errors with status codes in the message
+    error_msg = str(error)
+    for code in retry_status_codes:
+        if f"{code}" in error_msg or f"status {code}" in error_msg.lower():
+            return True
+
+    return False
+
+
+async def retry_with_backoff(
+    func: Callable[..., Any],
+    *args: Any,
+    max_retries: int = DEFAULT_MAX_RETRIES,
+    retry_delay: float = DEFAULT_RETRY_DELAY,
+    retry_status_codes: set[int] | None = None,
+    backoff_factor: float = DEFAULT_BACKOFF_FACTOR,
+    operation_name: str | None = None,
+    **kwargs: Any,
+) -> Any:
+    """
+    Execute an async function with retry logic and exponential backoff.
+
+    This matches the retry behavior in requests.py but can be applied
+    to any async function, particularly MCP client operations.
+
+    Args:
+        func: The async function to retry
+        *args: Positional arguments for the function
+        max_retries: Maximum number of retry attempts
+        retry_delay: Initial delay between retries in seconds
+        retry_status_codes: HTTP status codes to retry on
+        backoff_factor: Multiplier for exponential backoff
+        operation_name: Name of the operation for logging
+        **kwargs: Keyword arguments for the function
+
+    Returns:
+        The result of the function call
+
+    Raises:
+        The last exception if all retries are exhausted
+    """
+    if retry_status_codes is None:
+        retry_status_codes = DEFAULT_RETRY_STATUS_CODES
+
+    operation = operation_name or func.__name__
+    last_error = None
+
+    for attempt in range(max_retries + 1):
+        try:
+            result = await func(*args, **kwargs)
+            return result
+        except Exception as e:
+            last_error = e
+
+            # Check if this is a retryable error
+            if not is_retryable_error(e, retry_status_codes):
+                # Not retryable, raise immediately
+                raise
+
+            # Don't retry if we've exhausted attempts
+            if attempt >= max_retries:
+                logger.debug(
+                    "Operation '%s' failed after %d retries: %s",
+                    operation,
+                    max_retries,
+                    e,
+                )
+                raise
+
+            # Calculate backoff delay (exponential backoff)
+            delay = retry_delay * (backoff_factor**attempt)
+
+            logger.warning(
+                "Operation '%s' failed with retryable error, "
+                "retrying in %.2f seconds (attempt %d/%d): %s",
+                operation,
+                delay,
+                attempt + 1,
+                max_retries,
+                e,
+            )
+
+            await asyncio.sleep(delay)
+
+    # This should never be reached, but just in case
+    if last_error:
+        raise last_error
+    raise RuntimeError(f"Unexpected retry loop exit for operation '{operation}'")
+
+
+def with_retry(
+    max_retries: int = DEFAULT_MAX_RETRIES,
+    retry_delay: float = DEFAULT_RETRY_DELAY,
+    retry_status_codes: set[int] | None = None,
+    backoff_factor: float = DEFAULT_BACKOFF_FACTOR,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """
+    Decorator to add retry logic to async methods.
+
+    Usage:
+        @with_retry(max_retries=3)
+        async def my_method(self, ...):
+            ...
+
+    Args:
+        max_retries: Maximum number of retry attempts
+        retry_delay: Initial delay between retries
+        retry_status_codes: HTTP status codes to retry on
+        backoff_factor: Multiplier for exponential backoff
+
+    Returns:
+        Decorated function with retry logic
+    """
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return await retry_with_backoff(
+                func,
+                *args,
+                max_retries=max_retries,
+                retry_delay=retry_delay,
+                retry_status_codes=retry_status_codes,
+                backoff_factor=backoff_factor,
+                operation_name=func.__name__,
+                **kwargs,
+            )
+
+        return wrapper
+
+    return decorator

hud/datasets/execution/parallel.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import multiprocessing
 import os
 import traceback
 from concurrent.futures import ProcessPoolExecutor, as_completed
@@ -191,18 +192,15 @@ def _process_worker(
         # Run the async batch processing
         results = loop.run_until_complete(process_batch())
 
-        # CRITICAL: Ensure telemetry is fully sent before process exits
-        # Two things need to complete:
-        # 1. The trace context's __exit__ already called _update_task_status_sync (blocking)
-        # 2. But spans are buffered in BatchSpanProcessor and need explicit flush
+        # Ensure telemetry is fully sent before process exits
+        # Spans are buffered in BatchSpanProcessor and need explicit flush
 
         from opentelemetry import trace as otel_trace
 
         provider = otel_trace.get_tracer_provider()
         if provider and hasattr(provider, "force_flush"):
-            # This forces BatchSpanProcessor to export all buffered spans NOW
-            # The method returns True if successful, False if timeout
-            success = provider.force_flush(timeout_millis=5000)  # 5 second timeout # type: ignore
+            # Flush of buffered spans
+            success = provider.force_flush(timeout_millis=2000)  # type: ignore[arg-type]
             if not success:
                 logger.warning("Worker %s: Telemetry flush timed out", worker_id)
 
@@ -230,6 +228,18 @@ def _process_worker(
         logger.error("Worker %s batch processing failed: %s", worker_id, e)
         return [(idx, {"error": str(e), "isError": True}) for idx, _ in task_batch]
     finally:
+        try:
+            from opentelemetry import trace as otel_trace
+
+            provider = otel_trace.get_tracer_provider()
+            if provider and hasattr(provider, "force_flush"):
+                # Flush buffered spans with reasonable timeout
+                success = provider.force_flush(timeout_millis=2000)  # type: ignore[arg-type]
+                if not success:
+                    logger.warning("Worker %s: Telemetry flush timed out", worker_id)
+        except Exception as e:
+            logger.warning("Worker %s: Failed to flush telemetry: %s", worker_id, e)
+
         # Clean up the event loop
         try:
             loop.close()
@@ -342,6 +352,8 @@ async def run_dataset_parallel_manual(
         for task_dict in task_dicts:
             if "system_prompt" not in task_dict:
                 task_dict["system_prompt"] = custom_system_prompt
+            else:
+                task_dict["system_prompt"] += "\n" + custom_system_prompt
 
     # Prepare job metadata
     job_metadata = metadata or {}
@@ -366,6 +378,8 @@ async def run_dataset_parallel_manual(
         except Exception:
             logger.warning("Failed to extract dataset verification info")
 
+    # task_dicts = task_dicts[:10]
+
     # Create job context
     with hud.job(name, metadata=job_metadata, dataset_link=dataset_link) as job_obj:
         # Prepare agent class info for pickling
@@ -410,7 +424,10 @@ async def run_dataset_parallel_manual(
         )
 
         # Process batches in parallel using ProcessPoolExecutor
-        executor = ProcessPoolExecutor(max_workers=max_workers)
+        executor = ProcessPoolExecutor(
+            max_workers=max_workers,
+            mp_context=multiprocessing.get_context("spawn"),
+        )
         try:
             # Submit all batches to workers
             future_to_batch = {

hud/otel/config.py

@@ -97,7 +97,16 @@ def configure_telemetry(
         exporter = HudSpanExporter(
             telemetry_url=settings.hud_telemetry_url, api_key=settings.api_key
         )
-        provider.add_span_processor(BatchSpanProcessor(exporter))
+        # Export more continuously to avoid big end flushes
+        provider.add_span_processor(
+            BatchSpanProcessor(
+                exporter,
+                schedule_delay_millis=1000,
+                max_queue_size=8192,
+                max_export_batch_size=256,
+                export_timeout_millis=30000,
+            )
+        )
     elif settings.telemetry_enabled and not settings.api_key and not enable_otlp:
         # Error if no exporters are configured
         raise ValueError(
@@ -127,7 +136,15 @@ def configure_telemetry(
                 otlp_config["headers"] = otlp_headers
 
             otlp_exporter = OTLPSpanExporter(**otlp_config)
-            provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
+            provider.add_span_processor(
+                BatchSpanProcessor(
+                    otlp_exporter,
+                    schedule_delay_millis=1000,
+                    max_queue_size=8192,
+                    max_export_batch_size=256,
+                    export_timeout_millis=30000,
+                )
+            )
             logger.info("OTLP HTTP exporter enabled - endpoint: %s", otlp_config["endpoint"])
         except ImportError:
             logger.warning(

hud/utils/tests/test_version.py

@@ -5,4 +5,4 @@ def test_import():
     """Test that the package can be imported."""
     import hud
 
-    assert hud.__version__ == "0.4.23"
+    assert hud.__version__ == "0.4.25"

hud/version.py

@@ -4,4 +4,4 @@ Version information for the HUD SDK.
 
 from __future__ import annotations
 
-__version__ = "0.4.23"
+__version__ = "0.4.25"

{hud_python-0.4.23.dist-info → hud_python-0.4.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hud-python
-Version: 0.4.23
+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.23.dist-info → hud_python-0.4.25.dist-info}/RECORD

@@ -2,10 +2,10 @@ hud/__init__.py,sha256=KU7G-_Mj6Mjf7trXA6X0ufN6QUmqhVi19NKbnNIvD74,532
 hud/__main__.py,sha256=YR8Dq8OhINOsVfQ55PmRXXg4fEK84Rt_-rMtJ5rvhWo,145
 hud/settings.py,sha256=bgq_zstNGlan7UDZOiElY8Aw5ZU4xL7Ds5HP_Xzph1A,2535
 hud/types.py,sha256=jta4Hgj_rBdIMcf7mz0tsdgfA1pssp327bxcT6Mfp00,6107
-hud/version.py,sha256=5rH72_P2FBJH7JH3nuYeXSXzeAEZoA-JgfNAoljtscY,105
+hud/version.py,sha256=9XXnK9XmJJQ0fZ1VqYaKnR6Lp3ekC6Wp0RRaPTJwdJ8,105
 hud/agents/__init__.py,sha256=UoIkljWdbq4bM0LD-mSaw6w826EqdEjOk7r6glNYwYQ,286
-hud/agents/base.py,sha256=qkTj0bnlFDs9OA4A3g1kKyVaDaRtPSVB_HzOIV4DpVw,30980
-hud/agents/claude.py,sha256=Ixc05qg-r0H1OoC_DE5Ov6RtODWp7uJcTJJqpmoIob0,15478
+hud/agents/base.py,sha256=Ui4FC4FPCTrJ6UZBObuY2mzA99veKzbWFCsFFSVOIJc,32406
+hud/agents/claude.py,sha256=yTVbSNMwNjF9qZAnwDSsgrscaenJpMGnGdW_4rFW_pY,15603
 hud/agents/grounded_openai.py,sha256=U-FHjB2Nh1_o0gmlxY5F17lWJ3oHsNRIB2a7z-IKB64,11231
 hud/agents/langchain.py,sha256=1EgCy8jfjunsWxlPC5XfvfLS6_XZVrIF1ZjtHcrvhYw,9584
 hud/agents/openai.py,sha256=ovARRWNuHqKkZ2Q_OCYSVCIZckrh8XY2jUB2p2x1m88,14259
@@ -70,18 +70,21 @@ hud/clients/README.md,sha256=XNE3mch95ozDgVqfwCGcrhlHY9CwT1GKfNANNboowto,3826
 hud/clients/__init__.py,sha256=N5M_gZv4nP7dLRwpAiaqqaxyaLieGW6397FszeG7JGw,364
 hud/clients/base.py,sha256=vt9mRMZwZg2DtVxQmccR_VZZGamXhx3dvlFJPulbOd8,14131
 hud/clients/fastmcp.py,sha256=KJGi8bmds0Q6rHnkTXb_Hw9ZqWmSo0OfjW05SSuyEJU,9182
-hud/clients/mcp_use.py,sha256=fKG_yRM-GlucQbd7f4aYF2Go0kSjwApHqO3DhbIPG04,13034
+hud/clients/mcp_use.py,sha256=TsIIItqmt93hcdW1hj2HlMPJppVo4JDoUuCqQ-ZBoBI,13521
 hud/clients/tests/__init__.py,sha256=sKOtJFFa4mDIXh1U6O8ZUHjigE8CiRMQ2PzJTIBZuVE,33
 hud/clients/tests/test_client_integration.py,sha256=kohU6jfCNfwSnAushHeB1_CmDlRfQc7VBL0GEdJYSeI,4198
 hud/clients/tests/test_fastmcp.py,sha256=4q3TzDjuieTZa89taiNJIrzbUncNkYOG4MaubypA21k,13030
+hud/clients/tests/test_mcp_use_retry.py,sha256=v9vwKa9nrW0EYiWzkq_vpyju40NaOm3SBiJh8VJYX2E,13103
 hud/clients/tests/test_protocol.py,sha256=aK4CS4g3j1D5jPo83ykzZuHUvcZFAulYtIq9T9Hb_fQ,6640
-hud/clients/utils/__init__.py,sha256=ucYJqOVpEsN-D9OFE2YTNLG628MgxcZAzfYhnbzx02k,32
+hud/clients/utils/__init__.py,sha256=-zZjcKIWGj2tXbVDOW45UgoGghhLJzFQVZ6miKenuA4,595
+hud/clients/utils/mcp_use_retry.py,sha256=sBCjtgnAXiXASjzFF_AtBEtmizay0Fi0nPL6sVoooeI,6675
+hud/clients/utils/retry.py,sha256=mMs2T_mAlb8AYhSqMR4AmCw7838gqCC4mdG3zjMAYM4,5744
 hud/clients/utils/retry_transport.py,sha256=Rsq25eiKKt_pM1bas78QEZvO0illK97X_3opmaS3A3w,6809
 hud/datasets/__init__.py,sha256=74T4mrjELKtE04XkZKwU8QAJcg2wjqXLqRO9s4GlPr4,678
 hud/datasets/task.py,sha256=bDVLy4EOBfjiU4i8hrFqcQ3dc077vzVNRXIbyNXFnp8,3916
 hud/datasets/utils.py,sha256=3hKvZTkZuCRkTeITB86nNdA1dtHZAqFfAdSPMtcTUhs,4275
 hud/datasets/execution/__init__.py,sha256=4m1AEpMQaUSJFVN_iAXvY6zFttVgZKwE6oQtC0Rrk7U,330
-hud/datasets/execution/parallel.py,sha256=4aL1XpS3vOBqZjgs0vrMZJ4eAoi86Td8C-m5SUtVxMs,25231
+hud/datasets/execution/parallel.py,sha256=ZEMMmH018QXOL1QoD_AyspQkGs_41F-GB_mTgdXh6J4,25780
 hud/datasets/execution/runner.py,sha256=EEvb90vvAqFXXx8NyVKLfK5p-gtsfJqiFJAoqSjyfXg,4695
 hud/misc/__init__.py,sha256=m_pprQQ-G-Y0Sd0NEiR8MtAMbElnuFZ2OWT8TXrw7c4,43
 hud/misc/claude_plays_pokemon.py,sha256=IthAkjDVr2Q-GNvX-QLJyMzN7-0pHqqJbagGNv2m7yo,10453
@@ -92,7 +95,7 @@ hud/native/tests/test_comparator.py,sha256=x1gFLXEDRIiJhH8tg5Rd3ptY-modYaHgSm6-h
 hud/native/tests/test_native_init.py,sha256=Is8fcDZimp1Oi2Bv4zavqM3KrpS86_DUXFnqc0AsCH0,2736
 hud/otel/__init__.py,sha256=ii17ayoWiS5vAhA7UAmZ8TkmP52gs2pWyHsD46-uYbE,1003
 hud/otel/collector.py,sha256=jLZymZ8r7xt2VDuWexfbnT7PY1-0aiyLMgjBy8KDY1M,4497
-hud/otel/config.py,sha256=6np_C2UXhtKHHjY41HQxZElua2Eh_EUCBiRB_YuiSuc,6249
+hud/otel/config.py,sha256=BPEyINdlxjsTj-IOO7F_JXAqrbMEWNMuhi2GDGqB5Kk,6785
 hud/otel/context.py,sha256=C9MvO99cRSNNDEDC7ehO3eoTPnb6J7AemUYvEp57yEU,17774
 hud/otel/exporters.py,sha256=RLAjWa8b2DJEU21740Idq4fmeIuabLEqGGUspcFDcH4,14331
 hud/otel/instrumentation.py,sha256=bBWxQ5vkiP-2WZ_2ztR5LAw0Wu02SZNKUgs5f-rU-ro,3734
@@ -174,10 +177,10 @@ hud/utils/tests/test_init.py,sha256=2QLQSGgyP9wJhOvPCusm_zjJad0qApOZi1BXpxcdHXQ,
 hud/utils/tests/test_mcp.py,sha256=0pUa16mL-bqbZDXp5NHBnt1gO5o10BOg7zTMHZ1DNPM,4023
 hud/utils/tests/test_progress.py,sha256=QSF7Kpi03Ff_l3mAeqW9qs1nhK50j9vBiSobZq7T4f4,7394
 hud/utils/tests/test_telemetry.py,sha256=5jl7bEx8C8b-FfFUko5pf4UY-mPOR-9HaeL98dGtVHM,2781
-hud/utils/tests/test_version.py,sha256=IWI4m88C4T_FrgySB42WwPuaGnc3Smpe9QddpcNoOIc,160
+hud/utils/tests/test_version.py,sha256=hizll2edGZxX6BgcLfQx2WKR2sU8TTyUdgkCY0hpMMQ,160
 hud/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-hud_python-0.4.23.dist-info/METADATA,sha256=iQt7VOpj3uXIcsbAHXdv4Vu3KNgm194i5mi7jrbuvJg,20239
-hud_python-0.4.23.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hud_python-0.4.23.dist-info/entry_points.txt,sha256=jJbodNFg1m0-CDofe5AHvB4zKBq7sSdP97-ohaQ3ae4,63
-hud_python-0.4.23.dist-info/licenses/LICENSE,sha256=yIzBheVUf86FC1bztAcr7RYWWNxyd3B-UJQ3uddg1HA,1078
-hud_python-0.4.23.dist-info/RECORD,,
+hud_python-0.4.25.dist-info/METADATA,sha256=cV33nGpP96-h1DTaoqkI19o_9ow69_tiEOd6a3yL-Wo,20239
+hud_python-0.4.25.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hud_python-0.4.25.dist-info/entry_points.txt,sha256=jJbodNFg1m0-CDofe5AHvB4zKBq7sSdP97-ohaQ3ae4,63
+hud_python-0.4.25.dist-info/licenses/LICENSE,sha256=yIzBheVUf86FC1bztAcr7RYWWNxyd3B-UJQ3uddg1HA,1078
+hud_python-0.4.25.dist-info/RECORD,,