diff --git a/README.md b/README.md index b1bb099f5..00b78ef94 100644 --- a/README.md +++ b/README.md @@ -158,9 +158,9 @@ From the [platform UI](https://hud.ai) you can run batches, compare models on th Hosted Claude Code and Codex harnesses reach platform inference through an environment-owned, workspace-local endpoint. The endpoint is available only to -`bwrap` workspaces with network isolation; the workspace receives an opaque -per-session key, while platform credentials and trace attribution stay outside -its environment and manifest. +`bwrap` workspaces with network isolation and is bound to the exact CLI process +selected by the harness. Platform credentials stay in the environment-owned +relay rather than the CLI environment, workspace manifest, or child processes. → [Run & deploy](https://docs.hud.ai/v6/reference/runtime) diff --git a/hud/__init__.py b/hud/__init__.py index 6cde008be..0af253cd5 100644 --- a/hud/__init__.py +++ b/hud/__init__.py @@ -7,6 +7,7 @@ # Apply patches to third-party libraries early, before other imports from . import patches as _patches # noqa: F401 +from .capabilities import Connection from .clients import connect from .environment import Environment from .eval import ( @@ -16,7 +17,6 @@ Grade, HostedRuntime, HUDRuntime, - InferenceConnection, Job, LocalRuntime, Run, @@ -39,12 +39,12 @@ __all__ = [ "Chat", "ComposeProject", + "Connection", "DockerRuntime", "Environment", "Grade", "HUDRuntime", "HostedRuntime", - "InferenceConnection", "Job", "LocalRuntime", "Run", diff --git a/hud/agents/claude/sdk/agent.py b/hud/agents/claude/sdk/agent.py index abd1336f7..f77fafec2 100644 --- a/hud/agents/claude/sdk/agent.py +++ b/hud/agents/claude/sdk/agent.py @@ -33,8 +33,8 @@ from .events import ClaudeEvents if TYPE_CHECKING: - from hud.capabilities import SSHClient - from hud.eval.run import InferenceConnection, Run + from hud.capabilities import Connection, SSHClient + from hud.eval.run import Run logger = logging.getLogger(__name__) @@ -111,7 +111,7 @@ async def __call__(self, run: Run) -> None: mcp_servers=mcp_servers, prompt=run.prompt_text, executable=executable, - inference=run.inference, + connection=run.connections.get("inference"), ) async def _exec( @@ -123,7 +123,7 @@ async def _exec( mcp_servers: dict[str, dict[str, Any]], prompt: str, executable: str = "claude", - inference: InferenceConnection | None = None, + connection: Connection | None = None, ) -> None: mcp_config_path = await self._write_mcp_config(ssh, mcp_servers) input_text = ( @@ -147,7 +147,7 @@ async def _exec( shell=shell, mcp_config_path=mcp_config_path, executable=executable, - inference=inference, + connection=connection, ) if shell in WINDOWS_SHELLS: await ssh.write_text(RUN_SCRIPT_PATH, f"@echo off\r\n{command}\r\n") @@ -162,6 +162,7 @@ async def _exec( command, events.consume, input_text=None if shell in WINDOWS_SHELLS else input_text, + connections=(connection,) if connection is not None else (), ) logger.info("exit=%s stderr=%d", returncode, len(stderr)) events.finish(returncode=returncode, stderr=stderr) @@ -176,16 +177,16 @@ async def _exec( except (OSError, asyncssh.Error): logger.warning("Failed to remove Claude CLI runtime files") - def _build_env_vars(self, inference: InferenceConnection | None = None) -> dict[str, str]: + def _build_env_vars(self, connection: Connection | None = None) -> dict[str, str]: env: dict[str, str] = {} use_hud_gateway = self.config.use_hud_gateway if use_hud_gateway is None: - use_hud_gateway = inference is not None or settings.api_key is not None + use_hud_gateway = connection is not None or settings.api_key is not None if use_hud_gateway: - if inference is not None: - base_url = inference.base_url - api_key = inference.credential + if connection is not None: + base_url = connection.client_url + api_key = "hud-process-bound" elif settings.api_key: base_url = settings.hud_gateway_url api_key = settings.api_key @@ -195,7 +196,7 @@ def _build_env_vars(self, inference: InferenceConnection | None = None) -> dict[ env["ANTHROPIC_API_KEY"] = api_key env["CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS"] = "1" env["DISABLE_AUTO_COMPACT"] = "1" - if inference is None and (trace_id := get_current_trace_id()): + if connection is None and (trace_id := get_current_trace_id()): env["ANTHROPIC_CUSTOM_HEADERS"] = f"Trace-Id: {trace_id}" elif settings.anthropic_api_key: env["ANTHROPIC_API_KEY"] = settings.anthropic_api_key @@ -236,9 +237,9 @@ def _build_cli_command( shell: str, mcp_config_path: str | None = None, executable: str = "claude", - inference: InferenceConnection | None = None, + connection: Connection | None = None, ) -> str: - env_vars = self._build_env_vars(inference) + env_vars = self._build_env_vars(connection) is_win = shell in WINDOWS_SHELLS base_args: list[str] = [ executable, @@ -272,7 +273,10 @@ def _build_cli_command( cli_parts = [shlex.quote(a) for a in base_args] cli_cmd = " ".join(cli_parts) env_prefix = " ".join(f"{k}={shlex.quote(v)}" for k, v in env_vars.items()) - return f'export PATH="$HOME/.local/bin:$PATH"; {env_prefix} {cli_cmd}' + invocation = f"{env_prefix} {cli_cmd}" + if connection is not None: + invocation = f"exec env {env_prefix} {cli_cmd}" + return f'export PATH="$HOME/.local/bin:$PATH"; {invocation}' __all__ = ["ClaudeCLIAgent"] diff --git a/hud/agents/cli.py b/hud/agents/cli.py index d626c024e..92e8197f0 100644 --- a/hud/agents/cli.py +++ b/hud/agents/cli.py @@ -11,9 +11,9 @@ import asyncssh if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Sequence - from hud.capabilities import SSHClient + from hud.capabilities import Connection, SSHClient from hud.eval.runtime import RuntimeConfig WINDOWS_SHELLS = ("cmd", "powershell") @@ -135,9 +135,10 @@ async def run_jsonl( consume: Callable[[str], None], *, input_text: str | None = None, + connections: Sequence[Connection] = (), ) -> tuple[int, str]: """Stream one remote JSONL process and own its cancellation cleanup.""" - process = await ssh.create_process(command) + process = await ssh.create_process(command, connections=connections) stderr_task = asyncio.create_task(process.stderr.read()) try: if input_text is not None: diff --git a/hud/agents/codex/agent.py b/hud/agents/codex/agent.py index 345443c1f..c26447a65 100644 --- a/hud/agents/codex/agent.py +++ b/hud/agents/codex/agent.py @@ -24,8 +24,8 @@ from hud.utils.time import now_iso if TYPE_CHECKING: - from hud.capabilities import SSHClient - from hud.eval.run import InferenceConnection, Run + from hud.capabilities import Connection, SSHClient + from hud.eval.run import Run logger = logging.getLogger(__name__) @@ -212,7 +212,7 @@ def codex_command( config: CodexCLIConfig, shell: str, executable: str = "codex", - inference: InferenceConnection | None = None, + connection: Connection | None = None, ) -> str: env: dict[str, str] = {} args = [ @@ -231,12 +231,12 @@ def codex_command( use_hud_gateway = config.use_hud_gateway if use_hud_gateway is None: - use_hud_gateway = inference is not None or settings.api_key is not None + use_hud_gateway = connection is not None or settings.api_key is not None if use_hud_gateway: - if inference is not None: - base_url = inference.base_url - credential = inference.credential - credential_env = "HUD_RUNTIME_INFERENCE_TOKEN" + if connection is not None: + base_url = connection.client_url + credential = "hud-process-bound" + credential_env = "HUD_CONNECTION_CREDENTIAL" elif settings.api_key: base_url = settings.hud_gateway_url credential = settings.api_key @@ -253,7 +253,7 @@ def codex_command( } for key, value in overrides.items(): args.extend(["-c", f"{key}={json.dumps(value)}"]) - if inference is None and (trace_id := get_current_trace_id()): + if connection is None and (trace_id := get_current_trace_id()): args.extend( [ "-c", @@ -292,7 +292,14 @@ def codex_command( env_prefix = " ".join(f"{key}={shlex.quote(value)}" for key, value in env.items()) invocation = f"{env_prefix} {command}" if env_prefix else command statements = ['export PATH="$HOME/.local/bin:$PATH"', invocation] - if isolate_home: + if connection is not None: + statements = [ + 'codex_home=$(mktemp -d "${TMPDIR:-/tmp}/hud-codex.XXXXXX") || exit 1', + 'export CODEX_HOME="$codex_home"', + 'export PATH="$HOME/.local/bin:$PATH"', + f"exec env {env_prefix} {command}", + ] + elif isolate_home: statements = [ 'codex_home=$(mktemp -d "${TMPDIR:-/tmp}/hud-codex.XXXXXX") || exit 1', "trap 'rm -rf -- \"$codex_home\"' EXIT", @@ -310,12 +317,18 @@ async def run_codex( shell: str, prompt: str, executable: str = "codex", - inference: InferenceConnection | None = None, + connection: Connection | None = None, ) -> None: - command = codex_command(config, shell, executable, inference=inference) + command = codex_command(config, shell, executable, connection=connection) logger.info("SSH exec codex CLI (%d chars)", len(command)) events = CodexEvents(run, model=config.model, started_at=now_iso()) - returncode, stderr = await run_jsonl(ssh, command, events.consume, input_text=prompt) + returncode, stderr = await run_jsonl( + ssh, + command, + events.consume, + input_text=prompt, + connections=(connection,) if connection is not None else (), + ) logger.info("exit=%s stderr=%d", returncode, len(stderr)) events.finish(returncode=returncode, stderr=stderr) @@ -343,7 +356,7 @@ async def __call__(self, run: Run) -> None: shell=ssh.capability.params.get("shell", "bash"), prompt=run.prompt_text, executable=executable, - inference=run.inference, + connection=run.connections.get("inference"), ) diff --git a/hud/agents/tests/test_claude_cli_agent.py b/hud/agents/tests/test_claude_cli_agent.py index 07a793d73..7537aca2c 100644 --- a/hud/agents/tests/test_claude_cli_agent.py +++ b/hud/agents/tests/test_claude_cli_agent.py @@ -28,9 +28,8 @@ from hud.agents.tests.cli_fakes import FakeProcess as _FakeStreamProcess from hud.agents.tests.cli_fakes import fake_run as _fake_run from hud.agents.types import AgentStep, ClaudeCLIConfig, ToolStep -from hud.capabilities import Capability, SSHClient +from hud.capabilities import Capability, Connection, SSHClient from hud.capabilities.rfb import WebPScreenshotEncoding -from hud.eval import InferenceConnection from hud.settings import settings from hud.telemetry.context import set_trace_context from hud.types import MCPToolResult @@ -69,21 +68,25 @@ def test_command_follows_explicit_gateway_routing(monkeypatch: pytest.MonkeyPatc assert "ANTHROPIC_MODEL=claude-sonnet-5" in provider -def test_command_prefers_rollout_inference_connection() -> None: - inference = InferenceConnection( - base_url="https://inference.hud.so", - credential="scoped-runtime-token", +def test_command_uses_process_bound_connection_without_its_credential() -> None: + connection = Connection( + name="inference", + capability="ssh", + url="https://inference.hud.so", + headers={"Authorization": "Bearer scoped-runtime-token"}, ) gateway = ClaudeCLIAgent(ClaudeCLIConfig(use_hud_gateway=True))._build_cli_command( shell="bash", - inference=inference, + connection=connection, ) - assert "ANTHROPIC_BASE_URL=https://inference.hud.so" in gateway - assert "ANTHROPIC_API_KEY=scoped-runtime-token" in gateway + assert f"ANTHROPIC_BASE_URL={connection.client_url}" in gateway + assert "ANTHROPIC_API_KEY=hud-process-bound" in gateway + assert "scoped-runtime-token" not in gateway assert "HUD_API_KEY" not in gateway assert "Trace-Id" not in gateway + assert "exec env" in gateway for name in ( "ANTHROPIC_MODEL", "ANTHROPIC_SMALL_FAST_MODEL", @@ -500,7 +503,7 @@ async def open(self, ref: str) -> SSHClient: cast( "Any", SimpleNamespace( - client=Client(), prompt_text="call the tool", runtime_config=None, inference=None + client=Client(), prompt_text="call the tool", runtime_config=None, connections={} ), ) ) @@ -574,7 +577,7 @@ async def execute(*_args: Any, **_kwargs: Any) -> None: cast( "Any", SimpleNamespace( - client=Client(), prompt_text="use the computer", runtime_config=None, inference=None + client=Client(), prompt_text="use the computer", runtime_config=None, connections={} ), ) ) @@ -667,7 +670,7 @@ async def execute(*_args: Any, **kwargs: Any) -> None: client=Client(), prompt_text="use both screens", runtime_config=None, - inference=None, + connections={}, ), ) ) @@ -934,13 +937,13 @@ async def execute( agent = ClaudeCLIAgent() monkeypatch.setattr(agent, "_exec", execute) run_a = SimpleNamespace( - client=Client(shell_a, ssh_a), prompt_text="first", runtime_config=None, inference=None + client=Client(shell_a, ssh_a), prompt_text="first", runtime_config=None, connections={} ) run_b = SimpleNamespace( client=Client(shell_b, ssh_b), prompt_text="second", runtime_config=None, - inference=None, + connections={}, ) first = asyncio.create_task(agent(cast("Any", run_a))) diff --git a/hud/agents/tests/test_codex_cli_agent.py b/hud/agents/tests/test_codex_cli_agent.py index 3caeb021d..b896425a6 100644 --- a/hud/agents/tests/test_codex_cli_agent.py +++ b/hud/agents/tests/test_codex_cli_agent.py @@ -17,8 +17,7 @@ from hud.agents.tests.cli_fakes import FakeProcess as _FakeProcess from hud.agents.tests.cli_fakes import fake_run as _fake_run from hud.agents.types import AgentStep, CodexCLIConfig, ToolStep -from hud.capabilities import Capability, SSHClient -from hud.eval import InferenceConnection +from hud.capabilities import Capability, Connection, SSHClient from hud.eval.runtime import RuntimeConfig, RuntimeResources from hud.settings import settings from hud.telemetry.context import set_trace_context @@ -45,7 +44,9 @@ def __init__(self, process: _FakeProcess, *, shell: str = "bash") -> None: ) self.commands: list[str] = [] - async def create_process(self, command: str) -> _FakeProcess: + async def create_process( + self, command: str, *, connections: tuple[Connection, ...] = () + ) -> _FakeProcess: self.commands.append(command) return self.process @@ -105,19 +106,23 @@ def test_command_follows_explicit_gateway_routing(monkeypatch: pytest.MonkeyPatc assert command.endswith(" -") -def test_command_prefers_rollout_inference_connection() -> None: - inference = InferenceConnection( - base_url="https://inference.hud.so", - credential="scoped-runtime-token", +def test_command_uses_process_bound_connection_without_its_credential() -> None: + connection = Connection( + name="inference", + capability="ssh", + url="https://inference.hud.so", + headers={"Authorization": "Bearer scoped-runtime-token"}, ) - command = codex_command(CodexCLIConfig(use_hud_gateway=True), "bash", inference=inference) + command = codex_command(CodexCLIConfig(use_hud_gateway=True), "bash", connection=connection) - assert "HUD_RUNTIME_INFERENCE_TOKEN=scoped-runtime-token" in command - assert 'model_providers.hud.env_key="HUD_RUNTIME_INFERENCE_TOKEN"' in command + assert "scoped-runtime-token" not in command + assert "HUD_CONNECTION_CREDENTIAL=hud-process-bound" in command + assert 'model_providers.hud.env_key="HUD_CONNECTION_CREDENTIAL"' in command assert "HUD_API_KEY" not in command - assert 'model_providers.hud.base_url="https://inference.hud.so"' in command + assert f'model_providers.hud.base_url="{connection.client_url}"' in command assert "Trace-Id" not in command + assert "exec env" in command @pytest.mark.parametrize("shell", ["bash", "powershell"]) @@ -319,7 +324,7 @@ async def open(self, ref: str) -> _FakeSSH: execute = AsyncMock() monkeypatch.setattr("hud.agents.codex.agent.run_codex", execute) run = SimpleNamespace( - client=Client(), prompt_text="Fix it", runtime_config=None, inference=None + client=Client(), prompt_text="Fix it", runtime_config=None, connections={} ) await agent(cast("Any", run)) @@ -331,7 +336,7 @@ async def open(self, ref: str) -> _FakeSSH: shell="powershell", prompt="Fix it", executable="codex", - inference=None, + connection=None, ) diff --git a/hud/capabilities/__init__.py b/hud/capabilities/__init__.py index 714e061ab..075c6d0ff 100644 --- a/hud/capabilities/__init__.py +++ b/hud/capabilities/__init__.py @@ -6,6 +6,7 @@ from .base import Capability, CapabilityClient from .cdp import CDPClient +from .connection import Connection from .mcp import MCPClient from .rfb import RFBClient from .ssh import SSHClient @@ -28,6 +29,7 @@ def __getattr__(name: str) -> object: "CDPClient", "Capability", "CapabilityClient", + "Connection", "MCPClient", "RFBClient", "RobotClient", diff --git a/hud/capabilities/base.py b/hud/capabilities/base.py index 83808d160..64f8f2fbf 100644 --- a/hud/capabilities/base.py +++ b/hud/capabilities/base.py @@ -93,6 +93,7 @@ def ssh( shell: str | None = None, cwd: str | None = None, isolation: Literal["bwrap", "none"] | None = None, + process_connections: bool = False, ) -> Capability: """``ssh/2`` — SSH daemon with publickey auth. @@ -119,6 +120,8 @@ def ssh( params["cwd"] = cwd if isolation is not None: params["isolation"] = isolation + if process_connections: + params["process_connections"] = True return cls(name=name, protocol="ssh/2", url=normalized, params=params) @classmethod diff --git a/hud/capabilities/connection.py b/hud/capabilities/connection.py new file mode 100644 index 000000000..123d98906 --- /dev/null +++ b/hud/capabilities/connection.py @@ -0,0 +1,106 @@ +"""Controller-provided HTTP connections bound to trusted process executions.""" + +from __future__ import annotations + +import hashlib +import re +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import TYPE_CHECKING, cast +from urllib.parse import urlsplit, urlunsplit + +if TYPE_CHECKING: + from collections.abc import Mapping + +_NAME = re.compile(r"[a-z][a-z0-9_-]{0,62}") +_HEADER_NAME = re.compile(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]+") + + +@dataclass(frozen=True, slots=True) +class Connection: + """An authenticated HTTP endpoint available only to one bound process. + + ``url`` and ``headers`` describe the controller-selected upstream. The + process receives :attr:`client_url` instead and never receives the header + values. A runtime which cannot enforce process-bound connections must + reject the connection rather than expose it workspace-wide. + """ + + name: str + capability: str + url: str + headers: Mapping[str, str] = field(repr=False) + + def __post_init__(self) -> None: + if not _NAME.fullmatch(self.name): + raise ValueError( + "connection name must start with a lowercase letter and contain only " + "lowercase letters, digits, underscores, or hyphens" + ) + if not self.capability or self.capability.strip() != self.capability: + raise ValueError("connection capability must not be empty or padded") + parts = urlsplit(self.url) + if parts.scheme not in {"http", "https"} or parts.hostname is None: + raise ValueError("connection url must be an HTTP(S) URL with a hostname") + if parts.username is not None or parts.password is not None: + raise ValueError("connection url must not contain credentials") + if parts.query or parts.fragment: + raise ValueError("connection url must not contain a query or fragment") + if not self.headers: + raise ValueError("connection headers must not be empty") + for name, value in self.headers.items(): + if not _HEADER_NAME.fullmatch(name): + raise ValueError(f"invalid connection header name: {name!r}") + if not value or any(character in value for character in "\r\n"): + raise ValueError(f"invalid connection header value for {name!r}") + object.__setattr__(self, "headers", MappingProxyType(dict(self.headers))) + + @property + def host(self) -> str: + """Private hostname installed only inside the selected workspace.""" + digest = hashlib.sha256(f"{self.name}\0{self.url}".encode()).hexdigest()[:12] + return f"{self.name}-{digest}.hud.invalid" + + @property + def port(self) -> int: + return 80 + + @property + def client_url(self) -> str: + """Credential-free URL configured in the bound process.""" + parts = urlsplit(self.url) + return urlunsplit(("http", self.host, parts.path, "", "")) + + def to_wire(self) -> dict[str, object]: + return { + "name": self.name, + "capability": self.capability, + "url": self.url, + "headers": dict(self.headers), + } + + @classmethod + def from_wire(cls, value: object) -> Connection: + if not isinstance(value, dict): + raise ValueError("connections must be objects") + name = value.get("name") + capability = value.get("capability") + url = value.get("url") + raw_headers = value.get("headers") + if not all(isinstance(item, str) for item in (name, capability, url)): + raise ValueError("connection name, capability, and url must be strings") + if not isinstance(raw_headers, dict) or not all( + isinstance(key, str) and isinstance(header_value, str) + for key, header_value in raw_headers.items() + ): + raise ValueError("connection headers must map strings to strings") + assert isinstance(name, str) and isinstance(capability, str) and isinstance(url, str) + return cls( + name=name, + capability=capability, + url=url, + headers=dict(cast("dict[str, str]", raw_headers)), + ) + + +__all__ = ["Connection"] diff --git a/hud/capabilities/ssh.py b/hud/capabilities/ssh.py index 93b6f062f..1731e4b50 100644 --- a/hud/capabilities/ssh.py +++ b/hud/capabilities/ssh.py @@ -5,17 +5,24 @@ import asyncio import base64 import contextlib +import json import shlex -from typing import Any, ClassVar, Self +from typing import TYPE_CHECKING, Any, ClassVar, Self from urllib.parse import urlsplit import asyncssh +if TYPE_CHECKING: + from collections.abc import Sequence + + from .connection import Connection + from .base import Capability, CapabilityClient SSH_RECONNECT_ATTEMPTS = 3 SSH_RECONNECT_BASE_DELAY_S = 0.25 SSH_SESSION_CLOSE_TIMEOUT_S = 5.0 +PROCESS_CONNECTIONS_REQUEST = "HUD_PROCESS_CONNECTIONS" class SSHConnectionError(ConnectionError): @@ -75,11 +82,38 @@ def conn(self) -> asyncssh.SSHClientConnection: """Raw asyncssh connection for commands and port forwarding.""" return self._conn - async def create_process(self, command: str) -> asyncssh.SSHClientProcess[bytes]: + async def create_process( + self, + command: str, + *, + connections: Sequence[Connection] = (), + ) -> asyncssh.SSHClientProcess[bytes]: """Open a binary exec channel after restoring a dropped SSH transport.""" + if connections and self.capability.params.get("process_connections") is not True: + raise ValueError("SSH capability does not support process-bound connections") + capability_refs = { + self.capability.name, + self.capability.protocol, + self.capability.protocol.split("/", 1)[0], + } + if mismatched := [ + connection.name + for connection in connections + if connection.capability not in capability_refs + ]: + raise ValueError( + f"connections do not belong to SSH capability {self.capability.name!r}: " + f"{', '.join(mismatched)}" + ) + requested = list(dict.fromkeys(connection.name for connection in connections)) try: conn = await self._connection() - return await conn.create_process(command, encoding=None) + kwargs: dict[str, Any] = {"encoding": None} + if requested: + kwargs["env"] = { + PROCESS_CONNECTIONS_REQUEST: json.dumps(requested, separators=(",", ":")) + } + return await conn.create_process(command, **kwargs) except asyncssh.ChannelOpenError as exc: raise SSHConnectionError("SSH server rejected the session") from exc except asyncssh.ConnectionLost as exc: @@ -274,4 +308,4 @@ def _powershell_quote(value: str) -> str: return "'" + value.replace("'", "''") + "'" -__all__ = ["SSHClient"] +__all__ = ["PROCESS_CONNECTIONS_REQUEST", "SSHClient"] diff --git a/hud/capabilities/tests/test_ssh.py b/hud/capabilities/tests/test_ssh.py index 6499cf5a8..0c619dc28 100644 --- a/hud/capabilities/tests/test_ssh.py +++ b/hud/capabilities/tests/test_ssh.py @@ -8,7 +8,8 @@ import pytest from hud.capabilities.base import Capability -from hud.capabilities.ssh import SSHClient, SSHConnectionError +from hud.capabilities.connection import Connection +from hud.capabilities.ssh import PROCESS_CONNECTIONS_REQUEST, SSHClient, SSHConnectionError if TYPE_CHECKING: from collections.abc import Callable @@ -77,6 +78,7 @@ def __init__( self.open_error = open_error self.open_cancelled = False self.commands: list[str] = [] + self.process_kwargs: dict[str, Any] = {} def is_closed(self) -> bool: return self.closed @@ -97,7 +99,7 @@ async def run(self, command: str, **kwargs: Any) -> _Completed: return _Completed() async def create_process(self, *args: object, **kwargs: Any) -> _Process: - del kwargs + self.process_kwargs = kwargs self.commands.append(str(args[0])) if self.stall_open: try: @@ -196,6 +198,53 @@ async def test_create_process_reconnects_before_opening_channel( assert replacement.commands == ["bridge"] +async def test_create_process_sends_only_connection_names() -> None: + transport = _Connection() + capability = _capability() + capability.params["process_connections"] = True + client = SSHClient(capability, cast("Any", transport)) + connection = Connection( + name="inference", + capability="ssh", + url="https://inference.hud.so", + headers={"Authorization": "Bearer scoped-secret"}, + ) + + await client.create_process("agent", connections=(connection,)) + + assert transport.process_kwargs["env"] == {PROCESS_CONNECTIONS_REQUEST: '["inference"]'} + assert "scoped-secret" not in str(transport.process_kwargs) + + +async def test_create_process_rejects_connections_the_server_did_not_advertise() -> None: + client = _client(_Connection()) + connection = Connection( + name="inference", + capability="ssh", + url="https://inference.hud.so", + headers={"Authorization": "Bearer scoped-secret"}, + ) + + with pytest.raises(ValueError, match="does not support"): + await client.create_process("agent", connections=(connection,)) + + +async def test_create_process_rejects_connections_for_another_capability() -> None: + transport = _Connection() + capability = _capability() + capability.params["process_connections"] = True + client = SSHClient(capability, cast("Any", transport)) + connection = Connection( + name="inference", + capability="other-workspace", + url="https://inference.hud.so", + headers={"Authorization": "Bearer scoped-secret"}, + ) + + with pytest.raises(ValueError, match="do not belong"): + await client.create_process("agent", connections=(connection,)) + + async def test_create_process_preserves_reconnect_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/hud/clients/client.py b/hud/clients/client.py index 9ee17a92a..85b6abb83 100644 --- a/hud/clients/client.py +++ b/hud/clients/client.py @@ -23,6 +23,7 @@ Capability, CapabilityClient, CDPClient, + Connection, MCPClient, RFBClient, SSHClient, @@ -161,6 +162,7 @@ async def hello( session_id: str | None = None, *, workspace_routes: Sequence[WorkspaceRoute] = (), + connections: Sequence[Connection] = (), ) -> Manifest: """Send ``hello``; cache and return the parsed ``Manifest``. @@ -171,6 +173,8 @@ async def hello( params: dict[str, Any] = {} if workspace_routes: params["workspace_routes"] = [route.to_wire() for route in workspace_routes] + if connections: + params["connections"] = [connection.to_wire() for connection in connections] if session_id is not None: params["session_id"] = session_id result = await self._call("hello", params) @@ -384,6 +388,7 @@ async def _connect_ready( *, ready_timeout: float, workspace_routes: Sequence[WorkspaceRoute], + connections: Sequence[Connection], interval: float = 0.5, ) -> HudClient: """Connect and complete ``hello``, retrying until the env is ready. @@ -407,7 +412,7 @@ async def _connect_ready( client = HudClient(reader, writer, endpoint=(host, port)) try: - await client.hello(workspace_routes=workspace_routes) + await client.hello(workspace_routes=workspace_routes, connections=connections) except asyncio.CancelledError: client.abort() raise @@ -445,6 +450,7 @@ async def connect( *, ready_timeout: float = 240.0, workspace_routes: Sequence[WorkspaceRoute] = (), + connections: Sequence[Connection] = (), ) -> AsyncIterator[HudClient]: """Connect a :class:`HudClient` to a provisioned substrate's control channel. @@ -463,6 +469,7 @@ async def connect( parts.port or 0, ready_timeout=_runtime_ready_timeout(runtime, ready_timeout), workspace_routes=workspace_routes, + connections=connections, ) owner = asyncio.current_task() assert owner is not None @@ -477,6 +484,8 @@ async def heartbeat() -> None: params: dict[str, Any] = {"session_id": client.manifest.session_id} if workspace_routes: params["workspace_routes"] = [route.to_wire() for route in workspace_routes] + if connections: + params["connections"] = [connection.to_wire() for connection in connections] await client._call( "hello", params, diff --git a/hud/clients/tests/test_connect.py b/hud/clients/tests/test_connect.py index 555e78f1c..1daf55df0 100644 --- a/hud/clients/tests/test_connect.py +++ b/hud/clients/tests/test_connect.py @@ -17,7 +17,7 @@ import pytest import hud.clients.client as client_module -from hud.capabilities import Capability, CapabilityClient +from hud.capabilities import Capability, CapabilityClient, Connection from hud.clients import connect from hud.environment import WorkspaceRoute from hud.environment.utils import read_frame, send_frame @@ -69,6 +69,42 @@ async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> assert params == {"workspace_routes": [route.to_wire()]} +async def test_connect_sends_controller_connections_in_hello() -> None: + requests: list[dict[str, object]] = [] + + async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + hello = await read_frame(reader) + assert hello is not None + requests.append(hello) + await send_frame(writer, {"jsonrpc": "2.0", "id": hello["id"], "result": HELLO_RESULT}) + await read_frame(reader) + finally: + writer.close() + + server = await asyncio.start_server(handler, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + connection = Connection( + name="inference", + capability="ssh", + url="https://inference.hud.so", + headers={"Authorization": "Bearer secret"}, + ) + try: + async with connect( + Runtime(f"tcp://127.0.0.1:{port}"), + connections=(connection,), + ): + pass + finally: + server.close() + await server.wait_closed() + + params = requests[0]["params"] + assert isinstance(params, dict) + assert params == {"connections": [connection.to_wire()]} + + async def test_open_retries_transient_capability_connection_failures( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -333,12 +369,14 @@ async def fake_connect_ready( *, ready_timeout: float, workspace_routes: tuple[WorkspaceRoute, ...], + connections: tuple[Connection, ...], interval: float = 0.5, ) -> _FakeClient: seen["host"] = host seen["port"] = port seen["ready_timeout"] = ready_timeout assert workspace_routes == () + assert connections == () seen["interval"] = interval return _FakeClient() diff --git a/hud/environment/egress.py b/hud/environment/egress.py index 44cd93488..a37a9ae48 100644 --- a/hud/environment/egress.py +++ b/hud/environment/egress.py @@ -40,13 +40,15 @@ import threading import urllib.parse from dataclasses import dataclass -from http.server import BaseHTTPRequestHandler +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from collections.abc import Collection, Sequence + from hud.capabilities import Connection + LOGGER = logging.getLogger("hud.environment.egress") #: In an allowlist, the entry that permits everything. @@ -589,6 +591,172 @@ def get_request(self) -> tuple[socket.socket, tuple[str, int]]: return request, ("workspace", 0) +class _ConnectionProxy(BaseHTTPRequestHandler): + """Fixed-upstream reverse proxy which replaces controller-owned headers.""" + + protocol_version = "HTTP/1.1" + binding: Connection + + def log_message(self, format: str, *args: Any) -> None: + """The bound process's traffic is not the environment's log.""" + + def _fail(self, status: int, reason: str) -> None: + self.send_response(status) + self.send_header("X-Connection-Error", reason) + self.send_header("Content-Length", "0") + self.send_header("Connection", "close") + self.end_headers() + self.close_connection = True + + def _request_body(self) -> bytes | None: + transfer = self.headers.get("Transfer-Encoding") + length = self.headers.get("Content-Length") + if transfer is None: + if length is None: + return None + size = int(length) + if size < 0: + raise ValueError + body = self.rfile.read(size) + if len(body) != size: + raise ValueError + return body + if length is not None or transfer.strip().lower() != "chunked": + raise ValueError + + chunks: list[bytes] = [] + while True: + line = self.rfile.readline(65537) + if len(line) > 65536 or not line.endswith(b"\r\n"): + raise ValueError + size_text = line[:-2].split(b";", 1)[0].strip() + if not size_text or any(byte not in b"0123456789abcdefABCDEF" for byte in size_text): + raise ValueError + size = int(size_text, 16) + if size == 0: + while True: + trailer = self.rfile.readline(65537) + if len(trailer) > 65536 or not trailer.endswith(b"\r\n"): + raise ValueError + if trailer == b"\r\n": + return b"".join(chunks) + chunk = self.rfile.read(size) + if len(chunk) != size or self.rfile.read(2) != b"\r\n": + raise ValueError + chunks.append(chunk) + + def _forward(self) -> None: + request_target = urllib.parse.urlsplit(self.path) + if request_target.scheme or request_target.netloc: + self._fail(400, "absolute-request-target") + return + upstream = urllib.parse.urlsplit(self.binding.url) + base_path = upstream.path.rstrip("/") + request_path = request_target.path or "/" + if base_path and request_path != base_path and not request_path.startswith(f"{base_path}/"): + self._fail(403, "outside-connection-prefix") + return + try: + body = self._request_body() + except ValueError: + self._fail(400, "invalid-request-body") + return + + replacements = { + name.casefold(): (name, value) for name, value in self.binding.headers.items() + } + headers = { + key: value + for key, value in self.headers.items() + if key.casefold() not in _HOP_BY_HOP + and key.casefold() not in {"host", "content-length"} + and key.casefold() not in replacements + } + headers.update(dict(replacements.values())) + path = urllib.parse.urlunsplit(("", "", request_path, request_target.query, "")) + connection_type = ( + http.client.HTTPSConnection + if upstream.scheme == "https" + else http.client.HTTPConnection + ) + connection = connection_type( + upstream.hostname or "", + upstream.port, + timeout=300, + ) + response_started = False + try: + connection.request(self.command, path, body=body, headers=headers) + response = connection.getresponse() + relayed = [ + _field(key, value) + for key, value in response.getheaders() + if key.casefold() not in _HOP_BY_HOP and key.casefold() != "content-length" + ] + length = response.getheader("Content-Length") + framed = length is not None and length.strip().isdigit() + _field("Reason", response.reason or "") + response_started = True + self.send_response(response.status, response.reason) + for key, value in relayed: + self.send_header(key, value) + if framed: + assert length is not None + self.send_header("Content-Length", length.strip()) + else: + self.send_header("Connection", "close") + self.close_connection = True + self.end_headers() + while chunk := response.read1(65536): + self.wfile.write(chunk) + self.wfile.flush() + except _Unrelayable as error: + LOGGER.warning("refusing to relay connection response: %s", error) + self._fail(502, "unrelayable-upstream-header") + except (OSError, http.client.HTTPException): + if response_started: + self.close_connection = True + else: + self._fail(502, "upstream-failure") + finally: + connection.close() + + do_GET = _forward + do_HEAD = _forward + do_POST = _forward + do_PUT = _forward + do_DELETE = _forward + do_PATCH = _forward + do_OPTIONS = _forward + + +class ConnectionRelay: + """Credential-owning HTTP relay for one controller-provided connection.""" + + def __init__(self, connection: Connection) -> None: + self.connection = connection + handler = type( + f"_{connection.name.title().replace('_', '')}ConnectionProxy", + (_ConnectionProxy,), + {"binding": connection}, + ) + self._server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + self._server.daemon_threads = True + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + + @property + def port(self) -> int: + return int(self._server.server_address[1]) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> None: + self._server.shutdown() + self._server.server_close() + self._thread.join() + + class Egress: """A workspace's routes out, and the policy applied to them. @@ -716,6 +884,7 @@ def stop(self) -> None: "ANY_HOST", "BRIDGE_PORT", "VISITOR_PORT", + "ConnectionRelay", "Egress", "Peer", "WorkspaceRoute", diff --git a/hud/environment/env.py b/hud/environment/env.py index 90364a8c8..a6364b0d6 100644 --- a/hud/environment/env.py +++ b/hud/environment/env.py @@ -15,9 +15,9 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, create_model -from hud.capabilities import Capability +from hud.capabilities import Capability, Connection -from .egress import Peer, WorkspaceRoute +from .egress import ConnectionRelay, Peer, WorkspaceRoute from .workspace import Workspace if TYPE_CHECKING: @@ -165,6 +165,7 @@ def __init__( self._on_task_teardown: list[Callable[[], Awaitable[None]]] = [] self._workspaces: dict[str, Workspace] = {} self._workspace_routes: dict[WorkspaceRoute, tuple[Workspace, Peer | None]] = {} + self._connections: dict[str, tuple[Connection, Workspace, Peer, ConnectionRelay]] = {} # ─── task registration ─────────────────────────────────────────── @@ -359,9 +360,72 @@ async def stop(self) -> None: if peer is not None: workspace.remove_peer(peer) self._workspace_routes.clear() + for connection, workspace, peer, relay in reversed(self._connections.values()): + workspace.remove_process_connection(connection.name, peer) + workspace.remove_peer(peer) + relay.stop() + self._connections.clear() self._started = False self._hooks_done = False + def _workspace_for(self, capability: str) -> Workspace: + workspace = self._workspaces.get(capability) + if workspace is None and capability in {"ssh", "ssh/2"}: + if len(self._workspaces) > 1: + names = ", ".join(sorted(self._workspaces)) + raise RuntimeError(f"workspace capability {capability!r} is ambiguous: {names}") + workspace = next(iter(self._workspaces.values()), None) + if workspace is None: + raise RuntimeError(f"workspace capability {capability!r} does not exist") + return workspace + + def bind_connections(self, connections: Sequence[Connection]) -> None: + """Install controller connections before a workspace starts its sandbox.""" + if not self._started: + raise RuntimeError("environment must be started before connections are bound") + + bound: list[tuple[Connection, Workspace, Peer, ConnectionRelay]] = [] + try: + for connection in connections: + existing = self._connections.get(connection.name) + if existing is not None: + if existing[0] != connection: + raise RuntimeError(f"connection {connection.name!r} was already bound") + continue + workspace = self._workspace_for(connection.capability) + if not workspace.supports_process_connections: + raise RuntimeError( + f"workspace capability {connection.capability!r} does not support " + "process-bound connections" + ) + if any( + peer.name == connection.host and peer.port == connection.port + for peer in workspace.peers + ): + raise RuntimeError( + f"connection endpoint {connection.host}:{connection.port} conflicts with " + "an authored peer" + ) + relay = ConnectionRelay(connection) + relay.start() + peer = Peer( + connection.host, + connection.port, + target=("127.0.0.1", relay.port), + ) + workspace.add_peer(peer, first=True) + workspace.add_process_connection(connection.name, peer) + record = (connection, workspace, peer, relay) + self._connections[connection.name] = record + bound.append(record) + except BaseException: + for connection, workspace, peer, relay in reversed(bound): + workspace.remove_process_connection(connection.name, peer) + workspace.remove_peer(peer) + relay.stop() + self._connections.pop(connection.name, None) + raise + def bind_workspace_routes(self, routes: Sequence[WorkspaceRoute]) -> None: """Install controller routes before a workspace starts its sandbox.""" if not self._started: @@ -371,16 +435,7 @@ def bind_workspace_routes(self, routes: Sequence[WorkspaceRoute]) -> None: for route in dict.fromkeys(routes): if route in self._workspace_routes: continue - workspace = self._workspaces.get(route.capability) - if workspace is None and route.capability in {"ssh", "ssh/2"}: - if len(self._workspaces) > 1: - names = ", ".join(sorted(self._workspaces)) - raise RuntimeError( - f"workspace capability {route.capability!r} is ambiguous: {names}" - ) - workspace = next(iter(self._workspaces.values()), None) - if workspace is None: - raise RuntimeError(f"workspace capability {route.capability!r} does not exist") + workspace = self._workspace_for(route.capability) if not workspace.bwrap_available or not workspace.owns_netns: raise RuntimeError( f"workspace route for {route.capability!r} requires an isolated network" diff --git a/hud/environment/process_guard.py b/hud/environment/process_guard.py new file mode 100644 index 000000000..1e6e89e9b --- /dev/null +++ b/hud/environment/process_guard.py @@ -0,0 +1,397 @@ +"""Linux process-bound network connection enforcement.""" + +from __future__ import annotations + +import argparse +import array +import asyncio +import contextlib +import ctypes +import errno +import fcntl +import ipaddress +import os +import platform +import select +import socket +import struct +import subprocess +import sys +import threading +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Collection, Sequence + +_PR_SET_NO_NEW_PRIVS = 38 +_SECCOMP_SET_MODE_FILTER = 1 +_SECCOMP_FILTER_FLAG_NEW_LISTENER = 8 +_SECCOMP_RET_KILL_PROCESS = 0x80000000 +_SECCOMP_RET_USER_NOTIF = 0x7FC00000 +_SECCOMP_RET_ERRNO = 0x00050000 +_SECCOMP_RET_ALLOW = 0x7FFF0000 + +_BPF_LD_W_ABS = 0x20 +_BPF_JMP_JEQ_K = 0x15 +_BPF_RET_K = 0x06 + +_SECCOMP_IOCTL_NOTIF_RECV = 0xC0502100 +_SECCOMP_IOCTL_NOTIF_SEND = 0xC0182101 +_PIDFD_GETFD_SYSCALL = 438 +_REGISTER_ADDRESS = b"\0hud-process-connection-register" +_SANDBOX_SOCKET = "/tmp/.hud-process-connection/control.sock" # noqa: S108 +_READY_TIMEOUT_SECONDS = 10.0 + +_LIBC = ctypes.CDLL(None, use_errno=True) +_supported: bool | None = None + + +class _SockFilter(ctypes.Structure): + _fields_ = [ + ("code", ctypes.c_ushort), + ("jt", ctypes.c_ubyte), + ("jf", ctypes.c_ubyte), + ("k", ctypes.c_uint), + ] + + +class _SockFprog(ctypes.Structure): + _fields_ = [("length", ctypes.c_ushort), ("filters", ctypes.POINTER(_SockFilter))] + + +class _SeccompData(ctypes.Structure): + _fields_ = [ + ("nr", ctypes.c_int), + ("arch", ctypes.c_uint), + ("instruction_pointer", ctypes.c_ulonglong), + ("args", ctypes.c_ulonglong * 6), + ] + + +class _SeccompNotif(ctypes.Structure): + _fields_ = [ + ("id", ctypes.c_ulonglong), + ("pid", ctypes.c_uint), + ("flags", ctypes.c_uint), + ("data", _SeccompData), + ] + + +class _SeccompNotifResp(ctypes.Structure): + _fields_ = [ + ("id", ctypes.c_ulonglong), + ("val", ctypes.c_longlong), + ("error", ctypes.c_int), + ("flags", ctypes.c_uint), + ] + + +def _architecture() -> tuple[int, int, int, int, int]: + machine = platform.machine() + if machine == "x86_64": + return 0xC000003E, 317, 42, 425, 426 + if machine in {"aarch64", "arm64"}: + return 0xC00000B7, 277, 203, 425, 426 + raise RuntimeError(f"process-bound connections do not support {machine!r}") + + +def _install_connect_listener() -> int: + audit_arch, seccomp_syscall, connect_syscall, io_uring_setup, io_uring_enter = _architecture() + instructions = (_SockFilter * 11)( + _SockFilter(_BPF_LD_W_ABS, 0, 0, 4), + _SockFilter(_BPF_JMP_JEQ_K, 1, 0, audit_arch), + _SockFilter(_BPF_RET_K, 0, 0, _SECCOMP_RET_KILL_PROCESS), + _SockFilter(_BPF_LD_W_ABS, 0, 0, 0), + _SockFilter(_BPF_JMP_JEQ_K, 0, 1, connect_syscall), + _SockFilter(_BPF_RET_K, 0, 0, _SECCOMP_RET_USER_NOTIF), + _SockFilter(_BPF_JMP_JEQ_K, 0, 1, io_uring_setup), + _SockFilter(_BPF_RET_K, 0, 0, _SECCOMP_RET_ERRNO | errno.EPERM), + _SockFilter(_BPF_JMP_JEQ_K, 0, 1, io_uring_enter), + _SockFilter(_BPF_RET_K, 0, 0, _SECCOMP_RET_ERRNO | errno.EPERM), + _SockFilter(_BPF_RET_K, 0, 0, _SECCOMP_RET_ALLOW), + ) + program = _SockFprog(len(instructions), instructions) + if _LIBC.prctl(_PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "prctl(PR_SET_NO_NEW_PRIVS)") + listener = _LIBC.syscall( + seccomp_syscall, + _SECCOMP_SET_MODE_FILTER, + _SECCOMP_FILTER_FLAG_NEW_LISTENER, + ctypes.byref(program), + ) + if listener < 0: + raise OSError(ctypes.get_errno(), "seccomp(NEW_LISTENER)") + return int(listener) + + +def process_connections_supported() -> bool: + """Whether this substrate can install and broker seccomp notifications.""" + global _supported + if _supported is not None: + return _supported + if sys.platform != "linux" or not hasattr(os, "pidfd_open"): + _supported = False + return False + probe = subprocess.run( + [sys.executable, "-m", __name__, "--probe"], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=10, + ) + _supported = probe.returncode == 0 + return _supported + + +def _send_fd(channel: socket.socket, descriptor: int) -> None: + descriptors = array.array("i", [descriptor]) + channel.sendmsg([b"R"], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, descriptors)]) + + +def _receive_fd(channel: socket.socket) -> int: + _, ancillary, _, _ = channel.recvmsg(1, socket.CMSG_SPACE(array.array("i").itemsize)) + for level, kind, data in ancillary: + if level == socket.SOL_SOCKET and kind == socket.SCM_RIGHTS: + descriptors = array.array("i") + descriptors.frombytes(data[: descriptors.itemsize]) + return descriptors[0] + raise RuntimeError("guard process did not send its seccomp listener") + + +def _tgid(pid: int) -> int: + with Path(f"/proc/{pid}/status").open(encoding="ascii") as status: + for line in status: + if line.startswith("Tgid:"): + return int(line.split()[1]) + raise RuntimeError(f"process {pid} has no Tgid") + + +def _read_process(pid: int, address: int, length: int) -> bytes: + descriptor = os.open(f"/proc/{pid}/mem", os.O_RDONLY) + try: + data = os.pread(descriptor, length, address) + finally: + os.close(descriptor) + if len(data) != length: + raise OSError(errno.EFAULT, "short process memory read") + return data + + +def _destination(raw: bytes) -> tuple[str, int] | None: + if len(raw) < 2: + return None + family = struct.unpack_from("H", raw)[0] + if family == socket.AF_INET and len(raw) >= 8: + return socket.inet_ntop(socket.AF_INET, raw[4:8]), struct.unpack_from("!H", raw, 2)[0] + if family == socket.AF_INET6 and len(raw) >= 24: + address = ipaddress.IPv6Address(raw[8:24]) + host = str(address.ipv4_mapped or address) + return host, struct.unpack_from("!H", raw, 2)[0] + return None + + +def _emulate_connect(pid: int, descriptor: int, address: bytes) -> int: + pidfd_open = getattr(os, "pidfd_open", None) + if pidfd_open is None: + return -errno.ENOSYS + pidfd = int(pidfd_open(pid)) + try: + duplicate = _LIBC.syscall(_PIDFD_GETFD_SYSCALL, pidfd, descriptor, 0) + if duplicate < 0: + return -ctypes.get_errno() + finally: + os.close(pidfd) + try: + buffer = ctypes.create_string_buffer(address) + if _LIBC.connect(duplicate, ctypes.byref(buffer), len(address)) == 0: + return 0 + return -ctypes.get_errno() + finally: + os.close(duplicate) + + +class ProcessConnectionGuard: + """Broker connects for one trusted process while constraining descendants.""" + + def __init__( + self, + directory: Path, + protected: Collection[tuple[str, int]], + allowed: Collection[tuple[str, int]], + ) -> None: + self.directory = directory + self.socket_path = directory / "control.sock" + self.protected = frozenset(protected) + self.allowed = frozenset(allowed) + if not self.allowed <= self.protected: + raise ValueError("allowed process connections must be protected destinations") + self._server: socket.socket | None = None + self._listener: int | None = None + self._stop_read, self._stop_write = os.pipe() + self._ready = threading.Event() + self._error: BaseException | None = None + self._thread: threading.Thread | None = None + + @property + def sandbox_socket(self) -> str: + return _SANDBOX_SOCKET + + def start(self) -> None: + self.directory.mkdir(mode=0o700, parents=True, exist_ok=True) + self._server = socket.socket(socket.AF_UNIX) + self._server.bind(str(self.socket_path)) + self._server.listen(1) + self.socket_path.chmod(0o600) + self._thread = threading.Thread(target=self._serve, daemon=True) + self._thread.start() + + async def wait_ready(self) -> None: + ready = await asyncio.to_thread(self._ready.wait, _READY_TIMEOUT_SECONDS) + if not ready: + raise TimeoutError("process connection guard did not become ready") + if self._error is not None: + raise RuntimeError("process connection guard failed") from self._error + + def close(self) -> None: + with contextlib.suppress(OSError): + os.write(self._stop_write, b"x") + if self._thread is not None: + self._thread.join(timeout=2) + if self._thread.is_alive(): + if self._server is not None: + self._server.close() + if self._listener is not None: + os.close(self._listener) + self._listener = None + self._thread.join(timeout=2) + self._thread = None + if self._server is not None: + self._server.close() + self._server = None + if self._listener is not None: + os.close(self._listener) + self._listener = None + for descriptor in (self._stop_read, self._stop_write): + with contextlib.suppress(OSError): + os.close(descriptor) + with contextlib.suppress(FileNotFoundError): + self.socket_path.unlink() + with contextlib.suppress(OSError): + self.directory.rmdir() + + def _serve(self) -> None: + try: + assert self._server is not None + poller = select.poll() + poller.register(self._server, select.POLLIN) + poller.register(self._stop_read, select.POLLIN) + ready = {descriptor for descriptor, _ in poller.poll()} + if self._stop_read in ready: + return + channel, _ = self._server.accept() + self._server.close() + self._server = None + with channel: + self._listener = _receive_fd(channel) + with contextlib.suppress(FileNotFoundError): + self.socket_path.unlink() + self._broker() + except BaseException as exc: + self._error = exc + self._ready.set() + + def _broker(self) -> None: + assert self._listener is not None + poller = select.poll() + poller.register(self._listener, select.POLLIN) + poller.register(self._stop_read, select.POLLIN) + trusted_tgid: int | None = None + while True: + ready = {descriptor for descriptor, _ in poller.poll()} + if self._stop_read in ready: + return + notification = _SeccompNotif() + try: + fcntl.ioctl(self._listener, _SECCOMP_IOCTL_NOTIF_RECV, notification) + except OSError as exc: + if exc.errno in {errno.EINTR, errno.ENOENT}: + continue + raise + response = _SeccompNotifResp(id=notification.id) + try: + process_tgid = _tgid(notification.pid) + if trusted_tgid is None: + trusted_tgid = process_tgid + response.error = -errno.ECONNREFUSED + self._ready.set() + else: + address = _read_process( + notification.pid, + notification.data.args[1], + notification.data.args[2], + ) + target = _destination(address) + if target in self.protected and ( + process_tgid != trusted_tgid or target not in self.allowed + ): + response.error = -errno.EPERM + else: + result = _emulate_connect( + notification.pid, + notification.data.args[0], + address, + ) + response.error = result if result < 0 else 0 + response.val = result if result >= 0 else 0 + except (OSError, RuntimeError): + response.error = -errno.EPERM + try: + fcntl.ioctl(self._listener, _SECCOMP_IOCTL_NOTIF_SEND, response) + except OSError as exc: + if exc.errno != errno.ENOENT: + raise + + +def guarded_exec(socket_path: str, argv: Sequence[str]) -> None: + """Install the guard, register this process, and replace it with ``argv``.""" + if not argv: + raise ValueError("guarded execution requires a command") + channel = socket.socket(socket.AF_UNIX) + channel.connect(socket_path) + listener = _install_connect_listener() + try: + _send_fd(channel, listener) + finally: + os.close(listener) + channel.close() + registration = socket.socket(socket.AF_UNIX) + try: + registration.connect(_REGISTER_ADDRESS) + except ConnectionRefusedError: + pass + finally: + registration.close() + os.execvp(argv[0], list(argv)) # noqa: S606 - exact controller-built argv + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--probe", action="store_true") + parser.add_argument("socket", nargs="?") + parser.add_argument("argv", nargs=argparse.REMAINDER) + args = parser.parse_args() + if args.probe: + listener = _install_connect_listener() + os.close(listener) + return + if args.socket is None: + parser.error("socket is required") + argv = args.argv[1:] if args.argv[:1] == ["--"] else args.argv + guarded_exec(args.socket, argv) + + +if __name__ == "__main__": + main() + +__all__ = ["ProcessConnectionGuard", "guarded_exec", "process_connections_supported"] diff --git a/hud/environment/server.py b/hud/environment/server.py index 28d8b15aa..d8b00e731 100644 --- a/hud/environment/server.py +++ b/hud/environment/server.py @@ -26,6 +26,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError +from hud.capabilities import Connection from hud.graders.results import EvaluationResult from .egress import WorkspaceRoute @@ -351,6 +352,18 @@ async def error_to(msg_id: int | None, code: int, message: str) -> None: except ValueError as exc: await error_to(msg_id, -32602, f"hello: {exc}") continue + raw_connections = params.get("connections", []) + if not isinstance(raw_connections, list): + await error_to(msg_id, -32602, "hello: 'connections' must be a list") + continue + try: + connections = [ + Connection.from_wire(connection) for connection in raw_connections + ] + except ValueError as exc: + await error_to(msg_id, -32602, f"hello: {exc}") + continue + env.bind_connections(connections) env.bind_workspace_routes(workspace_routes) # env.start() ran before serving, so hook-published # capabilities (e.g. a workspace's ssh address) are diff --git a/hud/environment/tests/test_process_connection.py b/hud/environment/tests/test_process_connection.py new file mode 100644 index 000000000..7523fba1e --- /dev/null +++ b/hud/environment/tests/test_process_connection.py @@ -0,0 +1,201 @@ +"""Process-bound controller connection integration.""" + +from __future__ import annotations + +import shlex +import sys +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import TYPE_CHECKING, Any, cast + +import pytest + +from hud.capabilities import Connection, SSHClient +from hud.clients import connect +from hud.environment import Environment, Peer +from hud.environment.egress import ANY_HOST, BRIDGE_PORT +from hud.environment.process_guard import process_connections_supported +from hud.eval import LocalRuntime, Task + +if TYPE_CHECKING: + from pathlib import Path + +pytestmark = pytest.mark.skipif( + not process_connections_supported(), + reason="seccomp user notification is unavailable", +) + + +class _ProtectedUpstream(BaseHTTPRequestHandler): + authorization: str | None = None + + def do_GET(self) -> None: + type(self).authorization = self.headers.get("Authorization") + body = b"ok" + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: Any) -> None: + pass + + +class _OrdinaryUpstream(_ProtectedUpstream): + authorization: str | None = None + + +def _server(handler: type[BaseHTTPRequestHandler]) -> tuple[HTTPServer, threading.Thread]: + server = HTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return server, thread + + +def _fetch_source(url: str) -> str: + return ( + "import urllib.request;" + f"request=urllib.request.Request({url!r},headers={{'Authorization':'Bearer visible'}});" + "print(urllib.request.urlopen(request,timeout=5).read().decode())" + ) + + +def _fetch_script(url: str) -> str: + return f"exec {shlex.quote(sys.executable)} -c {shlex.quote(_fetch_source(url))}" + + +def _proxy_fetch_source(url: str) -> str: + return ( + "import http.client,sys;" + f"connection=http.client.HTTPConnection('127.0.0.1',{BRIDGE_PORT},timeout=5);" + f"connection.request('GET',{url!r});" + "response=connection.getresponse();body=response.read();" + "sys.exit(0 if response.status==200 and body==b'ok' else 9)" + ) + + +def _direct_fetch_source(host: str, port: int) -> str: + return ( + "import http.client;" + f"connection=http.client.HTTPConnection({host!r},{port},timeout=5);" + "connection.request('GET','/');" + "print(connection.getresponse().read().decode())" + ) + + +async def test_only_bound_process_reaches_controller_connection(tmp_path: Path) -> None: + protected, protected_thread = _server(_ProtectedUpstream) + ordinary, ordinary_thread = _server(_OrdinaryUpstream) + connection = Connection( + name="inference", + capability="ssh", + url=f"http://127.0.0.1:{protected.server_address[1]}", + headers={"Authorization": "Bearer scoped-runtime-token"}, + ) + env = Environment("process-connection") + env.workspace( + tmp_path / "root", + peers=( + Peer( + "ordinary.hud.invalid", + 80, + target=("127.0.0.1", ordinary.server_address[1]), + ), + ), + allowed_hosts={ANY_HOST}, + require_isolation=True, + track_files=False, + ) + task = Task(env=env.name, id="test") + try: + async with ( + LocalRuntime(env)(task) as runtime, + connect(runtime, connections=(connection,)) as client, + ): + ssh = cast("SSHClient", await client.open("ssh")) + assert ssh.capability.params["process_connections"] is True + + unbound = await ssh.run(_fetch_script(connection.client_url), check=False) + assert unbound.returncode != 0 + + bound = await ssh.create_process( + _fetch_script(connection.client_url), + connections=(connection,), + ) + completed = await bound.wait() + assert completed.returncode == 0 + assert completed.stdout == b"ok\n" + + child_source = ( + "import subprocess,sys;" + f"result=subprocess.run([sys.executable,'-c',{_fetch_source(connection.client_url)!r}]);" + "print(result.returncode)" + ) + child = await ssh.create_process( + f"exec {shlex.quote(sys.executable)} -c {shlex.quote(child_source)}", + connections=(connection,), + ) + child_result = await child.wait() + assert child_result.returncode == 0 + assert isinstance(child_result.stdout, bytes) + assert child_result.stdout.strip() != b"0" + + mapped_fetch = _direct_fetch_source("::ffff:127.0.0.1", connection.port) + mapped_child_source = ( + "import subprocess,sys;" + f"result=subprocess.run([sys.executable,'-c',{mapped_fetch!r}]);" + "print(result.returncode)" + ) + mapped_child = await ssh.create_process( + f"exec {shlex.quote(sys.executable)} -c {shlex.quote(mapped_child_source)}", + connections=(connection,), + ) + mapped_child_result = await mapped_child.wait() + assert mapped_child_result.returncode == 0 + assert isinstance(mapped_child_result.stdout, bytes) + assert mapped_child_result.stdout.strip() != b"0" + + proxy_child_source = ( + "import subprocess,sys;" + f"result=subprocess.run([sys.executable,'-c',{_proxy_fetch_source(connection.client_url)!r}]);" + "print(result.returncode)" + ) + proxy_child = await ssh.create_process( + f"exec {shlex.quote(sys.executable)} -c {shlex.quote(proxy_child_source)}", + connections=(connection,), + ) + proxy_child_result = await proxy_child.wait() + assert proxy_child_result.returncode == 0 + assert isinstance(proxy_child_result.stdout, bytes) + assert proxy_child_result.stdout.strip() != b"0" + + io_uring_source = ( + "import ctypes,errno,sys;" + "libc=ctypes.CDLL(None,use_errno=True);" + "result=libc.syscall(425,0,None);" + "sys.exit(0 if result == -1 and ctypes.get_errno() == errno.EPERM else 1)" + ) + io_uring = await ssh.create_process( + f"exec {shlex.quote(sys.executable)} -c {shlex.quote(io_uring_source)}", + connections=(connection,), + ) + io_uring_result = await io_uring.wait() + assert io_uring_result.returncode == 0 + + ordinary_result = await ssh.run( + _fetch_script("http://ordinary.hud.invalid"), + check=False, + ) + assert ordinary_result.returncode == 0 + assert ordinary_result.stdout == "ok\n" + finally: + for server, thread in ( + (protected, protected_thread), + (ordinary, ordinary_thread), + ): + server.shutdown() + server.server_close() + thread.join() + + assert _ProtectedUpstream.authorization == "Bearer scoped-runtime-token" + assert _OrdinaryUpstream.authorization == "Bearer visible" diff --git a/hud/environment/tests/test_workspace.py b/hud/environment/tests/test_workspace.py index 451873137..3c929808e 100644 --- a/hud/environment/tests/test_workspace.py +++ b/hud/environment/tests/test_workspace.py @@ -27,10 +27,17 @@ import asyncssh import pytest -from hud.capabilities import SSHClient +from hud.capabilities import Connection, SSHClient from hud.environment import namespace as namespace_mod from hud.environment import workspace as workspace_mod -from hud.environment.egress import Peer, WorkspaceRoute, _field, _UnixServer, _Unrelayable +from hud.environment.egress import ( + ConnectionRelay, + Peer, + WorkspaceRoute, + _field, + _UnixServer, + _Unrelayable, +) from hud.environment.workspace import Bubblewrap, Mount, Workspace from hud.utils.process import ProcessGroup, ProcessResult @@ -1600,6 +1607,65 @@ def log_message(self, format: str, *args: Any) -> None: assert Chunked.request_transfer is None +def test_connection_relay_replaces_credentials_and_preserves_streaming() -> None: + import http.client + from http.server import BaseHTTPRequestHandler, HTTPServer + + class Upstream(BaseHTTPRequestHandler): + authorization: str | None = None + api_key: str | None = None + body = b"" + + def do_POST(self) -> None: + type(self).authorization = self.headers.get("Authorization") + type(self).api_key = self.headers.get("X-Api-Key") + type(self).body = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + self.send_response(200) + self.send_header("Transfer-Encoding", "chunked") + self.end_headers() + self.wfile.write(b"5\r\nfirst\r\n6\r\nsecond\r\n0\r\n\r\n") + + def log_message(self, format: str, *args: Any) -> None: + pass + + upstream = HTTPServer(("127.0.0.1", 0), Upstream) + upstream_thread = threading.Thread(target=upstream.serve_forever, daemon=True) + upstream_thread.start() + token = "scoped-runtime-token" + connection = Connection( + name="inference", + capability="ssh", + url=f"http://127.0.0.1:{upstream.server_address[1]}/v1", + headers={"Authorization": f"Bearer {token}", "X-Api-Key": token}, + ) + relay = ConnectionRelay(connection) + relay.start() + try: + client = http.client.HTTPConnection("127.0.0.1", relay.port, timeout=5) + client.request( + "POST", + "/v1/messages", + body=b"request", + headers={ + "Authorization": "Bearer model-visible", + "X-Api-Key": "model-visible", + }, + ) + response = client.getresponse() + assert response.status == 200 + assert response.read() == b"firstsecond" + client.close() + finally: + relay.stop() + upstream.shutdown() + upstream.server_close() + upstream_thread.join() + + assert Upstream.authorization == f"Bearer {token}" + assert Upstream.api_key == token + assert Upstream.body == b"request" + + @pytest.mark.parametrize( "payload", [ @@ -1835,7 +1901,7 @@ async def capture_spawn(*_args: str, **kwargs: Any) -> None: ws = Workspace(tmp_path / "root") monkeypatch.setattr(ws, "sandbox_pid", AsyncMock(return_value=7)) ws._namespace = cast("Any", SimpleNamespace(spawn=capture_spawn)) - process = SimpleNamespace(term_type=None, command="true") + process = SimpleNamespace(term_type=None, command="true", env={}) with pytest.raises(SpawnCaptured) as captured: await ws._handle_process(cast("Any", process)) diff --git a/hud/environment/workspace.py b/hud/environment/workspace.py index c0c93d028..2a174166f 100644 --- a/hud/environment/workspace.py +++ b/hud/environment/workspace.py @@ -21,8 +21,17 @@ import asyncssh -from hud.environment.egress import VISITOR_PORT, Egress, Peer, hosts_text, proxy_environment +from hud.capabilities.ssh import PROCESS_CONNECTIONS_REQUEST +from hud.environment.egress import ( + VISITOR_PORT, + Egress, + Peer, + bind_addresses, + hosts_text, + proxy_environment, +) from hud.environment.namespace import NamespaceHost, NamespaceProcess, install_identity_map +from hud.environment.process_guard import ProcessConnectionGuard, process_connections_supported from hud.utils.process import ProcessGroup, ProcessResult, create_process_group_exec if sys.platform != "win32": # the pty a session runs on has no Windows analogue @@ -516,6 +525,7 @@ def __init__( #: named here to exist for it. Nothing to do where sessions share the #: substrate's network: the services are already at those addresses. self.peers: tuple[Peer, ...] = tuple(peers) + self._process_connections: dict[str, Peer] = {} self.local_aliases = frozenset(local_aliases) self.ports = frozenset(ports) self._egress: Egress | None = None @@ -649,6 +659,22 @@ def remove_peer(self, peer: Peer) -> None: if self._hosts_path is not None: self._hosts_path = self._write_hosts() + @property + def supports_process_connections(self) -> bool: + return self.bwrap_available and self.owns_netns and process_connections_supported() + + def add_process_connection(self, name: str, peer: Peer) -> None: + if name in self._process_connections: + raise RuntimeError(f"process connection {name!r} was already bound") + if peer not in self.peers: + raise RuntimeError("process connection peer must be installed first") + self._process_connections[name] = peer + + def remove_process_connection(self, name: str, peer: Peer) -> None: + if self._process_connections.get(name) != peer: + raise RuntimeError(f"process connection {name!r} is not bound to that peer") + del self._process_connections[name] + def _setpriv(self) -> str | None: """Absolute path to ``setpriv``, resolved via the *server's* PATH. @@ -846,6 +872,7 @@ def capability(self, name: str = "shell") -> Capability: client_key_path=key_path, cwd=self._guest_path, isolation="bwrap" if self.bwrap_available else "none", + process_connections=self.supports_process_connections, ) @property @@ -1613,7 +1640,35 @@ def _session_env(self) -> dict[str, str] | None: return {**base, **self.env} return {**os.environ, **self.env} if self.env else None + def _requested_process_connections( + self, + process: asyncssh.SSHServerProcess[bytes], + ) -> tuple[str, ...]: + raw = getattr(process, "env", {}).get(PROCESS_CONNECTIONS_REQUEST) + if raw is None: + return () + try: + requested = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError("invalid process connection request") from exc + if not isinstance(requested, list) or not all(isinstance(name, str) for name in requested): + raise ValueError("process connection request must be a list of names") + names = tuple(dict.fromkeys(str(name) for name in requested)) + missing = [name for name in names if name not in self._process_connections] + if missing: + raise ValueError(f"process connection is not bound: {', '.join(missing)}") + return names + + def _process_connection_targets(self, names: Sequence[str]) -> frozenset[tuple[str, int]]: + addresses = bind_addresses(self.peers, reserved_ports=self.ports) + return frozenset( + (addresses[peer.name], peer.port) + for name in names + for peer in (self._process_connections[name],) + ) + async def _handle_process(self, process: asyncssh.SSHServerProcess[bytes]) -> None: + guard: ProcessConnectionGuard | None = None try: pid = await self.sandbox_pid() # Sessions start from an exact environment, so a terminal's TERM has to @@ -1622,11 +1677,53 @@ async def _handle_process(self, process: asyncssh.SSHServerProcess[bytes]) -> No term_type = process.term_type wants_tty = bool(term_type) session_env = {"TERM": term_type} if term_type else None - argv = ( - self.shell_argv(process.command, env=session_env, tty=wants_tty) - if pid is None - else self.session_argv(process.command, env=session_env, tty=wants_tty) - ) + requested_connections = self._requested_process_connections(process) + if self._process_connections: + if pid is None or not self.supports_process_connections: + raise RuntimeError("process-bound connections require an isolated workspace") + guard_directory = Path( + tempfile.mkdtemp(prefix="process-connection-", dir=self._credentials_dir()) + ) + guard = ProcessConnectionGuard( + guard_directory, + self._process_connection_targets(tuple(self._process_connections)), + self._process_connection_targets(requested_connections), + ) + guard.start() + sandbox_directory = str(Path(guard.sandbox_socket).parent) + shell_command = ( + ["bash", "-lc", process.command] + if process.command is not None + else ["bash", "-l"] + ) + guarded_command = [ + sys.executable, + "-m", + "hud.environment.process_guard", + guard.sandbox_socket, + "--", + *self._drop_argv(), + *shell_command, + ] + argv = self.bwrap_argv( + guarded_command, + env=session_env, + network=True, + isolate_processes=False, + isolate_users=False, + bind_devices=True, + mounts=( + *self.mounts, + Mount("ro", src=str(guard.directory), dst=sandbox_directory), + ), + tty=wants_tty, + ) + else: + argv = ( + self.shell_argv(process.command, env=session_env, tty=wants_tty) + if pid is None + else self.session_argv(process.command, env=session_env, tty=wants_tty) + ) if sys.platform != "win32": # Namespace/process wrappers must not receive caller-controlled # loader variables or server secrets. The inner payload injects @@ -1637,6 +1734,8 @@ async def _handle_process(self, process: asyncssh.SSHServerProcess[bytes]) -> No else: proc_env = self._session_env() except Exception as exc: + if guard is not None: + guard.close() LOGGER.warning("workspace session setup failed: %s", exc) if not process.channel.is_closing(): process.stderr.write(f"workspace: cannot prepare shell: {exc}\n".encode()) @@ -1760,6 +1859,7 @@ def spawn() -> tuple[_subprocess.Popen[bytes], _WindowsJob]: return pty_pair = _open_pty(process) if wants_tty and pid is None else None + sub: ProcessGroup | NamespaceProcess | None = None try: if pid is None: child_fds: dict[str, Any] = ( @@ -1771,7 +1871,7 @@ def spawn() -> tuple[_subprocess.Popen[bytes], _WindowsJob]: if pty_pair is None else {"stdin": pty_pair[1], "stdout": pty_pair[1], "stderr": pty_pair[1]} ) - sub: ProcessGroup | NamespaceProcess = await create_process_group_exec( + sub = await create_process_group_exec( *argv, **child_fds, cwd=str(self.root), env=proc_env ) else: @@ -1780,18 +1880,35 @@ def spawn() -> tuple[_subprocess.Popen[bytes], _WindowsJob]: argv, cwd=self.root, env=proc_env, + mount_view="host" if guard is not None else "workspace", tty=wants_tty, terminal_size=process.get_terminal_size() if wants_tty else (80, 24, 0, 0), - persistent=True, + persistent=guard is None, ) + if guard is not None: + await guard.wait_ready() except FileNotFoundError as exc: + if guard is not None: + guard.close() if pty_pair is not None: os.close(pty_pair[0]) os.close(pty_pair[1]) process.stderr.write(f"workspace: cannot spawn shell: {exc}\n".encode()) process.exit(127) return + except (OSError, RuntimeError) as exc: + if sub is not None: + await sub.terminate() + if guard is not None: + guard.close() + if pty_pair is not None: + os.close(pty_pair[0]) + os.close(pty_pair[1]) + process.stderr.write(f"workspace: cannot spawn shell: {exc}\n".encode()) + process.exit(1) + return + assert sub is not None if pty_pair is not None: # The child holds the terminal now; this side keeps only the master. os.close(pty_pair[1]) @@ -1909,6 +2026,8 @@ async def relay_output( for task in output_pending: task.cancel() await asyncio.gather(*output_tasks, return_exceptions=True) + if guard is not None: + guard.close() if process.channel.is_closing(): return diff --git a/hud/eval/__init__.py b/hud/eval/__init__.py index 7411f9105..0bce06d61 100644 --- a/hud/eval/__init__.py +++ b/hud/eval/__init__.py @@ -32,7 +32,7 @@ from .chat import Chat from .job import Job -from .run import Grade, InferenceConnection, Run, rollout +from .run import Grade, Run, rollout from .runtime import ( ComposeProject, DaytonaRuntime, @@ -63,7 +63,6 @@ "Grade", "HUDRuntime", "HostedRuntime", - "InferenceConnection", "Job", "LocalRuntime", "ModalRuntime", diff --git a/hud/eval/run.py b/hud/eval/run.py index 1be0b8b1e..c60abe045 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -27,7 +27,6 @@ import uuid from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, Self, cast -from urllib.parse import urlsplit import mcp.types as mcp_types @@ -45,6 +44,7 @@ from types import TracebackType from hud.agents.base import Agent + from hud.capabilities import Connection from hud.clients.client import HudClient from hud.environment import WorkspaceRoute @@ -55,25 +55,6 @@ logger = logging.getLogger("hud.eval.run") -@dataclass(frozen=True, slots=True) -class InferenceConnection: - """Execution-scoped inference connection exposed to a live agent.""" - - base_url: str - credential: str = field(repr=False) - - def __post_init__(self) -> None: - parts = urlsplit(self.base_url) - if parts.scheme not in {"http", "https"} or parts.hostname is None: - raise ValueError("inference base_url must be an HTTP(S) URL with a hostname") - if parts.username is not None or parts.password is not None: - raise ValueError("inference base_url must not contain credentials") - if parts.query or parts.fragment: - raise ValueError("inference base_url must not contain a query or fragment") - if not self.credential: - raise ValueError("inference credential must not be empty") - - def validate_rollout_timeouts( task: Task, agent: Agent, @@ -222,14 +203,14 @@ def __init__( *, best_effort_grade: bool = False, runtime_config: RuntimeConfig | None = None, - inference: InferenceConnection | None = None, + connections: Sequence[Connection] = (), ) -> None: self._client = client self._task_id = task_id self._args = args self._best_effort_grade = best_effort_grade self.runtime_config = runtime_config - self.inference = inference + self.connections = {connection.name: connection for connection in connections} #: The task's opening prompt as ``tasks.start`` returned it: plain #: text, or a list of message dicts (``{"role", "content"}``) for #: chat-style / multi-turn prompts. Agents consume the normalized @@ -469,7 +450,7 @@ async def rollout( group_id: str | None = None, trace_id: str | None = None, rollout_timeout: float | None = None, - inference: InferenceConnection | None = None, + connections: Sequence[Connection] = (), workspace_routes: Sequence[WorkspaceRoute] = (), ) -> Run: """Drive one task to a graded :class:`Run` here, against ``runtime``'s channel. @@ -561,7 +542,11 @@ async def close_actor() -> None: scope.push_async_callback(close_actor) addr = await actor.enter_async_context(runtime(task)) _phase = "starting task" - async with connect(addr, workspace_routes=workspace_routes) as actor_client: + async with connect( + addr, + workspace_routes=workspace_routes, + connections=connections, + ) as actor_client: client = actor_client live = Run( actor_client, @@ -569,7 +554,7 @@ async def close_actor() -> None: task.args, best_effort_grade=task.verifier is not None, runtime_config=addr.config or actor_runtime_config, - inference=inference, + connections=connections, ) live._runtime = addr.url # the placement record for the receipt async with live: # start on enter; complete on exit @@ -601,7 +586,7 @@ async def close_actor() -> None: run.trace.status = "error" run.record(Step(source="system", error=f"[{_phase}] {detail}")) finally: - run.inference = None + run.connections.clear() _phase = "grading" if verifier is not None: @@ -701,7 +686,7 @@ async def close_actor() -> None: run.trace.status = "error" run.record(Step(source="system", error=f"[{_phase}] {detail}")) assert run is not None # the body bound it, or the handler synthesized it - run.inference = None + run.connections.clear() run.trace.trace_id = trace_id run.job_id = job_id run.group_id = group_id @@ -715,4 +700,4 @@ def _consume_task_result(task: asyncio.Future[Any]) -> None: task.result() -__all__ = ["Grade", "InferenceConnection", "Run", "rollout"] +__all__ = ["Grade", "Run", "rollout"] diff --git a/hud/eval/tests/test_rollout.py b/hud/eval/tests/test_rollout.py index 0b031b6ec..d0880b8b4 100644 --- a/hud/eval/tests/test_rollout.py +++ b/hud/eval/tests/test_rollout.py @@ -31,9 +31,9 @@ from hud.agents.base import Agent from hud.agents.openai_compatible import OpenAIChatAgent from hud.agents.types import OpenAIChatConfig +from hud.capabilities import Connection from hud.environment import Answer, Environment from hud.eval import ( - InferenceConnection, Job, LocalRuntime, Runtime, @@ -174,27 +174,17 @@ async def test_rollout_returns_graded_run_with_trace_id(env_file: Path) -> None: assert run.runtime.startswith("tcp://127.0.0.1:") -async def test_inference_connection_exists_only_during_agent_execution(env_file: Path) -> None: - connection = InferenceConnection( - base_url="https://inference.hud.so", - credential="scoped-runtime-token", +def test_run_owns_named_controller_connections() -> None: + connection = Connection( + name="inference", + capability="ssh", + url="https://inference.hud.so", + headers={"Authorization": "Bearer scoped-runtime-token"}, ) - observed: list[InferenceConnection | None] = [] + run = Run(None, "task", {}, connections=(connection,)) - class InspectingAgent(Agent): - async def __call__(self, run: Run) -> None: - observed.append(run.inference) - run.trace.content = _solve_add(run.prompt_text) - - run = await rollout( - _add_task(2, 3), - InspectingAgent(), - runtime=SubprocessRuntime(env_file), - inference=connection, - ) - - assert observed == [connection] - assert run.inference is None + assert run.connections == {"inference": connection} + assert "scoped-runtime-token" not in repr(connection) async def test_verifier_task_replaces_the_actor_grade_in_the_same_runtime() -> None: diff --git a/hud/tests/test_init_module.py b/hud/tests/test_init_module.py index 0281d756d..83437c458 100644 --- a/hud/tests/test_init_module.py +++ b/hud/tests/test_init_module.py @@ -20,13 +20,13 @@ def test_all_exports(self): expected = [ "Chat", "ComposeProject", + "Connection", "DockerRuntime", "Environment", "Grade", "Job", "HUDRuntime", "HostedRuntime", - "InferenceConnection", "Run", "Runtime", "RuntimeConfig",