hud-python 0.5.1__py3-none-any.whl → 0.5.13__py3-none-any.whl

hud/environment/tests/test_connectors.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from typing import Any
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 from hud.environment.connection import ConnectionType, Connector
 
@@ -180,39 +180,26 @@ class TestRemoteConnectorMixin:
         conn = env._connections["example"]
         assert conn._auth == "Bearer my-token"
 
-    @patch("httpx.Client")
-    def test_connect_hub_fetches_config(self, mock_httpx_cls: MagicMock) -> None:
-        """connect_hub fetches mcp_config from API."""
+    def test_connect_hub_creates_connection(self) -> None:
+        """connect_hub creates connection with correct config."""
         from hud.environment.connectors.remote import RemoteConnectorMixin
 
         class TestEnv(RemoteConnectorMixin):
             def __init__(self) -> None:
                 self._connections: dict[str, Connector] = {}
+                self._hub_config: dict[str, Any] | None = None
 
             def mount(self, server: Any, *, prefix: str | None = None) -> None:
                 pass
 
-        # Mock httpx response
-        mock_response = MagicMock()
-        mock_response.json.return_value = {
-            "mcp_config": {
-                "browser": {"url": "https://mcp.hud.ai/browser"},
-            }
-        }
-        mock_response.raise_for_status = MagicMock()
-
-        mock_client = MagicMock()
-        mock_client.get.return_value = mock_response
-        mock_client.__enter__ = MagicMock(return_value=mock_client)
-        mock_client.__exit__ = MagicMock(return_value=None)
-        mock_httpx_cls.return_value = mock_client
-
         env = TestEnv()
         with patch("hud.settings.settings") as mock_settings:
-            mock_settings.hud_api_url = "https://api.hud.so"
-            mock_settings.api_key = "test-key"
+            mock_settings.hud_mcp_url = "https://mcp.hud.ai"
+            mock_settings.client_timeout = 300  # Used in connect_mcp for sse_read_timeout
 
-            env.connect_hub("hud/browser")
+            env.connect_hub("browser")
 
-        # connect_hub creates a connection named "hud" (the server name)
+        # connect_hub creates a connection named "hud" (from mcp_config key)
         assert "hud" in env._connections
+        # Verify hub config is stored for serialization
+        assert env._hub_config == {"name": "browser"}

hud/environment/tests/test_environment.py

@@ -159,3 +159,435 @@ class TestEnvironmentSetupEvaluate:
         )
 
         assert len(env._setup_calls) == 2
+
+
+class TestEnvironmentMCPProtocol:
+    """Tests for MCP protocol overrides - Environment._env_list_tools and _env_call_tool.
+
+    These test that Environment properly exposes connector tools via MCP handlers.
+    """
+
+    @pytest.mark.asyncio
+    async def test_env_list_tools_includes_local_tools(self) -> None:
+        """_env_list_tools returns local tools after routing is built."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def my_tool(x: int) -> int:
+            """A test tool."""
+            return x * 2
+
+        # Build routing (simulates what __aenter__ does)
+        await env._build_routing()
+
+        # Call the handler that MCP will call
+        tools = await env._env_list_tools()
+
+        assert len(tools) == 1
+        assert tools[0].name == "my_tool"
+
+    @pytest.mark.asyncio
+    async def test_env_list_tools_includes_connector_tools(self) -> None:
+        """_env_list_tools returns tools from connectors (the key feature)."""
+        import mcp.types as mcp_types
+
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        # Create a mock connector with cached tools
+        mock_tools = [
+            mcp_types.Tool(
+                name="remote_tool",
+                description="A remote tool",
+                inputSchema={"type": "object"},
+            )
+        ]
+
+        class MockConnector:
+            is_connected = True
+            _tools_cache = mock_tools
+
+            @property
+            def cached_tools(self) -> list[mcp_types.Tool]:
+                return self._tools_cache
+
+            @property
+            def cached_prompts(self) -> list[mcp_types.Prompt]:
+                return []
+
+            @property
+            def cached_resources(self) -> list[mcp_types.Resource]:
+                return []
+
+            async def connect(self) -> None:
+                pass
+
+            async def disconnect(self) -> None:
+                pass
+
+            async def list_tools(self) -> list[mcp_types.Tool]:
+                return self._tools_cache
+
+        # Add the mock connector
+        env._connections["mock"] = MockConnector()  # type: ignore
+
+        # Build routing
+        await env._build_routing()
+
+        # Call the handler that MCP will call
+        tools = await env._env_list_tools()
+
+        # Should include the remote tool
+        tool_names = [t.name for t in tools]
+        assert "remote_tool" in tool_names
+
+    @pytest.mark.asyncio
+    async def test_env_call_tool_routes_to_local(self) -> None:
+        """_env_call_tool routes local tool calls correctly."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+        called_with: list[int] = []
+
+        @env.tool()
+        def my_tool(x: int) -> str:
+            """A test tool."""
+            called_with.append(x)
+            return f"result: {x}"
+
+        # Build routing
+        await env._build_routing()
+
+        # Call the handler that MCP will call
+        result = await env._env_call_tool("my_tool", {"x": 42})
+
+        assert called_with == [42]
+        assert len(result) == 1
+
+    @pytest.mark.asyncio
+    async def test_env_call_tool_routes_to_connector(self) -> None:
+        """_env_call_tool routes connector tool calls correctly."""
+        from unittest.mock import AsyncMock
+
+        import mcp.types as mcp_types
+
+        from hud.environment import Environment
+        from hud.types import MCPToolResult
+
+        env = Environment("test")
+
+        # Create a mock connector
+        mock_tools = [
+            mcp_types.Tool(
+                name="remote_tool",
+                description="A remote tool",
+                inputSchema={"type": "object"},
+            )
+        ]
+
+        class MockConnector:
+            is_connected = True
+            _tools_cache = mock_tools
+            call_tool = AsyncMock(
+                return_value=MCPToolResult(
+                    content=[mcp_types.TextContent(type="text", text="remote result")],
+                    isError=False,
+                )
+            )
+
+            @property
+            def cached_tools(self) -> list[mcp_types.Tool]:
+                return self._tools_cache
+
+            @property
+            def cached_prompts(self) -> list[mcp_types.Prompt]:
+                return []
+
+            @property
+            def cached_resources(self) -> list[mcp_types.Resource]:
+                return []
+
+            async def connect(self) -> None:
+                pass
+
+            async def disconnect(self) -> None:
+                pass
+
+            async def list_tools(self) -> list[mcp_types.Tool]:
+                return self._tools_cache
+
+        mock_conn = MockConnector()
+        env._connections["mock"] = mock_conn  # type: ignore
+
+        # Build routing
+        await env._build_routing()
+
+        # Call the handler that MCP will call
+        result = await env._env_call_tool("remote_tool", {"arg": "value"})
+
+        # Verify the connector was called
+        mock_conn.call_tool.assert_called_once_with("remote_tool", {"arg": "value"})
+        assert len(result) == 1
+
+    def test_setup_handlers_registers_custom_handlers(self) -> None:
+        """Verify _setup_handlers registers our _env_list_tools and _env_call_tool."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        # Verify the custom handlers exist
+        assert hasattr(env, "_env_list_tools")
+        assert hasattr(env, "_env_call_tool")
+        assert callable(env._env_list_tools)
+        assert callable(env._env_call_tool)
+
+
+class TestEnvironmentToolFiltering:
+    """Tests for agent-level tool filtering with wildcard support (v4 backwards compat)."""
+
+    @pytest.mark.asyncio
+    async def test_as_tools_no_filter(self) -> None:
+        """as_tools returns all tools when no filter is set."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def tool_a() -> str:
+            """Tool A."""
+            return "a"
+
+        @env.tool()
+        def tool_b() -> str:
+            """Tool B."""
+            return "b"
+
+        await env._build_routing()
+
+        tools = env.as_tools()
+        tool_names = [t.name for t in tools]
+
+        assert "tool_a" in tool_names
+        assert "tool_b" in tool_names
+
+    @pytest.mark.asyncio
+    async def test_as_tools_exact_include(self) -> None:
+        """as_tools filters with exact include list."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def tool_a() -> str:
+            """Tool A."""
+            return "a"
+
+        @env.tool()
+        def tool_b() -> str:
+            """Tool B."""
+            return "b"
+
+        env._agent_include = ["tool_a"]
+        await env._build_routing()
+
+        tools = env.as_tools()
+        tool_names = [t.name for t in tools]
+
+        assert "tool_a" in tool_names
+        assert "tool_b" not in tool_names
+
+    @pytest.mark.asyncio
+    async def test_as_tools_exact_exclude(self) -> None:
+        """as_tools filters with exact exclude list."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def tool_a() -> str:
+            """Tool A."""
+            return "a"
+
+        @env.tool()
+        def tool_b() -> str:
+            """Tool B."""
+            return "b"
+
+        env._agent_exclude = ["tool_a"]
+        await env._build_routing()
+
+        tools = env.as_tools()
+        tool_names = [t.name for t in tools]
+
+        assert "tool_a" not in tool_names
+        assert "tool_b" in tool_names
+
+    @pytest.mark.asyncio
+    async def test_as_tools_wildcard_exclude_prefix(self) -> None:
+        """as_tools filters with wildcard prefix pattern (e.g., 'setup_*')."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def setup_database() -> str:
+            """Setup tool."""
+            return "setup"
+
+        @env.tool()
+        def setup_user() -> str:
+            """Another setup tool."""
+            return "setup"
+
+        @env.tool()
+        def run_query() -> str:
+            """Regular tool."""
+            return "query"
+
+        env._agent_exclude = ["setup_*"]
+        await env._build_routing()
+
+        tools = env.as_tools()
+        tool_names = [t.name for t in tools]
+
+        assert "setup_database" not in tool_names
+        assert "setup_user" not in tool_names
+        assert "run_query" in tool_names
+
+    @pytest.mark.asyncio
+    async def test_as_tools_wildcard_exclude_contains(self) -> None:
+        """as_tools filters with wildcard contains pattern (e.g., '*setup*')."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def hud_setup() -> str:
+            """Contains setup."""
+            return "setup"
+
+        @env.tool()
+        def setup_env() -> str:
+            """Starts with setup."""
+            return "setup"
+
+        @env.tool()
+        def my_setup_tool() -> str:
+            """Contains setup in middle."""
+            return "setup"
+
+        @env.tool()
+        def run_query() -> str:
+            """No setup in name."""
+            return "query"
+
+        env._agent_exclude = ["*setup*"]
+        await env._build_routing()
+
+        tools = env.as_tools()
+        tool_names = [t.name for t in tools]
+
+        assert "hud_setup" not in tool_names
+        assert "setup_env" not in tool_names
+        assert "my_setup_tool" not in tool_names
+        assert "run_query" in tool_names
+
+    @pytest.mark.asyncio
+    async def test_as_tools_multiple_wildcard_patterns(self) -> None:
+        """as_tools filters with multiple wildcard patterns."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def setup_db() -> str:
+            """Setup tool."""
+            return "setup"
+
+        @env.tool()
+        def evaluate_result() -> str:
+            """Evaluate tool."""
+            return "evaluate"
+
+        @env.tool()
+        def checkout_branch() -> str:
+            """Checkout tool."""
+            return "checkout"
+
+        @env.tool()
+        def run_query() -> str:
+            """Regular tool."""
+            return "query"
+
+        env._agent_exclude = ["*setup*", "*evaluate*", "checkout_branch"]
+        await env._build_routing()
+
+        tools = env.as_tools()
+        tool_names = [t.name for t in tools]
+
+        assert "setup_db" not in tool_names
+        assert "evaluate_result" not in tool_names
+        assert "checkout_branch" not in tool_names
+        assert "run_query" in tool_names
+
+    @pytest.mark.asyncio
+    async def test_as_tools_wildcard_include_all(self) -> None:
+        """as_tools with ['*'] include pattern matches all tools."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def tool_a() -> str:
+            """Tool A."""
+            return "a"
+
+        @env.tool()
+        def tool_b() -> str:
+            """Tool B."""
+            return "b"
+
+        env._agent_include = ["*"]
+        await env._build_routing()
+
+        tools = env.as_tools()
+        tool_names = [t.name for t in tools]
+
+        assert "tool_a" in tool_names
+        assert "tool_b" in tool_names
+
+    @pytest.mark.asyncio
+    async def test_as_tools_include_and_exclude_combined(self) -> None:
+        """as_tools applies both include and exclude filters."""
+        from hud.environment import Environment
+
+        env = Environment("test")
+
+        @env.tool()
+        def browser_navigate() -> str:
+            """Browser tool."""
+            return "nav"
+
+        @env.tool()
+        def browser_setup() -> str:
+            """Browser setup - should be excluded."""
+            return "setup"
+
+        @env.tool()
+        def file_read() -> str:
+            """File tool - not included."""
+            return "read"
+
+        env._agent_include = ["browser_*"]
+        env._agent_exclude = ["*setup*"]
+        await env._build_routing()
+
+        tools = env.as_tools()
+        tool_names = [t.name for t in tools]
+
+        assert "browser_navigate" in tool_names
+        assert "browser_setup" not in tool_names  # Excluded by *setup*
+        assert "file_read" not in tool_names  # Not included by browser_*

hud/environment/tests/test_local_connectors.py

@@ -11,13 +11,10 @@ from hud.environment.connection import ConnectionType, Connector
 class TestConnectImage:
     """Tests for LocalConnectorMixin.connect_image."""
 
-    @patch("hud.cli.utils.docker.create_docker_run_command")
-    def test_connect_image_creates_local_connection(self, mock_docker_cmd: MagicMock) -> None:
+    def test_connect_image_creates_local_connection(self) -> None:
         """connect_image creates LOCAL connection with docker command."""
         from hud.environment.connectors.local import LocalConnectorMixin
 
-        mock_docker_cmd.return_value = ["docker", "run", "-i", "--rm", "mcp/fetch"]
-
         class TestEnv(LocalConnectorMixin):
             def __init__(self) -> None:
                 self._connections: dict[str, Connector] = {}
@@ -25,21 +22,32 @@ class TestConnectImage:
             def mount(self, server: Any, *, prefix: str | None = None) -> None:
                 pass
 
-        env = TestEnv()
-        env.connect_image("mcp/fetch")
-
-        assert "mcp/fetch" in env._connections
-        conn = env._connections["mcp/fetch"]
-        assert conn.connection_type == ConnectionType.LOCAL
-        mock_docker_cmd.assert_called_once()
-
-    @patch("hud.cli.utils.docker.create_docker_run_command")
-    def test_connect_image_with_alias(self, mock_docker_cmd: MagicMock) -> None:
+        # Mock the import that happens inside connect_image
+        mock_docker_utils = MagicMock()
+        mock_docker_utils.create_docker_run_command.return_value = [
+            "docker",
+            "run",
+            "-i",
+            "--rm",
+            "mcp/fetch",
+        ]
+
+        with patch.dict(
+            "sys.modules",
+            {"hud.cli.utils.docker": mock_docker_utils},
+        ):
+            env = TestEnv()
+            env.connect_image("mcp/fetch")
+
+            assert "mcp/fetch" in env._connections
+            conn = env._connections["mcp/fetch"]
+            assert conn.connection_type == ConnectionType.LOCAL
+            mock_docker_utils.create_docker_run_command.assert_called_once()
+
+    def test_connect_image_with_alias(self) -> None:
         """connect_image uses alias for connection name."""
         from hud.environment.connectors.local import LocalConnectorMixin
 
-        mock_docker_cmd.return_value = ["docker", "run", "-i", "--rm", "mcp/fetch"]
-
         class TestEnv(LocalConnectorMixin):
             def __init__(self) -> None:
                 self._connections: dict[str, Connector] = {}
@@ -47,19 +55,29 @@ class TestConnectImage:
             def mount(self, server: Any, *, prefix: str | None = None) -> None:
                 pass
 
-        env = TestEnv()
-        env.connect_image("mcp/fetch", alias="fetcher")
-
-        assert "fetcher" in env._connections
-        assert "mcp/fetch" not in env._connections
-
-    @patch("hud.cli.utils.docker.create_docker_run_command")
-    def test_connect_image_with_prefix(self, mock_docker_cmd: MagicMock) -> None:
+        mock_docker_utils = MagicMock()
+        mock_docker_utils.create_docker_run_command.return_value = [
+            "docker",
+            "run",
+            "-i",
+            "--rm",
+            "mcp/fetch",
+        ]
+
+        with patch.dict(
+            "sys.modules",
+            {"hud.cli.utils.docker": mock_docker_utils},
+        ):
+            env = TestEnv()
+            env.connect_image("mcp/fetch", alias="fetcher")
+
+            assert "fetcher" in env._connections
+            assert "mcp/fetch" not in env._connections
+
+    def test_connect_image_with_prefix(self) -> None:
         """connect_image passes prefix to config."""
         from hud.environment.connectors.local import LocalConnectorMixin
 
-        mock_docker_cmd.return_value = ["docker", "run", "-i", "--rm", "mcp/fetch"]
-
         class TestEnv(LocalConnectorMixin):
             def __init__(self) -> None:
                 self._connections: dict[str, Connector] = {}
@@ -67,19 +85,29 @@ class TestConnectImage:
             def mount(self, server: Any, *, prefix: str | None = None) -> None:
                 pass
 
-        env = TestEnv()
-        env.connect_image("mcp/fetch", prefix="fetch")
-
-        conn = env._connections["mcp/fetch"]
-        assert conn.config.prefix == "fetch"
-
-    @patch("hud.cli.utils.docker.create_docker_run_command")
-    def test_connect_image_returns_self(self, mock_docker_cmd: MagicMock) -> None:
+        mock_docker_utils = MagicMock()
+        mock_docker_utils.create_docker_run_command.return_value = [
+            "docker",
+            "run",
+            "-i",
+            "--rm",
+            "mcp/fetch",
+        ]
+
+        with patch.dict(
+            "sys.modules",
+            {"hud.cli.utils.docker": mock_docker_utils},
+        ):
+            env = TestEnv()
+            env.connect_image("mcp/fetch", prefix="fetch")
+
+            conn = env._connections["mcp/fetch"]
+            assert conn.config.prefix == "fetch"
+
+    def test_connect_image_returns_self(self) -> None:
         """connect_image returns self for chaining."""
         from hud.environment.connectors.local import LocalConnectorMixin
 
-        mock_docker_cmd.return_value = ["docker", "run", "-i", "--rm", "mcp/fetch"]
-
         class TestEnv(LocalConnectorMixin):
             def __init__(self) -> None:
                 self._connections: dict[str, Connector] = {}
@@ -87,10 +115,23 @@ class TestConnectImage:
             def mount(self, server: Any, *, prefix: str | None = None) -> None:
                 pass
 
-        env = TestEnv()
-        result = env.connect_image("mcp/fetch")
-
-        assert result is env
+        mock_docker_utils = MagicMock()
+        mock_docker_utils.create_docker_run_command.return_value = [
+            "docker",
+            "run",
+            "-i",
+            "--rm",
+            "mcp/fetch",
+        ]
+
+        with patch.dict(
+            "sys.modules",
+            {"hud.cli.utils.docker": mock_docker_utils},
+        ):
+            env = TestEnv()
+            result = env.connect_image("mcp/fetch")
+
+            assert result is env
 
 
 class TestConnectServer: