diff --git a/README.md b/README.md index f154e6089..621cdd101 100644 --- a/README.md +++ b/README.md @@ -148,6 +148,17 @@ 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 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. +Codex therefore uses the environment sandbox instead of starting a nested +Codex sandbox for process-bound hosted execution. +The workspace probes its substrate and uses either seccomp notification or a +ptrace-backed seccomp guard; it does not advertise process-bound connections +when neither enforcement backend is available. + → [Run & deploy](https://docs.hud.ai/v6/reference/runtime) ## Train on rewards diff --git a/docs/v6/guides/running-an-eval.mdx b/docs/v6/guides/running-an-eval.mdx index 0fe86a019..150a87a9d 100644 --- a/docs/v6/guides/running-an-eval.mdx +++ b/docs/v6/guides/running-an-eval.mdx @@ -38,9 +38,11 @@ and launch - no CLI required. See [evaluations on the platform](/platform/evalua ### Choosing an agent -The agent name (`claude`, `openai`, `gemini`) selects a built-in harness and routes calls through the -[HUD gateway](/v6/reference/agents), where one `HUD_API_KEY` covers every provider. Switching models is a -single flag, and `hud models list` shows every model the gateway knows. +The agent name (`claude`, `openai`, `gemini`) selects a provider harness and routes calls through the +[HUD gateway](/v6/reference/agents), where one `HUD_API_KEY` covers every provider. The `claude_cli` and +`codex_cli` harnesses instead run an installed CLI inside the environment; pass `--gateway` to route +their model calls through HUD. Switching models is a single flag, and `hud models list` shows every +model the gateway knows. ```bash hud eval "My Taskset" claude --model claude-haiku-4-5 # a cheaper model for fast iteration diff --git a/docs/v6/reference/runtime.mdx b/docs/v6/reference/runtime.mdx index 46dbe96a7..2f021e69e 100644 --- a/docs/v6/reference/runtime.mdx +++ b/docs/v6/reference/runtime.mdx @@ -252,9 +252,9 @@ provisioning. Where `HUDRuntime` runs the agent loop locally against a tunneled env, `HostedRuntime` runs the **whole rollout** remotely, with the agent running alongside the task environment. This process -only submits the rollout and polls its trace to completion. It supports gateway agents from -[`create_agent`](/v6/reference/agents#create-agent); agents with a custom `model_client` must use -`HUDRuntime` or `LocalRuntime`. +only submits the rollout and polls its trace to completion. It supports registered built-in agents; +[`create_agent`](/v6/reference/agents#create-agent) supplies the gateway-backed ones. Agents with a +custom `model_client` must use `HUDRuntime` or `LocalRuntime`. ### `Runtime` diff --git a/environments/argument-hints/.hud_eval.toml b/environments/argument-hints/.hud_eval.toml index 2315a717f..559160271 100644 --- a/environments/argument-hints/.hud_eval.toml +++ b/environments/argument-hints/.hud_eval.toml @@ -4,7 +4,7 @@ [eval] # source = "tasks.py" # tasks file, or a platform taskset name / id -# agent = "claude" # claude, openai, gemini, openai_compatible +# agent = "claude" # claude, claude_cli, codex_cli, openai, gemini, openai_compatible # model = "claude-sonnet-4-6" # all = false # run every task instead of just the first # task_ids = ["fix_bug", "0"] # slugs or 0-based indices @@ -21,6 +21,10 @@ [claude] # max_tokens = 16384 +[claude_cli] + +[codex_cli] + [openai] # temperature = 0.7 diff --git a/environments/blank/.hud_eval.toml b/environments/blank/.hud_eval.toml index 2315a717f..559160271 100644 --- a/environments/blank/.hud_eval.toml +++ b/environments/blank/.hud_eval.toml @@ -4,7 +4,7 @@ [eval] # source = "tasks.py" # tasks file, or a platform taskset name / id -# agent = "claude" # claude, openai, gemini, openai_compatible +# agent = "claude" # claude, claude_cli, codex_cli, openai, gemini, openai_compatible # model = "claude-sonnet-4-6" # all = false # run every task instead of just the first # task_ids = ["fix_bug", "0"] # slugs or 0-based indices @@ -21,6 +21,10 @@ [claude] # max_tokens = 16384 +[claude_cli] + +[codex_cli] + [openai] # temperature = 0.7 diff --git a/environments/coding/.hud_eval.toml b/environments/coding/.hud_eval.toml index 2315a717f..559160271 100644 --- a/environments/coding/.hud_eval.toml +++ b/environments/coding/.hud_eval.toml @@ -4,7 +4,7 @@ [eval] # source = "tasks.py" # tasks file, or a platform taskset name / id -# agent = "claude" # claude, openai, gemini, openai_compatible +# agent = "claude" # claude, claude_cli, codex_cli, openai, gemini, openai_compatible # model = "claude-sonnet-4-6" # all = false # run every task instead of just the first # task_ids = ["fix_bug", "0"] # slugs or 0-based indices @@ -21,6 +21,10 @@ [claude] # max_tokens = 16384 +[claude_cli] + +[codex_cli] + [openai] # temperature = 0.7 diff --git a/environments/cua/.hud_eval.toml b/environments/cua/.hud_eval.toml index 2315a717f..559160271 100644 --- a/environments/cua/.hud_eval.toml +++ b/environments/cua/.hud_eval.toml @@ -4,7 +4,7 @@ [eval] # source = "tasks.py" # tasks file, or a platform taskset name / id -# agent = "claude" # claude, openai, gemini, openai_compatible +# agent = "claude" # claude, claude_cli, codex_cli, openai, gemini, openai_compatible # model = "claude-sonnet-4-6" # all = false # run every task instead of just the first # task_ids = ["fix_bug", "0"] # slugs or 0-based indices @@ -21,6 +21,10 @@ [claude] # max_tokens = 16384 +[claude_cli] + +[codex_cli] + [openai] # temperature = 0.7 diff --git a/hud/__init__.py b/hud/__init__.py index fdcdc7b0f..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 ( @@ -38,6 +39,7 @@ __all__ = [ "Chat", "ComposeProject", + "Connection", "DockerRuntime", "Environment", "Grade", diff --git a/hud/agents/__init__.py b/hud/agents/__init__.py index ecb02b7e5..3545a11f7 100644 --- a/hud/agents/__init__.py +++ b/hud/agents/__init__.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, cast +from hud.agents.registry import dump_agent, load_agent from hud.settings import settings from hud.types import AgentType from hud.utils.exceptions import HudAuthenticationError @@ -15,7 +16,8 @@ if TYPE_CHECKING: from typing import TypeAlias - from hud.agents.claude import ClaudeAgent, ClaudeSDKAgent, ClaudeSDKConfig + from hud.agents.claude import ClaudeAgent, ClaudeCLIAgent, ClaudeCLIConfig + from hud.agents.codex import CodexCLIAgent, CodexCLIConfig from hud.agents.gemini import GeminiAgent from hud.agents.openai import OpenAIAgent from hud.agents.openai_compatible import OpenAIChatAgent @@ -47,11 +49,15 @@ def create_agent(model: str, **kwargs: Any) -> GatewayAgent: raise HudAuthenticationError("HUD_API_KEY is required to create a gateway agent") agent_type, model_id = resolve_agent_model(model) + if agent_type.is_cli: + raise ValueError( + f"create_agent only constructs provider API agents; instantiate " + f"{agent_type.cls.__name__} directly" + ) kwargs.setdefault("model", model_id) kwargs["gateway"] = True - # cls/config_cls are matched unions; the pairing is correct by construction. config = agent_type.config_cls(**kwargs) - return agent_type.cls(cast("Any", config)) + return cast("GatewayAgent", agent_type.instantiate(config)) def resolve_agent_model(model: str) -> tuple[AgentType, str]: @@ -82,8 +88,10 @@ def resolve_agent_model(model: str) -> tuple[AgentType, str]: _LAZY_EXPORTS = { "ClaudeAgent": ("hud.agents.claude", "ClaudeAgent"), - "ClaudeSDKAgent": ("hud.agents.claude", "ClaudeSDKAgent"), - "ClaudeSDKConfig": ("hud.agents.claude", "ClaudeSDKConfig"), + "ClaudeCLIAgent": ("hud.agents.claude", "ClaudeCLIAgent"), + "ClaudeCLIConfig": ("hud.agents.claude", "ClaudeCLIConfig"), + "CodexCLIAgent": ("hud.agents.codex", "CodexCLIAgent"), + "CodexCLIConfig": ("hud.agents.codex", "CodexCLIConfig"), "GeminiAgent": ("hud.agents.gemini", "GeminiAgent"), "MCPAgent": ("hud.agents.tool_agent", "ToolAgent"), "OpenAIAgent": ("hud.agents.openai", "OpenAIAgent"), @@ -92,13 +100,17 @@ def resolve_agent_model(model: str) -> tuple[AgentType, str]: __all__ = [ "ClaudeAgent", - "ClaudeSDKAgent", - "ClaudeSDKConfig", + "ClaudeCLIAgent", + "ClaudeCLIConfig", + "CodexCLIAgent", + "CodexCLIConfig", "GeminiAgent", "MCPAgent", "OpenAIAgent", "OpenAIChatAgent", "create_agent", + "dump_agent", + "load_agent", "resolve_agent_model", ] diff --git a/hud/agents/base.py b/hud/agents/base.py index 49bdbb499..d9d44814b 100644 --- a/hud/agents/base.py +++ b/hud/agents/base.py @@ -10,13 +10,13 @@ class Agent(ABC): - """Drives a live ``Run`` to completion by filling ``run.trace`` in place. + """Drives a live ``Run`` by recording its trajectory and final answer. Subclasses implement ``__call__(run)``; callers do ``await agent(run)``. Stateless per run — everything comes from ``run`` — so one instance drives many concurrent - rollouts. + rollouts. The caller owns lifecycle status, cancellation, and grading. """ @abstractmethod async def __call__(self, run: Run) -> None: - """Drive ``run`` to completion, filling ``run.trace`` (answer is ``trace.content``).""" + """Fill ``run.trace`` with the trajectory and final answer.""" diff --git a/hud/agents/browser_use/agent.py b/hud/agents/browser_use/agent.py index c9d3c8869..43cae6d37 100644 --- a/hud/agents/browser_use/agent.py +++ b/hud/agents/browser_use/agent.py @@ -24,7 +24,6 @@ from hud.agents.base import Agent from hud.agents.types import AgentStep, BrowserUseConfig from hud.settings import settings -from hud.types import Step if TYPE_CHECKING: from hud.eval.run import Run @@ -107,18 +106,12 @@ async def __call__(self, run: Run) -> None: try: history = await sdk_agent.run(max_steps=self.config.max_steps) - except Exception as exc: - LOGGER.exception("browser-use run failed") - trace.status = "error" - run.record(Step(source="system", error=str(exc))) - return finally: with contextlib.suppress(Exception): await browser.stop() successful = history.is_successful() content = history.final_result() or "" - trace.status = "error" if successful is False else "completed" trace.content = content trace.extra.update( { @@ -133,7 +126,6 @@ async def __call__(self, run: Run) -> None: AgentStep( content=content, done=history.is_done(), - error=content if successful is False else None, ), ) diff --git a/hud/agents/claude/__init__.py b/hud/agents/claude/__init__.py index f5c727565..f61f0add8 100644 --- a/hud/agents/claude/__init__.py +++ b/hud/agents/claude/__init__.py @@ -7,15 +7,15 @@ AsyncAnthropicBedrock, ClaudeAgent, ) -from .sdk import ClaudeSDKAgent, ClaudeSDKConfig +from .sdk import ClaudeCLIAgent, ClaudeCLIConfig from .tools import ClaudeToolSearchTool, ClaudeWebFetchTool, ClaudeWebSearchTool __all__ = [ "AsyncAnthropic", "AsyncAnthropicBedrock", "ClaudeAgent", - "ClaudeSDKAgent", - "ClaudeSDKConfig", + "ClaudeCLIAgent", + "ClaudeCLIConfig", "ClaudeToolSearchTool", "ClaudeWebFetchTool", "ClaudeWebSearchTool", diff --git a/hud/agents/claude/agent.py b/hud/agents/claude/agent.py index 864ef1a2c..8e749973e 100644 --- a/hud/agents/claude/agent.py +++ b/hud/agents/claude/agent.py @@ -360,26 +360,36 @@ async def get_response( if response is None: raise ValueError("Claude response missing after retries") - result = AgentStep(content="", done=True) - result.model = response.model - result.usage = Usage( - prompt_tokens=response.usage.input_tokens, - completion_tokens=response.usage.output_tokens, - cached_tokens=response.usage.cache_read_input_tokens, + return self.message_to_agent_step(response, citations_enabled=citations_enabled) + + @classmethod + def message_to_agent_step( + cls, + response: BetaMessage, + *, + citations_enabled: bool = False, + ) -> AgentStep: + result = AgentStep( + content="", + done=True, + model=response.model, + usage=Usage( + prompt_tokens=response.usage.input_tokens, + completion_tokens=response.usage.output_tokens, + cached_tokens=response.usage.cache_read_input_tokens, + ), ) text_parts: list[str] = [] thinking_parts: list[str] = [] - citations: list[Citation] = [] for block in response.content: match block.type: case "tool_use": - arguments = dict(block.input) if block.input else {} result.tool_calls.append( MCPToolCall( id=block.id, name=block.name, - arguments=arguments, + arguments=dict(block.input) if block.input else {}, _meta=mcp_types.RequestParams.Meta.model_validate( {"citations_enabled": citations_enabled}, ), @@ -387,9 +397,8 @@ async def get_response( ) result.done = False case "text": - text_block = block - text_parts.append(text_block.text) - citations.extend(self._citation(c) for c in (text_block.citations or [])) + text_parts.append(block.text) + result.citations.extend(cls._citation(c) for c in (block.citations or [])) case "thinking": if block.thinking: thinking_parts.append(block.thinking) @@ -397,7 +406,6 @@ async def get_response( pass result.content = "".join(text_parts) - result.citations = citations if thinking_parts: result.reasoning = "\n".join(thinking_parts) result.finish_reason = response.stop_reason diff --git a/hud/agents/claude/sdk/__init__.py b/hud/agents/claude/sdk/__init__.py index 57fd2773c..4511023b2 100644 --- a/hud/agents/claude/sdk/__init__.py +++ b/hud/agents/claude/sdk/__init__.py @@ -1,5 +1,7 @@ -"""Claude Agent SDK agent.""" +"""Agent that runs the ``claude`` CLI over SSH.""" -from .agent import ClaudeSDKAgent, ClaudeSDKConfig +from hud.agents.types import ClaudeCLIConfig -__all__ = ["ClaudeSDKAgent", "ClaudeSDKConfig"] +from .agent import ClaudeCLIAgent + +__all__ = ["ClaudeCLIAgent", "ClaudeCLIConfig"] diff --git a/hud/agents/claude/sdk/agent.py b/hud/agents/claude/sdk/agent.py index 68c3c712a..b8a8e85cd 100644 --- a/hud/agents/claude/sdk/agent.py +++ b/hud/agents/claude/sdk/agent.py @@ -1,76 +1,54 @@ -"""ClaudeSDKAgent — runs ``claude`` CLI over SSH inside the env workspace. +"""ClaudeCLIAgent — runs ``claude`` CLI over SSH inside the env workspace. SSH-execs the ``claude`` CLI on the remote workspace so all built-in tools (Bash, Read, Write, Edit, Glob, Grep) operate on the env's filesystem. MCP capabilities from the manifest are written as MCP server config so the CLI can call env-hosted MCP tools too. - -Inspired by harbor-framework/harbor's ClaudeCode agent. """ from __future__ import annotations -import base64 import json import logging import shlex from contextlib import AsyncExitStack -from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast +import asyncssh + from hud.agents.base import Agent -from hud.agents.types import AgentStep, ClaudeSDKConfig, Usage +from hud.agents.cli import ( + WINDOWS_SHELLS, + powershell, + powershell_quote, + resolve_executable, + run_jsonl, +) +from hud.agents.types import ClaudeCLIConfig from hud.settings import settings -from hud.telemetry.context import get_current_trace_id, get_trace_headers -from hud.types import Step +from hud.telemetry.context import get_trace_headers +from hud.utils.time import now_iso + +from . import computer_mcp +from .events import ClaudeEvents if TYPE_CHECKING: - from hud.capabilities import SSHClient + from hud.capabilities import Connection, SSHClient from hud.eval.run import Run logger = logging.getLogger(__name__) -WINDOWS_SHELLS = ("cmd", "powershell") -#: Bare ``claude`` install bootstrap for POSIX shells (no-op when already present). -_POSIX_INSTALL_CHECK = ( - "command -v claude >/dev/null 2>&1 || " - "{ curl -fsSL https://claude.ai/install.sh | bash -s -- 2>/dev/null; " - 'export PATH="$HOME/.local/bin:$PATH"; }' -) - - -@dataclass(slots=True) -class RemoteInvocation: - """How to run an assembled CLI command on the remote workspace shell. - - ``command`` is what gets exec'd over SSH. When ``script_name`` is set, that - file must be written (with ``script_body``) before exec'ing ``command``. - """ - - command: str - script_name: str | None = None - script_body: str | None = None +INPUT_PATH = ".hud_input.jsonl" +MCP_CONFIG_PATH = ".hud_mcp_config.json" +RUN_SCRIPT_PATH = ".hud_run.bat" +_MANAGED_CLAUDE_PATHS = { + "linux-x64": "/usr/local/lib/agents/claude/linux-x64/claude", + "linux-x64-musl": "/usr/local/lib/agents/claude/linux-x64-musl/claude", +} -def build_remote_invocation(shell: str, run_cmd: str) -> RemoteInvocation: - """Build the remote exec command for ``run_cmd`` under the given login shell. - Windows shells can't take the assembled command inline — ``cmd.exe`` mangles - the quotes — so it is written to a batch file and invoked through ``cmd /c``. - A bare ``.hud_run.bat`` is rejected as an unknown command, and silently fails - to run under a PowerShell default shell, so ``cmd /c`` is required for both. - POSIX shells take the command inline, prefixed with a one-shot install check. - """ - if shell in WINDOWS_SHELLS: - return RemoteInvocation( - command="cmd /c .hud_run.bat", - script_name=".hud_run.bat", - script_body=f"@echo off\r\n{run_cmd}\r\n", - ) - return RemoteInvocation(command=f"{_POSIX_INSTALL_CHECK} && {run_cmd}") - - -class ClaudeSDKAgent(Agent): +class ClaudeCLIAgent(Agent): """Runs ``claude`` CLI over SSH inside the env workspace. Stateless w.r.t. the env: driven by ``await agent(run)``. SSH is opened @@ -78,22 +56,24 @@ class ClaudeSDKAgent(Agent): servers are bridged over the run's SSH connection. """ - config: ClaudeSDKConfig + config: ClaudeCLIConfig - def __init__(self, config: ClaudeSDKConfig | None = None) -> None: - self.config = config or ClaudeSDKConfig() + def __init__(self, config: ClaudeCLIConfig | None = None) -> None: + self.config = config or ClaudeCLIConfig() async def __call__(self, run: Run) -> None: mcp_servers: dict[str, dict[str, Any]] = {} - trace_id = get_current_trace_id() - manifest = run.client.manifest - bindings = manifest.bindings if manifest is not None else [] - families = {c.protocol.split("/", 1)[0] for c in bindings} - - if "ssh" not in families: - raise RuntimeError("ClaudeSDKAgent requires an SSH capability") ssh = cast("SSHClient", await run.client.open("ssh")) + manifest = run.client.manifest + assert manifest is not None + bindings = manifest.bindings shell = ssh.capability.params.get("shell", "bash") + executable = await resolve_executable( + ssh, + "claude", + _MANAGED_CLAUDE_PATHS, + run.runtime_config, + ) rfb_bindings = [cap for cap in bindings if cap.protocol.split("/", 1)[0] == "rfb"] async with AsyncExitStack() as resources: @@ -103,17 +83,12 @@ async def __call__(self, run: Run) -> None: token = cap.params.get("auth_token") transport = "http" if cap.params["transport"] == "streamable-http" else "sse" server_config: dict[str, Any] = {"type": transport, "url": cap.url} - headers = {"Trace-Id": trace_id} if trace_id is not None else {} if token: - headers["Authorization"] = f"Bearer {token}" - if headers: - server_config["headers"] = headers + server_config["headers"] = {"Authorization": f"Bearer {token}"} if cap.name in mcp_servers: raise RuntimeError(f"duplicate MCP server name {cap.name!r}") mcp_servers[cap.name] = server_config elif family == "rfb": - from hud.agents.claude.sdk.computer_mcp import bridge_computer_mcp - server_name = ( "computer-use" if len(rfb_bindings) == 1 else f"computer-use-{cap.name}" ) @@ -121,7 +96,7 @@ async def __call__(self, run: Run) -> None: raise RuntimeError(f"duplicate MCP server name {server_name!r}") routed = run.client.binding(cap.name) mcp_servers[server_name] = await resources.enter_async_context( - bridge_computer_mcp( + computer_mcp.bridge_computer_mcp( ssh, routed, self.config.screenshot_encoding, @@ -135,8 +110,8 @@ async def __call__(self, run: Run) -> None: shell=shell, mcp_servers=mcp_servers, prompt=run.prompt_text, - max_steps=self.config.max_steps, - system_prompt=self.config.system_prompt, + executable=executable, + connection=run.connections.get("inference"), ) async def _exec( @@ -147,58 +122,86 @@ async def _exec( shell: str, mcp_servers: dict[str, dict[str, Any]], prompt: str, - max_steps: int = -1, - system_prompt: str | None = None, + executable: str = "claude", + connection: Connection | None = None, ) -> None: mcp_config_path = await self._write_mcp_config(ssh, mcp_servers) + input_text = ( + json.dumps( + { + "type": "user", + "message": { + "role": "user", + "content": [{"type": "text", "text": prompt}], + }, + } + ) + + "\n" + ) + files = [mcp_config_path] if mcp_config_path else [] + if shell in WINDOWS_SHELLS: + await ssh.write_text(INPUT_PATH, input_text) + files.append(INPUT_PATH) - await ssh.write_text(".hud_prompt.txt", prompt) - - run_cmd = self._build_cli_command( + command = self._build_cli_command( shell=shell, - prompt=prompt, - max_steps=max_steps, - system_prompt=system_prompt, mcp_config_path=mcp_config_path, + executable=executable, + connection=connection, ) - - invocation = build_remote_invocation(shell, run_cmd) - if invocation.script_name is not None: - assert invocation.script_body is not None - # cmd.exe mangles inline quotes, so the command rides a batch file. - await ssh.write_text(invocation.script_name, invocation.script_body) - - full_cmd = invocation.command - logger.info("SSH exec claude CLI (%d chars)", len(full_cmd)) - logger.info("Full command: %s", full_cmd) - - completed = await ssh.run(full_cmd, check=False) - stdout = completed.stdout if isinstance(completed.stdout, str) else "" - stderr = completed.stderr if isinstance(completed.stderr, str) else "" - returncode = completed.returncode - - logger.info("returncode=%s stdout=%d stderr=%d", returncode, len(stdout), len(stderr)) - - if returncode != 0 and not stdout.strip(): - error = stderr or f"claude CLI exited with return code {returncode}" - run.trace.status = "error" - run.trace.extra.update({"returncode": returncode, "stderr": stderr}) - run.record(Step(source="system", error=error)) - return - - self._parse_stream_json(run, stdout, stderr) - - def _build_env_vars(self) -> dict[str, str]: + if shell in WINDOWS_SHELLS: + await ssh.write_text(RUN_SCRIPT_PATH, f"@echo off\r\n{command}\r\n") + files.append(RUN_SCRIPT_PATH) + command = f"cmd /c {RUN_SCRIPT_PATH}" + + try: + logger.info("SSH exec claude CLI (%d chars)", len(command)) + events = ClaudeEvents(run, started_at=now_iso()) + returncode, stderr = await run_jsonl( + ssh, + 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) + finally: + if files: + if shell in WINDOWS_SHELLS: + cleanup = f"cmd /c del /f /q {' '.join(files)} 2>nul" + else: + cleanup = "rm -f -- " + " ".join(shlex.quote(path) for path in files) + try: + await ssh.run(cleanup, check=False) + except (OSError, asyncssh.Error): + logger.warning("Failed to remove Claude CLI runtime files") + + def _build_env_vars(self, connection: Connection | None = None) -> dict[str, str]: env: dict[str, str] = {} - - if settings.api_key: - env["ANTHROPIC_BASE_URL"] = settings.hud_gateway_url - env["ANTHROPIC_API_KEY"] = settings.api_key - trace_headers = get_trace_headers() - if trace_headers: - env["ANTHROPIC_CUSTOM_HEADERS"] = "\n".join( - f"{name}: {value}" for name, value in trace_headers.items() - ) + use_hud_gateway = self.config.use_hud_gateway + if use_hud_gateway is None: + use_hud_gateway = connection is not None or settings.api_key is not None + + if use_hud_gateway: + 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 + else: + raise ValueError("HUD_API_KEY is required for HUD gateway routing") + 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 connection is None: + trace_headers = get_trace_headers() + if trace_headers: + env["ANTHROPIC_CUSTOM_HEADERS"] = "\n".join( + f"{name}: {value}" for name, value in trace_headers.items() + ) elif settings.anthropic_api_key: env["ANTHROPIC_API_KEY"] = settings.anthropic_api_key @@ -214,8 +217,8 @@ def _build_env_vars(self) -> dict[str, str]: env["CLAUDE_CODE_SUBAGENT_MODEL"] = self.config.model env["CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"] = "1" + env["DISABLE_AUTOUPDATER"] = "1" env["IS_SANDBOX"] = "1" - return env async def _write_mcp_config( @@ -227,7 +230,7 @@ async def _write_mcp_config( if not mcp_servers: return None mcp_json = json.dumps({"mcpServers": mcp_servers}, indent=2) - path = ".hud_mcp_config.json" + path = MCP_CONFIG_PATH await ssh.write_text(path, mcp_json) logger.info("Wrote MCP config") return path @@ -236,120 +239,48 @@ def _build_cli_command( self, *, shell: str, - prompt: str, - max_steps: int, - system_prompt: str | None, mcp_config_path: str | None = None, + executable: str = "claude", + connection: Connection | None = None, ) -> str: - env_vars = self._build_env_vars() + env_vars = self._build_env_vars(connection) is_win = shell in WINDOWS_SHELLS - - # Raw args list (no shell quoting) — used directly for Windows Python launcher. base_args: list[str] = [ - "claude", + executable, "--verbose", + "--input-format=stream-json", "--output-format=stream-json", "--print", f"--permission-mode={self.config.permission_mode}", ] - if max_steps > 0: - base_args.append(f"--max-turns={max_steps}") - if system_prompt: - base_args.extend(["--system-prompt", system_prompt]) + if self.config.max_steps > 0: + base_args.append(f"--max-turns={self.config.max_steps}") + if self.config.system_prompt: + base_args.extend(["--system-prompt", self.config.system_prompt]) for tool in self.config.allowed_tools: base_args.extend(["--allowedTools", tool]) if mcp_config_path: base_args.extend(["--mcp-config", mcp_config_path]) if is_win: - # On Windows, two problems combine: - # 1. claude is installed as claude.cmd (Node.js wrapper) — Python's - # subprocess.run can't execute .cmd files via CreateProcess directly. - # 2. Embedding the prompt inline in the bat file breaks — cmd.exe parses - # line-by-line, so newlines inside quoted strings split the command. - # Solution: use `cmd /c claude [args]` (no inline prompt) and feed the - # prompt via stdin from .hud_prompt.txt. claude --print reads stdin as - # the initial message when no -- argument is provided. - cmd_args = ["cmd", "/c", "claude"] + base_args[1:] # noqa: RUF005 - py_args_repr = "[" + ",".join(f"'{a}'" for a in cmd_args) + "]" - encoded_env = base64.b64encode(json.dumps(env_vars).encode()).decode() - python_launcher = ( - 'python -c "' - "import base64,json,os,subprocess,sys;" - "env=os.environ.copy();" - f"env.update(json.loads(base64.b64decode('{encoded_env}')));" - f"r=subprocess.run({py_args_repr},stdin=open('.hud_prompt.txt','rb'),env=env);" - 'sys.exit(r.returncode)"' + script = ";".join( + [ + *(f"$env:{key}={powershell_quote(value)}" for key, value in env_vars.items()), + f"Get-Content -Raw -Encoding UTF8 {powershell_quote(INPUT_PATH)}" + f" | & {powershell_quote(executable)} " + f"{' '.join(powershell_quote(arg) for arg in base_args[1:])}", + "exit $LASTEXITCODE", + ] ) - return python_launcher + return powershell(script) - # POSIX path: shell-quote everything and embed prompt inline. cli_parts = [shlex.quote(a) for a in base_args] - cli_parts.extend(["--", shlex.quote(prompt)]) 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}' - - def _parse_stream_json(self, run: Run, stdout: str, stderr: str) -> None: - messages: list[dict[str, Any]] = [] - content_parts: list[str] = [] - is_error = False - info: dict[str, Any] = {} - cost_usd: float | None = None - num_turns: int | None = None - - for line in stdout.splitlines(): - line = line.strip() - if not line: - continue - try: - msg = json.loads(line) - except json.JSONDecodeError: - continue - - messages.append(msg) - msg_type = msg.get("type") - - if msg_type == "assistant" and isinstance(msg.get("message"), dict): - for raw_block in msg["message"].get("content", []): - if not isinstance(raw_block, dict): - continue - block = cast("dict[str, Any]", raw_block) - if block.get("type") == "text" and block.get("text"): - content_parts.append(str(block["text"])) - - elif msg_type == "result": - is_error = msg.get("is_error", False) - result_text = msg.get("result") - if result_text: - content_parts.append(result_text) - info["session_id"] = msg.get("session_id") - info["duration_ms"] = msg.get("duration_ms") - info["stop_reason"] = msg.get("stop_reason") - num_turns = msg.get("num_turns") - cost_usd = msg.get("total_cost_usd") - - content = "\n".join(content_parts) - trace = run.trace - trace.status = "error" if is_error else "completed" - trace.content = content - # Raw CLI stream kept locally; a claude-native serializer can take over - # per-turn fidelity later (the CLI session is its own span vocabulary). - trace.extra["messages"] = messages - if stderr: - trace.extra["stderr"] = stderr - - # The CLI run collapses to one coarse agent step with aggregate usage. - run.record( - AgentStep( - content=content, - done=True, - model=self.config.model, - usage=Usage(cost_usd=cost_usd, llm_call_count=num_turns), - error=content if is_error else None, - extra={k: v for k, v in info.items() if v is not None}, - ), - ) + 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__ = ["ClaudeSDKAgent", "ClaudeSDKConfig", "RemoteInvocation", "build_remote_invocation"] +__all__ = ["ClaudeCLIAgent"] diff --git a/hud/agents/claude/sdk/computer_mcp.py b/hud/agents/claude/sdk/computer_mcp.py index 84d0cc8a0..ad51496e4 100644 --- a/hud/agents/claude/sdk/computer_mcp.py +++ b/hud/agents/claude/sdk/computer_mcp.py @@ -19,8 +19,11 @@ import asyncssh import fastmcp +from fastmcp.exceptions import ToolError from pydantic import TypeAdapter +from hud.agents.claude.tools.computer import ClaudeComputerTool +from hud.agents.tools.base import AgentToolSpec, result_text from hud.capabilities import Capability from hud.capabilities.rfb import RFBClient, ScreenshotEncoding, WebPScreenshotEncoding @@ -46,18 +49,23 @@ def create_computer_mcp( """Build a FastMCP server with one ``computer`` tool backed by ``rfb``.""" mcp = fastmcp.FastMCP("computer-use") + tool = ClaudeComputerTool( + spec=AgentToolSpec(api_type="computer", api_name="computer"), + client=rfb, + screenshot_encoding=screenshot_encoding, + ) @mcp.tool() async def computer( action: str, - coordinate: str | None = None, + coordinate: list[int] | None = None, text: str | None = None, scroll_direction: str | None = None, scroll_amount: int | None = None, - start_coordinate: str | None = None, + start_coordinate: list[int] | None = None, duration: float | None = None, repeat: int | None = None, - region: str | None = None, + region: list[int] | None = None, ) -> list[Any]: """Control a remote screen — screenshot, click, type, key, scroll, move, drag, wait, zoom. @@ -67,85 +75,34 @@ async def computer( Returns the resulting screenshot image so you can see the screen state. """ - import mcp.types as mcp_types - - from hud.agents.claude.tools.computer import ClaudeComputerTool - from hud.agents.tools.base import AgentToolSpec - - arguments: dict[str, Any] = {"action": action} - if coordinate is not None: - try: - arguments["coordinate"] = json.loads(coordinate) - except json.JSONDecodeError: - arguments["coordinate"] = coordinate - if text is not None: - arguments["text"] = text - if scroll_direction is not None: - arguments["scroll_direction"] = scroll_direction - if scroll_amount is not None: - arguments["scroll_amount"] = scroll_amount - if start_coordinate is not None: - try: - arguments["start_coordinate"] = json.loads(start_coordinate) - except json.JSONDecodeError: - arguments["start_coordinate"] = start_coordinate - if duration is not None: - arguments["duration"] = duration - if repeat is not None: - arguments["repeat"] = repeat - if region is not None: - try: - arguments["region"] = json.loads(region) - except json.JSONDecodeError: - arguments["region"] = region - - spec = AgentToolSpec(api_type="computer", api_name="computer") - tool = ClaudeComputerTool( - spec=spec, - client=rfb, - screenshot_encoding=screenshot_encoding, - ) + arguments = { + name: value + for name, value in { + "action": action, + "coordinate": coordinate, + "text": text, + "scroll_direction": scroll_direction, + "scroll_amount": scroll_amount, + "start_coordinate": start_coordinate, + "duration": duration, + "repeat": repeat, + "region": region, + }.items() + if value is not None + } result = await tool.execute(arguments) - - # Return content blocks directly so the CLI/model sees real images. - blocks: list[Any] = [] - for block in result.content: - if isinstance(block, mcp_types.ImageContent): - blocks.append( - mcp_types.ImageContent( - type="image", - data=block.data, - mimeType=block.mimeType, - ), - ) - elif isinstance(block, mcp_types.TextContent): - blocks.append(mcp_types.TextContent(type="text", text=block.text)) - if not blocks: - blocks.append(mcp_types.TextContent(type="text", text="ok")) if result.isError: - blocks.insert(0, mcp_types.TextContent(type="text", text="ERROR")) - return blocks + raise ToolError(result_text(result) or "computer action failed") + return result.content return mcp -def _required_env(environ: Mapping[str, str], name: str) -> str: - try: - return environ[name] - except KeyError as exc: - raise RuntimeError(f"missing required environment variable {name}") from exc - - async def run_computer_mcp(environ: Mapping[str, str] = os.environ) -> None: """Run computer-use over stdio in a controller-side child process.""" - raw_manifest = json.loads(_required_env(environ, RFB_CAPABILITY_ENV)) - if not isinstance(raw_manifest, dict): - raise ValueError(f"{RFB_CAPABILITY_ENV} must contain a JSON object") - capability = Capability.from_manifest(raw_manifest) - if capability.protocol.split("/", 1)[0] != "rfb": - raise ValueError(f"{RFB_CAPABILITY_ENV} must describe an RFB capability") + capability = Capability.from_manifest(json.loads(environ[RFB_CAPABILITY_ENV])) screenshot_encoding = TypeAdapter(ScreenshotEncoding).validate_json( - _required_env(environ, SCREENSHOT_ENCODING_ENV) + environ[SCREENSHOT_ENCODING_ENV] ) rfb = await RFBClient.connect(capability) @@ -168,7 +125,7 @@ async def bridge_computer_mcp( ) -> AsyncIterator[dict[str, Any]]: """Bridge a controller-side computer MCP process into a remote POSIX shell.""" if shell in {"cmd", "powershell"}: - raise RuntimeError("ClaudeSDKAgent computer use requires a POSIX workspace") + raise RuntimeError("ClaudeCLIAgent computer use requires a POSIX workspace") token = secrets.token_hex(16) request_path = str(_REMOTE_TMP / f"hud-computer-{token}.request") diff --git a/hud/agents/claude/sdk/events.py b/hud/agents/claude/sdk/events.py new file mode 100644 index 000000000..737cb8776 --- /dev/null +++ b/hud/agents/claude/sdk/events.py @@ -0,0 +1,135 @@ +"""Claude CLI stream translation.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +import mcp.types as mcp_types +from anthropic.types.beta import BetaMessage + +from hud.agents.claude.agent import ClaudeAgent +from hud.agents.types import ToolStep +from hud.types import MCPToolCall, MCPToolResult +from hud.utils.time import now_iso + +if TYPE_CHECKING: + from hud.eval.run import Run + + +class ClaudeEvents: + """Translate Claude CLI stream messages into canonical HUD steps.""" + + def __init__(self, run: Run, *, started_at: str) -> None: + self.run = run + self.agent_started_at = started_at + self.pending_calls: dict[str, tuple[MCPToolCall, str]] = {} + self.saw_result = False + self.error: str | None = None + + def consume(self, line: str) -> None: + line = line.strip() + if not line: + return + message = json.loads(line) + if not isinstance(message, dict): + raise ValueError("Claude stream event must be an object") + received_at = now_iso() + match message.get("type"): + case "system" if message.get("subtype") == "init": + self.agent_started_at = received_at + case "assistant": + step = ClaudeAgent.message_to_agent_step( + BetaMessage.model_validate(message["message"]) + ) + step.started_at = self.agent_started_at + step.ended_at = received_at + if step.content: + self.run.trace.content = step.content + self.run.record(step) + for call in step.tool_calls: + self.pending_calls[call.id] = (call, received_at) + case "user": + saw_result = False + for block in message["message"]["content"]: + if block["type"] != "tool_result": + continue + call_id = block["tool_use_id"] + try: + call, started_at = self.pending_calls.pop(call_id) + except KeyError: + raise ValueError( + f"Claude returned a result for unknown tool call {call_id!r}" + ) from None + + raw_result = block.get("content") + raw_items = raw_result if isinstance(raw_result, list) else [raw_result] + content: list[mcp_types.ContentBlock] = [] + for item in raw_items: + if isinstance(item, str): + content.append(mcp_types.TextContent(type="text", text=item)) + elif item["type"] == "text": + content.append(mcp_types.TextContent(type="text", text=item["text"])) + elif item["type"] == "image": + source = item["source"] + content.append( + mcp_types.ImageContent( + type="image", + data=source["data"], + mimeType=source["media_type"], + ) + ) + else: + raise ValueError(f"unsupported Claude tool result block: {item!r}") + + self.run.record( + ToolStep( + call=call, + result=MCPToolResult( + call_id=call_id, + content=content, + isError=block.get("is_error") is True, + ), + started_at=started_at, + ended_at=received_at, + ) + ) + saw_result = True + if saw_result: + self.agent_started_at = received_at + case "result": + self.saw_result = True + trace = self.run.trace + result = message.get("result") + if isinstance(result, str): + trace.content = result + if message.get("is_error") is True: + self.error = trace.content or "claude CLI reported an error" + for key in ( + "subtype", + "session_id", + "duration_ms", + "duration_api_ms", + "stop_reason", + "num_turns", + "total_cost_usd", + ): + if (value := message.get(key)) is not None: + trace.extra[key] = value + + def finish(self, *, returncode: int, stderr: str) -> None: + trace = self.run.trace + error = self.error + if returncode != 0: + trace.extra["returncode"] = returncode + error = stderr.strip() or f"claude CLI exited with return code {returncode}" + elif not self.saw_result: + error = "claude CLI exited without a result event" + elif self.pending_calls: + missing = ", ".join(sorted(self.pending_calls)) + error = f"claude CLI exited without results for tool calls: {missing}" + + if error is not None and stderr: + trace.extra["stderr"] = stderr + if error is not None: + raise RuntimeError(error) diff --git a/hud/agents/cli.py b/hud/agents/cli.py new file mode 100644 index 000000000..92e8197f0 --- /dev/null +++ b/hud/agents/cli.py @@ -0,0 +1,164 @@ +"""Process boundary for JSONL CLI agents.""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import shlex +from typing import TYPE_CHECKING + +import asyncssh + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from hud.capabilities import Connection, SSHClient + from hud.eval.runtime import RuntimeConfig + +WINDOWS_SHELLS = ("cmd", "powershell") +PROCESS_CLOSE_TIMEOUT_S = 5.0 + + +async def resolve_executable( + ssh: SSHClient, + command: str, + managed_paths: dict[str, str], + runtime_config: RuntimeConfig | None, +) -> str: + """Resolve a CLI against the live SSH target and its declared runtime config.""" + platform = await _runtime_platform(ssh) + _validate_runtime_os(runtime_config, platform.partition("-")[0]) + + managed = managed_paths.get(platform) + if managed is not None: + result = await ssh.run( + f"test -x {shlex.quote(managed)}", + check=False, + encoding=None, + ) + if result.returncode == 0: + return managed + + if platform.startswith("windows-"): + result = await ssh.run(f"where.exe {command}", check=False, encoding=None) + else: + result = await ssh.run(f"command -v -- {command}", check=False, encoding=None) + if result.returncode == 0: + stdout = _output_text(result.stdout) + if path := stdout.splitlines()[0].strip(): + return path + + raise RuntimeError( + f"{command} is unavailable for runtime platform {platform}; " + "install it in the environment or provide a managed runtime bundle" + ) + + +async def _runtime_platform(ssh: SSHClient) -> str: + shell = ssh.capability.params.get("shell", "bash") + if shell in WINDOWS_SHELLS: + result = await ssh.run( + powershell("[System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture"), + check=True, + encoding=None, + ) + arch = _output_text(result.stdout).strip().lower() + return f"windows-{_normalize_arch(arch)}" + + result = await ssh.run( + "uname -s; uname -m; " + "if ls /lib/ld-musl-*.so.1 >/dev/null 2>&1; then echo musl; else echo gnu; fi", + check=True, + encoding=None, + ) + lines = _output_text(result.stdout).splitlines() + if len(lines) != 3: + raise RuntimeError("SSH runtime platform probe returned an invalid response") + system, machine, libc = (line.strip().lower() for line in lines) + os_name = {"darwin": "darwin", "linux": "linux"}.get(system) + if os_name is None: + raise RuntimeError(f"unsupported SSH runtime operating system {system!r}") + platform = f"{os_name}-{_normalize_arch(machine)}" + return f"{platform}-musl" if os_name == "linux" and libc == "musl" else platform + + +def _normalize_arch(value: str) -> str: + normalized = { + "amd64": "x64", + "x86_64": "x64", + "arm64": "arm64", + "aarch64": "arm64", + }.get(value) + if normalized is None: + raise RuntimeError(f"unsupported SSH runtime architecture {value!r}") + return normalized + + +def _output_text(value: bytes | str | None) -> str: + if isinstance(value, bytes): + return value.decode(errors="replace") + return value or "" + + +def _validate_runtime_os(runtime_config: RuntimeConfig | None, actual: str) -> None: + if runtime_config is None or runtime_config.resources is None: + return + declared = runtime_config.resources.os + if declared is None: + return + normalized = { + "darwin": "darwin", + "linux": "linux", + "macos": "darwin", + "windows": "windows", + }.get(declared.lower()) + if normalized is not None and normalized != actual: + raise RuntimeError( + f"runtime_config.resources.os requested {declared!r}, " + f"but the SSH runtime reports {actual!r}" + ) + + +def powershell_quote(value: str) -> str: + return "'" + value.replace("'", "''") + "'" + + +def powershell(script: str) -> str: + encoded = base64.b64encode(script.encode("utf-16-le")).decode("ascii") + return f"powershell -NoProfile -NonInteractive -EncodedCommand {encoded}" + + +async def run_jsonl( + ssh: SSHClient, + command: str, + 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, connections=connections) + stderr_task = asyncio.create_task(process.stderr.read()) + try: + if input_text is not None: + process.stdin.write(input_text.encode()) + await process.stdin.drain() + process.stdin.write_eof() + while line := await process.stdout.readline(): + consume(line.decode(errors="replace")) + await process.wait_closed() + stderr = (await stderr_task).decode(errors="replace") + except BaseException: + process.close() + if not stderr_task.done(): + stderr_task.cancel() + await asyncio.gather(stderr_task, return_exceptions=True) + with contextlib.suppress(OSError, TimeoutError, asyncssh.Error): + async with asyncio.timeout(PROCESS_CLOSE_TIMEOUT_S): + await process.wait_closed() + raise + + if process.returncode is None: + raise RuntimeError("CLI process closed without an exit status") + return process.returncode, stderr diff --git a/hud/agents/codex/__init__.py b/hud/agents/codex/__init__.py new file mode 100644 index 000000000..0b5e99611 --- /dev/null +++ b/hud/agents/codex/__init__.py @@ -0,0 +1,7 @@ +"""Codex CLI agent.""" + +from hud.agents.types import CodexCLIConfig + +from .agent import CodexCLIAgent + +__all__ = ["CodexCLIAgent", "CodexCLIConfig"] diff --git a/hud/agents/codex/agent.py b/hud/agents/codex/agent.py new file mode 100644 index 000000000..26e4eddf6 --- /dev/null +++ b/hud/agents/codex/agent.py @@ -0,0 +1,364 @@ +"""Codex CLI harness over a workspace SSH capability.""" + +from __future__ import annotations + +import json +import logging +import shlex +from typing import TYPE_CHECKING, Any, cast + +import mcp.types as mcp_types + +from hud.agents.base import Agent +from hud.agents.cli import ( + WINDOWS_SHELLS, + powershell, + powershell_quote, + resolve_executable, + run_jsonl, +) +from hud.agents.types import AgentStep, CodexCLIConfig, ToolStep +from hud.settings import settings +from hud.telemetry.context import get_current_trace_id +from hud.types import MCPToolCall, MCPToolResult, Step +from hud.utils.time import now_iso + +if TYPE_CHECKING: + from hud.capabilities import Connection, SSHClient + from hud.eval.run import Run + +logger = logging.getLogger(__name__) + +_MANAGED_CODEX_PATHS = { + "linux-x64": "/usr/local/lib/agents/codex/bin/codex", + "linux-x64-musl": "/usr/local/lib/agents/codex/bin/codex", +} + + +class CodexEvents: + """Translate ``codex exec --json`` events into canonical HUD steps.""" + + def __init__(self, run: Run, *, model: str, started_at: str) -> None: + self.run = run + self.model = model + self.agent_started_at = started_at + self.item_started_at: dict[str, str] = {} + self.saw_completion = False + self.error: str | None = None + + def consume(self, line: str) -> None: + line = line.strip() + if not line: + return + event = json.loads(line) + if not isinstance(event, dict): + raise ValueError("Codex stream event must be an object") + + received_at = now_iso() + match event.get("type"): + case "thread.started": + self.run.trace.extra["codex_thread_id"] = event["thread_id"] + case "turn.started": + self.agent_started_at = received_at + case "item.started": + item = event["item"] + self.item_started_at[item["id"]] = received_at + case "item.completed": + self.record(event["item"], received_at) + case "turn.completed": + self.run.trace.extra["usage"] = event["usage"] + self.saw_completion = True + case "turn.failed": + self.error = event["error"]["message"] + case "error": + self.error = event["message"] + + def finish(self, *, returncode: int, stderr: str) -> None: + trace = self.run.trace + error = self.error + if returncode != 0: + trace.extra["returncode"] = returncode + error = error or stderr.strip() or f"codex CLI exited with return code {returncode}" + elif error is None and not self.saw_completion: + error = "codex CLI exited without a turn.completed event" + + if error is not None and stderr and self.error is None: + trace.extra["stderr"] = stderr + if error is not None: + raise RuntimeError(error) + + def record(self, item: dict[str, Any], received_at: str) -> None: + item_id = item["id"] + started_at = self.item_started_at.pop(item_id, self.agent_started_at) + match item["type"]: + case "agent_message": + text = item["text"] + self.run.trace.content = text + self.run.record( + AgentStep( + content=text, + model=self.model, + raw=item, + started_at=started_at, + ended_at=received_at, + ) + ) + case "reasoning": + self.run.record( + AgentStep( + reasoning=item["text"], + model=self.model, + raw=item, + started_at=started_at, + ended_at=received_at, + ) + ) + case "command_execution" | "file_change" | "mcp_tool_call" | "web_search": + self.record_tool(item, started_at, received_at) + case _: + self.run.record( + Step( + source="agent", + extra={"codex_item": item}, + started_at=started_at, + ended_at=received_at, + ) + ) + self.agent_started_at = received_at + + def record_tool(self, item: dict[str, Any], started_at: str, ended_at: str) -> None: + call_id = item["id"] + match item["type"]: + case "command_execution": + call = MCPToolCall( + id=call_id, + name="shell", + arguments={"command": item["command"]}, + ) + result = MCPToolResult( + call_id=call_id, + content=[mcp_types.TextContent(type="text", text=item["aggregated_output"])], + isError=item["status"] != "completed" or item["exit_code"] not in (None, 0), + ) + case "file_change": + changes = item["changes"] + call = MCPToolCall( + id=call_id, + name="apply_patch", + arguments={"changes": changes}, + ) + result = MCPToolResult( + call_id=call_id, + content=[ + mcp_types.TextContent( + type="text", + text="\n".join( + f"{change['kind']}: {change['path']}" for change in changes + ), + ) + ], + isError=item["status"] != "completed", + ) + case "mcp_tool_call": + call = MCPToolCall( + id=call_id, + name=item["tool"], + provider_name=f"{item['server']}.{item['tool']}", + arguments=item["arguments"], + ) + raw_result = item.get("result") or {} + error = item.get("error") + result = MCPToolResult.model_validate( + { + "call_id": call_id, + "content": raw_result.get("content") + or ([{"type": "text", "text": error["message"]}] if error else []), + "structuredContent": raw_result.get("structured_content"), + "_meta": raw_result.get("_meta"), + "isError": item["status"] == "failed", + } + ) + case "web_search": + call = MCPToolCall( + id=call_id, + name="web_search", + arguments={"query": item["query"], "action": item["action"]}, + ) + result = MCPToolResult( + call_id=call_id, + content=[ + mcp_types.TextContent( + type="text", + text=json.dumps(item["action"], separators=(",", ":")), + ) + ], + isError=False, + ) + case _: + raise ValueError(f"unsupported Codex tool item {item['type']!r}") + + self.run.record( + ToolStep( + call=call, + result=result, + extra={"codex_item": item}, + started_at=started_at, + ended_at=ended_at, + ) + ) + + +def codex_command( + config: CodexCLIConfig, + shell: str, + executable: str = "codex", + connection: Connection | None = None, +) -> str: + env: dict[str, str] = {} + sandbox = "danger-full-access" if connection is not None else config.sandbox + args = [ + executable, + "exec", + "--json", + "--ephemeral", + "--skip-git-repo-check", + "--color", + "never", + "--sandbox", + sandbox, + "--model", + config.model, + ] + + use_hud_gateway = config.use_hud_gateway + if use_hud_gateway is None: + use_hud_gateway = connection is not None or settings.api_key is not None + if use_hud_gateway: + 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 + credential_env = "HUD_API_KEY" + else: + raise ValueError("HUD_API_KEY is required for HUD gateway routing") + env[credential_env] = credential + overrides = { + "model_provider": "hud", + "model_providers.hud.name": "HUD", + "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 connection is None and (trace_id := get_current_trace_id()): + args.extend( + [ + "-c", + f'model_providers.hud.http_headers={{"Trace-Id"={json.dumps(trace_id)}}}', + ] + ) + elif settings.openai_api_key: + env["CODEX_API_KEY"] = settings.openai_api_key + + args.append("-") + isolate_home = bool(env) + if shell in WINDOWS_SHELLS: + 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 {{ {invocation} }} finally {{ Remove-Item -Recurse -Force $codexHome }}", + "exit $hudExitCode", + ] + 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 + statements = ['export PATH="$HOME/.local/bin:$PATH"', invocation] + 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", + 'export CODEX_HOME="$codex_home"', + *statements, + ] + return "; ".join(statements) + + +async def run_codex( + config: CodexCLIConfig, + run: Run, + *, + ssh: SSHClient, + shell: str, + prompt: str, + executable: str = "codex", + connection: Connection | None = None, +) -> None: + 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, + connections=(connection,) if connection is not None else (), + ) + logger.info("exit=%s stderr=%d", returncode, len(stderr)) + events.finish(returncode=returncode, stderr=stderr) + + +class CodexCLIAgent(Agent): + """Runs ``codex exec`` over SSH inside the environment workspace.""" + + config: CodexCLIConfig + + def __init__(self, config: CodexCLIConfig | None = None) -> None: + self.config = config or CodexCLIConfig() + + async def __call__(self, run: Run) -> None: + ssh = cast("SSHClient", await run.client.open("ssh")) + executable = await resolve_executable( + ssh, + "codex", + _MANAGED_CODEX_PATHS, + run.runtime_config, + ) + await run_codex( + self.config, + run, + ssh=ssh, + shell=ssh.capability.params.get("shell", "bash"), + prompt=run.prompt_text, + executable=executable, + connection=run.connections.get("inference"), + ) + + +__all__ = ["CodexCLIAgent"] diff --git a/hud/agents/openai_compatible/agent.py b/hud/agents/openai_compatible/agent.py index 3b1413e8d..65c1caa32 100644 --- a/hud/agents/openai_compatible/agent.py +++ b/hud/agents/openai_compatible/agent.py @@ -149,23 +149,16 @@ async def get_response( if return_token_ids: request_kwargs.setdefault("logprobs", True) - try: - response: ChatCompletion = await self.oai.chat.completions.create( - model=self.config.model, - messages=( - [{"role": "system", "content": system_prompt}, *messages] - if system_prompt is not None - else messages - ), - stream=False, - **request_kwargs, - ) - except Exception as e: - error_content = f"Error getting response {e}" - if "Invalid JSON" in str(e): - error_content = "Invalid JSON, response was truncated" - logger.warning(error_content) - return AgentStep(error=error_content, done=True) + response: ChatCompletion = await self.oai.chat.completions.create( + model=self.config.model, + messages=( + [{"role": "system", "content": system_prompt}, *messages] + if system_prompt is not None + else messages + ), + stream=False, + **request_kwargs, + ) choice = response.choices[0] message = choice.message diff --git a/hud/agents/registry.py b/hud/agents/registry.py new file mode 100644 index 000000000..347cc2500 --- /dev/null +++ b/hud/agents/registry.py @@ -0,0 +1,56 @@ +"""Serialization and reconstruction for built-in agents.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +from hud.agents.types import AgentConfig, ToolAgentConfig +from hud.types import AgentType + +if TYPE_CHECKING: + from hud.agents.base import Agent + +_RUNTIME_ONLY_CONFIG_FIELDS = { + "model_client", + "gateway", + "api_key", + "base_url", + "hosted_tools", +} + + +def dump_agent(agent: Agent) -> dict[str, Any]: + """Serialize a registered agent without credentials or live clients.""" + agent_type = AgentType.of(agent) + config = getattr(agent, "config", None) + if agent_type is None or not isinstance(config, AgentConfig): + raise ValueError( + f"agent must be one of the registered types " + f"({', '.join(member.value for member in AgentType)}); " + f"got {type(agent).__name__}" + ) + if isinstance(config, ToolAgentConfig) and config.model_client is not None: + raise ValueError( + "agents with a custom model_client cannot run remotely; use HUDRuntime or LocalRuntime" + ) + + payload = config.model_dump( + mode="json", + exclude=_RUNTIME_ONLY_CONFIG_FIELDS, + ) + return {"type": agent_type.value, "config": payload} + + +def load_agent(data: Mapping[str, Any]) -> Agent: + """Reconstruct a registered agent from :func:`dump_agent` output.""" + try: + agent_type = AgentType(data["type"]) + except (KeyError, TypeError, ValueError): + raise ValueError(f"unsupported agent type {data.get('type')!r}") from None + + raw_config = data.get("config") + if not isinstance(raw_config, Mapping): + raise ValueError("agent config must be an object") + config = agent_type.config_cls.model_validate(dict(raw_config)) + return agent_type.instantiate(config) diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index c9f27da10..e10454bcc 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -114,7 +114,6 @@ async def __call__(self, run: Run, *, max_steps: int | None = None) -> None: writer.end_episode() finally: await robot.close() - run.trace.status = "completed" run.trace.content = "done" async def _loop( diff --git a/hud/agents/tests/cli_fakes.py b/hud/agents/tests/cli_fakes.py new file mode 100644 index 000000000..6cd59bab4 --- /dev/null +++ b/hud/agents/tests/cli_fakes.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any + + +class FakeReader: + def __init__(self, value: str, *, pause_after: int | None = None) -> None: + self._raw = value.encode() + self._lines = self._raw.splitlines(keepends=True) + self._pause_after = pause_after + self._index = 0 + self.blocked = asyncio.Event() + self.release = asyncio.Event() + + async def readline(self) -> bytes: + if self._pause_after == self._index: + self.blocked.set() + await self.release.wait() + self._pause_after = None + if self._index == len(self._lines): + return b"" + line = self._lines[self._index] + self._index += 1 + return line + + async def read(self) -> bytes: + return self._raw + + +class FakeWriter: + def __init__(self) -> None: + self.data = bytearray() + self.eof = False + + def write(self, data: bytes) -> None: + self.data.extend(data) + + async def drain(self) -> None: + pass + + def write_eof(self) -> None: + self.eof = True + + +class FakeProcess: + def __init__( + self, + stdout: str, + *, + stderr: str = "", + exit_status: int | None = 0, + returncode: int | None = None, + pause_after: int | None = None, + ) -> None: + self.stdin = FakeWriter() + self.stdout = FakeReader(stdout, pause_after=pause_after) + self.stderr = FakeReader(stderr) + self.exit_status = exit_status + self.returncode = exit_status if returncode is None else returncode + self.closed = False + self.terminated = False + + def terminate(self) -> None: + self.terminated = True + + def close(self) -> None: + self.closed = True + + async def wait_closed(self) -> None: + pass + + +def fake_run() -> Any: + trace = SimpleNamespace(status=None, content="", extra={}) + steps: list[Any] = [] + return SimpleNamespace(trace=trace, record=steps.append, steps=steps) diff --git a/hud/agents/tests/test_base.py b/hud/agents/tests/test_base.py index 80269dd66..4b9051ece 100644 --- a/hud/agents/tests/test_base.py +++ b/hud/agents/tests/test_base.py @@ -12,7 +12,15 @@ import pytest -from hud.agents import OpenAIAgent, OpenAIChatAgent, create_agent +from hud.agents import ( + ClaudeCLIAgent, + CodexCLIAgent, + OpenAIAgent, + OpenAIChatAgent, + create_agent, + dump_agent, + load_agent, +) from hud.agents.base import Agent from hud.agents.types import OpenAIConfig from hud.types import AgentType @@ -24,9 +32,6 @@ async def __call__(self, run: Any) -> None: run.trace.content = "done" -# ─── the ABC contract ───────────────────────────────────────────────── - - def test_agent_requires_call_implementation() -> None: with pytest.raises(TypeError): Agent() @@ -40,15 +45,48 @@ async def test_agent_call_fills_trace() -> None: assert run.trace.content == "done" -# ─── AgentType resolution ───────────────────────────────────────────── - - def test_agent_type_maps_value_to_class_and_provider() -> None: assert AgentType("openai").cls is OpenAIAgent assert AgentType("openai_compatible").cls is OpenAIChatAgent assert isinstance(AgentType("openai").gateway_provider, str) +def test_agent_type_registers_cli_agent() -> None: + assert AgentType("claude_cli").cls is ClaudeCLIAgent + assert AgentType.of(ClaudeCLIAgent()) == AgentType.CLAUDE_CLI + assert AgentType("codex_cli").cls is CodexCLIAgent + assert AgentType.of(CodexCLIAgent()) == AgentType.CODEX_CLI + + +def test_cli_agent_round_trips_through_registered_wire_format() -> None: + agent = ClaudeCLIAgent() + spec = dump_agent(agent) + + loaded = load_agent(spec) + + assert spec["type"] == "claude_cli" + assert spec["config"]["model"] == "claude-sonnet-5" + assert isinstance(loaded, ClaudeCLIAgent) + assert loaded.config == agent.config + + +def test_codex_cli_agent_round_trips_through_registry() -> None: + agent = CodexCLIAgent() + spec = dump_agent(agent) + + loaded = load_agent(spec) + + assert spec["type"] == "codex_cli" + assert spec["config"]["model"] == "gpt-5.6-sol" + assert isinstance(loaded, CodexCLIAgent) + assert loaded.config == agent.config + + +def test_dump_agent_rejects_unregistered_agent() -> None: + with pytest.raises(ValueError, match="registered types"): + dump_agent(_FillingAgent()) + + def test_missing_provider_dependency_points_at_agents_extra( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -77,9 +115,6 @@ def find_spec(self, fullname: str, path: Any = None, target: Any = None) -> None _ = AgentType.CLAUDE.cls -# ─── create_agent routing ───────────────────────────────────────────── - - @pytest.fixture(autouse=True) def gateway_api_key(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("hud.agents.settings.api_key", "test-key") @@ -107,7 +142,7 @@ def _build_client(_provider: str) -> object: assert isinstance(agent, OpenAIAgent) assert agent.config.model_client is None assert agent.openai_client is sentinel - assert agent.hosted_spec()["config"]["prompt_cache_key"] == agent.config.prompt_cache_key + assert dump_agent(agent)["config"]["prompt_cache_key"] == agent.config.prompt_cache_key def test_create_agent_uses_the_gateway_even_with_a_provider_key( @@ -127,7 +162,7 @@ def test_create_agent_uses_the_gateway_even_with_a_provider_key( # The same config built directly honours the provider key. assert OpenAIAgent(OpenAIConfig()).openai_client is direct.return_value # Routing is config, not a client, so the agent stays hosted-serializable. - assert "gateway" not in agent.hosted_spec()["config"] + assert "gateway" not in dump_agent(agent)["config"] def test_create_agent_resolves_gateway_model_metadata( diff --git a/hud/agents/tests/test_claude_sdk_agent.py b/hud/agents/tests/test_claude_cli_agent.py similarity index 53% rename from hud/agents/tests/test_claude_sdk_agent.py rename to hud/agents/tests/test_claude_cli_agent.py index 4e8af68bf..7537aca2c 100644 --- a/hud/agents/tests/test_claude_sdk_agent.py +++ b/hud/agents/tests/test_claude_cli_agent.py @@ -1,4 +1,4 @@ -"""ClaudeSDKAgent remote-command construction over the workspace SSH. +"""ClaudeCLIAgent remote-command construction over the workspace SSH. The agent runs the ``claude`` CLI on the remote workspace. These cover how the command is assembled per login shell — especially the Windows path, where the @@ -19,127 +19,200 @@ from typing import TYPE_CHECKING, Any, Literal, cast from unittest.mock import AsyncMock, Mock +import fastmcp import pytest +from mcp.types import ImageContent, TextContent from hud.agents.claude.sdk import computer_mcp -from hud.agents.claude.sdk.agent import ClaudeSDKAgent, build_remote_invocation -from hud.agents.types import ClaudeSDKConfig -from hud.capabilities import Capability, SSHClient +from hud.agents.claude.sdk.agent import ClaudeCLIAgent +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, Connection, SSHClient from hud.capabilities.rfb import WebPScreenshotEncoding from hud.settings import settings from hud.telemetry.context import set_trace_context +from hud.types import MCPToolResult if TYPE_CHECKING: from pathlib import Path -# ─── build_remote_invocation (pure) ─────────────────────────────────── - - -@pytest.mark.parametrize("shell", ["cmd", "powershell"]) -def test_windows_shell_runs_batch_file_via_cmd(shell: str) -> None: - inv = build_remote_invocation(shell, "claude --print -- hi") - - # The bare filename is rejected by the remote shell; cmd /c runs it. - assert inv.command == "cmd /c .hud_run.bat" - assert inv.script_name == ".hud_run.bat" - assert inv.script_body == "@echo off\r\nclaude --print -- hi\r\n" +@pytest.fixture(autouse=True) +def _clear_api_keys(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "api_key", None) + monkeypatch.setattr(settings, "anthropic_api_key", None) + monkeypatch.setattr( + "hud.agents.claude.sdk.agent.resolve_executable", + AsyncMock(return_value="claude"), + ) -def test_posix_shell_runs_inline_with_install_check() -> None: - inv = build_remote_invocation("bash", "claude --print -- hi") - assert inv.script_name is None - assert inv.script_body is None - assert "install.sh" in inv.command # one-shot bootstrap prefix - assert inv.command.endswith(" && claude --print -- hi") +def test_command_follows_explicit_gateway_routing(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "api_key", "hud-key") + monkeypatch.setattr(settings, "anthropic_api_key", "anthropic-key") + + gateway_agent = ClaudeCLIAgent(ClaudeCLIConfig(use_hud_gateway=True)) + provider_agent = ClaudeCLIAgent(ClaudeCLIConfig(use_hud_gateway=False)) + gateway = gateway_agent._build_cli_command(shell="bash") + provider = provider_agent._build_cli_command(shell="bash") + + assert f"ANTHROPIC_BASE_URL={settings.hud_gateway_url}" in gateway + assert "ANTHROPIC_API_KEY=hud-key" in gateway + assert "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS=1" in gateway + assert "DISABLE_AUTO_COMPACT=1" in gateway + assert "ANTHROPIC_API_KEY=anthropic-key" in provider + assert "ANTHROPIC_BASE_URL" not in provider + assert "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS" not in provider + assert "DISABLE_AUTO_COMPACT" not in provider + assert "ANTHROPIC_MODEL=claude-sonnet-5" in provider + + +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", + connection=connection, + ) -def test_gateway_trace_headers_are_forwarded_to_claude_cli( + 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", + "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: - monkeypatch.setattr(settings, "api_key", "test-key") - agent = ClaudeSDKAgent() + monkeypatch.setattr(settings, "api_key", "hud&key's") + config = ClaudeCLIConfig( + use_hud_gateway=True, + max_steps=3, + system_prompt="don't $expand", + ) + agent = ClaudeCLIAgent(config) + command = agent._build_cli_command(shell="powershell") - with set_trace_context("child-trace", parent_trace_id="parent-trace"): - env = agent._build_env_vars() + encoded = command.rsplit(" ", 1)[1] + script = base64.b64decode(encoded).decode("utf-16-le") + assert "$env:ANTHROPIC_API_KEY='hud&key''s'" in script + assert "$env:CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS='1'" in script + assert "$env:DISABLE_AUTO_COMPACT='1'" in script + assert "'--system-prompt' 'don''t $expand'" in script + assert "Get-Content -Raw -Encoding UTF8 '.hud_input.jsonl' | & 'claude'" in script + assert "'--input-format=stream-json'" in script + assert "python" not in script - assert env["ANTHROPIC_CUSTOM_HEADERS"] == ( - "Trace-Id: child-trace\nX-HUD-Parent-Trace-Id: parent-trace" - ) +class _FakeCompletedProcess: + async def wait(self, *, check: bool, **kwargs: Any) -> Any: + del check + assert kwargs == {"timeout": None} + return SimpleNamespace(stdout=b"", stderr=b"", exit_status=0, returncode=0) -# ─── _exec end-to-end over a fake SSH workspace ──────────────────────── + def terminate(self) -> None: + pass + + def close(self) -> None: + pass + + async def wait_closed(self) -> None: + pass class _FakeConn: - def __init__(self, sink: dict[str, bytes], result: Any) -> None: + def __init__(self, sink: dict[str, bytes], process: _FakeStreamProcess) -> None: self._sink = sink - self._result = result + self._process = process self.ran: list[str] = [] self.write_commands: list[str] = [] + self.written: dict[str, bytes] = {} + self.deleted: list[str] = [] def is_closed(self) -> bool: return False - async def run( - self, - cmd: str, - *, - input: str | None = None, - check: bool = True, - encoding: str | None = "utf-8", - ) -> Any: - if input is not None or cmd.startswith("powershell "): + async def create_process(self, cmd: str, **kwargs: Any) -> Any: + input_value = kwargs.get("input") + if cmd.startswith(("rm -f -- ", "cmd /c del /f /q ")): + paths = [path for path in self._sink if path in cmd] + for path in paths: + self._sink.pop(path) + self.deleted.extend(paths) + return _FakeCompletedProcess() + if input_value is not None or cmd.startswith("powershell "): self.write_commands.append(cmd) script = cmd if match := re.search(r"-EncodedCommand (\S+)", cmd): script = base64.b64decode(match.group(1)).decode("utf-16-le") name = next( path - for path in (".hud_prompt.txt", ".hud_run.bat", ".hud_mcp_config.json") + for path in (".hud_input.jsonl", ".hud_run.bat", ".hud_mcp_config.json") if path in script ) - if input is not None: - self._sink[name] = input.encode() + if input_value is not None: + self._sink[name] = str(input_value).encode() elif match := re.search(r"FromBase64String\('([^']+)'\)", script): self._sink[name] += base64.b64decode(match.group(1)) else: self._sink[name] = b"" - return SimpleNamespace(stdout="", stderr="", exit_status=0, returncode=0) + self.written[name] = self._sink[name] + return _FakeCompletedProcess() + assert kwargs == {"encoding": None} self.ran.append(cmd) - return self._result - - async def create_process(self, cmd: str, **kwargs: Any) -> _FakeProcess: - return _FakeProcess(await self.run(cmd, **kwargs)) - - -class _FakeProcess: - def __init__(self, result: Any) -> None: - self._result = result - - async def wait(self, *, check: bool, **kwargs: Any) -> Any: - del check - assert kwargs == {"timeout": None} - return self._result - - def terminate(self) -> None: - pass - - def close(self) -> None: - pass - - async def wait_closed(self) -> None: - pass + return self._process -def _fake_run() -> Any: - trace = SimpleNamespace(status="", content="", extra={}) - steps: list[Any] = [] - return SimpleNamespace(trace=trace, record=steps.append, steps=steps) +async def run_claude( + config: ClaudeCLIConfig, + run: Any, + *, + ssh: SSHClient, + shell: str, + mcp_servers: dict[str, dict[str, Any]], + prompt: str, +) -> None: + agent = ClaudeCLIAgent(config) + await agent._exec( + run, + ssh=ssh, + shell=shell, + mcp_servers=mcp_servers, + prompt=prompt, + ) _STREAM_JSON = ( - '{"type":"assistant","message":{"content":[{"type":"text","text":"working"}]}}\n' + '{"type":"assistant","message":{"id":"msg-1","type":"message",' + '"role":"assistant","model":"claude-test","content":[{"type":"text",' + '"text":"editing"},{"type":"tool_use","id":"tool-1","name":"Write","input":{}}],' + '"stop_reason":"tool_use","stop_sequence":null,"usage":{"input_tokens":11,' + '"output_tokens":7,"cache_read_input_tokens":3}}}\n' + '{"type":"user","message":{"content":[{"type":"tool_result",' + '"tool_use_id":"tool-1","content":[{"type":"text","text":"wrote a.txt"},' + '{"type":"image","source":{"type":"base64","media_type":"image/png",' + '"data":"aW1hZ2U="}}],"is_error":false}]}}\n' + '{"type":"assistant","message":{"id":"msg-2","type":"message",' + '"role":"assistant","model":"claude-test","content":[{"type":"text",' + '"text":"done"}],"stop_reason":"end_turn","stop_sequence":null,' + '"usage":{"input_tokens":11,"output_tokens":7,"cache_read_input_tokens":3}}}\n' '{"type":"result","is_error":false,"result":"done","session_id":"s",' '"duration_ms":5,"num_turns":2,"total_cost_usd":0.01}\n' ) @@ -157,76 +230,239 @@ def _ssh_with_conn(shell: str, conn: _FakeConn) -> SSHClient: async def test_exec_on_windows_writes_batch_and_execs_via_cmd() -> None: sink: dict[str, bytes] = {} - conn = _FakeConn( - sink, - SimpleNamespace(stdout=_STREAM_JSON, stderr="", exit_status=0, returncode=0), - ) - agent = ClaudeSDKAgent() + conn = _FakeConn(sink, _FakeStreamProcess(_STREAM_JSON)) ssh = _ssh_with_conn("cmd", conn) run = _fake_run() - await agent._exec(run, ssh=ssh, shell="cmd", mcp_servers={}, prompt="build it", max_steps=5) + await run_claude( + ClaudeCLIConfig(), run, ssh=ssh, shell="cmd", mcp_servers={}, prompt="build it" + ) assert conn.ran == ["cmd /c .hud_run.bat"] assert all(command.startswith("powershell ") for command in conn.write_commands) - assert sink[".hud_run.bat"].startswith(b"@echo off\r\n") - assert sink[".hud_prompt.txt"] == b"build it" - assert run.trace.status == "completed" - assert "done" in run.trace.content + assert conn.written[".hud_run.bat"].startswith(b"@echo off\r\n") + assert conn.written[".hud_input.jsonl"].endswith(b"\n") + assert json.loads(conn.written[".hud_input.jsonl"]) == { + "type": "user", + "message": { + "role": "user", + "content": [{"type": "text", "text": "build it"}], + }, + } + assert sink == {} + assert set(conn.deleted) == {".hud_input.jsonl", ".hud_run.bat"} + assert run.trace.status is None + assert run.trace.content == "done" + assert "messages" not in run.trace.extra async def test_exec_on_bash_runs_inline_without_batch() -> None: sink: dict[str, bytes] = {} - conn = _FakeConn( - sink, - SimpleNamespace(stdout=_STREAM_JSON, stderr="", exit_status=0, returncode=0), - ) - agent = ClaudeSDKAgent() + process = _FakeStreamProcess(_STREAM_JSON) + conn = _FakeConn(sink, process) ssh = _ssh_with_conn("bash", conn) run = _fake_run() - await agent._exec(run, ssh=ssh, shell="bash", mcp_servers={}, prompt="build it", max_steps=5) + await run_claude( + ClaudeCLIConfig(), run, ssh=ssh, shell="bash", mcp_servers={}, prompt="build it" + ) - assert ".hud_run.bat" not in sink - assert conn.write_commands == ["cat > .hud_prompt.txt"] + assert sink == {} + assert conn.write_commands == [] + assert conn.deleted == [] assert len(conn.ran) == 1 - assert "install.sh" in conn.ran[0] assert "claude" in conn.ran[0] - assert run.trace.status == "completed" + assert "--input-format=stream-json" in conn.ran[0] + assert "build it" not in conn.ran[0] + assert process.stdin.data.endswith(b"\n") + assert json.loads(process.stdin.data) == { + "type": "user", + "message": { + "role": "user", + "content": [{"type": "text", "text": "build it"}], + }, + } + assert process.stdin.eof is True + assert run.trace.status is None + assert run.trace.content == "done" + assert "messages" not in run.trace.extra -async def test_exec_nonzero_exit_with_no_stdout_records_system_error() -> None: +async def test_exec_removes_mcp_config_after_run() -> None: sink: dict[str, bytes] = {} - conn = _FakeConn( - sink, - SimpleNamespace(stdout="", stderr="boom", exit_status=1, returncode=1), + conn = _FakeConn(sink, _FakeStreamProcess(_STREAM_JSON)) + await run_claude( + ClaudeCLIConfig(), + _fake_run(), + ssh=_ssh_with_conn("bash", conn), + shell="bash", + mcp_servers={"database": {"type": "http", "url": "http://db/mcp"}}, + prompt="build it", + ) + + config = json.loads(conn.written[".hud_mcp_config.json"]) + assert config == {"mcpServers": {"database": {"type": "http", "url": "http://db/mcp"}}} + assert sink == {} + assert conn.deleted == [".hud_mcp_config.json"] + assert "--mcp-config .hud_mcp_config.json" in conn.ran[0] + + +async def test_exec_records_steps_before_process_exit() -> None: + process = _FakeStreamProcess(_STREAM_JSON, pause_after=1) + conn = _FakeConn({}, process) + ssh = _ssh_with_conn("bash", conn) + run = _fake_run() + + execution = asyncio.create_task( + run_claude( + ClaudeCLIConfig(), + run, + ssh=ssh, + shell="bash", + mcp_servers={}, + prompt="edit it", + ) ) - agent = ClaudeSDKAgent() + await process.stdout.blocked.wait() + + assert not execution.done() + assert len(run.steps) == 1 + first = run.steps[0] + assert isinstance(first, AgentStep) + assert first.content == "editing" + assert first.tool_calls[0].id == "tool-1" + + process.stdout.release.set() + await execution + + assert [type(step) for step in run.steps] == [AgentStep, ToolStep, AgentStep] + tool = cast("ToolStep", run.steps[1]) + assert tool.started_at == first.ended_at + assert tool.result is not None + text = tool.result.content[0] + assert isinstance(text, TextContent) + assert text.text == "wrote a.txt" + image = tool.result.content[1] + assert isinstance(image, ImageContent) + assert image.mimeType == "image/png" + assert image.data == "aW1hZ2U=" + final = cast("AgentStep", run.steps[2]) + assert final.started_at == tool.ended_at + assert run.trace.status is None + assert run.trace.content == "done" + + +async def test_exec_forwards_trace_id_only_to_hud_gateway( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(settings, "api_key", "hud-key") + monkeypatch.setattr(settings, "anthropic_api_key", "anthropic-key") + + gateway_conn = _FakeConn({}, _FakeStreamProcess(_STREAM_JSON)) + gateway = ClaudeCLIConfig(use_hud_gateway=True) + with set_trace_context("trace-123"): + await run_claude( + gateway, + _fake_run(), + ssh=_ssh_with_conn("bash", gateway_conn), + shell="bash", + mcp_servers={}, + prompt="build it", + ) + assert "ANTHROPIC_CUSTOM_HEADERS='Trace-Id: trace-123'" in gateway_conn.ran[0] + + provider_conn = _FakeConn({}, _FakeStreamProcess(_STREAM_JSON)) + provider = ClaudeCLIConfig(use_hud_gateway=False) + with set_trace_context("trace-123"): + await run_claude( + provider, + _fake_run(), + ssh=_ssh_with_conn("bash", provider_conn), + shell="bash", + mcp_servers={}, + prompt="build it", + ) + assert "ANTHROPIC_CUSTOM_HEADERS" not in provider_conn.ran[0] + + +async def test_exec_closes_streaming_process_when_cancelled() -> None: + process = _FakeStreamProcess(_STREAM_JSON, pause_after=0) + conn = _FakeConn({}, process) + execution = asyncio.create_task( + run_claude( + ClaudeCLIConfig(), + _fake_run(), + ssh=_ssh_with_conn("bash", conn), + shell="bash", + mcp_servers={}, + prompt="build it", + ) + ) + await process.stdout.blocked.wait() + execution.cancel() + + with pytest.raises(asyncio.CancelledError): + await execution + + assert process.closed + + +async def test_exec_nonzero_exit_with_no_stdout_raises() -> None: + sink: dict[str, bytes] = {} + conn = _FakeConn(sink, _FakeStreamProcess("", stderr="boom", exit_status=1)) ssh = _ssh_with_conn("cmd", conn) run = _fake_run() - await agent._exec(run, ssh=ssh, shell="cmd", mcp_servers={}, prompt="x", max_steps=1) + with pytest.raises(RuntimeError, match="boom"): + await run_claude(ClaudeCLIConfig(), run, ssh=ssh, shell="cmd", mcp_servers={}, prompt="x") - assert run.trace.status == "error" assert run.trace.extra["returncode"] == 1 - assert run.steps[0].error == "boom" async def test_exec_signal_exit_records_the_returncode() -> None: sink: dict[str, bytes] = {} conn = _FakeConn( sink, - SimpleNamespace(stdout="", stderr="", exit_status=None, returncode=-15), + _FakeStreamProcess("", exit_status=None, returncode=-15), ) - agent = ClaudeSDKAgent() ssh = _ssh_with_conn("bash", conn) run = _fake_run() - await agent._exec(run, ssh=ssh, shell="bash", mcp_servers={}, prompt="x", max_steps=1) + with pytest.raises(RuntimeError, match="return code -15"): + await run_claude(ClaudeCLIConfig(), run, ssh=ssh, shell="bash", mcp_servers={}, prompt="x") - assert run.trace.status == "error" assert run.trace.extra["returncode"] == -15 - assert run.steps[0].error == "claude CLI exited with return code -15" + + +async def test_exec_nonzero_exit_with_result_stream_remains_an_error() -> None: + sink: dict[str, bytes] = {} + conn = _FakeConn( + sink, + _FakeStreamProcess(_STREAM_JSON, stderr="transport failed", exit_status=1), + ) + ssh = _ssh_with_conn("bash", conn) + + run = _fake_run() + with pytest.raises(RuntimeError, match="transport failed"): + await run_claude(ClaudeCLIConfig(), run, ssh=ssh, shell="bash", mcp_servers={}, prompt="x") + + assert run.trace.content == "done" + assert run.trace.extra["returncode"] == 1 + assert run.trace.extra["stderr"] == "transport failed" + assert "messages" not in run.trace.extra + + +async def test_exec_zero_exit_without_result_event_is_an_error() -> None: + sink: dict[str, bytes] = {} + stdout = _STREAM_JSON.rsplit('{"type":"result"', 1)[0] + conn = _FakeConn(sink, _FakeStreamProcess(stdout)) + ssh = _ssh_with_conn("bash", conn) + + run = _fake_run() + with pytest.raises(RuntimeError, match="without a result event"): + await run_claude(ClaudeCLIConfig(), run, ssh=ssh, shell="bash", mcp_servers={}, prompt="x") + + assert run.trace.content == "done" @pytest.mark.parametrize( @@ -252,32 +488,30 @@ 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: assert ref == "ssh" return ssh - agent = ClaudeSDKAgent() + agent = ClaudeCLIAgent() execute = AsyncMock() monkeypatch.setattr(agent, "_exec", execute) - with set_trace_context("child-trace"): - await agent( - cast( - "Any", - SimpleNamespace(client=Client(), prompt_text="call the tool"), - ) + await agent( + cast( + "Any", + SimpleNamespace( + client=Client(), prompt_text="call the tool", runtime_config=None, connections={} + ), ) + ) await_args = execute.await_args assert await_args is not None assert await_args.kwargs["mcp_servers"] == { - "database": { - "type": claude_type, - "url": "http://database:8000/mcp", - "headers": {"Trace-Id": "child-trace"}, - } + "database": {"type": claude_type, "url": "http://database:8000/mcp"} } execute.assert_awaited_once() @@ -298,6 +532,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: @@ -329,7 +564,7 @@ async def bridge( bridge_active = False encoding = WebPScreenshotEncoding(quality=42) - agent = ClaudeSDKAgent(ClaudeSDKConfig(screenshot_encoding=encoding)) + agent = ClaudeCLIAgent(ClaudeCLIConfig(screenshot_encoding=encoding)) async def execute(*_args: Any, **_kwargs: Any) -> None: assert bridge_active @@ -341,7 +576,9 @@ async def execute(*_args: Any, **_kwargs: Any) -> None: await agent( cast( "Any", - SimpleNamespace(client=Client(), prompt_text="use the computer"), + SimpleNamespace( + client=Client(), prompt_text="use the computer", runtime_config=None, connections={} + ), ) ) @@ -382,6 +619,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: @@ -421,11 +659,21 @@ async def execute(*_args: Any, **kwargs: Any) -> None: }, } - agent = ClaudeSDKAgent() + agent = ClaudeCLIAgent() monkeypatch.setattr(computer_mcp, "bridge_computer_mcp", bridge) monkeypatch.setattr(agent, "_exec", execute) - await agent(cast("Any", SimpleNamespace(client=Client(), prompt_text="use both screens"))) + await agent( + cast( + "Any", + SimpleNamespace( + client=Client(), + prompt_text="use both screens", + runtime_config=None, + connections={}, + ), + ) + ) assert bridged == [] @@ -455,6 +703,26 @@ async def test_computer_mcp_stdio_owns_rfb_lifetime( rfb.close.assert_awaited_once() +async def test_computer_mcp_preserves_tool_result(monkeypatch: pytest.MonkeyPatch) -> None: + result = MCPToolResult( + content=[TextContent(type="text", text="failed")], + isError=True, + ) + execute = AsyncMock(return_value=result) + monkeypatch.setattr(computer_mcp.ClaudeComputerTool, "execute", execute) + server = computer_mcp.create_computer_mcp(cast("Any", object())) + + async with fastmcp.Client(server) as client: + received = await client.call_tool_mcp( + "computer", + {"action": "left_click", "coordinate": [10, 20]}, + ) + + execute.assert_awaited_once_with({"action": "left_click", "coordinate": [10, 20]}) + assert received.isError is True + assert received.content == result.content + + class _ByteWriter: def __init__(self) -> None: self.closed = False @@ -640,6 +908,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 @@ -665,10 +934,17 @@ async def execute( first_entered.set() await release_first.wait() - agent = ClaudeSDKAgent() + agent = ClaudeCLIAgent() monkeypatch.setattr(agent, "_exec", execute) - run_a = SimpleNamespace(client=Client(shell_a, ssh_a), prompt_text="first") - run_b = SimpleNamespace(client=Client(shell_b, ssh_b), prompt_text="second") + run_a = SimpleNamespace( + 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, + connections={}, + ) first = asyncio.create_task(agent(cast("Any", run_a))) await first_entered.wait() diff --git a/hud/agents/tests/test_codex_cli_agent.py b/hud/agents/tests/test_codex_cli_agent.py new file mode 100644 index 000000000..395e3422e --- /dev/null +++ b/hud/agents/tests/test_codex_cli_agent.py @@ -0,0 +1,387 @@ +"""CodexCLIAgent command construction and JSONL trajectory mapping.""" + +from __future__ import annotations + +import asyncio +import base64 +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import AsyncMock + +import pytest +from mcp.types import TextContent + +from hud.agents.cli import resolve_executable +from hud.agents.codex import CodexCLIAgent +from hud.agents.codex.agent import codex_command, run_codex +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, Connection, SSHClient +from hud.eval.runtime import RuntimeConfig, RuntimeResources +from hud.settings import settings +from hud.telemetry.context import set_trace_context + + +@pytest.fixture(autouse=True) +def _clear_api_keys(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "api_key", None) + monkeypatch.setattr(settings, "openai_api_key", None) + monkeypatch.setattr( + "hud.agents.codex.agent.resolve_executable", + AsyncMock(return_value="codex"), + ) + + +class _FakeSSH: + def __init__(self, process: _FakeProcess, *, shell: str = "bash") -> None: + self.process = process + self.capability = Capability( + name="shell", + protocol="ssh/2", + url="ssh://localhost:22", + params={"shell": shell}, + ) + self.commands: list[str] = [] + + async def create_process( + self, command: str, *, connections: tuple[Connection, ...] = () + ) -> _FakeProcess: + self.commands.append(command) + return self.process + + +_STREAM_JSON = ( + '{"type":"thread.started","thread_id":"thread-1"}\n' + '{"type":"turn.started"}\n' + '{"type":"item.started","item":{"id":"cmd-1","type":"command_execution",' + '"command":"pytest -q","aggregated_output":"","exit_code":null,' + '"status":"in_progress"}}\n' + '{"type":"item.completed","item":{"id":"cmd-1","type":"command_execution",' + '"command":"pytest -q","aggregated_output":"1 passed\\n","exit_code":0,' + '"status":"completed"}}\n' + '{"type":"item.completed","item":{"id":"patch-1","type":"file_change",' + '"changes":[{"path":"calc.py","kind":"update"}],"status":"completed"}}\n' + '{"type":"item.started","item":{"id":"mcp-1","type":"mcp_tool_call",' + '"server":"db","tool":"query","arguments":{"sql":"select 42"},' + '"result":null,"error":null,"status":"in_progress"}}\n' + '{"type":"item.completed","item":{"id":"mcp-1","type":"mcp_tool_call",' + '"server":"db","tool":"query","arguments":{"sql":"select 42"},' + '"result":{"content":[{"type":"text","text":"42"}],' + '"structured_content":{"answer":42}},"error":null,"status":"completed"}}\n' + '{"type":"item.completed","item":{"id":"search-1","type":"web_search",' + '"query":"HUD evals","action":{"type":"search"}}}\n' + '{"type":"item.completed","item":{"id":"reason-1","type":"reasoning",' + '"text":"The test now passes."}}\n' + '{"type":"item.completed","item":{"id":"message-1","type":"agent_message",' + '"text":"Implemented and verified."}}\n' + '{"type":"turn.completed","usage":{"input_tokens":20,"cached_input_tokens":5,' + '"output_tokens":8,"reasoning_output_tokens":3}}\n' +) + + +def test_command_follows_explicit_gateway_routing(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "api_key", "hud-key") + monkeypatch.setattr(settings, "openai_api_key", "openai-key") + + with set_trace_context("trace-123"): + gateway = codex_command(CodexCLIConfig(use_hud_gateway=True), "bash") + provider = codex_command(CodexCLIConfig(use_hud_gateway=False), "bash") + + assert "HUD_API_KEY=hud-key" in gateway + assert 'model_provider="hud"' in gateway + assert f'model_providers.hud.base_url="{settings.hud_gateway_url}"' in gateway + assert "Trace-Id" in gateway + assert "CODEX_API_KEY=openai-key" in provider + assert "model_provider" not in provider + for command in (gateway, provider): + assert "codex exec" in command + assert "--json" in command + assert "--ephemeral" in command + assert "--ignore-user-config" not in command + assert "mktemp -d" in command + assert 'export CODEX_HOME="$codex_home"' in command + assert "--sandbox workspace-write" in command + assert "--model gpt-5.6-sol" in command + assert command.endswith(" -") + + +@pytest.mark.parametrize("sandbox", ["read-only", "workspace-write", "danger-full-access"]) +def test_command_uses_process_bound_connection_without_its_credential( + sandbox: str, +) -> None: + connection = Connection( + name="inference", + capability="ssh", + url="https://inference.hud.so", + headers={"Authorization": "Bearer scoped-runtime-token"}, + ) + + command = codex_command( + CodexCLIConfig.model_validate({"use_hud_gateway": True, "sandbox": sandbox}), + "bash", + connection=connection, + ) + + 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 f'model_providers.hud.base_url="{connection.client_url}"' in command + assert "Trace-Id" not in command + assert "exec env" in command + assert "--sandbox danger-full-access" in command + assert "--sandbox workspace-write" not in command + assert "--sandbox read-only" 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: + monkeypatch.setattr(settings, "openai_api_key", "key&value's") + config = CodexCLIConfig(use_hud_gateway=False, sandbox="danger-full-access") + command = codex_command(config, "powershell") + + script = base64.b64decode(command.rsplit(" ", 1)[1]).decode("utf-16-le") + assert "$env:CODEX_API_KEY='key&value''s'" in script + assert "$env:CODEX_HOME=$codexHome" in script + assert "[System.Guid]::NewGuid()" in script + assert "Remove-Item -Recurse -Force $codexHome" in script + assert "--ignore-user-config" not in script + assert "'--sandbox' 'danger-full-access'" in script + assert "& 'codex' 'exec'" in script + assert script.endswith(";exit $hudExitCode") + + +async def test_exec_streams_prompt_and_records_codex_items() -> None: + process = _FakeProcess(_STREAM_JSON) + ssh = _FakeSSH(process) + run = _fake_run() + + await run_codex( + CodexCLIConfig(), + run, + ssh=cast("SSHClient", ssh), + shell="bash", + prompt="Fix the failing test", + ) + + assert process.stdin.data == b"Fix the failing test" + assert process.stdin.eof + assert [type(step) for step in run.steps] == [ + ToolStep, + ToolStep, + ToolStep, + ToolStep, + AgentStep, + AgentStep, + ] + command = cast("ToolStep", run.steps[0]) + assert command.call is not None + assert command.call.name == "shell" + assert command.call.arguments == {"command": "pytest -q"} + assert command.result is not None + assert command.result.isError is False + output = command.result.content[0] + assert isinstance(output, TextContent) + assert output.text == "1 passed\n" + patch = cast("ToolStep", run.steps[1]) + assert patch.call is not None + assert patch.call.name == "apply_patch" + mcp = cast("ToolStep", run.steps[2]) + assert mcp.call is not None + assert mcp.call.name == "query" + assert mcp.call.provider_name == "db.query" + assert mcp.result is not None + assert mcp.result.structuredContent == {"answer": 42} + search = cast("ToolStep", run.steps[3]) + assert search.call is not None + assert search.call.name == "web_search" + assert cast("AgentStep", run.steps[4]).reasoning == "The test now passes." + assert cast("AgentStep", run.steps[5]).content == "Implemented and verified." + assert run.trace.content == "Implemented and verified." + assert run.trace.extra["codex_thread_id"] == "thread-1" + assert run.trace.extra["usage"]["cached_input_tokens"] == 5 + assert run.trace.status is None + + +async def test_exec_records_completed_items_before_process_exit() -> None: + process = _FakeProcess(_STREAM_JSON, pause_after=4) + ssh = _FakeSSH(process) + run = _fake_run() + execution = asyncio.create_task( + run_codex( + CodexCLIConfig(), + run, + ssh=cast("SSHClient", ssh), + shell="bash", + prompt="Fix it", + ) + ) + await process.stdout.blocked.wait() + + assert not execution.done() + assert len(run.steps) == 1 + assert isinstance(run.steps[0], ToolStep) + + process.stdout.release.set() + await execution + + +async def test_exec_turn_failure_raises() -> None: + stream = ( + '{"type":"thread.started","thread_id":"thread-1"}\n' + '{"type":"turn.started"}\n' + '{"type":"turn.failed","error":{"message":"model unavailable"}}\n' + ) + run = _fake_run() + + with pytest.raises(RuntimeError, match="model unavailable"): + await run_codex( + CodexCLIConfig(), + run, + ssh=cast("SSHClient", _FakeSSH(_FakeProcess(stream))), + shell="bash", + prompt="Fix it", + ) + + +async def test_exec_nonzero_exit_raises_stderr() -> None: + run = _fake_run() + + with pytest.raises(RuntimeError, match="authentication failed"): + await run_codex( + CodexCLIConfig(), + run, + ssh=cast( + "SSHClient", + _FakeSSH(_FakeProcess("", stderr="authentication failed", returncode=1)), + ), + shell="bash", + prompt="Fix it", + ) + + assert run.trace.extra["returncode"] == 1 + + +async def test_exec_nonzero_exit_prefers_structured_error() -> None: + run = _fake_run() + stream = '{"type":"error","message":"gateway rejected streaming"}\n' + + with pytest.raises(RuntimeError, match="gateway rejected streaming"): + await run_codex( + CodexCLIConfig(), + run, + ssh=cast( + "SSHClient", + _FakeSSH(_FakeProcess(stream, stderr="noisy warning", returncode=1)), + ), + shell="bash", + prompt="Fix it", + ) + + assert "stderr" not in run.trace.extra + + +async def test_exec_closes_process_when_cancelled() -> None: + process = _FakeProcess(_STREAM_JSON, pause_after=0) + execution = asyncio.create_task( + run_codex( + CodexCLIConfig(), + _fake_run(), + ssh=cast("SSHClient", _FakeSSH(process)), + shell="bash", + prompt="Fix it", + ) + ) + await process.stdout.blocked.wait() + execution.cancel() + + with pytest.raises(asyncio.CancelledError): + await execution + + assert process.closed + + +async def test_agent_opens_ssh_and_uses_workspace_prompt(monkeypatch: pytest.MonkeyPatch) -> None: + ssh = _FakeSSH(_FakeProcess(_STREAM_JSON), shell="powershell") + + class Client: + inference = None + + async def open(self, ref: str) -> _FakeSSH: + assert ref == "ssh" + return ssh + + agent = CodexCLIAgent() + execute = AsyncMock() + monkeypatch.setattr("hud.agents.codex.agent.run_codex", execute) + run = SimpleNamespace( + client=Client(), prompt_text="Fix it", runtime_config=None, connections={} + ) + + await agent(cast("Any", run)) + + execute.assert_awaited_once_with( + agent.config, + run, + ssh=ssh, + shell="powershell", + prompt="Fix it", + executable="codex", + connection=None, + ) + + +async def test_executable_resolution_prefers_matching_managed_bundle() -> None: + ssh = SimpleNamespace( + capability=Capability.ssh(url="ssh://localhost:22", host_pubkey="key", shell="bash"), + run=AsyncMock( + side_effect=[ + SimpleNamespace(returncode=0, stdout=b"Linux\nx86_64\ngnu\n"), + SimpleNamespace(returncode=0, stdout=b""), + ] + ), + ) + + executable = await resolve_executable( + cast("Any", ssh), + "codex", + {"linux-x64": "/usr/local/lib/agents/codex/bin/codex"}, + RuntimeConfig(resources=RuntimeResources(os="linux")), + ) + + assert executable == "/usr/local/lib/agents/codex/bin/codex" + assert ssh.run.await_count == 2 + + +async def test_executable_resolution_rejects_runtime_os_mismatch() -> None: + ssh = SimpleNamespace( + capability=Capability.ssh(url="ssh://localhost:22", host_pubkey="key", shell="bash"), + run=AsyncMock(return_value=SimpleNamespace(returncode=0, stdout=b"Linux\nx86_64\ngnu\n")), + ) + + with pytest.raises(RuntimeError, match=r"requested 'windows'.*reports 'linux'"): + await resolve_executable( + cast("Any", ssh), + "codex", + {}, + RuntimeConfig(resources=RuntimeResources(os="windows")), + ) diff --git a/hud/agents/tests/test_openai_compatible_agent.py b/hud/agents/tests/test_openai_compatible_agent.py index aa8d09114..ea607d773 100644 --- a/hud/agents/tests/test_openai_compatible_agent.py +++ b/hud/agents/tests/test_openai_compatible_agent.py @@ -5,6 +5,8 @@ from types import SimpleNamespace from typing import Any, cast +import pytest + from hud.agents.openai_compatible.agent import OpenAIChatAgent, OpenAIChatRunState from hud.agents.types import OpenAIChatConfig @@ -87,9 +89,8 @@ async def test_get_response_with_tool_call() -> None: async def test_get_response_error_path() -> None: agent = _agent(None, error=RuntimeError("boom")) - result = await agent.get_response(_state(agent)) - assert result.done is True - assert result.error is not None and "boom" in result.error + with pytest.raises(RuntimeError, match="boom"): + await agent.get_response(_state(agent)) async def test_get_response_malformed_tool_args() -> None: diff --git a/hud/agents/tests/test_tool_agent.py b/hud/agents/tests/test_tool_agent.py index 2fa7f9fc3..e5b3b747b 100644 --- a/hud/agents/tests/test_tool_agent.py +++ b/hud/agents/tests/test_tool_agent.py @@ -25,7 +25,13 @@ from hud.agents.tools.base import AgentToolSpec, result_text from hud.agents.tools.rfb import RFBTool from hud.agents.tools.ssh import SSHInfrastructureErrorResult -from hud.agents.types import AgentConfig, AgentStep, ClaudeConfig, ClaudeSDKConfig, ToolStep +from hud.agents.types import ( + AgentStep, + ClaudeCLIConfig, + ClaudeConfig, + ToolAgentConfig, + ToolStep, +) from hud.capabilities import ( Capability, CapabilityClient, @@ -55,11 +61,11 @@ def record(self, step: Step) -> None: self.trace.record(step) -class DictAgent(ToolAgent[_Msg, AgentConfig]): +class DictAgent(ToolAgent[_Msg, ToolAgentConfig]): """Minimal concrete ToolAgent over plain-dict messages.""" def __init__(self, turns: list[AgentStep], **config: Any) -> None: - self.config = AgentConfig(model="test-model", **config) + self.config = ToolAgentConfig(model="test-model", **config) self._turns = list(turns) async def _initialize_state(self, *, prompt: Any) -> RunState[_Msg]: @@ -90,9 +96,9 @@ class WithCatalog(DictAgent): def test_claude_defaults_to_configurable_webp_screenshots() -> None: - assert AgentConfig().screenshot_encoding == PngScreenshotEncoding() + assert ToolAgentConfig().screenshot_encoding == PngScreenshotEncoding() assert ClaudeConfig().screenshot_encoding == WebPScreenshotEncoding() - assert ClaudeSDKConfig().screenshot_encoding == WebPScreenshotEncoding() + assert ClaudeCLIConfig().screenshot_encoding == WebPScreenshotEncoding() configured = ClaudeConfig.model_validate( {"screenshot_encoding": {"mime_type": "image/webp", "quality": 42}}, @@ -111,8 +117,7 @@ def test_only_claude_provider_has_a_default_tool_timeout() -> None: assert config.timeout_seconds == 600 assert config.tool_timeout_seconds == 120 assert ClaudeConfig(tool_timeout_seconds=None).tool_timeout_seconds is None - assert AgentConfig().tool_timeout_seconds is None - assert ClaudeSDKConfig().tool_timeout_seconds is None + assert ToolAgentConfig().tool_timeout_seconds is None async def test_agent_passes_screenshot_encoding_to_rfb_tools() -> None: @@ -521,7 +526,8 @@ async def wait(*, check: bool, timeout: None) -> None: # noqa: ASYNC109 state = RunState(messages=[], tools={"bash": tool}) run = cast("Run", _FakeRun()) - await agent._loop(run, state, max_steps=3) + agent.config.max_steps = 3 + await agent._loop(run, state) assert started.is_set() process.terminate.assert_called_once_with() @@ -576,7 +582,8 @@ async def test_loop_finishes_on_done_response() -> None: agent = DictAgent([AgentStep(content="final answer", done=True)]) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=3) + agent.config.max_steps = 3 + await agent._loop(run, RunState()) assert run.trace.status == "completed" assert run.trace.content == "final answer" @@ -618,7 +625,8 @@ async def test_loop_discards_degenerate_turn_and_resamples() -> None: ) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=3) + agent.config.max_steps = 3 + await agent._loop(run, RunState()) # The degenerate turn is discarded (consuming a step) and never recorded. assert run.trace.status == "completed" @@ -630,7 +638,8 @@ async def test_loop_fails_when_degenerate_turn_exhausts_steps() -> None: agent = _DegenerateDictAgent([DegenerateTurnError("empty shell_call")] * 2) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=2) + agent.config.max_steps = 2 + await agent._loop(run, RunState()) assert run.trace.status == "error" assert run.trace.error == "empty shell_call" @@ -645,7 +654,8 @@ async def test_loop_dispatches_tool_calls_then_finishes() -> None: ) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=3) + agent.config.max_steps = 3 + await agent._loop(run, RunState()) assert run.trace.content == "done now" assert [step.source for step in run.trace.steps] == ["agent", "tool", "agent"] @@ -679,7 +689,8 @@ async def test_loop_resets_infrastructure_error_count_after_other_result( monkeypatch.setattr(agent, "_dispatch_call", dispatch) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=10) + agent.config.max_steps = 10 + await agent._loop(run, RunState()) assert dispatch.await_count == 5 assert run.trace.status == "error" @@ -697,7 +708,8 @@ async def test_loop_max_steps_is_normal_termination() -> None: agent = DictAgent(never_done) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=2) + agent.config.max_steps = 2 + await agent._loop(run, RunState()) assert run.trace.is_error is False assert run.trace.status == "completed" @@ -715,7 +727,8 @@ async def test_loop_marks_length_finish_as_truncated() -> None: agent = DictAgent([AgentStep(content="partial", done=True, finish_reason=finish_reason)]) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=3) + agent.config.max_steps = 3 + await agent._loop(run, RunState()) assert run.trace.status == "completed" assert run.trace.stop_reason == "length" @@ -736,7 +749,8 @@ async def test_loop_answers_malformed_call_by_default() -> None: ) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=3) + agent.config.max_steps = 3 + await agent._loop(run, RunState()) assert run.trace.content == "recovered" tool_step = run.trace.steps[1] @@ -760,7 +774,8 @@ async def test_loop_stops_on_malformed_call_when_configured() -> None: ) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=3) + agent.config.max_steps = 3 + await agent._loop(run, RunState()) assert run.trace.status == "completed" assert run.trace.stop_reason == "malformed_tool_call" @@ -783,7 +798,8 @@ async def test_loop_stops_on_length_when_configured() -> None: ) run = cast("Run", _FakeRun()) - await agent._loop(run, RunState(), max_steps=3) + agent.config.max_steps = 3 + await agent._loop(run, RunState()) assert run.trace.stop_reason == "length" assert run.trace.is_truncated is True diff --git a/hud/agents/tool_agent.py b/hud/agents/tool_agent.py index a9e744c09..541ead57d 100644 --- a/hud/agents/tool_agent.py +++ b/hud/agents/tool_agent.py @@ -1,19 +1,4 @@ -"""ToolAgent: catalog-driven provider tool-call loop. - -Subclass contract:: - - class ClaudeAgent(ToolAgent[BetaMessageParam, ClaudeConfig]): - tool_catalog = (ClaudeBashTool, ClaudeTextEditorTool, ClaudeMCPProxyTool) - - async def _initialize_state(self, *, prompt) -> RunState[BetaMessageParam]: ... - async def get_response(self, state, *, system_prompt, citations_enabled): ... - def _format_message(self, role, text) -> BetaMessageParam: ... - def _format_result(self, call, result) -> BetaMessageParam | None: ... - -``RunState`` carries the messages *and* the tools/params built for one run, so a -single agent instance can drive many concurrent ``rollout`` calls with no shared -mutable state. -""" +"""Catalog-driven provider tool-call agents.""" from __future__ import annotations @@ -34,11 +19,11 @@ def _format_result(self, call, result) -> BetaMessageParam | None: ... from hud.agents.types import AgentStep, ToolStep from hud.capabilities import MCPClient, RFBClient from hud.capabilities.ssh import SSHConnectionError -from hud.types import AgentType, MCPToolCall, MCPToolResult, Step, StopCondition +from hud.types import MCPToolCall, MCPToolResult, Step, StopCondition from hud.utils.time import now_iso if TYPE_CHECKING: - from hud.agents.types import AgentConfig + from hud.agents.types import ToolAgentConfig from hud.capabilities import CapabilityClient from hud.eval.run import Run @@ -50,7 +35,7 @@ def _format_result(self, call, result) -> BetaMessageParam | None: ... MAX_CONSECUTIVE_SSH_FAILURES = 3 MessageT = TypeVar("MessageT") -ConfigT = TypeVar("ConfigT", bound="AgentConfig") +ConfigT = TypeVar("ConfigT", bound="ToolAgentConfig") class DegenerateTurnError(Exception): @@ -64,20 +49,13 @@ class DegenerateTurnError(Exception): def _message_text(message: mcp_types.PromptMessage) -> str: - """Best-effort plain text for a prompt message (text content only for now).""" content = message.content - if isinstance(content, mcp_types.TextContent): - return content.text - return getattr(content, "text", "") or "" + return content.text if isinstance(content, mcp_types.TextContent) else "" @dataclass class RunState(Generic[MessageT]): - """Mutable per-run state: messages + the tools/params built for this run. - - Created fresh per ``rollout`` (or ``run``) call, so one agent instance can - drive many concurrent rollouts without shared mutable state. - """ + """Provider messages and tools for one run.""" messages: list[MessageT] = field(default_factory=list[MessageT]) tools: dict[str, AgentTool[Any]] = field(default_factory=dict[str, AgentTool[Any]]) @@ -88,10 +66,7 @@ class ToolAgent(Agent, Generic[MessageT, ConfigT]): """Catalog-driven provider tool-call loop.""" tool_catalog: ClassVar[tuple[type[AgentTool[Any]], ...]] = () - #: Capability-client types this agent can drive (derived from the catalog). clients: ClassVar[tuple[type[CapabilityClient], ...]] = () - - #: The agent's typed config; set by subclass __init__. config: ConfigT def __init_subclass__(cls, **kwargs: Any) -> None: @@ -102,47 +77,12 @@ def __init_subclass__(cls, **kwargs: Any) -> None: seen.setdefault(t.client_type, None) cls.clients = tuple(seen.keys()) - def hosted_spec(self) -> dict[str, Any]: - """HUD-hosted execution runs the agent remotely, so it is - reconstructed there from this identity (type, model, step budget, system - prompt, provider kwargs) with the model resolved through the HUD gateway. - """ - if self.config.model_client is not None: - raise ValueError( - "hosted execution cannot serialize a custom model_client; " - "use create_agent(model, ...) so the hosted runner rebuilds the " - "gateway client, or run the agent loop locally with HUDRuntime() " - "/ LocalRuntime (recommended for TrainingClient workflows that " - "attach a BYOK client)" - ) - agent_type = AgentType.of(self) - if agent_type is None: - raise ValueError( - f"hosted execution supports the gateway agent types " - f"({', '.join(at.value for at in AgentType)}); got {type(self).__name__}" - ) - config = self.config.model_dump( - mode="json", - exclude={"model_client", "gateway", "api_key", "base_url", "hosted_tools"}, - ) - return {"type": agent_type.value, "config": config} - async def __call__(self, run: Run) -> None: - """Drive this (stateless) agent over a live ``Run``, filling ``run.trace``. - - Opens the capabilities this agent's catalog supports off the connection, - builds the tools into a fresh ``RunState``, - then runs the loop against ``run.prompt_messages``, accumulating the - trajectory onto ``run.trace``. Loop budget and prompting come from the agent's config - (``max_steps``, ``system_prompt``, ``citations_enabled``). No per-rollout - state is stored on ``self``, so one instance may drive many concurrent - rollouts. - """ connections: dict[str, CapabilityClient] = {} opened_protocols: set[str] = set() manifest = run.client.manifest if manifest is not None: - wanted = {cls.protocol for cls in type(self).clients} + wanted = {client.protocol for client in type(self).clients} for cap in manifest.bindings: if cap.protocol not in wanted: continue @@ -152,13 +92,7 @@ async def __call__(self, run: Run) -> None: opened_protocols.add(cap.protocol) state = await self._initialize_state(prompt=run.prompt_messages) state.tools, state.params = await self._build_tools(connections) - await self._loop( - run, - state, - max_steps=self.config.max_steps, - system_prompt=self.config.system_prompt, - citations_enabled=self.config.citations_enabled, - ) + await self._loop(run, state) async def _build_tools( self, @@ -168,12 +102,11 @@ async def _build_tools( tools: dict[str, AgentTool[Any]] = {} params: list[Any] = [] model = self.config.model - hosted_tools = self.config.hosted_tools mcp_clients = [c for c in connections.values() if isinstance(c, MCPClient)] mcp_lists = await asyncio.gather(*(c.list_tools() for c in mcp_clients)) mcp_by_client: dict[MCPClient, list[mcp_types.Tool]] = dict( - zip(mcp_clients, mcp_lists, strict=False), + zip(mcp_clients, mcp_lists, strict=True), ) qualify_mcp_names = len(mcp_clients) > 1 @@ -216,7 +149,11 @@ async def _build_tools( tools[tool.provider_name] = tool params.append(tool.to_params()) - params.extend(hosted.to_params() for hosted in hosted_tools if hosted.supports_model(model)) + params.extend( + hosted.to_params() + for hosted in self.config.hosted_tools + if hosted.supports_model(model) + ) return tools, params @@ -224,10 +161,6 @@ async def _loop( self, run: Run, state: RunState[MessageT], - *, - max_steps: int = 10, - system_prompt: str | None = None, - citations_enabled: bool = False, ) -> None: trace = run.trace try: @@ -236,17 +169,17 @@ async def _loop( stopped: StopCondition | None = None consecutive_ssh_failures = 0 - for turn in range(1, max_steps + 1): - logger.info("step %d/%d", turn, max_steps) + for turn in range(1, self.config.max_steps + 1): + logger.info("step %d/%d", turn, self.config.max_steps) started_at = now_iso() try: step = await self.get_response( state, - system_prompt=system_prompt, - citations_enabled=citations_enabled, + system_prompt=self.config.system_prompt, + citations_enabled=self.config.citations_enabled, ) except DegenerateTurnError as exc: - if turn == max_steps: + if turn == self.config.max_steps: raise logger.warning("Discarded degenerate turn: %s", exc) continue @@ -297,7 +230,7 @@ async def _loop( run.record(Step(source="system", error=error)) return - if turn == max_steps: + if turn == self.config.max_steps: hit_max = True trace.content = step.content if step else None @@ -392,12 +325,14 @@ async def _dispatch_call( isError=True, ) - # ─── provider hooks ─────────────────────────────────────────────── - def _initial_messages(self, prompt: list[mcp_types.PromptMessage]) -> list[MessageT]: """Map normalized prompt turns onto provider messages.""" return [self._format_message(message.role, _message_text(message)) for message in prompt] + def _format_user_text(self, text: str) -> MessageT: + """Wrap a plain text string as a provider user message.""" + return self._format_message("user", text) + @abstractmethod async def _initialize_state( self, *, prompt: list[mcp_types.PromptMessage] @@ -414,14 +349,10 @@ async def get_response( ) -> AgentStep: """Call the provider API and return the model's turn as an ``AgentStep``. - The loop stamps ``started_at``/``model`` fallbacks and records it; - a failed call is an ``AgentStep`` with ``error`` set and ``done=True``. + The loop stamps ``started_at``/``model`` fallbacks, records the step, + and raises its error if present. """ - def _format_user_text(self, text: str) -> MessageT: - """Wrap a plain text string as a provider user message.""" - return self._format_message("user", text) - @abstractmethod def _format_message(self, role: str, text: str) -> MessageT: """Wrap text as a provider message of the given role (``user``/``assistant``).""" diff --git a/hud/agents/types.py b/hud/agents/types.py index 56789fdfd..41f0d2814 100644 --- a/hud/agents/types.py +++ b/hud/agents/types.py @@ -51,11 +51,16 @@ class AgentConfig(BaseModel): + timeout_seconds: float | None = Field(default=None, gt=0, allow_inf_nan=False) + model_name: str = "Agent" + model: str = Field(default="unknown", validation_alias=_model_alias) + + +class ToolAgentConfig(AgentConfig): model_config = ConfigDict(arbitrary_types_allowed=True) auto_respond: bool = False max_steps: int = 10 - timeout_seconds: float | None = Field(default=None, gt=0, allow_inf_nan=False) tool_timeout_seconds: float | None = Field(default=None, gt=0, allow_inf_nan=False) system_prompt: str | None = None citations_enabled: bool = False @@ -65,8 +70,6 @@ class AgentConfig(BaseModel): hosted_tools: list[HostedTool[object]] = Field(default_factory=list[HostedTool[object]]) screenshot_encoding: ScreenshotEncoding = Field(default_factory=PngScreenshotEncoding) - model_name: str = "Agent" - model: str = Field(default="unknown", validation_alias=_model_alias) #: Provider client (AsyncAnthropic, AsyncOpenAI, genai.Client, ...). When unset, #: agents build one: the provider's own key when set, otherwise the HUD gateway. model_client: Any = None @@ -80,7 +83,7 @@ class AgentConfig(BaseModel): # ----------------------------------------------------------------------------- -class ClaudeConfig(AgentConfig): +class ClaudeConfig(ToolAgentConfig): model_name: str = "Claude" model: str = Field(default="claude-sonnet-4-6", validation_alias=_model_alias) tool_timeout_seconds: float | None = Field(default=120, gt=0, allow_inf_nan=False) @@ -94,7 +97,7 @@ class ClaudeConfig(AgentConfig): # ----------------------------------------------------------------------------- -class GeminiConfig(AgentConfig): +class GeminiConfig(ToolAgentConfig): """Configuration for GeminiAgent.""" model_name: str = "Gemini" @@ -113,7 +116,7 @@ class GeminiConfig(AgentConfig): # ----------------------------------------------------------------------------- -class OpenAIConfig(AgentConfig): +class OpenAIConfig(ToolAgentConfig): """Configuration for OpenAIAgent.""" model_name: str = "OpenAI" @@ -132,7 +135,7 @@ class OpenAIConfig(AgentConfig): ) -class OpenAIChatConfig(AgentConfig): +class OpenAIChatConfig(ToolAgentConfig): """Configuration for OpenAIChatAgent.""" model_name: str = "OpenAI Chat" @@ -150,19 +153,21 @@ class OpenAIChatConfig(AgentConfig): # ----------------------------------------------------------------------------- -# Claude Code (CLI over SSH) +# Claude CLI (over SSH) # ----------------------------------------------------------------------------- -class ClaudeSDKConfig(AgentConfig): - """Configuration for ClaudeSDKAgent (runs the ``claude`` CLI over SSH). +class ClaudeCLIConfig(AgentConfig): + """Configuration for ClaudeCLIAgent (runs the ``claude`` CLI over SSH). - ``system_prompt`` is inherited from ``AgentConfig``. ``max_steps`` maps to the - CLI's ``--max-turns``; values <= 0 leave the turn budget to the CLI (unlimited). + ``max_steps`` maps to the CLI's ``--max-turns``; values <= 0 leave the turn + budget to the CLI (unlimited). """ - model_name: str = "Claude Code" - model: str = Field(default="claude-sonnet-4-6", validation_alias=_model_alias) + system_prompt: str | None = None + model_name: str = "Claude CLI" + model: str = Field(default="claude-sonnet-5", validation_alias=_model_alias) + use_hud_gateway: bool | None = None permission_mode: str = "bypassPermissions" max_steps: int = -1 screenshot_encoding: ScreenshotEncoding = Field(default_factory=WebPScreenshotEncoding) @@ -180,6 +185,26 @@ class ClaudeSDKConfig(AgentConfig): ) +# ----------------------------------------------------------------------------- +# Codex CLI (over SSH) +# ----------------------------------------------------------------------------- + + +class CodexCLIConfig(AgentConfig): + """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. + A process-bound inference connection runs Codex without its inner sandbox; + the connection is available only inside the environment's isolated workspace. + """ + + model_name: str = "Codex CLI" + model: str = Field(default="gpt-5.6-sol", validation_alias=_model_alias) + use_hud_gateway: bool | None = None + sandbox: Literal["read-only", "workspace-write", "danger-full-access"] = "workspace-write" + + # ----------------------------------------------------------------------------- # Browser Use # ----------------------------------------------------------------------------- @@ -189,9 +214,7 @@ class BrowserUseConfig(AgentConfig): """Configuration for BrowserUseAgent. Lives here (not in the agent module) so it can be imported and serialized - without the optional ``browser-use`` dependency installed. The ``auto_respond`` - / ``system_prompt`` / ``hosted_tools`` fields from ``AgentConfig`` do not apply - — browser-use runs its own agent loop. + without the optional ``browser-use`` dependency installed. """ model_name: str = "Browser Use" 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/cli/eval.py b/hud/cli/eval.py index b5d754d5e..26a4b03a8 100644 --- a/hud/cli/eval.py +++ b/hud/cli/eval.py @@ -465,10 +465,13 @@ def spawn(task: Task) -> AbstractAsyncContextManager[Runtime]: if cfg.model: agent_kwargs["model"] = cfg.model agent_kwargs["max_steps"] = cfg.max_steps - if cfg.auto_respond: + if cfg.auto_respond and not agent_type.is_cli: agent_kwargs["auto_respond"] = True if cfg.gateway: - agent_kwargs["gateway"] = True + if agent_type.is_cli: + agent_kwargs["use_hud_gateway"] = True + else: + agent_kwargs["gateway"] = True table = Table(title="Evaluation Settings", title_style="bold cyan", box=box.ROUNDED) table.add_column("Setting", style="yellow") diff --git a/hud/cli/tests/test_eval.py b/hud/cli/tests/test_eval.py index 557e2a02e..6fce49afc 100644 --- a/hud/cli/tests/test_eval.py +++ b/hud/cli/tests/test_eval.py @@ -12,7 +12,7 @@ import pytest from typer.testing import CliRunner -from hud.agents import OpenAIAgent +from hud.agents import OpenAIAgent, dump_agent from hud.agents.types import ClaudeConfig from hud.cli import eval as eval_mod from hud.cli.__main__ import app @@ -518,7 +518,7 @@ def test_hosted_agent_keeps_client_out_of_serialized_config( monkeypatch.setattr("hud.utils.gateway.build_gateway_client", MagicMock(return_value=object())) eval_cli.invoke("tasks.py", "openai", "--remote", "--yes") assert eval_cli.agent.config.model_client is None - assert "model_client" not in eval_cli.agent.hosted_spec()["config"] + assert "model_client" not in dump_agent(eval_cli.agent)["config"] def test_openai_compatible_routes_through_gateway_despite_openai_key( diff --git a/hud/clients/client.py b/hud/clients/client.py index d878ffc66..33cd9dce2 100644 --- a/hud/clients/client.py +++ b/hud/clients/client.py @@ -23,6 +23,7 @@ Capability, CapabilityClient, CDPClient, + Connection, MCPClient, RFBClient, SSHClient, @@ -36,8 +37,9 @@ ) 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") @@ -161,14 +163,26 @@ 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] = (), + connections: Sequence[Connection] = (), + ) -> 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 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) env = result["env"] bindings = [Capability.from_manifest(binding) for binding in result["bindings"]] @@ -383,6 +397,8 @@ async def _connect_ready( port: int, *, 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. @@ -406,7 +422,7 @@ async def _connect_ready( client = HudClient(reader, writer, endpoint=(host, port)) try: - await client.hello() + await client.hello(workspace_routes=workspace_routes, connections=connections) except asyncio.CancelledError: client.abort() raise @@ -439,7 +455,13 @@ 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] = (), + connections: Sequence[Connection] = (), +) -> AsyncIterator[HudClient]: """Connect a :class:`HudClient` to a provisioned substrate's control channel. Takes the :class:`~hud.eval.runtime.Runtime` a provider yielded (or @@ -456,6 +478,8 @@ 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, + connections=connections, ) owner = asyncio.current_task() assert owner is not None @@ -467,9 +491,14 @@ 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] + if connections: + params["connections"] = [connection.to_wire() for connection in connections] await client._call( "hello", - {"session_id": client.manifest.session_id}, + params, reply_timeout=_CONTROL_HEARTBEAT_TIMEOUT_SECONDS, ) except HudProtocolError as exc: @@ -495,4 +524,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 2134d7bbf..7aa42455c 100644 --- a/hud/clients/tests/test_connect.py +++ b/hud/clients/tests/test_connect.py @@ -17,14 +17,94 @@ 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 HudProtocolError, 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_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: @@ -288,11 +368,15 @@ async def fake_connect_ready( port: int, *, 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/__init__.py b/hud/environment/__init__.py index c1920a9f9..baa2f36e7 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 @@ -114,5 +114,6 @@ def load_environment( "Peer", "PromptArg", "Workspace", + "WorkspaceRoute", "load_environment", ] diff --git a/hud/environment/egress.py b/hud/environment/egress.py index 836cdcb81..0e0f46240 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. @@ -185,6 +187,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], *, @@ -538,6 +593,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. @@ -665,8 +886,10 @@ def stop(self) -> None: "ANY_HOST", "BRIDGE_PORT", "VISITOR_PORT", + "ConnectionRelay", "Egress", "Peer", + "WorkspaceRoute", "bind_addresses", "hosts_text", "permitted", diff --git a/hud/environment/env.py b/hud/environment/env.py index 86023f3da..a6364b0d6 100644 --- a/hud/environment/env.py +++ b/hud/environment/env.py @@ -15,8 +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 ConnectionRelay, Peer, WorkspaceRoute from .workspace import Workspace if TYPE_CHECKING: @@ -162,6 +163,9 @@ 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]] = {} + self._connections: dict[str, tuple[Connection, Workspace, Peer, ConnectionRelay]] = {} # ─── task registration ─────────────────────────────────────────── @@ -284,7 +288,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 +356,120 @@ 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() + 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: + 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._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" + ) + 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/process_guard.py b/hud/environment/process_guard.py new file mode 100644 index 000000000..fdc1a429f --- /dev/null +++ b/hud/environment/process_guard.py @@ -0,0 +1,723 @@ +"""Linux process-bound network connection enforcement.""" + +# ruff: noqa: UP045 + +from __future__ import annotations + +import argparse +import array +import asyncio +import contextlib +import ctypes +import errno +import fcntl +import ipaddress +import json +import os +import platform +import select +import shutil +import signal +import socket +import struct +import subprocess +import sys +import threading +from pathlib import Path +from typing import TYPE_CHECKING, Literal, NoReturn, Optional, cast + +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_TRACE = 0x7FF00000 +_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_OPEN_SYSCALL = 434 +_PIDFD_GETFD_SYSCALL = 438 +_REGISTER_ADDRESS = b"\0hud-process-connection-register" +_SANDBOX_SOCKET = "/tmp/.hud-process-connection/control.sock" # noqa: S108 +_SANDBOX_HELPER = "/tmp/.hud-process-connection/process_guard.py" # noqa: S108 +_READY_TIMEOUT_SECONDS = 10.0 + +_PTRACE_TRACEME = 0 +_PTRACE_CONT = 7 +_PTRACE_GETREGS = 12 +_PTRACE_SETREGS = 13 +_PTRACE_SETOPTIONS = 0x4200 +_PTRACE_O_TRACEFORK = 0x00000002 +_PTRACE_O_TRACEVFORK = 0x00000004 +_PTRACE_O_TRACECLONE = 0x00000008 +_PTRACE_O_TRACEEXEC = 0x00000010 +_PTRACE_O_TRACESECCOMP = 0x00000080 +_PTRACE_O_EXITKILL = 0x00100000 +_PTRACE_EVENT_SECCOMP = 7 +_WAIT_ALL = getattr(os, "WALL", 0x40000000) +_MAX_SOCKADDR_BYTES = 128 + +_LIBC = ctypes.CDLL(None, use_errno=True) +_LIBC.ptrace.restype = ctypes.c_long +GuardBackend = Literal["notify", "ptrace"] +_backend: Optional[GuardBackend] = None +_backend_probed = False + + +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), + ] + + +class _UserRegsStruct(ctypes.Structure): + _fields_ = [ + ("r15", ctypes.c_ulonglong), + ("r14", ctypes.c_ulonglong), + ("r13", ctypes.c_ulonglong), + ("r12", ctypes.c_ulonglong), + ("rbp", ctypes.c_ulonglong), + ("rbx", ctypes.c_ulonglong), + ("r11", ctypes.c_ulonglong), + ("r10", ctypes.c_ulonglong), + ("r9", ctypes.c_ulonglong), + ("r8", ctypes.c_ulonglong), + ("rax", ctypes.c_ulonglong), + ("rcx", ctypes.c_ulonglong), + ("rdx", ctypes.c_ulonglong), + ("rsi", ctypes.c_ulonglong), + ("rdi", ctypes.c_ulonglong), + ("orig_rax", ctypes.c_ulonglong), + ("rip", ctypes.c_ulonglong), + ("cs", ctypes.c_ulonglong), + ("eflags", ctypes.c_ulonglong), + ("rsp", ctypes.c_ulonglong), + ("ss", ctypes.c_ulonglong), + ("fs_base", ctypes.c_ulonglong), + ("gs_base", ctypes.c_ulonglong), + ("ds", ctypes.c_ulonglong), + ("es", ctypes.c_ulonglong), + ("fs", ctypes.c_ulonglong), + ("gs", ctypes.c_ulonglong), + ] + + +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 _connect_filter(connect_action: int) -> tuple[_SockFprog, object]: + audit_arch, _, 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, connect_action), + _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), + ) + return _SockFprog(len(instructions), instructions), instructions + + +def _install_connect_listener() -> int: + _, seccomp_syscall, _, _, _ = _architecture() + program, instructions = _connect_filter(_SECCOMP_RET_USER_NOTIF) + 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), + ) + del instructions + if listener < 0: + raise OSError(ctypes.get_errno(), "seccomp(NEW_LISTENER)") + return int(listener) + + +def _install_connect_trace() -> None: + _, seccomp_syscall, _, _, _ = _architecture() + program, instructions = _connect_filter(_SECCOMP_RET_TRACE) + if _LIBC.prctl(_PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "prctl(PR_SET_NO_NEW_PRIVS)") + result = _LIBC.syscall( + seccomp_syscall, + _SECCOMP_SET_MODE_FILTER, + 0, + ctypes.byref(program), + ) + del instructions + if result != 0: + raise OSError(ctypes.get_errno(), "seccomp(TRACE)") + + +def _detected_backend() -> Optional[GuardBackend]: + global _backend, _backend_probed + if _backend_probed: + return _backend + _backend_probed = True + if sys.platform != "linux": + return None + try: + probe = subprocess.run( + [sys.executable, str(Path(__file__).resolve()), "--probe"], + check=False, + capture_output=True, + text=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError): + return None + candidate = probe.stdout.strip() + if probe.returncode == 0 and candidate in {"notify", "ptrace"}: + _backend = cast("GuardBackend", candidate) + return _backend + + +def process_connections_supported() -> bool: + """Whether this substrate has a race-free process connection guard.""" + return _detected_backend() is not None + + +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) -> Optional[tuple[str, int]]: + 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(tgid: int, descriptor: int, address: bytes) -> int: + pidfd = _LIBC.syscall(_PIDFD_OPEN_SYSCALL, tgid, 0) + if pidfd < 0: + return -ctypes.get_errno() + 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) + + +def _ptrace(request: int, pid: int, address: object = 0, data: object = 0) -> int: + ctypes.set_errno(0) + result = _LIBC.ptrace(request, pid, address, data) + if result == -1: + error = ctypes.get_errno() + if error: + raise OSError(error, f"ptrace({request})") + return int(result) + + +def _trace_options(pid: int) -> None: + options = ( + _PTRACE_O_TRACEFORK + | _PTRACE_O_TRACEVFORK + | _PTRACE_O_TRACECLONE + | _PTRACE_O_TRACEEXEC + | _PTRACE_O_TRACESECCOMP + | _PTRACE_O_EXITKILL + ) + _ptrace(_PTRACE_SETOPTIONS, pid, 0, options) + + +def _trace_registers(pid: int) -> _UserRegsStruct: + registers = _UserRegsStruct() + _ptrace(_PTRACE_GETREGS, pid, 0, ctypes.byref(registers)) + return registers + + +def _complete_traced_connect( + pid: int, + trusted_tgid: int, + protected: Collection[tuple[str, int]], + allowed: Collection[tuple[str, int]], +) -> None: + registers = _trace_registers(pid) + process_tgid = _tgid(pid) + length = int(registers.rdx) + if length < 0 or length > _MAX_SOCKADDR_BYTES: + result = -errno.EFAULT + else: + try: + address = _read_process(pid, int(registers.rsi), length) + target = _destination(address) + if target in protected and (process_tgid != trusted_tgid or target not in allowed): + result = -errno.EPERM + else: + result = _emulate_connect(process_tgid, int(registers.rdi), address) + except (OSError, RuntimeError): + result = -errno.EPERM + registers.orig_rax = ctypes.c_ulonglong(-1).value + registers.rax = ctypes.c_ulonglong(result).value + _ptrace(_PTRACE_SETREGS, pid, 0, ctypes.byref(registers)) + + +def _trace_loop( + original: int, + protected: Collection[tuple[str, int]], + allowed: Collection[tuple[str, int]], +) -> int: + while True: + try: + pid, status = os.waitpid(-1, _WAIT_ALL) + except InterruptedError: + continue + if os.WIFEXITED(status) or os.WIFSIGNALED(status): + if pid == original: + return status + continue + if not os.WIFSTOPPED(status): + continue + event = status >> 16 + stop_signal = os.WSTOPSIG(status) + if event == _PTRACE_EVENT_SECCOMP: + _complete_traced_connect(pid, original, protected, allowed) + deliver = 0 + elif event or stop_signal in {signal.SIGSTOP, signal.SIGTRAP}: + deliver = 0 + else: + deliver = stop_signal + try: + _ptrace(_PTRACE_CONT, pid, 0, deliver) + except OSError as exc: + if exc.errno != errno.ESRCH: + raise + + +def _forward_signals(original: int) -> None: + def forward(signum: int, _frame: object) -> None: + with contextlib.suppress(ProcessLookupError): + os.kill(original, signum) + + for name in ("SIGHUP", "SIGINT", "SIGQUIT", "SIGTERM", "SIGTSTP", "SIGCONT", "SIGWINCH"): + if signum := getattr(signal, name, None): + signal.signal(signum, forward) + + +def _exit_from_wait_status(status: int) -> NoReturn: + if os.WIFEXITED(status): + raise SystemExit(os.WEXITSTATUS(status)) + signum = os.WTERMSIG(status) + signal.signal(signum, signal.SIG_DFL) + os.kill(os.getpid(), signum) + raise RuntimeError("failed to reproduce traced process signal exit") + + +def _trace_exec( + channel: socket.socket, + protected: Collection[tuple[str, int]], + allowed: Collection[tuple[str, int]], + argv: Sequence[str], +) -> NoReturn: + if platform.machine() != "x86_64": + raise RuntimeError("ptrace process connections currently require x86_64") + original = os.fork() + if original == 0: + channel.close() + try: + _ptrace(_PTRACE_TRACEME, 0) + os.kill(os.getpid(), signal.SIGSTOP) + _install_connect_trace() + os.execvp(argv[0], list(argv)) # noqa: S606 - exact controller-built argv + except BaseException: + os._exit(127) + try: + _, status = os.waitpid(original, 0) + if not os.WIFSTOPPED(status): + raise RuntimeError("guarded process did not stop for ptrace") + _trace_options(original) + _forward_signals(original) + channel.sendall(b"R") + channel.close() + _ptrace(_PTRACE_CONT, original) + _exit_from_wait_status(_trace_loop(original, protected, allowed)) + except BaseException: + with contextlib.suppress(ProcessLookupError): + os.kill(original, signal.SIGKILL) + with contextlib.suppress(ChildProcessError): + os.waitpid(original, 0) + raise + + +def _receive_policy(channel: socket.socket) -> tuple[frozenset[tuple[str, int]], ...]: + raw = bytearray() + while not raw.endswith(b"\n"): + chunk = channel.recv(65536 - len(raw)) + if not chunk: + raise RuntimeError("guard broker closed before sending its policy") + raw.extend(chunk) + if len(raw) >= 65536: + raise RuntimeError("guard policy exceeded 64 KiB") + document = json.loads(raw) + protected = frozenset((str(host), int(port)) for host, port in document["protected"]) + allowed = frozenset((str(host), int(port)) for host, port in document["allowed"]) + if not allowed <= protected: + raise RuntimeError("guard policy allowed an unprotected destination") + return protected, allowed + + +def _probe_ptrace() -> bool: + if platform.machine() != "x86_64": + return False + original = os.fork() + if original == 0: + try: + _ptrace(_PTRACE_TRACEME, 0) + os.kill(os.getpid(), signal.SIGSTOP) + _install_connect_trace() + descriptor = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + descriptor.connect(("127.0.0.1", 9)) + except OSError as exc: + os._exit(0 if exc.errno == errno.EPERM else 1) + os._exit(1) + except BaseException: + os._exit(1) + try: + _, status = os.waitpid(original, 0) + if not os.WIFSTOPPED(status): + return False + _trace_options(original) + _ptrace(_PTRACE_CONT, original) + status = _trace_loop(original, {("127.0.0.1", 9)}, set()) + return os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0 + except (OSError, RuntimeError): + with contextlib.suppress(ProcessLookupError): + os.kill(original, signal.SIGKILL) + with contextlib.suppress(ChildProcessError): + os.waitpid(original, 0) + return False + + +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]], + *, + backend: Optional[GuardBackend] = None, + ) -> None: + self.directory = directory + self.socket_path = directory / "control.sock" + self.helper_path = directory / "process_guard.py" + self.protected = frozenset(protected) + self.allowed = frozenset(allowed) + if not self.allowed <= self.protected: + raise ValueError("allowed process connections must be protected destinations") + selected_backend = backend or _detected_backend() + if selected_backend is None: + raise RuntimeError("process connection guards are unavailable on this substrate") + self.backend: GuardBackend = selected_backend + self._server: Optional[socket.socket] = None + self._listener: Optional[int] = None + self._stop_read, self._stop_write = os.pipe() + self._ready = threading.Event() + self._error: Optional[BaseException] = None + self._thread: Optional[threading.Thread] = None + + @property + def sandbox_socket(self) -> str: + return _SANDBOX_SOCKET + + @property + def sandbox_helper(self) -> str: + return _SANDBOX_HELPER + + def start(self) -> None: + self.directory.mkdir(mode=0o700, parents=True, exist_ok=True) + shutil.copyfile(Path(__file__), self.helper_path) + self.helper_path.chmod(0o500) + 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(FileNotFoundError): + self.helper_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: + if self.backend == "notify": + self._listener = _receive_fd(channel) + else: + policy = { + "protected": sorted([host, port] for host, port in self.protected), + "allowed": sorted([host, port] for host, port in self.allowed), + } + channel.sendall(json.dumps(policy, separators=(",", ":")).encode() + b"\n") + if channel.recv(1) != b"R": + raise RuntimeError("ptrace guard did not acknowledge its policy") + with contextlib.suppress(FileNotFoundError): + self.socket_path.unlink() + if self.backend == "notify": + self._broker() + else: + self._ready.set() + 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: Optional[int] = 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_length = int(notification.data.args[2]) + if address_length < 0 or address_length > _MAX_SOCKADDR_BYTES: + raise OSError(errno.EFAULT, "invalid sockaddr length") + address = _read_process( + notification.pid, + notification.data.args[1], + address_length, + ) + 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( + process_tgid, + 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(backend: GuardBackend, 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) + if backend == "ptrace": + protected, allowed = _receive_policy(channel) + _trace_exec(channel, protected, allowed, argv) + 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", choices=("auto", "notify", "ptrace"), nargs="?", const="auto") + parser.add_argument("--backend", choices=("notify", "ptrace")) + parser.add_argument("socket", nargs="?") + parser.add_argument("argv", nargs=argparse.REMAINDER) + args = parser.parse_args() + if args.probe: + if args.probe in {"auto", "notify"}: + try: + listener = _install_connect_listener() + except OSError: + if args.probe == "notify": + raise SystemExit(1) from None + else: + os.close(listener) + sys.stdout.write("notify\n") + return + if args.probe in {"auto", "ptrace"} and _probe_ptrace(): + sys.stdout.write("ptrace\n") + return + raise SystemExit(1) + if args.backend is None: + parser.error("--backend is required") + if args.socket is None: + parser.error("socket is required") + argv = args.argv[1:] if args.argv[:1] == ["--"] else args.argv + guarded_exec(args.backend, 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 ed320b929..65259c57a 100644 --- a/hud/environment/server.py +++ b/hud/environment/server.py @@ -26,8 +26,10 @@ from pydantic import BaseModel, TypeAdapter, ValidationError +from hud.capabilities import Connection from hud.graders.results import EvaluationResult +from .egress import WorkspaceRoute from .env import Answer, current_session_id from .utils import ( CONTROL_FRAME_LIMIT_BYTES, @@ -243,7 +245,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: @@ -278,7 +280,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 @@ -289,6 +291,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): @@ -343,6 +348,32 @@ async def error_to(msg_id: int | str | 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 + 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 # already concrete here. diff --git a/hud/environment/tests/test_process_connection.py b/hud/environment/tests/test_process_connection.py new file mode 100644 index 000000000..87428d856 --- /dev/null +++ b/hud/environment/tests/test_process_connection.py @@ -0,0 +1,324 @@ +"""Process-bound controller connection integration.""" + +from __future__ import annotations + +import asyncio +import shlex +import subprocess +import sys +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer +from pathlib import Path +from typing import 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 ProcessConnectionGuard, process_connections_supported +from hud.environment.workspace import usable_bwrap +from hud.eval import LocalRuntime, Task + +pytestmark = pytest.mark.skipif( + not process_connections_supported(), + reason="process connection guards are unavailable", +) + +_GUARD_PATH = Path(__file__).parents[1] / "process_guard.py" +_PTRACE_SUPPORTED = ( + sys.platform == "linux" + and subprocess.run( + [sys.executable, str(_GUARD_PATH), "--probe", "ptrace"], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode + == 0 +) + + +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 _threaded_fetch_source(url: str) -> str: + return ( + "import threading,urllib.request;" + "result=[];" + f"request=urllib.request.Request({url!r});" + "thread=threading.Thread(target=lambda:result.append(" + "urllib.request.urlopen(request,timeout=5).read().decode()));" + "thread.start();thread.join();print(result[0])" + ) + + +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())" + ) + + +def _proxy_environment_source() -> str: + return ( + "import os;" + "print(os.environ.get('http_proxy',''));" + "print(os.environ.get('https_proxy',''));" + "print(os.environ.get('no_proxy',''))" + ) + + +def test_guard_projects_a_standalone_helper(tmp_path: Path) -> None: + guard = ProcessConnectionGuard(tmp_path / "guard", set(), set(), backend="notify") + try: + guard.start() + assert guard.helper_path.read_bytes() == _GUARD_PATH.read_bytes() + assert guard.helper_path.stat().st_mode & 0o777 == 0o500 + finally: + guard.close() + + +@pytest.mark.skipif(not _PTRACE_SUPPORTED, reason="ptrace guard backend is unavailable") +@pytest.mark.asyncio +async def test_ptrace_backend_emulates_connects_and_blocks_descendants(tmp_path: Path) -> None: + protected, protected_thread = _server(_ProtectedUpstream) + ordinary, ordinary_thread = _server(_OrdinaryUpstream) + protected_url = f"http://127.0.0.1:{protected.server_address[1]}" + ordinary_url = f"http://127.0.0.1:{ordinary.server_address[1]}" + child = ( + "import subprocess,sys,threading,urllib.request;" + f"protected={_fetch_source(protected_url)!r};" + f"ordinary={_fetch_source(ordinary_url)!r};" + f"print(urllib.request.urlopen({protected_url!r},timeout=5).read().decode());" + "threaded=[];" + f"thread=threading.Thread(target=lambda:threaded.append(urllib.request.urlopen({protected_url!r},timeout=5).read().decode()));" + "thread.start();thread.join();print(threaded[0]);" + "blocked=subprocess.run([sys.executable,'-c',protected]);" + "print(f'blocked={blocked.returncode}',flush=True);" + "permitted=subprocess.run([sys.executable,'-c',ordinary]);" + "print(f'ordinary={permitted.returncode}')" + ) + target = ("127.0.0.1", protected.server_address[1]) + guard = ProcessConnectionGuard(tmp_path / "guard", {target}, {target}, backend="ptrace") + process = None + try: + guard.start() + process = await asyncio.create_subprocess_exec( + sys.executable, + str(_GUARD_PATH), + "--backend", + "ptrace", + str(guard.socket_path), + "--", + sys.executable, + "-c", + child, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await guard.wait_ready() + stdout, stderr = await asyncio.wait_for(process.communicate(), 20) + assert process.returncode == 0, stderr.decode(errors="replace") + lines = stdout.decode().splitlines() + assert lines[:2] == ["ok", "ok"] + assert lines[2] != "blocked=0" + assert lines[3:] == ["ok", "ordinary=0"] + finally: + if process is not None and process.returncode is None: + process.kill() + await process.wait() + guard.close() + for server, thread in ( + (protected, protected_thread), + (ordinary, ordinary_thread), + ): + server.shutdown() + server.server_close() + thread.join() + + +@pytest.mark.skipif(usable_bwrap() is None, reason="bubblewrap isolation is unavailable") +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" + + proxy_environment = await ssh.create_process( + f"exec {shlex.quote(sys.executable)} -c {shlex.quote(_proxy_environment_source())}", + connections=(connection,), + ) + proxy_environment_result = await proxy_environment.wait() + assert proxy_environment_result.returncode == 0 + assert isinstance(proxy_environment_result.stdout, bytes) + proxy_values = proxy_environment_result.stdout.decode().splitlines() + assert proxy_values[:2] == [f"http://127.0.0.1:{BRIDGE_PORT}"] * 2 + assert "ordinary.hud.invalid" in proxy_values[2].split(",") + + threaded = await ssh.create_process( + f"exec {shlex.quote(sys.executable)} -c " + f"{shlex.quote(_threaded_fetch_source(connection.client_url))}", + connections=(connection,), + ) + threaded_result = await threaded.wait() + assert threaded_result.returncode == 0 + assert threaded_result.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 2db2d2fe7..7cbf52916 100644 --- a/hud/environment/tests/test_workspace.py +++ b/hud/environment/tests/test_workspace.py @@ -27,10 +27,18 @@ import asyncssh import pytest -from hud.capabilities import SSHClient +from hud.capabilities import Connection, SSHClient +from hud.capabilities.ssh import PROCESS_CONNECTIONS_REQUEST from hud.environment import namespace as namespace_mod +from hud.environment import process_guard as process_guard_mod from hud.environment import workspace as workspace_mod -from hud.environment.egress import Peer, _field, _Unrelayable +from hud.environment.egress import ( + ConnectionRelay, + Peer, + WorkspaceRoute, + _field, + _Unrelayable, +) from hud.environment.workspace import Bubblewrap, Mount, Workspace from hud.utils.process import ProcessGroup, ProcessResult @@ -932,6 +940,54 @@ async def test_namespace_management_does_not_share_process_connections( management.wait_closed.assert_awaited_once_with() +def test_process_guard_executes_projected_file_without_module_reentry() -> None: + guard = cast( + "Any", + SimpleNamespace( + backend="ptrace", + sandbox_helper="/tmp/.hud-process-connection/process_guard.py", + sandbox_socket="/tmp/.hud-process-connection/control.sock", + ), + ) + argv = workspace_mod._guarded_process_argv( + "/usr/bin/python3", + guard, + ["bash", "-lc", "true"], + ) + + assert argv == [ + "/usr/bin/python3", + "/tmp/.hud-process-connection/process_guard.py", + "--backend", + "ptrace", + "/tmp/.hud-process-connection/control.sock", + "--", + "bash", + "-lc", + "true", + ] + result = subprocess.run( + [sys.executable, str(Path(workspace_mod.__file__).with_name("process_guard.py")), "--help"], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert result.stderr == "" + + +def test_process_guard_selects_probed_ptrace_backend(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(process_guard_mod, "_backend", None) + monkeypatch.setattr(process_guard_mod, "_backend_probed", False) + monkeypatch.setattr(process_guard_mod.sys, "platform", "linux") + run = Mock(return_value=SimpleNamespace(returncode=0, stdout="ptrace\n")) + monkeypatch.setattr(process_guard_mod.subprocess, "run", run) + + assert process_guard_mod._detected_backend() == "ptrace" + assert process_guard_mod._detected_backend() == "ptrace" + run.assert_called_once() + + @pytest.mark.asyncio async def test_namespace_host_only_terminates_a_used_session_holder( tmp_path: Path, @@ -1165,6 +1221,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 @@ -1607,6 +1685,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", [ @@ -1842,7 +1979,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)) @@ -1850,6 +1987,66 @@ async def capture_spawn(*_args: str, **kwargs: Any) -> None: assert captured.value.env == {"PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")} +@pytest.mark.asyncio +async def test_process_guard_preserves_session_proxy_environment( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + captured: list[dict[str, Any]] = [] + + class Guard: + sandbox_helper = "/run/hud/process_guard.py" + sandbox_socket = "/run/hud/process_guard.sock" + backend = "notify" + + def __init__(self, directory: Path, *_args: object) -> None: + self.directory = directory + + def start(self) -> None: + pass + + def close(self) -> None: + pass + + def capture_bwrap(*_args: object, **kwargs: Any) -> list[str]: + captured.append(kwargs) + raise RuntimeError("captured") + + proxy = { + "http_proxy": "http://127.0.0.1:3128", + "https_proxy": "http://127.0.0.1:3128", + "no_proxy": "127.0.0.1,localhost,main", + } + monkeypatch.setenv("HUD_API_KEY", "server-secret") + monkeypatch.setattr(workspace_mod, "_process_guard_interpreter", lambda: "/usr/bin/python3") + monkeypatch.setattr(workspace_mod, "ProcessConnectionGuard", Guard) + monkeypatch.setattr(Workspace, "supports_process_connections", True) + ws = Workspace(tmp_path / "root", env={"WORKSPACE_ENV": "present"}) + ws._process_connections["inference"] = cast("Any", object()) + ws._egress = cast("Any", SimpleNamespace(environment=lambda: proxy)) + monkeypatch.setattr(ws, "sandbox_pid", AsyncMock(return_value=7)) + monkeypatch.setattr(ws, "_process_connection_targets", lambda _names: frozenset()) + monkeypatch.setattr(ws, "bwrap_argv", capture_bwrap) + process = SimpleNamespace( + term_type="xterm-256color", + command="true", + env={PROCESS_CONNECTIONS_REQUEST: '["inference"]'}, + channel=SimpleNamespace(is_closing=Mock(return_value=False)), + stderr=SimpleNamespace(write=Mock()), + exit=Mock(), + ) + + await ws._handle_process(cast("Any", process)) + + assert len(captured) == 1 + session_env = captured[0]["env"] + assert {name: session_env[name] for name in proxy} == proxy + assert session_env["WORKSPACE_ENV"] == "present" + assert session_env["TERM"] == "xterm-256color" + assert "HUD_API_KEY" not in session_env + assert captured[0]["inherit_host_env"] is False + assert captured[0]["inherit_workspace_env"] is False + + @pytest.mark.asyncio async def test_namespace_wait_status_is_forwarded_to_ssh_client( tmp_path: Path, monkeypatch: pytest.MonkeyPatch diff --git a/hud/environment/workspace.py b/hud/environment/workspace.py index 10be8d7e3..8d7e8b0db 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 @@ -48,6 +57,40 @@ _INVALID_DWORD = 0xFFFFFFFF +def _process_guard_interpreter() -> str | None: + candidates = ( + Path("/usr/bin/python3"), + Path("/usr/local/bin/python3"), + Path(sys.executable).resolve(), + ) + return next( + ( + str(candidate) + for candidate in dict.fromkeys(candidates) + if candidate.is_file() + and os.access(candidate, os.X_OK) + and candidate.is_relative_to("/usr") + ), + None, + ) + + +def _guarded_process_argv( + interpreter: str, + guard: ProcessConnectionGuard, + argv: Sequence[str], +) -> list[str]: + return [ + interpreter, + guard.sandbox_helper, + "--backend", + guard.backend, + guard.sandbox_socket, + "--", + *argv, + ] + + class _WindowsJob: """Windows Job Object which owns a subprocess and all of its descendants.""" @@ -516,6 +559,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 @@ -633,6 +677,43 @@ 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() + + @property + def supports_process_connections(self) -> bool: + return ( + self.bwrap_available + and self.owns_netns + and _process_guard_interpreter() is not None + 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. @@ -830,6 +911,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 @@ -1599,7 +1681,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 @@ -1608,11 +1718,59 @@ 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: + guard_interpreter = _process_guard_interpreter() + if ( + pid is None + or not self.supports_process_connections + or guard_interpreter is None + ): + 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 = _guarded_process_argv( + guard_interpreter, + guard, + [ + *self._drop_argv(), + *shell_command, + ], + ) + argv = self.bwrap_argv( + guarded_command, + env=self._full_env(session_env), + inherit_host_env=False, + inherit_workspace_env=False, + 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 @@ -1623,6 +1781,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()) @@ -1746,6 +1906,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] = ( @@ -1757,7 +1918,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: @@ -1766,18 +1927,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]) @@ -1895,6 +2073,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/run.py b/hud/eval/run.py index 123b260e2..572cef395 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -42,10 +42,13 @@ 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.capabilities import Connection from hud.clients.client import HudClient + from hud.environment import WorkspaceRoute from .runtime import Provider from .runtime.core import RuntimeConfig @@ -63,9 +66,10 @@ def validate_rollout_timeouts( verifier_runtime_config: RuntimeConfig | None, ) -> float | None: """Validate configured phase limits and return the effective agent timeout.""" - from hud.agents.tool_agent import ToolAgent + from hud.agents.types import AgentConfig - agent_timeout = agent.config.timeout_seconds if isinstance(agent, ToolAgent) else None + config = getattr(agent, "config", None) + agent_timeout = config.timeout_seconds if isinstance(config, AgentConfig) else None if task.agent_config is not None: agent_timeout = task.agent_config.get("timeout_seconds", agent_timeout) @@ -200,11 +204,15 @@ def __init__( args: dict[str, Any], *, best_effort_grade: bool = False, + runtime_config: RuntimeConfig | 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.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 @@ -444,6 +452,8 @@ async def rollout( group_id: str | None = None, trace_id: str | None = None, rollout_timeout: float | None = None, + connections: Sequence[Connection] = (), + workspace_routes: Sequence[WorkspaceRoute] = (), ) -> Run: """Drive one task to a graded :class:`Run` here, against ``runtime``'s channel. @@ -477,11 +487,12 @@ async def rollout( """ from .runtime.core import resolve_runtime_config + actor_runtime_config = resolve_runtime_config(runtime, task) agent_timeout = validate_rollout_timeouts( task, agent, rollout_timeout, - actor_runtime_config=resolve_runtime_config(runtime, task), + actor_runtime_config=actor_runtime_config, verifier_runtime_config=( resolve_runtime_config(runtime, task.verifier) if task.verifier is not None else None ), @@ -499,9 +510,10 @@ async def rollout( # trace to it on enter. Only LLM tool agents carry an inference-model slug # (``config.model``); robot/other agents have none. Local import avoids an # eval<->agents import cycle. - from hud.agents.tool_agent import ToolAgent + from hud.agents.types import AgentConfig - agent_model = agent.config.model if isinstance(agent, ToolAgent) else None + config = getattr(agent, "config", None) + agent_model = config.model if isinstance(config, AgentConfig) else None with set_trace_context(trace_id, parent_trace_id=parent_trace_id): await trace_enter( trace_id, @@ -535,42 +547,51 @@ 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, + connections=connections, + ) as actor_client: client = actor_client live = Run( actor_client, task.id, task.args, best_effort_grade=task.verifier is not None, + runtime_config=addr.config or actor_runtime_config, + connections=connections, ) 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.connections.clear() _phase = "grading" if verifier is not None: @@ -670,6 +691,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.connections.clear() run.trace.trace_id = trace_id run.job_id = job_id run.group_id = group_id diff --git a/hud/eval/runtime/__init__.py b/hud/eval/runtime/__init__.py index 46d8ab207..ba815d4a1 100644 --- a/hud/eval/runtime/__init__.py +++ b/hud/eval/runtime/__init__.py @@ -1,6 +1,6 @@ """Runtime placement and provider configuration.""" -from .compose import ComposeProject +from .compose import ComposeProject, DockerBindMount from .core import ( Provider, Runtime, @@ -21,6 +21,7 @@ __all__ = [ "ComposeProject", "DaytonaRuntime", + "DockerBindMount", "DockerRuntime", "HUDRuntime", "HostedRuntime", diff --git a/hud/eval/runtime/compose.py b/hud/eval/runtime/compose.py index e1ec7e4c1..8506c76f0 100644 --- a/hud/eval/runtime/compose.py +++ b/hud/eval/runtime/compose.py @@ -11,7 +11,7 @@ import tarfile import tempfile from dataclasses import dataclass -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING, Any import yaml @@ -29,7 +29,7 @@ from yaml.nodes import MappingNode, Node, ScalarNode, SequenceNode if TYPE_CHECKING: - from collections.abc import Iterator, Mapping + from collections.abc import Iterator, Mapping, Sequence _COMPOSE_VARIABLE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") @@ -443,6 +443,47 @@ class ComposeLaunchFiles: archive: Path | None +@dataclass(frozen=True, slots=True) +class DockerBindMount: + """A provider-owned bind mount injected into a Docker environment.""" + + source: Path + target: PurePosixPath + read_only: bool = True + + def __init__( + self, + source: str | Path, + target: str | PurePosixPath, + read_only: bool = True, + ) -> None: + resolved_source = Path(source) + resolved_target = PurePosixPath(target) + if not resolved_source.is_absolute(): + raise ValueError("Docker bind mount source must be absolute") + if not resolved_target.is_absolute(): + raise ValueError("Docker bind mount target must be absolute") + if "," in str(resolved_source) or "," in str(resolved_target): + raise ValueError("Docker bind mount paths cannot contain commas") + object.__setattr__(self, "source", resolved_source) + object.__setattr__(self, "target", resolved_target) + object.__setattr__(self, "read_only", read_only) + + def docker_argument(self) -> str: + argument = f"type=bind,source={self.source},target={self.target}" + return f"{argument},readonly" if self.read_only else argument + + def compose_volume(self) -> dict[str, str | bool]: + volume: dict[str, str | bool] = { + "type": "bind", + "source": str(self.source), + "target": str(self.target), + } + if self.read_only: + volume["read_only"] = True + return volume + + class ComposeProject(BaseModel): """A Compose recipe and the project data it may need at runtime.""" @@ -512,6 +553,7 @@ def stage( port_service: str = "main", seccomp: str | Path, service_socket: str | None = None, + bind_mounts: Sequence[DockerBindMount] = (), env_vars: Mapping[str, str] | None = None, cpu: float | None = None, memory_mb: int | None = None, @@ -528,14 +570,17 @@ def stage( "apparmor=unconfined", ], } + volumes = [mount.compose_volume() for mount in bind_mounts] if service_socket is not None: - main["volumes"] = [ + volumes.append( { "type": "bind", "source": service_socket, "target": "/media/hud/docker.sock", } - ] + ) + if volumes: + main["volumes"] = volumes if env_vars: main["environment"] = dict(env_vars) if cpu is not None: diff --git a/hud/eval/runtime/docker.py b/hud/eval/runtime/docker.py index 4582e2df4..75c92ee88 100644 --- a/hud/eval/runtime/docker.py +++ b/hud/eval/runtime/docker.py @@ -18,7 +18,7 @@ from hud.utils.docker import docker as _docker from hud.utils.process import create_process_group_exec, finish_output, stream_output -from .compose import ComposeConfig +from .compose import ComposeConfig, DockerBindMount from .core import Runtime, RuntimeConfig, validate_session_id if TYPE_CHECKING: @@ -134,12 +134,14 @@ def __init__( *, port: int = 8765, run_args: Sequence[str] = (), + bind_mounts: Sequence[DockerBindMount] = (), compose_service_socket: str | Path | None = None, runtime_config: RuntimeConfig | dict[str, Any] | None = None, env_vars: Mapping[str, str] | None = None, ) -> None: self.port = port self.run_args = tuple(run_args) + self.bind_mounts = tuple(bind_mounts) self.env_vars = dict(env_vars or {}) self.compose_service_socket = ( str(Path(compose_service_socket)) if compose_service_socket is not None else None @@ -204,6 +206,7 @@ async def __call__(self, task: Task) -> AsyncIterator[Runtime]: port_service=port_service, seccomp=_DOCKER_SECCOMP_PROFILE, service_socket=service_socket, + bind_mounts=self.bind_mounts, env_vars=self.env_vars, cpu=resources.cpu if resources is not None else None, memory_mb=resources.memory_mb if resources is not None else None, @@ -296,12 +299,16 @@ async def __call__(self, task: Task) -> AsyncIterator[Runtime]: env_args: list[str] = [] for key, value in self.env_vars.items(): env_args.extend(("--env", f"{key}={value}")) + mount_args: list[str] = [] + for mount in self.bind_mounts: + mount_args.extend(("--mount", mount.docker_argument())) out, _ = await _docker( "run", "--detach", *self.run_args, *env_args, *resource_args, + *mount_args, *_DOCKER_SECURITY_ARGS, "--publish", f"127.0.0.1::{self.port}", diff --git a/hud/eval/runtime/hosted.py b/hud/eval/runtime/hosted.py index cfe261e07..c0e4569da 100644 --- a/hud/eval/runtime/hosted.py +++ b/hud/eval/runtime/hosted.py @@ -32,7 +32,7 @@ class HostedRuntime: agent runs alongside the task environment. This process only submits the rollout and polls the trace to completion, folding the result into a :class:`~hud.eval.run.Run`. Because the agent runs remotely, its identity - travels via :func:`_agent_spec`. + travels via :func:`hud.agents.dump_agent`. ``run_timeout`` is a deprecated constructor alias for ``rollout_timeout``. A local cancel (Ctrl-C) requests remote cancellation before propagating. @@ -143,14 +143,9 @@ async def _submit_and_await( trace_id: str, parent_trace_id: str | None, ) -> dict[str, Any]: - from hud.agents.tool_agent import ToolAgent + from hud.agents.registry import dump_agent - if not isinstance(agent, ToolAgent): - raise ValueError( - f"hosted execution requires a gateway agent that can serialize its " - f"identity (Claude/OpenAI/Gemini/OpenAIChat); got {type(agent).__name__}" - ) - spec = agent.hosted_spec() + spec = dump_agent(agent) if task.agent_config: spec = { **spec, diff --git a/hud/eval/tests/test_docker_provider.py b/hud/eval/tests/test_docker_provider.py index 3e86c8bef..054416fd3 100644 --- a/hud/eval/tests/test_docker_provider.py +++ b/hud/eval/tests/test_docker_provider.py @@ -26,6 +26,7 @@ import hud.utils.process as process_module from hud.eval.runtime import ( DaytonaRuntime, + DockerBindMount, DockerRuntime, ModalRuntime, RuntimeConfig, @@ -603,6 +604,46 @@ async def test_acquisition_publishes_ephemeral_port_and_removes_container( assert capsys.readouterr().out == "ImportError: boom\n" +async def test_docker_runtime_injects_provider_bind_mount( + tmp_path: Path, + docker_log: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_fake_docker(tmp_path, port_behavior="echo 127.0.0.1:43210", monkeypatch=monkeypatch) + bundle = tmp_path / "agents" / "codex" + bundle.mkdir(parents=True) + + provider = DockerRuntime( + "img:tag", + bind_mounts=(DockerBindMount(bundle, "/usr/local/lib/agents/codex"),), + ) + async with provider(_row()): + pass + + assert (await _docker_calls(docker_log))[0] == ( + f"run --detach --mount type=bind,source={bundle}," + "target=/usr/local/lib/agents/codex,readonly " + f"{_docker_security_args()} --publish 127.0.0.1::8765 img:tag" + ) + + +@pytest.mark.parametrize( + ("source", "target", "message"), + [ + ("relative", "/opt/agents", "source must be absolute"), + ("/opt/agents", "relative", "target must be absolute"), + ("/opt/agents,old", "/opt/agents", "paths cannot contain commas"), + ], +) +def test_docker_bind_mount_rejects_ambiguous_paths( + source: str, + target: str, + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + DockerBindMount(source, target) + + async def test_docker_session_archives_inside_the_container( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -995,6 +1036,53 @@ async def fake_docker(*args: str, **_kwargs: Any) -> tuple[str, str]: ] +async def test_docker_runtime_stages_provider_mount_with_service_socket( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + compose = tmp_path / "compose.yaml" + compose.write_text("services:\n main:\n image: hud-env:one\n", encoding="utf-8") + bundle = tmp_path / "agents" / "claude" + bundle.mkdir(parents=True) + rendered: dict[str, Any] = {} + + async def fake_docker(*args: str, **_kwargs: Any) -> tuple[str, str]: + if args[-4:] == ("up", "--detach", "--build", "--remove-orphans"): + files = [Path(args[index + 1]) for index, value in enumerate(args) if value == "--file"] + rendered.update(json.loads(files[1].read_text("utf-8"))) + if args[-3:] == ("port", "main", "8765"): + return "127.0.0.1:43210\n", "" + return "", "" + + monkeypatch.setattr(runtime_module, "_docker", fake_docker) + task = Task( + env="any-env", + id="t", + runtime_config=RuntimeConfig(compose=ComposeProject(document=compose, service_access=True)), + ) + provider = DockerRuntime( + compose_service_socket="/vm/run/docker.sock", + bind_mounts=(DockerBindMount(bundle, "/usr/local/lib/agents/claude"),), + ) + + async with provider(task): + pass + + assert rendered["services"]["main"]["volumes"] == [ + { + "type": "bind", + "source": str(bundle), + "target": "/usr/local/lib/agents/claude", + "read_only": True, + }, + { + "type": "bind", + "source": "/vm/run/docker.sock", + "target": "/media/hud/docker.sock", + }, + ] + + def test_docker_runtime_accepts_only_one_environment_definition(tmp_path: Path) -> None: with pytest.raises(ValueError, match="either image or compose"): RuntimeConfig( diff --git a/hud/eval/tests/test_hosted.py b/hud/eval/tests/test_hosted.py index 5f95654f7..24d22666b 100644 --- a/hud/eval/tests/test_hosted.py +++ b/hud/eval/tests/test_hosted.py @@ -18,6 +18,8 @@ import pytest +from hud.agents import dump_agent +from hud.agents.claude import ClaudeCLIAgent, ClaudeCLIConfig from hud.agents.openai_compatible import OpenAIChatAgent from hud.agents.types import OpenAIChatConfig from hud.eval.job import Job @@ -92,14 +94,14 @@ def test_runtime_constructor_timeout_is_a_deprecated_alias(runtime_type: type[An assert runtime.run_timeout == 90.0 -def test_hosted_spec_serializes_full_config() -> None: +def test_dump_agent_serializes_full_config() -> None: agent = _agent() agent.config.system_prompt = "be brief" agent.config.max_steps = 7 agent.config.timeout_seconds = 3600 agent.config.tool_timeout_seconds = 1800 - spec = agent.hosted_spec() + spec = dump_agent(agent) assert spec["type"] == "openai_compatible" config = spec["config"] @@ -116,7 +118,7 @@ def test_hosted_spec_serializes_full_config() -> None: assert "hosted_tools" not in config -def test_create_agent_hosted_spec_preserves_training_config( +def test_dump_agent_preserves_training_config( monkeypatch: pytest.MonkeyPatch, ) -> None: """The constructor builds the runtime client without putting it in config.""" @@ -151,7 +153,7 @@ class _GatewayStub: assert agent.config.model_client is None assert agent.oai is client - spec = agent.hosted_spec() + spec = dump_agent(agent) config = spec["config"] assert spec["type"] == "openai_compatible" assert config["model"] == "arith-rl" @@ -163,17 +165,15 @@ class _GatewayStub: assert "model_client" not in config -def test_hosted_spec_rejects_custom_model_client() -> None: +def test_dump_agent_rejects_custom_model_client() -> None: agent = _agent() agent.config = OpenAIChatConfig(model="m", model_client=object()) - with pytest.raises(ValueError, match="custom model_client"): - agent.hosted_spec() - with pytest.raises(ValueError, match="HUDRuntime"): - agent.hosted_spec() + with pytest.raises(ValueError, match=r"custom model_client.*HUDRuntime"): + dump_agent(agent) @pytest.mark.asyncio -async def test_run_rejects_non_gateway_agent() -> None: +async def test_run_rejects_unregistered_agent() -> None: """An agent that can't serialize its identity yields a failed Run, not a crash.""" run = await HostedRuntime(poll_interval=0.0).run( Task(env="e", id="x"), @@ -181,7 +181,7 @@ async def test_run_rejects_non_gateway_agent() -> None: job_id="j", ) assert run.trace.is_error - assert "gateway agent" in (run.trace.error or "") + assert "registered types" in (run.trace.error or "") @pytest.mark.asyncio @@ -259,6 +259,35 @@ async def test_run_submits_and_polls_to_terminal(monkeypatch: pytest.MonkeyPatch assert payload["agent"]["config"]["timeout_seconds"] == 45.0 +@pytest.mark.asyncio +async def test_run_submits_registered_cli_agent(monkeypatch: pytest.MonkeyPatch) -> None: + platform = _FakePlatform([{"status": "completed", "reward": 1.0}]) + monkeypatch.setattr( + "hud.eval.runtime.hosted.PlatformClient.from_settings", classmethod(lambda cls: platform) + ) + agent = ClaudeCLIAgent( + ClaudeCLIConfig( + model="claude-sonnet-4-6", + max_steps=23, + use_hud_gateway=True, + ) + ) + + run = await HostedRuntime(poll_interval=0.0).run( + Task(env="coding", id="solve"), + agent, + job_id=uuid.uuid4().hex, + trace_id=uuid.uuid4().hex, + ) + + assert run.reward == 1.0 + submitted = platform.posts[0][1]["agent"] + assert submitted["type"] == "claude_cli" + assert submitted["config"]["model"] == "claude-sonnet-4-6" + assert submitted["config"]["max_steps"] == 23 + assert submitted["config"]["use_hud_gateway"] is True + + @pytest.mark.asyncio async def test_run_preserves_runtime_config_null_override( monkeypatch: pytest.MonkeyPatch, diff --git a/hud/eval/tests/test_rollout.py b/hud/eval/tests/test_rollout.py index 3942d097e..4db19fe3f 100644 --- a/hud/eval/tests/test_rollout.py +++ b/hud/eval/tests/test_rollout.py @@ -32,9 +32,17 @@ 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.clients.client import HudClient from hud.environment import Answer, Environment -from hud.eval import Job, LocalRuntime, Runtime, SubprocessRuntime, Task, Taskset +from hud.eval import ( + Job, + LocalRuntime, + Runtime, + SubprocessRuntime, + Task, + Taskset, +) from hud.eval.run import Run, rollout from hud.telemetry.context import get_current_trace_id, get_trace_headers, set_trace_context @@ -175,6 +183,19 @@ async def test_rollout_returns_graded_run_with_trace_id(env_file: Path) -> None: assert run.runtime.startswith("tcp://127.0.0.1:") +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"}, + ) + run = Run(None, "task", {}, connections=(connection,)) + + 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: env = Environment("reviewed") completed: list[str] = [] diff --git a/hud/tests/test_init_module.py b/hud/tests/test_init_module.py index 62bd04723..83437c458 100644 --- a/hud/tests/test_init_module.py +++ b/hud/tests/test_init_module.py @@ -20,6 +20,7 @@ def test_all_exports(self): expected = [ "Chat", "ComposeProject", + "Connection", "DockerRuntime", "Environment", "Grade", diff --git a/hud/types.py b/hud/types.py index 3b530a6e0..05982985d 100644 --- a/hud/types.py +++ b/hud/types.py @@ -44,33 +44,39 @@ if TYPE_CHECKING: from collections.abc import Callable - from hud.agents.claude import ClaudeAgent - from hud.agents.gemini import GeminiAgent - from hud.agents.openai import OpenAIAgent - from hud.agents.openai_compatible import OpenAIChatAgent - from hud.agents.types import ClaudeConfig, GeminiConfig, OpenAIChatConfig, OpenAIConfig - - AgentClass: TypeAlias = type[ClaudeAgent | GeminiAgent | OpenAIAgent | OpenAIChatAgent] - AgentConfigClass: TypeAlias = type[ - ClaudeConfig | GeminiConfig | OpenAIConfig | OpenAIChatConfig - ] + from hud.agents.base import Agent + from hud.agents.types import AgentConfig T = TypeVar("T") class AgentType(StrEnum): CLAUDE = "claude" + CLAUDE_CLI = "claude_cli" + CODEX_CLI = "codex_cli" OPENAI = "openai" GEMINI = "gemini" OPENAI_COMPATIBLE = "openai_compatible" @property - def cls(self) -> AgentClass: + def is_cli(self) -> bool: + return self in (AgentType.CLAUDE_CLI, AgentType.CODEX_CLI) + + @property + def cls(self) -> type[Agent]: match self: case AgentType.CLAUDE: from hud.agents import ClaudeAgent return ClaudeAgent + case AgentType.CLAUDE_CLI: + from hud.agents import ClaudeCLIAgent + + return ClaudeCLIAgent + case AgentType.CODEX_CLI: + from hud.agents import CodexCLIAgent + + return CodexCLIAgent case AgentType.OPENAI: from hud.agents import OpenAIAgent @@ -85,13 +91,24 @@ def cls(self) -> AgentClass: return OpenAIChatAgent @property - def config_cls(self) -> AgentConfigClass: + def config_cls(self) -> type[AgentConfig]: """Get config class without importing agent (avoids SDK dependency).""" - from hud.agents.types import ClaudeConfig, GeminiConfig, OpenAIChatConfig, OpenAIConfig + from hud.agents.types import ( + ClaudeCLIConfig, + ClaudeConfig, + CodexCLIConfig, + GeminiConfig, + OpenAIChatConfig, + OpenAIConfig, + ) match self: case AgentType.CLAUDE: return ClaudeConfig + case AgentType.CLAUDE_CLI: + return ClaudeCLIConfig + case AgentType.CODEX_CLI: + return CodexCLIConfig case AgentType.OPENAI: return OpenAIConfig case AgentType.GEMINI: @@ -99,12 +116,19 @@ def config_cls(self) -> AgentConfigClass: case AgentType.OPENAI_COMPATIBLE: return OpenAIChatConfig + def instantiate(self, config: AgentConfig) -> Agent: + return cast("Any", self.cls)(config) + @property def gateway_provider(self) -> str: """Default provider client used when this agent type is a gateway shortcut.""" match self: case AgentType.CLAUDE: return "anthropic" + case AgentType.CLAUDE_CLI: + return "anthropic" + case AgentType.CODEX_CLI: + return "openai" case AgentType.OPENAI: return "openai" case AgentType.GEMINI: @@ -114,15 +138,15 @@ def gateway_provider(self) -> str: @classmethod def of(cls, agent: object) -> AgentType | None: - """The gateway agent type *agent* is an instance of, or ``None``. + """The registered agent type *agent* is an instance of, or ``None``. - Reverse of :attr:`cls`. Provider extras (anthropic, google-genai, ...) + Reverse of :attr:`cls`. Agent extras (anthropic, google-genai, ...) may be uninstalled, so importing a type's agent class can fail; that - simply means *agent* is not that type. ``None`` for a custom ``Agent`` - subclass that is not one of the gateway shortcuts. + simply means *agent* is not that type. ``None`` means the ``Agent`` + implementation is not registered for reconstruction. """ for agent_type in cls: - with contextlib.suppress(Exception): + with contextlib.suppress(ImportError): if isinstance(agent, agent_type.cls): return agent_type return None @@ -302,7 +326,7 @@ def emit(self, *, trace_id: str | None = None) -> None: #: Why the rollout stopped; anything but "done" means a limit cut it off. StopReason: TypeAlias = Literal["done", "max_steps", "length", "timeout", "malformed_tool_call"] -#: The configurable subset of stop reasons (``AgentConfig.stop_on``): policy +#: The configurable subset of stop reasons (``ToolAgentConfig.stop_on``): policy #: conditions the loop may either stop on or answer with an error result. StopCondition: TypeAlias = Literal["length", "malformed_tool_call"]