diff --git a/README.md b/README.md index 4a06edc15..b1bb099f5 100644 --- a/README.md +++ b/README.md @@ -156,6 +156,12 @@ A **capability** is a connection the environment exposes; a **harness** attaches From the [platform UI](https://hud.ai) you can run batches, compare models on the same taskset, and inspect every trace. +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. + → [Run & deploy](https://docs.hud.ai/v6/reference/runtime) ## Train on rewards diff --git a/hud/__init__.py b/hud/__init__.py index fdcdc7b0f..6cde008be 100644 --- a/hud/__init__.py +++ b/hud/__init__.py @@ -16,6 +16,7 @@ Grade, HostedRuntime, HUDRuntime, + InferenceConnection, Job, LocalRuntime, Run, @@ -43,6 +44,7 @@ "Grade", "HUDRuntime", "HostedRuntime", + "InferenceConnection", "Job", "LocalRuntime", "Run", diff --git a/hud/agents/claude/sdk/agent.py b/hud/agents/claude/sdk/agent.py index 76fd85c8e..abd1336f7 100644 --- a/hud/agents/claude/sdk/agent.py +++ b/hud/agents/claude/sdk/agent.py @@ -34,7 +34,7 @@ if TYPE_CHECKING: from hud.capabilities import SSHClient - from hud.eval.run import Run + from hud.eval.run import InferenceConnection, Run logger = logging.getLogger(__name__) @@ -43,8 +43,8 @@ RUN_SCRIPT_PATH = ".hud_run.bat" _MANAGED_CLAUDE_PATHS = { - "linux-x64": "/media/hud/bin/claude/linux-x64/claude", - "linux-x64-musl": "/media/hud/bin/claude/linux-x64-musl/claude", + "linux-x64": "/usr/local/lib/agents/claude/linux-x64/claude", + "linux-x64-musl": "/usr/local/lib/agents/claude/linux-x64-musl/claude", } @@ -111,6 +111,7 @@ async def __call__(self, run: Run) -> None: mcp_servers=mcp_servers, prompt=run.prompt_text, executable=executable, + inference=run.inference, ) async def _exec( @@ -122,6 +123,7 @@ async def _exec( mcp_servers: dict[str, dict[str, Any]], prompt: str, executable: str = "claude", + inference: InferenceConnection | None = None, ) -> None: mcp_config_path = await self._write_mcp_config(ssh, mcp_servers) input_text = ( @@ -145,6 +147,7 @@ async def _exec( shell=shell, mcp_config_path=mcp_config_path, executable=executable, + inference=inference, ) if shell in WINDOWS_SHELLS: await ssh.write_text(RUN_SCRIPT_PATH, f"@echo off\r\n{command}\r\n") @@ -173,20 +176,26 @@ async def _exec( except (OSError, asyncssh.Error): logger.warning("Failed to remove Claude CLI runtime files") - def _build_env_vars(self) -> dict[str, str]: + def _build_env_vars(self, inference: InferenceConnection | 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 = settings.api_key is not None + use_hud_gateway = inference is not None or settings.api_key is not None if use_hud_gateway: - if not settings.api_key: + if inference is not None: + base_url = inference.base_url + api_key = inference.credential + elif settings.api_key: + base_url = settings.hud_gateway_url + api_key = settings.api_key + else: raise ValueError("HUD_API_KEY is required for HUD gateway routing") - env["ANTHROPIC_BASE_URL"] = settings.hud_gateway_url - env["ANTHROPIC_API_KEY"] = settings.api_key + env["ANTHROPIC_BASE_URL"] = base_url + env["ANTHROPIC_API_KEY"] = api_key env["CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS"] = "1" env["DISABLE_AUTO_COMPACT"] = "1" - if trace_id := get_current_trace_id(): + if inference 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 @@ -227,8 +236,9 @@ def _build_cli_command( shell: str, mcp_config_path: str | None = None, executable: str = "claude", + inference: InferenceConnection | None = None, ) -> str: - env_vars = self._build_env_vars() + env_vars = self._build_env_vars(inference) is_win = shell in WINDOWS_SHELLS base_args: list[str] = [ executable, diff --git a/hud/agents/codex/agent.py b/hud/agents/codex/agent.py index ed4f6946b..345443c1f 100644 --- a/hud/agents/codex/agent.py +++ b/hud/agents/codex/agent.py @@ -25,13 +25,13 @@ if TYPE_CHECKING: from hud.capabilities import SSHClient - from hud.eval.run import Run + from hud.eval.run import InferenceConnection, Run logger = logging.getLogger(__name__) _MANAGED_CODEX_PATHS = { - "linux-x64": "/media/hud/bin/codex/bin/codex", - "linux-x64-musl": "/media/hud/bin/codex/bin/codex", + "linux-x64": "/usr/local/lib/agents/codex/bin/codex", + "linux-x64-musl": "/usr/local/lib/agents/codex/bin/codex", } @@ -208,7 +208,12 @@ def record_tool(self, item: dict[str, Any], started_at: str, ended_at: str) -> N ) -def codex_command(config: CodexCLIConfig, shell: str, executable: str = "codex") -> str: +def codex_command( + config: CodexCLIConfig, + shell: str, + executable: str = "codex", + inference: InferenceConnection | None = None, +) -> str: env: dict[str, str] = {} args = [ executable, @@ -226,21 +231,29 @@ def codex_command(config: CodexCLIConfig, shell: str, executable: str = "codex") use_hud_gateway = config.use_hud_gateway if use_hud_gateway is None: - use_hud_gateway = settings.api_key is not None + use_hud_gateway = inference is not None or settings.api_key is not None if use_hud_gateway: - if not settings.api_key: + if inference is not None: + base_url = inference.base_url + credential = inference.credential + credential_env = "HUD_RUNTIME_INFERENCE_TOKEN" + elif settings.api_key: + base_url = settings.hud_gateway_url + credential = settings.api_key + credential_env = "HUD_API_KEY" + else: raise ValueError("HUD_API_KEY is required for HUD gateway routing") - env["HUD_API_KEY"] = settings.api_key + env[credential_env] = credential overrides = { "model_provider": "hud", "model_providers.hud.name": "HUD", - "model_providers.hud.base_url": settings.hud_gateway_url, - "model_providers.hud.env_key": "HUD_API_KEY", + "model_providers.hud.base_url": base_url, + "model_providers.hud.env_key": credential_env, "model_providers.hud.wire_api": "responses", } for key, value in overrides.items(): args.extend(["-c", f"{key}={json.dumps(value)}"]) - if trace_id := get_current_trace_id(): + if inference is None and (trace_id := get_current_trace_id()): args.extend( [ "-c", @@ -251,35 +264,42 @@ def codex_command(config: CodexCLIConfig, shell: str, executable: str = "codex") env["CODEX_API_KEY"] = settings.openai_api_key args.append("-") + isolate_home = bool(env) if shell in WINDOWS_SHELLS: - script = ";".join( - [ + invocation = ( + f"& {powershell_quote(executable)} " + f"{' '.join(powershell_quote(arg) for arg in args[1:])}; " + "$hudExitCode=$LASTEXITCODE" + ) + statements = [ + *(f"$env:{key}={powershell_quote(value)}" for key, value in env.items()), + invocation, + "exit $hudExitCode", + ] + if isolate_home: + statements = [ "$codexHome=Join-Path ([System.IO.Path]::GetTempPath()) " "('hud-codex-' + [System.Guid]::NewGuid())", "New-Item -ItemType Directory -Force -Path $codexHome | Out-Null", "$env:CODEX_HOME=$codexHome", *(f"$env:{key}={powershell_quote(value)}" for key, value in env.items()), - f"try {{ & {powershell_quote(executable)} " - f"{' '.join(powershell_quote(arg) for arg in args[1:])}; " - "$hudExitCode=$LASTEXITCODE } finally { Remove-Item -Recurse -Force " - "$codexHome }", + f"try {{ {invocation} }} finally {{ Remove-Item -Recurse -Force $codexHome }}", "exit $hudExitCode", ] - ) - return powershell(script) + return powershell(";".join(statements)) command = " ".join(shlex.quote(arg) for arg in args) env_prefix = " ".join(f"{key}={shlex.quote(value)}" for key, value in env.items()) invocation = f"{env_prefix} {command}" if env_prefix else command - return "; ".join( - [ + statements = ['export PATH="$HOME/.local/bin:$PATH"', invocation] + if isolate_home: + statements = [ 'codex_home=$(mktemp -d "${TMPDIR:-/tmp}/hud-codex.XXXXXX") || exit 1', "trap 'rm -rf -- \"$codex_home\"' EXIT", 'export CODEX_HOME="$codex_home"', - 'export PATH="$HOME/.local/bin:$PATH"', - invocation, + *statements, ] - ) + return "; ".join(statements) async def run_codex( @@ -290,8 +310,9 @@ async def run_codex( shell: str, prompt: str, executable: str = "codex", + inference: InferenceConnection | None = None, ) -> None: - command = codex_command(config, shell, executable) + command = codex_command(config, shell, executable, inference=inference) 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) @@ -322,6 +343,7 @@ async def __call__(self, run: Run) -> None: shell=ssh.capability.params.get("shell", "bash"), prompt=run.prompt_text, executable=executable, + inference=run.inference, ) diff --git a/hud/agents/tests/test_claude_cli_agent.py b/hud/agents/tests/test_claude_cli_agent.py index b6528b24b..07a793d73 100644 --- a/hud/agents/tests/test_claude_cli_agent.py +++ b/hud/agents/tests/test_claude_cli_agent.py @@ -30,6 +30,7 @@ from hud.agents.types import AgentStep, ClaudeCLIConfig, ToolStep from hud.capabilities import Capability, 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 @@ -68,6 +69,32 @@ 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", + ) + + gateway = ClaudeCLIAgent(ClaudeCLIConfig(use_hud_gateway=True))._build_cli_command( + shell="bash", + inference=inference, + ) + + assert "ANTHROPIC_BASE_URL=https://inference.hud.so" in gateway + assert "ANTHROPIC_API_KEY=scoped-runtime-token" in gateway + assert "HUD_API_KEY" not in gateway + assert "Trace-Id" not in gateway + for name in ( + "ANTHROPIC_MODEL", + "ANTHROPIC_SMALL_FAST_MODEL", + "ANTHROPIC_DEFAULT_SONNET_MODEL", + "ANTHROPIC_DEFAULT_OPUS_MODEL", + "ANTHROPIC_DEFAULT_HAIKU_MODEL", + "CLAUDE_CODE_SUBAGENT_MODEL", + ): + assert f"{name}=claude-sonnet-5" in gateway + + def test_windows_command_encodes_environment_and_arguments( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -458,6 +485,7 @@ async def test_manifest_mcp_capability_is_written_for_remote_claude( ssh = SSHClient(shell, cast("Any", object())) class Client: + inference = None manifest = SimpleNamespace(bindings=[shell, mcp]) async def open(self, ref: str) -> SSHClient: @@ -471,7 +499,9 @@ async def open(self, ref: str) -> SSHClient: await agent( cast( "Any", - SimpleNamespace(client=Client(), prompt_text="call the tool", runtime_config=None), + SimpleNamespace( + client=Client(), prompt_text="call the tool", runtime_config=None, inference=None + ), ) ) @@ -499,6 +529,7 @@ async def test_remote_claude_passes_screenshot_encoding_to_computer_mcp( bridge_active = False class Client: + inference = None manifest = SimpleNamespace(bindings=[shell, screen]) async def open(self, ref: str) -> SSHClient: @@ -542,7 +573,9 @@ async def execute(*_args: Any, **_kwargs: Any) -> None: await agent( cast( "Any", - SimpleNamespace(client=Client(), prompt_text="use the computer", runtime_config=None), + SimpleNamespace( + client=Client(), prompt_text="use the computer", runtime_config=None, inference=None + ), ) ) @@ -583,6 +616,7 @@ async def test_remote_claude_preserves_multiple_rfb_bindings( bridged: list[str] = [] class Client: + inference = None manifest = SimpleNamespace(bindings=[shell, *screens]) async def open(self, ref: str) -> SSHClient: @@ -633,6 +667,7 @@ async def execute(*_args: Any, **kwargs: Any) -> None: client=Client(), prompt_text="use both screens", runtime_config=None, + inference=None, ), ) ) @@ -870,6 +905,7 @@ async def test_concurrent_runs_keep_their_ssh_state_isolated( class Client: def __init__(self, shell: Capability, ssh: SSHClient) -> None: + self.inference = None self.manifest = SimpleNamespace(bindings=[shell]) self.ssh = ssh @@ -897,9 +933,14 @@ 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) + run_a = SimpleNamespace( + client=Client(shell_a, ssh_a), prompt_text="first", runtime_config=None, inference=None + ) run_b = SimpleNamespace( - client=Client(shell_b, ssh_b), prompt_text="second", runtime_config=None + client=Client(shell_b, ssh_b), + prompt_text="second", + runtime_config=None, + inference=None, ) 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 9819c1827..3caeb021d 100644 --- a/hud/agents/tests/test_codex_cli_agent.py +++ b/hud/agents/tests/test_codex_cli_agent.py @@ -18,6 +18,7 @@ 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.eval.runtime import RuntimeConfig, RuntimeResources from hud.settings import settings from hud.telemetry.context import set_trace_context @@ -104,6 +105,38 @@ 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", + ) + + command = codex_command(CodexCLIConfig(use_hud_gateway=True), "bash", inference=inference) + + assert "HUD_RUNTIME_INFERENCE_TOKEN=scoped-runtime-token" in command + assert 'model_providers.hud.env_key="HUD_RUNTIME_INFERENCE_TOKEN"' in command + assert "HUD_API_KEY" not in command + assert 'model_providers.hud.base_url="https://inference.hud.so"' in command + assert "Trace-Id" not in command + + +@pytest.mark.parametrize("shell", ["bash", "powershell"]) +def test_command_preserves_ambient_codex_login_without_explicit_credentials(shell: str) -> None: + command = codex_command(CodexCLIConfig(use_hud_gateway=False), shell) + script = ( + base64.b64decode(command.rsplit(" ", 1)[1]).decode("utf-16-le") + if shell == "powershell" + else command + ) + + assert "CODEX_HOME" not in script + assert "CODEX_API_KEY" not in script + assert "HUD_API_KEY" not in script + assert "HUD_RUNTIME_INFERENCE_TOKEN" not in script + assert "mktemp" not in script + assert "codex exec" in script or "& 'codex' 'exec'" in script + + def test_windows_command_encodes_environment_and_arguments( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -276,6 +309,8 @@ async def test_agent_opens_ssh_and_uses_workspace_prompt(monkeypatch: pytest.Mon ssh = _FakeSSH(_FakeProcess(_STREAM_JSON), shell="powershell") class Client: + inference = None + async def open(self, ref: str) -> _FakeSSH: assert ref == "ssh" return ssh @@ -283,7 +318,9 @@ async def open(self, ref: str) -> _FakeSSH: agent = CodexCLIAgent() execute = AsyncMock() monkeypatch.setattr("hud.agents.codex.agent.run_codex", execute) - run = SimpleNamespace(client=Client(), prompt_text="Fix it", runtime_config=None) + run = SimpleNamespace( + client=Client(), prompt_text="Fix it", runtime_config=None, inference=None + ) await agent(cast("Any", run)) @@ -294,6 +331,7 @@ async def open(self, ref: str) -> _FakeSSH: shell="powershell", prompt="Fix it", executable="codex", + inference=None, ) @@ -311,11 +349,11 @@ async def test_executable_resolution_prefers_matching_managed_bundle() -> None: executable = await resolve_executable( cast("Any", ssh), "codex", - {"linux-x64": "/media/hud/bin/codex/bin/codex"}, + {"linux-x64": "/usr/local/lib/agents/codex/bin/codex"}, RuntimeConfig(resources=RuntimeResources(os="linux")), ) - assert executable == "/media/hud/bin/codex/bin/codex" + assert executable == "/usr/local/lib/agents/codex/bin/codex" assert ssh.run.await_count == 2 diff --git a/hud/agents/types.py b/hud/agents/types.py index 4dbf3e3bd..be1956597 100644 --- a/hud/agents/types.py +++ b/hud/agents/types.py @@ -182,7 +182,11 @@ class ClaudeCLIConfig(AgentConfig): class CodexCLIConfig(AgentConfig): - """Configuration for CodexCLIAgent (runs ``codex exec`` over SSH).""" + """Configuration for CodexCLIAgent (runs ``codex exec`` over SSH). + + Without an explicit inference connection or API key, the agent leaves + ``CODEX_HOME`` unchanged so a login in that execution environment can apply. + """ model_name: str = "Codex CLI" model: str = Field(default="gpt-5.6-sol", validation_alias=_model_alias) diff --git a/hud/clients/client.py b/hud/clients/client.py index 755c6f542..9ee17a92a 100644 --- a/hud/clients/client.py +++ b/hud/clients/client.py @@ -30,8 +30,9 @@ from hud.environment.utils import read_frame, send_frame, splice if TYPE_CHECKING: - from collections.abc import AsyncIterator + from collections.abc import AsyncIterator, Sequence + from hud.environment.egress import WorkspaceRoute from hud.eval.runtime import Runtime LOGGER = logging.getLogger("hud.clients") @@ -155,14 +156,23 @@ def abort(self) -> None: # ─── handshake ──────────────────────────────────────────────────── - async def hello(self, session_id: str | None = None) -> Manifest: + async def hello( + self, + session_id: str | None = None, + *, + workspace_routes: Sequence[WorkspaceRoute] = (), + ) -> Manifest: """Send ``hello``; cache and return the parsed ``Manifest``. ``session_id`` resumes that parked session on the env — its suspended task, e.g. one a prior connection started — instead of minting a fresh session. """ - params: dict[str, Any] = {} if session_id is None else {"session_id": session_id} + params: dict[str, Any] = {} + if workspace_routes: + params["workspace_routes"] = [route.to_wire() for route in workspace_routes] + if session_id is not None: + params["session_id"] = session_id result = await self._call("hello", params) env = result["env"] bindings = [Capability.from_manifest(binding) for binding in result["bindings"]] @@ -373,6 +383,7 @@ async def _connect_ready( port: int, *, ready_timeout: float, + workspace_routes: Sequence[WorkspaceRoute], interval: float = 0.5, ) -> HudClient: """Connect and complete ``hello``, retrying until the env is ready. @@ -396,7 +407,7 @@ async def _connect_ready( client = HudClient(reader, writer, endpoint=(host, port)) try: - await client.hello() + await client.hello(workspace_routes=workspace_routes) except asyncio.CancelledError: client.abort() raise @@ -429,7 +440,12 @@ def _runtime_ready_timeout(runtime: Runtime, default: float) -> float: @asynccontextmanager -async def connect(runtime: Runtime, *, ready_timeout: float = 240.0) -> AsyncIterator[HudClient]: +async def connect( + runtime: Runtime, + *, + ready_timeout: float = 240.0, + workspace_routes: Sequence[WorkspaceRoute] = (), +) -> AsyncIterator[HudClient]: """Connect a :class:`HudClient` to a provisioned substrate's control channel. Takes the :class:`~hud.eval.runtime.Runtime` a provider yielded (or @@ -446,6 +462,7 @@ async def connect(runtime: Runtime, *, ready_timeout: float = 240.0) -> AsyncIte parts.hostname or "127.0.0.1", parts.port or 0, ready_timeout=_runtime_ready_timeout(runtime, ready_timeout), + workspace_routes=workspace_routes, ) owner = asyncio.current_task() assert owner is not None @@ -457,9 +474,12 @@ async def heartbeat() -> None: await asyncio.sleep(_CONTROL_HEARTBEAT_INTERVAL_SECONDS) assert client.manifest is not None try: + params: dict[str, Any] = {"session_id": client.manifest.session_id} + if workspace_routes: + params["workspace_routes"] = [route.to_wire() for route in workspace_routes] await client._call( "hello", - {"session_id": client.manifest.session_id}, + params, reply_timeout=_CONTROL_HEARTBEAT_TIMEOUT_SECONDS, ) except HudProtocolError as exc: @@ -485,4 +505,10 @@ async def heartbeat() -> None: raise -__all__ = ["HudClient", "HudProtocolError", "Manifest", "ServerInfo", "connect"] +__all__ = [ + "HudClient", + "HudProtocolError", + "Manifest", + "ServerInfo", + "connect", +] diff --git a/hud/clients/tests/test_connect.py b/hud/clients/tests/test_connect.py index a5f921d4c..555e78f1c 100644 --- a/hud/clients/tests/test_connect.py +++ b/hud/clients/tests/test_connect.py @@ -19,12 +19,56 @@ import hud.clients.client as client_module from hud.capabilities import Capability, CapabilityClient from hud.clients import connect +from hud.environment import WorkspaceRoute from hud.environment.utils import read_frame, send_frame from hud.eval.runtime import Runtime HELLO_RESULT = {"session_id": "s-1", "env": {"name": "stub", "version": "1.0"}, "bindings": []} +def test_workspace_route_from_url_extracts_transport_address() -> None: + assert WorkspaceRoute.from_url("ssh", "https://inference.hud.so/v1") == WorkspaceRoute( + "ssh", + "inference.hud.so", + 443, + ) + assert WorkspaceRoute.from_url("shell", "http://gateway.test:8080") == WorkspaceRoute( + "shell", + "gateway.test", + 8080, + ) + + +async def test_connect_sends_workspace_routes_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] + runtime = Runtime(f"tcp://127.0.0.1:{port}") + route = WorkspaceRoute("ssh", "inference.hud.so", 443) + try: + async with connect(runtime, workspace_routes=(route,)): + pass + finally: + server.close() + await server.wait_closed() + + assert [request["method"] for request in requests] == ["hello"] + params = requests[0]["params"] + assert isinstance(params, dict) + assert params == {"workspace_routes": [route.to_wire()]} + + async def test_open_retries_transient_capability_connection_failures( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -288,11 +332,13 @@ async def fake_connect_ready( port: int, *, ready_timeout: float, + workspace_routes: tuple[WorkspaceRoute, ...], interval: float = 0.5, ) -> _FakeClient: seen["host"] = host seen["port"] = port seen["ready_timeout"] = ready_timeout + assert workspace_routes == () seen["interval"] = interval return _FakeClient() diff --git a/hud/environment/__init__.py b/hud/environment/__init__.py index 1beb9bfb9..acbbb96b5 100644 --- a/hud/environment/__init__.py +++ b/hud/environment/__init__.py @@ -23,7 +23,7 @@ from hud.utils.modules import iter_modules from .arguments import DataFileArg, DataFileRef, DataFilesArg, GradingArg, PromptArg -from .egress import Peer +from .egress import Peer, WorkspaceRoute from .env import Answer, Environment from .workspace import DEFAULT_SYSTEM_MOUNTS, Mount, MountKind, Workspace @@ -100,5 +100,6 @@ def load_environment( "Peer", "PromptArg", "Workspace", + "WorkspaceRoute", "load_environment", ] diff --git a/hud/environment/egress.py b/hud/environment/egress.py index 37ddc3732..44cd93488 100644 --- a/hud/environment/egress.py +++ b/hud/environment/egress.py @@ -185,6 +185,59 @@ def address(self) -> tuple[str, int]: return self.target or ("127.0.0.1", self.port) +@dataclass(frozen=True, slots=True) +class WorkspaceRoute: + """A controller-provided host route exposed through one workspace capability.""" + + capability: str + host: str + port: int + + def __post_init__(self) -> None: + if not self.capability or self.capability.strip() != self.capability: + raise ValueError("workspace route capability must not be empty or padded") + if not self.host or self.host.strip() != self.host or any(c.isspace() for c in self.host): + raise ValueError("workspace route host must be a hostname without whitespace") + try: + ipaddress.ip_address(self.host) + except ValueError: + pass + else: + raise ValueError("workspace routes require a hostname, not an IP address") + if not 1 <= self.port <= 65535: + raise ValueError("workspace route port must be between 1 and 65535") + + def to_wire(self) -> dict[str, str | int]: + return {"capability": self.capability, "host": self.host, "port": self.port} + + @classmethod + def from_url(cls, capability: str, url: str) -> WorkspaceRoute: + """Build a host route for one HTTP(S) endpoint.""" + parts = urllib.parse.urlsplit(url) + if parts.scheme not in {"http", "https"} or parts.hostname is None: + raise ValueError("workspace route URL must be HTTP(S) with a hostname") + if parts.username is not None or parts.password is not None: + raise ValueError("workspace route URL must not contain credentials") + return cls( + capability=capability, + host=parts.hostname, + port=parts.port or (443 if parts.scheme == "https" else 80), + ) + + @classmethod + def from_wire(cls, value: object) -> WorkspaceRoute: + if not isinstance(value, dict): + raise ValueError("workspace routes must be objects") + capability = value.get("capability") + host = value.get("host") + port = value.get("port") + if not isinstance(capability, str) or not isinstance(host, str): + raise ValueError("workspace route capability and host must be strings") + if isinstance(port, bool) or not isinstance(port, int): + raise ValueError("workspace route port must be an integer") + return cls(capability=capability, host=host, port=port) + + def bind_addresses( peers: Sequence[Peer], *, @@ -665,6 +718,7 @@ def stop(self) -> None: "VISITOR_PORT", "Egress", "Peer", + "WorkspaceRoute", "bind_addresses", "hosts_text", "permitted", diff --git a/hud/environment/env.py b/hud/environment/env.py index 86023f3da..90364a8c8 100644 --- a/hud/environment/env.py +++ b/hud/environment/env.py @@ -17,6 +17,7 @@ from hud.capabilities import Capability +from .egress import Peer, WorkspaceRoute from .workspace import Workspace if TYPE_CHECKING: @@ -162,6 +163,8 @@ def __init__( self._on_stop: list[Callable[[], Awaitable[None]]] = [] # Per task-session end (cancel / bye / post-grade cleanup). self._on_task_teardown: list[Callable[[], Awaitable[None]]] = [] + self._workspaces: dict[str, Workspace] = {} + self._workspace_routes: dict[WorkspaceRoute, tuple[Workspace, Peer | None]] = {} # ─── task registration ─────────────────────────────────────────── @@ -284,7 +287,10 @@ def workspace( from hud.settings import settings track_files = settings.file_tracking_enabled + if name in self._workspaces: + raise ValueError(f"workspace capability {name!r} is already attached") ws = Workspace(root, track_files=track_files, **kwargs) + self._workspaces[name] = ws @self.initialize async def _up() -> None: @@ -349,5 +355,66 @@ async def stop(self) -> None: for hook in reversed(self._on_stop): with contextlib.suppress(Exception): await hook() + for workspace, peer in reversed(self._workspace_routes.values()): + if peer is not None: + workspace.remove_peer(peer) + self._workspace_routes.clear() self._started = False self._hooks_done = False + + def bind_workspace_routes(self, routes: Sequence[WorkspaceRoute]) -> None: + """Install controller routes before a workspace starts its sandbox.""" + if not self._started: + raise RuntimeError("environment must be started before workspace routes are bound") + + planned: list[tuple[WorkspaceRoute, Workspace, Peer | 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") + if not workspace.bwrap_available or not workspace.owns_netns: + raise RuntimeError( + f"workspace route for {route.capability!r} requires an isolated network" + ) + matching = [ + peer + for peer in workspace.peers + if peer.name == route.host and peer.port == route.port + ] + if matching: + if any(peer.address != (route.host, route.port) for peer in matching): + raise RuntimeError( + f"workspace route {route.host}:{route.port} conflicts with an authored peer" + ) + planned.append((route, workspace, None)) + continue + planned.append( + ( + route, + workspace, + Peer(route.host, route.port, target=(route.host, route.port)), + ) + ) + + bound: list[tuple[WorkspaceRoute, Workspace, Peer | None]] = [] + try: + for route, workspace, peer in planned: + if peer is not None: + workspace.add_peer(peer, first=True) + self._workspace_routes[route] = (workspace, peer) + bound.append((route, workspace, peer)) + except BaseException: + for route, workspace, peer in reversed(bound): + if peer is not None: + workspace.remove_peer(peer) + self._workspace_routes.pop(route, None) + raise diff --git a/hud/environment/server.py b/hud/environment/server.py index 36ad47818..28d8b15aa 100644 --- a/hud/environment/server.py +++ b/hud/environment/server.py @@ -28,6 +28,7 @@ from hud.graders.results import EvaluationResult +from .egress import WorkspaceRoute from .env import Answer, current_session_id from .utils import error, read_frame, reply, send_frame, splice @@ -234,7 +235,7 @@ def __init__(self, env: Environment) -> None: self._live: set[str] = set() async def start(self, session_id: str, task_id: str, args: dict[str, Any]) -> dict[str, Any]: - await self.cancel(session_id) + await self._cancel_runner(session_id) runner = TaskRunner(self.env.tasks[task_id], args) self._runners[session_id] = runner try: @@ -269,7 +270,7 @@ def _adopt_parked(self) -> tuple[str, TaskRunner]: sid = parked[0] return sid, self._runners.pop(sid) - async def cancel(self, session_id: str) -> None: + async def _cancel_runner(self, session_id: str) -> None: runner = self._runners.pop(session_id, None) if runner is None: return @@ -280,6 +281,9 @@ async def cancel(self, session_id: str) -> None: finally: current_session_id.reset(token) + async def cancel(self, session_id: str) -> None: + await self._cancel_runner(session_id) + async def cancel_all(self) -> None: """Tear down every suspended/live task (server shutdown).""" for session_id in list(self._runners): @@ -334,6 +338,20 @@ async def error_to(msg_id: int | None, code: int, message: str) -> None: self._live.add(session_id) current_session_id.reset(session_token) session_token = current_session_id.set(session_id) + raw_routes = params.get("workspace_routes", []) + if not isinstance(raw_routes, list): + await error_to( + msg_id, -32602, "hello: 'workspace_routes' must be a list" + ) + continue + try: + workspace_routes = [ + WorkspaceRoute.from_wire(route) for route in raw_routes + ] + except ValueError as exc: + await error_to(msg_id, -32602, f"hello: {exc}") + continue + env.bind_workspace_routes(workspace_routes) # env.start() ran before serving, so hook-published # capabilities (e.g. a workspace's ssh address) are # already concrete here. diff --git a/hud/environment/tests/test_workspace.py b/hud/environment/tests/test_workspace.py index 5b857ab59..451873137 100644 --- a/hud/environment/tests/test_workspace.py +++ b/hud/environment/tests/test_workspace.py @@ -30,7 +30,7 @@ from hud.capabilities import SSHClient from hud.environment import namespace as namespace_mod from hud.environment import workspace as workspace_mod -from hud.environment.egress import Peer, _field, _UnixServer, _Unrelayable +from hud.environment.egress import Peer, WorkspaceRoute, _field, _UnixServer, _Unrelayable from hud.environment.workspace import Bubblewrap, Mount, Workspace from hud.utils.process import ProcessGroup, ProcessResult @@ -1136,6 +1136,28 @@ def test_a_peer_answers_at_the_address_the_task_expects() -> None: bind_addresses([Peer("db", 5432), Peer("db", 5432)]) +async def test_workspace_route_is_bound_once_and_removed_on_stop(tmp_path: Path) -> None: + from hud.environment import Environment + + env = Environment() + workspace = env.workspace(tmp_path / "root", track_files=False) + workspace._bwrap = cast("Any", object()) + env._started = True + route = WorkspaceRoute("ssh", "inference.hud.so", 443) + + env.bind_workspace_routes([route, route]) + env.bind_workspace_routes([route]) + + assert workspace.peers == (Peer("inference.hud.so", 443, target=("inference.hud.so", 443)),) + await env.stop() + assert workspace.peers == () + + +def test_workspace_route_rejects_ip_literals() -> None: + with pytest.raises(ValueError, match="hostname"): + WorkspaceRoute("ssh", "127.0.0.1", 443) + + def test_workspace_names_are_added_to_the_substrates_hosts_rather_than_replacing_it() -> None: """Dropping the substrate's entries would cost the workspace localhost.""" from hud.environment.egress import Peer, hosts_text diff --git a/hud/environment/workspace.py b/hud/environment/workspace.py index 8dc7784ae..c0c93d028 100644 --- a/hud/environment/workspace.py +++ b/hud/environment/workspace.py @@ -633,6 +633,22 @@ def owns_netns(self) -> bool: """ return not self.network or self.allowed_hosts is not None + def add_peer(self, peer: Peer, *, first: bool = False) -> None: + """Add a substrate service before the workspace accepts sessions.""" + if self._sandbox is not None: + raise RuntimeError("workspace peers must be bound before its sandbox starts") + self.peers = (peer, *self.peers) if first else (*self.peers, peer) + if self._hosts_path is not None: + self._hosts_path = self._write_hosts() + + def remove_peer(self, peer: Peer) -> None: + """Remove a substrate service after the workspace has stopped.""" + if self._sandbox is not None: + raise RuntimeError("workspace peers must be unbound after its sandbox stops") + self.peers = tuple(candidate for candidate in self.peers if candidate != peer) + if self._hosts_path is not None: + self._hosts_path = self._write_hosts() + def _setpriv(self) -> str | None: """Absolute path to ``setpriv``, resolved via the *server's* PATH. diff --git a/hud/eval/__init__.py b/hud/eval/__init__.py index 0bce06d61..7411f9105 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, Run, rollout +from .run import Grade, InferenceConnection, Run, rollout from .runtime import ( ComposeProject, DaytonaRuntime, @@ -63,6 +63,7 @@ "Grade", "HUDRuntime", "HostedRuntime", + "InferenceConnection", "Job", "LocalRuntime", "ModalRuntime", diff --git a/hud/eval/run.py b/hud/eval/run.py index 839d106e9..1be0b8b1e 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -27,6 +27,7 @@ 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 @@ -40,10 +41,12 @@ from .job import job_enter, trace_enter, trace_exit if TYPE_CHECKING: + from collections.abc import Sequence from types import TracebackType from hud.agents.base import Agent from hud.clients.client import HudClient + from hud.environment import WorkspaceRoute from .runtime import Provider from .runtime.core import RuntimeConfig @@ -52,6 +55,25 @@ 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, @@ -200,12 +222,14 @@ def __init__( *, best_effort_grade: bool = False, runtime_config: RuntimeConfig | None = None, + inference: InferenceConnection | None = None, ) -> 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 #: 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 @@ -445,6 +469,8 @@ async def rollout( group_id: str | None = None, trace_id: str | None = None, rollout_timeout: float | None = None, + inference: InferenceConnection | None = None, + workspace_routes: Sequence[WorkspaceRoute] = (), ) -> Run: """Drive one task to a graded :class:`Run` here, against ``runtime``'s channel. @@ -535,7 +561,7 @@ 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) as actor_client: + async with connect(addr, workspace_routes=workspace_routes) as actor_client: client = actor_client live = Run( actor_client, @@ -543,35 +569,39 @@ 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, ) live._runtime = addr.url # the placement record for the receipt async with live: # start on enter; complete on exit run = live # bound only once live: an earlier failure synthesizes _phase = "agent loop" try: - async with file_tracking_observer(actor_client): - if agent_timeout is None: - await agent(run) - else: - deadline = asyncio.timeout(agent_timeout) - try: - async with deadline: - await agent(run) - except TimeoutError: - if not deadline.expired(): - raise - detail = f"agent timed out after {agent_timeout:g}s" - logger.warning(detail) - run.trace.status = "error" - run.trace.stop_reason = "timeout" - run.record(Step(source="system", error=detail)) - except Exception as exc: - if task.verifier is None: - raise - detail = "".join(traceback.format_exception_only(exc)).strip() - logger.warning("rollout failed mid-run (%s): %s", _phase, detail) - run.trace.status = "error" - run.record(Step(source="system", error=f"[{_phase}] {detail}")) + try: + async with file_tracking_observer(actor_client): + if agent_timeout is None: + await agent(run) + else: + deadline = asyncio.timeout(agent_timeout) + try: + async with deadline: + await agent(run) + except TimeoutError: + if not deadline.expired(): + raise + detail = f"agent timed out after {agent_timeout:g}s" + logger.warning(detail) + run.trace.status = "error" + run.trace.stop_reason = "timeout" + run.record(Step(source="system", error=detail)) + except Exception as exc: + if task.verifier is None: + raise + detail = "".join(traceback.format_exception_only(exc)).strip() + logger.warning("rollout failed mid-run (%s): %s", _phase, detail) + run.trace.status = "error" + run.record(Step(source="system", error=f"[{_phase}] {detail}")) + finally: + run.inference = None _phase = "grading" if verifier is not None: @@ -671,6 +701,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.trace.trace_id = trace_id run.job_id = job_id run.group_id = group_id @@ -684,4 +715,4 @@ def _consume_task_result(task: asyncio.Future[Any]) -> None: task.result() -__all__ = ["Grade", "Run", "rollout"] +__all__ = ["Grade", "InferenceConnection", "Run", "rollout"] diff --git a/hud/eval/tests/test_rollout.py b/hud/eval/tests/test_rollout.py index 8210e2078..0b031b6ec 100644 --- a/hud/eval/tests/test_rollout.py +++ b/hud/eval/tests/test_rollout.py @@ -32,7 +32,15 @@ from hud.agents.openai_compatible import OpenAIChatAgent from hud.agents.types import OpenAIChatConfig from hud.environment import Answer, Environment -from hud.eval import Job, LocalRuntime, Runtime, SubprocessRuntime, Task, Taskset +from hud.eval import ( + InferenceConnection, + Job, + LocalRuntime, + Runtime, + SubprocessRuntime, + Task, + Taskset, +) from hud.eval.run import Run, rollout if TYPE_CHECKING: @@ -166,6 +174,29 @@ 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", + ) + observed: list[InferenceConnection | None] = [] + + 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 + + async def test_verifier_task_replaces_the_actor_grade_in_the_same_runtime() -> None: env = Environment("reviewed") completed: list[str] = [] diff --git a/hud/integrations/harbor/env.py b/hud/integrations/harbor/env.py index 40f2fb973..df8137e83 100644 --- a/hud/integrations/harbor/env.py +++ b/hud/integrations/harbor/env.py @@ -260,7 +260,6 @@ def home(user_id: int | None, *, root: Path | None = None) -> str | None: ) agent_mounts = ( *harness_mounts, - Mount("ro", src=str(ROOT / "bin"), dst=str(ROOT / "bin")), Mount("tmpfs", dst=str(TESTS)), Mount("tmpfs", dst=str(VERIFIER_LOGS)), Mount("ro", src="/dev/null", dst=str(AGENT_ANSWER)), diff --git a/hud/integrations/harbor/tests/test_contract.py b/hud/integrations/harbor/tests/test_contract.py index 531bf4c55..58124a04f 100644 --- a/hud/integrations/harbor/tests/test_contract.py +++ b/hud/integrations/harbor/tests/test_contract.py @@ -130,7 +130,6 @@ def test_adapt_packages_an_image_task_as_a_compose_project(tmp_path: Path) -> No served = (context / "env.py").read_text(encoding="utf-8") assert f'Environment("{context.name}")' in served assert 'Environment(CONFIG["name"])' not in served - assert 'Mount("ro", src=str(ROOT / "bin"), dst=str(ROOT / "bin"))' in served project_root = context / "compose-project" assert _tree_snapshot(project_root / "environment") == authored_environment payload = project_root / "hud" diff --git a/hud/tests/test_init_module.py b/hud/tests/test_init_module.py index 62bd04723..0281d756d 100644 --- a/hud/tests/test_init_module.py +++ b/hud/tests/test_init_module.py @@ -26,6 +26,7 @@ def test_all_exports(self): "Job", "HUDRuntime", "HostedRuntime", + "InferenceConnection", "Run", "Runtime", "RuntimeConfig",