hud-python 0.4.24__py3-none-any.whl → 0.4.26__py3-none-any.whl

hud/agents/base.py

@@ -556,26 +556,11 @@ class MCPAgent(ABC):
 
         all_tools = await self.mcp_client.list_tools()
 
-        # Filter tools
-        self._available_tools = []
-        self._tool_map = {}
 
-        # Track response tools by server
         response_tools_by_server: dict[str, str] = {}  # server_name -> tool_name
-
         for tool in all_tools:
-            # Check if tool should be included
-            if self.allowed_tools and tool.name not in self.allowed_tools:
-                continue
-            if tool.name in self.disallowed_tools:
-                continue
-
-            self._available_tools.append(tool)
-            # Simplified mapping - just tool name to tool
-            self._tool_map[tool.name] = tool
-
-            # Track response tools
             if "response" in tool.name or tool.name == "response":
+                self.console.debug(f"Found response tool: '{tool.name}'")
                 # Extract server name from tool name (e.g., "grader_response" -> "grader")
                 if "_" in tool.name:
                     server_name = tool.name.split("_", 1)[0]
@@ -583,27 +568,61 @@ class MCPAgent(ABC):
                 else:
                     response_tools_by_server["_default"] = tool.name
 
-        # Find the response tool to use (prioritize last server in config)
+        # Add response tool to lifecycle tools BEFORE filtering
         if response_tools_by_server and hasattr(self.mcp_client, "mcp_config"):
             # Get server names in order from mcp_config
             server_names = list(self.mcp_client.mcp_config.keys())
+            self.console.debug(f"Server names: {server_names}")
 
             # Try to find response tool from last server first
             response_tool_name = None
             for server_name in reversed(server_names):
                 if server_name in response_tools_by_server:
                     response_tool_name = response_tools_by_server[server_name]
+                    self.console.debug(f"Found response tool '{response_tool_name}' from server '{server_name}'")
                     break
 
             # Fallback to any response tool
             if not response_tool_name and response_tools_by_server:
                 response_tool_name = next(iter(response_tools_by_server.values()))
+                self.console.debug(f"Using fallback response tool '{response_tool_name}'")
 
             # Add to lifecycle tools if found
             if response_tool_name and response_tool_name not in self.lifecycle_tools:
                 self.console.debug(f"Auto-detected '{response_tool_name}' tool as a lifecycle tool")
                 self.response_tool_name = response_tool_name
                 self.lifecycle_tools.append(response_tool_name)
+            elif response_tool_name:
+                self.console.debug(f"Response tool '{response_tool_name}' already in lifecycle_tools")
+                self.response_tool_name = response_tool_name
+        else:
+            self.console.debug(f"No response tools found or no mcp_config")
+
+        # Filter tools
+        self._available_tools = []
+        self._tool_map = {}
+
+        self.console.debug(f"All tools: {[t.name for t in all_tools]}")
+        self.console.debug(f"Allowed tools: {self.allowed_tools}")
+        self.console.debug(f"Disallowed tools: {self.disallowed_tools}")
+        self.console.debug(f"Lifecycle tools: {self.lifecycle_tools}")
+
+        for tool in all_tools:
+            # Lifecycle tools (setup, evaluate, response) should always be included
+            is_lifecycle = tool.name in self.lifecycle_tools
+            
+            # Check if tool should be included
+            if not is_lifecycle:
+                if self.allowed_tools and tool.name not in self.allowed_tools:
+                    self.console.debug(f"Skipping tool '{tool.name}' - not in allowed_tools")
+                    continue
+                if tool.name in self.disallowed_tools:
+                    self.console.debug(f"Skipping tool '{tool.name}' - in disallowed_tools")
+                    continue
+
+            self.console.debug(f"Adding tool '{tool.name}' to available tools (lifecycle={is_lifecycle})")
+            self._available_tools.append(tool)
+            self._tool_map[tool.name] = tool
 
         # Check if all required tools are available
         if self.required_tools:

hud/agents/misc/response_agent.py

@@ -34,7 +34,8 @@ class ResponseAgent:
         
         - STOP: If the agent indicates it has successfully completed a task, even if phrased as a question
           like "I have entered the right values into this form. Would you like me to do anything else?"
-          or "Here is the website. Is there any other information you need?"
+          or "Here is the website. Is there any other information you need?" or if the agent has 
+          strongly determined it wants to stop the task.
         
         - CONTINUE: If the agent is asking for clarification before proceeding with a task
           like "I'm about to clear cookies from this website. Would you like me to proceed?"

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/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/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/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
@@ -114,58 +115,36 @@ def _process_worker(
                     task_name = task_dict.get("prompt") or f"Task {index}"
 
                     # Use the job_id to group all tasks under the same job
-                    with hud.trace(
-                        task_name, job_id=job_id, task_id=task_dict.get("id")
-                    ):
-                        try:
-                            # Convert dict to Task
-                            task = Task(**task_dict)
-
-                            # Create agent instance
-                            agent = agent_class(**(agent_config or {}))
-
-                            if auto_respond:
-                                agent.response_agent = ResponseAgent()
-
-                            # Run the task - this should ALWAYS return a result, even on error
-                            result = await agent.run(task, max_steps=max_steps)
-
-                            # Extract and print evaluation score for visibility
-                            reward = getattr(result, "reward", "N/A")
-                            logger.info(
-                                "[Worker %s] Task %s: ✓ Completed (reward: %s)",
-                                worker_id,
-                                index,
-                                reward,
-                            )
-
-                            logger.info(
-                                "[Worker %s] Completed task %s (reward: %s)",
-                                worker_id,
-                                index,
-                                reward,
-                            )
-
-                            return (index, result)
-                        except Exception as e:
-                            # Even if there's an exception, ensure we have a proper result
-                            logger.error(
-                                "[Worker %s] Task %s failed during execution: %s",
-                                worker_id,
-                                index,
-                                str(e)[:200],
-                            )
-                            # Create a proper Trace result for errors
-                            from hud.types import Trace
-
-                            error_result = Trace(
-                                reward=0.0,
-                                done=True,
-                                content=f"Task execution failed: {e}",
-                                isError=True,
-                                info={"error": str(e), "traceback": traceback.format_exc()},
-                            )
-                            return (index, error_result)
+                    with hud.trace(task_name, job_id=job_id, task_id=task_dict.get("id")):
+                        # Convert dict to Task
+                        task = Task(**task_dict)
+
+                        # Create agent instance
+                        agent = agent_class(**(agent_config or {}))
+
+                        if auto_respond:
+                            agent.response_agent = ResponseAgent()
+
+                        # Run the task
+                        result = await agent.run(task, max_steps=max_steps)
+
+                        # Extract and print evaluation score for visibility
+                        reward = getattr(result, "reward", "N/A")
+                        logger.info(
+                            "[Worker %s] Task %s: ✓ Completed (reward: %s)",
+                            worker_id,
+                            index,
+                            reward,
+                        )
+
+                        logger.info(
+                            "[Worker %s] Completed task %s (reward: %s)",
+                            worker_id,
+                            index,
+                            reward,
+                        )
+
+                        return (index, result)
 
                 except Exception as e:
                     error_msg = f"Worker {worker_id}: Task {index} failed: {e}"
@@ -212,6 +191,19 @@ def _process_worker(
     try:
         # Run the async batch processing
         results = loop.run_until_complete(process_batch())
+
+        # 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"):
+            # 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)
+
         return results
     except KeyboardInterrupt:
         logger.info("Worker %s: Interrupted by user, stopping gracefully", worker_id)
@@ -236,24 +228,17 @@ 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:
-        # CRITICAL: Always ensure telemetry is fully sent before process exits
-        # This must happen in finally block to ensure it runs even on errors
         try:
             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=10000
-                )  # 10 second timeout # type: ignore
+                # 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)
-                else:
-                    logger.debug("Worker %s: Telemetry flushed successfully", worker_id)
-        except Exception as flush_error:
-            logger.error("Worker %s: Failed to flush telemetry: %s", worker_id, flush_error)
+        except Exception as e:
+            logger.warning("Worker %s: Failed to flush telemetry: %s", worker_id, e)
 
         # Clean up the event loop
         try:
@@ -367,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 {}
@@ -391,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
@@ -435,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.24"
+    assert hud.__version__ == "0.4.26"

hud/version.py

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

{hud_python-0.4.24.dist-info → hud_python-0.4.26.dist-info}/METADATA

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

@@ -2,16 +2,16 @@ 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=CXij8SS3KpB4bx6E_ofkWx1WyFRHLlZUr_NQ4Sp8q8w,105
+hud/version.py,sha256=EbQzTU9GpvP2361rMenfZ5YkA5new-xa4K7myfCFFh4,105
 hud/agents/__init__.py,sha256=UoIkljWdbq4bM0LD-mSaw6w826EqdEjOk7r6glNYwYQ,286
-hud/agents/base.py,sha256=Ui4FC4FPCTrJ6UZBObuY2mzA99veKzbWFCsFFSVOIJc,32406
+hud/agents/base.py,sha256=rlHnzRxz1jER97pYKkpg3rDBkbt94C3lLgwuxopFQ8k,33711
 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
 hud/agents/openai_chat_generic.py,sha256=BDtMQnf9ddYN4fKy-CN9IotLyYvHGdZqjju5PODEnU4,10683
 hud/agents/misc/__init__.py,sha256=BYi4Ytp9b_vycpZFXnr5Oyw6ncKLNNGml8Jrb7bWUb4,136
-hud/agents/misc/response_agent.py,sha256=tdyRI-4A1iECW5NBA-dIJwDHhXgTi6W7FHhe--SHvMQ,3150
+hud/agents/misc/response_agent.py,sha256=Kx7RRF2KE3i_t82bDkcyp49z7cmqtRoY48Vv0UCDFXs,3228
 hud/agents/tests/__init__.py,sha256=W-O-_4i34d9TTyEHV-O_q1Ai1gLhzwDaaPo02_TWQIY,34
 hud/agents/tests/test_base.py,sha256=F39ajSqASGUbPyPoWSY9KARFav62qNTK74W11Tr1Tg4,28970
 hud/agents/tests/test_claude.py,sha256=wqEKlzEvx8obz1sSm4NY0j-Zyt1qWNfDOmRqYIuAEd0,13069
@@ -70,19 +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=zXXhu0ofrXiRA0zcpGTZ7AmXjOVfMuvOPYr7VQ5NREw,13694
+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=-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=zReqsk6orWOQ8RiCzvAQ-nVT_H-6tMWhNMxc0ZM1LMc,26630
+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
@@ -93,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
@@ -175,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=zEqEy_FWrqyHU5diD54OLOypKJ8z_hQzHIWtqs0-Swo,160
+hud/utils/tests/test_version.py,sha256=GYpR9LSDJ5o46gRfZeSnywz0JY_LafuZBJObj7oU0y8,160
 hud/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-hud_python-0.4.24.dist-info/METADATA,sha256=6Nb5a8gPvq9MIrhxfs2RII-0ql0jYZygWm7MMRCS5wE,20239
-hud_python-0.4.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hud_python-0.4.24.dist-info/entry_points.txt,sha256=jJbodNFg1m0-CDofe5AHvB4zKBq7sSdP97-ohaQ3ae4,63
-hud_python-0.4.24.dist-info/licenses/LICENSE,sha256=yIzBheVUf86FC1bztAcr7RYWWNxyd3B-UJQ3uddg1HA,1078
-hud_python-0.4.24.dist-info/RECORD,,
+hud_python-0.4.26.dist-info/METADATA,sha256=dYgepch672ATiMJJXpFsqse1nuw9Zp-T9gWVyvKLBHM,20239
+hud_python-0.4.26.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hud_python-0.4.26.dist-info/entry_points.txt,sha256=jJbodNFg1m0-CDofe5AHvB4zKBq7sSdP97-ohaQ3ae4,63
+hud_python-0.4.26.dist-info/licenses/LICENSE,sha256=yIzBheVUf86FC1bztAcr7RYWWNxyd3B-UJQ3uddg1HA,1078
+hud_python-0.4.26.dist-info/RECORD,,