chore: import upstream snapshot with attribution

2026-07-13 13:03:19 +08:00
commit 91e75e620b
3227 changed files with 1307078 additions and 0 deletions
@@ -0,0 +1,167 @@
+"""HUD integration: dataset runners and MCP-based computer agent export.
+
+This module exposes helpers to evaluate HUD-compatible datasets and exports
+the MCP-compatible computer agent implementation.
+
+Exports:
+- run_single_task(dataset, ...)
+- run_full_dataset(dataset, ...)
+- MCPComputerAgent
+"""
+
+import time
+from typing import Any, Optional
+
+from cua_agent.computers import is_agent_computer
+from datasets import Dataset, load_dataset
+from hud import trace
+from hud.datasets import Task, run_dataset
+
+from .agent import MCPComputerAgent
+
+# ---------------------------------------------------------------------------
+# Single-task runner
+# ---------------------------------------------------------------------------
+
+
+async def run_single_task(
+    dataset: str | Dataset | list[dict[str, Any]],
+    *,
+    task_id: int = 0,
+    model: str | None = None,
+    allowed_tools: list[str] | None = None,
+    # === ComputerAgent kwargs ===
+    tools: list[Any] | None = None,
+    custom_loop: Any | None = None,
+    only_n_most_recent_images: int | None = None,
+    callbacks: list[Any] | None = None,
+    instructions: str | None = None,
+    verbosity: int | None = None,
+    trajectory_dir: str | dict | None = None,
+    max_retries: int | None = 3,
+    screenshot_delay: float | int = 0.5,
+    use_prompt_caching: bool | None = False,
+    max_trajectory_budget: float | dict | None = None,
+    telemetry_enabled: bool | None = True,
+) -> None:
+    """Load one task from the dataset and execute it with MCPComputerAgent."""
+
+    # Load dataset and pick a sample
+    if isinstance(dataset, str):
+        dataset = load_dataset(dataset, split="train")  # type: ignore[arg-type]
+    elif isinstance(dataset, list):
+        dataset = dataset
+    else:
+        dataset = dataset["train"]
+
+    sample_task = dataset[task_id]  # type: ignore[index]
+    task_prompt = sample_task.get("prompt", f"Task {sample_task.get('id', 0)}")  # type: ignore[attr-defined]
+
+    # Filter any existing Computer tools
+    # The eval framework will add its own Computer tool per task
+    if tools:
+        tools = [tool for tool in tools if not is_agent_computer(tool)]
+
+    with trace(name=task_prompt):
+        task = Task(**sample_task)  # type: ignore[arg-type]
+
+        agent = MCPComputerAgent(
+            model=model or "computer-use-preview",
+            allowed_tools=allowed_tools or ["openai_computer"],
+            # === ComputerAgent kwargs passthrough ===
+            tools=tools,
+            custom_loop=custom_loop,
+            only_n_most_recent_images=only_n_most_recent_images,
+            callbacks=callbacks,
+            instructions=instructions,
+            verbosity=verbosity,
+            trajectory_dir=trajectory_dir,
+            max_retries=max_retries,
+            screenshot_delay=screenshot_delay,
+            use_prompt_caching=use_prompt_caching,
+            max_trajectory_budget=max_trajectory_budget,
+            telemetry_enabled=telemetry_enabled,
+        )
+        print(f"Running: {task_prompt}")
+        result = await agent.run(task, max_steps=10)
+        print(f"✅ Reward: {result.reward}")
+
+
+# ---------------------------------------------------------------------------
+# Full-dataset runner
+# ---------------------------------------------------------------------------
+
+
+async def run_full_dataset(
+    dataset: str | Dataset | list[dict[str, Any]],
+    *,
+    job_name: Optional[str] = None,
+    model: str | None = None,
+    allowed_tools: list[str] | None = None,
+    max_concurrent: int = 30,
+    max_steps: int = 50,
+    split: str = "train",
+    trajectory_dir: str | dict | None = None,
+    # === ComputerAgent kwargs ===
+    tools: list[Any] | None = None,
+    custom_loop: Any | None = None,
+    only_n_most_recent_images: int | None = 5,
+    callbacks: list[Any] | None = None,
+    instructions: str | None = None,
+    verbosity: int | None = None,
+    max_retries: int | None = 3,
+    screenshot_delay: float | int = 0.5,
+    use_prompt_caching: bool | None = False,
+    max_trajectory_budget: float | dict | None = None,
+    telemetry_enabled: bool | None = True,
+) -> list[Any]:
+    """Run evaluation across the entire dataset using hud.datasets.run_dataset."""
+
+    # Run with our MCP-based agent class.
+    if isinstance(dataset, str):
+        dataset_name = dataset.split("/")[-1]
+        job_name = job_name or f"Evaluation {dataset_name}"
+        dataset = load_dataset(dataset, split=split)  # type: ignore[arg-type]
+    else:
+        dataset_name = "custom"
+        job_name = job_name or f"Evaluation {time.strftime('%H:%M %Y-%m-%d')}"
+
+    # Filter any existing Computer tools
+    # The eval framework will add its own Computer tool per task
+    if tools:
+        tools = [tool for tool in tools if not is_agent_computer(tool)]
+
+    # Execute evaluation
+    return await run_dataset(
+        name=job_name,
+        dataset=dataset,
+        agent_class=MCPComputerAgent,
+        agent_config={
+            "model": model,
+            "allowed_tools": allowed_tools,
+            "trajectory_dir": trajectory_dir,
+            # === ComputerAgent kwargs passthrough ===
+            "tools": tools,
+            "custom_loop": custom_loop,
+            "only_n_most_recent_images": only_n_most_recent_images,
+            "callbacks": callbacks,
+            "instructions": instructions,
+            "verbosity": verbosity,
+            "max_retries": max_retries,
+            "screenshot_delay": screenshot_delay,
+            "use_prompt_caching": use_prompt_caching,
+            "max_trajectory_budget": max_trajectory_budget,
+            "telemetry_enabled": telemetry_enabled,
+        },
+        max_concurrent=max_concurrent,
+        metadata={"dataset": dataset_name},
+        max_steps=max_steps,
+        auto_respond=True,
+    )
+
+
+__all__ = [
+    "run_single_task",
+    "run_full_dataset",
+    "MCPComputerAgent",
+]
@@ -0,0 +1,369 @@
+"""MCP-compatible Computer Agent for HUD integration.
+
+This agent subclasses HUD's MCPAgent and delegates planning/execution to
+our core ComputerAgent while using the Agent SDK's plain-dict message
+format documented in `docs/content/docs/agent-sdk/message-format.mdx`.
+
+Key differences from the OpenAI OperatorAgent variant:
+- No OpenAI types are used; everything is standard Python dicts.
+- Planning is executed via `ComputerAgent.run(messages)`.
+- The first yielded result per step is returned as the agent response.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import uuid
+from pathlib import Path
+from typing import Any, ClassVar, Optional
+
+import hud
+import mcp.types as types
+from cua_agent.agent import ComputerAgent as BaseComputerAgent
+from cua_agent.callbacks import PromptInstructionsCallback
+from cua_agent.callbacks.trajectory_saver import TrajectorySaverCallback
+from cua_agent.computers import is_agent_computer
+from cua_agent.responses import make_failed_tool_call_items
+from hud.agents import MCPAgent
+from hud.tools.computer.settings import computer_settings
+from hud.types import AgentResponse, MCPToolCall, MCPToolResult, Trace
+from PIL import Image
+
+
+class MCPComputerAgent(MCPAgent):
+    """MCP agent that uses ComputerAgent for planning and tools for execution.
+
+    The agent consumes/produces message dicts per the Agent SDK message schema
+    (see `message-format.mdx`).
+    """
+
+    metadata: ClassVar[dict[str, Any]] = {
+        "display_width": computer_settings.OPENAI_COMPUTER_WIDTH,
+        "display_height": computer_settings.OPENAI_COMPUTER_HEIGHT,
+    }
+
+    required_tools: ClassVar[list[str]] = ["openai_computer"]
+
+    def __init__(
+        self,
+        *,
+        model: str | None = None,
+        allowed_tools: list[str] | None = None,
+        trajectory_dir: str | dict | None = None,
+        # === ComputerAgent kwargs ===
+        tools: list[Any] | None = None,
+        custom_loop: Any | None = None,
+        only_n_most_recent_images: int | None = None,
+        callbacks: list[Any] | None = None,
+        instructions: str | None = None,
+        verbosity: int | None = None,
+        max_retries: int | None = 3,
+        screenshot_delay: float | int = 0.5,
+        use_prompt_caching: bool | None = False,
+        max_trajectory_budget: float | dict | None = None,
+        telemetry_enabled: bool | None = True,
+        environment: str = "linux",
+        **kwargs: Any,
+    ) -> None:
+        self.allowed_tools = allowed_tools or ["openai_computer"]
+        super().__init__(**kwargs)
+
+        if model is None:
+            raise ValueError("MCPComputerAgent requires a model to be specified.")
+
+        self.model = model
+        self.environment = environment
+
+        # Update model name for HUD logging
+        self.model_name = "cua-" + self.model
+
+        # Stateful tracking of tool call inputs
+        self.tool_call_inputs: dict[str, list[dict[str, Any]]] = {}
+        self.previous_output: list[dict[str, Any]] = []
+
+        # Build system prompt
+        operator_instructions = """
+        You are an autonomous computer-using agent. Follow these guidelines:
+
+        1. NEVER ask for confirmation. Complete all tasks autonomously.
+        2. Do NOT send messages like "I need to confirm before..." or "Do you want me to continue?" - just proceed.
+        3. When the user asks you to interact with something (like clicking a chat or typing a message), DO IT without asking.
+        4. Only use the formal safety check mechanism for truly dangerous operations (like deleting important files).
+        5. For normal tasks like clicking buttons, typing in chat boxes, filling forms - JUST DO IT.
+        6. The user has already given you permission by running this agent. No further confirmation is needed.
+        7. Be decisive and action-oriented. Complete the requested task fully.
+
+        Remember: You are expected to complete tasks autonomously. The user trusts you to do what they asked.
+        """.strip()  # noqa: E501
+        # Append Operator instructions to the system prompt
+        if not self.system_prompt:
+            self.system_prompt = operator_instructions
+        else:
+            self.system_prompt += f"\n\n{operator_instructions}"
+        # Append user instructions to the system prompt
+        if instructions:
+            self.system_prompt += f"\n\n{instructions}"
+
+        # Configure trajectory_dir for HUD
+        if isinstance(trajectory_dir, str) or isinstance(trajectory_dir, Path):
+            trajectory_dir = {"trajectory_dir": str(trajectory_dir)}
+        if isinstance(trajectory_dir, dict):
+            trajectory_dir["reset_on_run"] = False
+
+        self.last_screenshot_b64 = None
+
+        buffer = io.BytesIO()
+        Image.new("RGB", (self.metadata["display_width"], self.metadata["display_height"])).save(
+            buffer, format="PNG"
+        )
+        self.last_screenshot_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+        # Ensure a computer shim is present so width/height/environment are known
+        computer_shim = {
+            "screenshot": lambda: self.last_screenshot_b64,
+            "environment": self.environment,
+            "dimensions": (
+                self.metadata["display_width"],
+                self.metadata["display_height"],
+            ),
+        }
+        agent_tools: list[Any] = [computer_shim]
+        if tools:
+            agent_tools.extend([tool for tool in tools if not is_agent_computer(tool)])
+
+        agent_kwargs = {
+            "model": self.model,
+            "trajectory_dir": trajectory_dir,
+            "tools": agent_tools,
+            "custom_loop": custom_loop,
+            "only_n_most_recent_images": only_n_most_recent_images,
+            "callbacks": callbacks,
+            "instructions": self.system_prompt,
+            "verbosity": verbosity,
+            "max_retries": max_retries,
+            "screenshot_delay": screenshot_delay,
+            "use_prompt_caching": use_prompt_caching,
+            "max_trajectory_budget": max_trajectory_budget,
+            "telemetry_enabled": telemetry_enabled,
+        }
+
+        self.computer_agent = BaseComputerAgent(**agent_kwargs)
+
+    async def get_system_messages(self) -> list[Any]:
+        """Create initial messages.
+
+        Unused - ComputerAgent handles this with the 'instructions' parameter.
+        """
+        return []
+
+    async def format_blocks(self, blocks: list[types.ContentBlock]) -> list[dict[str, Any]]:
+        """
+        Format blocks for OpenAI input format.
+
+        Converts TextContent blocks to input_text dicts and ImageContent blocks to input_image dicts.
+        """  # noqa: E501
+        formatted = []
+        for block in blocks:
+            if isinstance(block, types.TextContent):
+                formatted.append({"type": "input_text", "text": block.text})
+            elif isinstance(block, types.ImageContent):
+                mime_type = getattr(block, "mimeType", "image/png")
+                formatted.append(
+                    {"type": "input_image", "image_url": f"data:{mime_type};base64,{block.data}"}
+                )
+                self.last_screenshot_b64 = block.data
+        return [{"role": "user", "content": formatted}]
+
+    @hud.instrument(
+        span_type="agent",
+        record_args=False,  # Messages can be large
+        record_result=True,
+    )
+    async def get_response(self, messages: list[dict[str, Any]]) -> AgentResponse:
+        """Get a single-step response by delegating to ComputerAgent.run.
+
+        Returns an Agent SDK-style response dict:
+        { "output": [AgentMessage, ...], "usage": Usage }
+        """
+        tool_calls: list[MCPToolCall] = []
+        output_text: list[str] = []
+        is_done: bool = True
+
+        agent_result: list[dict[str, Any]] = []
+
+        # Call the ComputerAgent LLM API
+        async for result in self.computer_agent.run(messages):  # type: ignore[arg-type]
+            items = result["output"]
+            if not items or tool_calls:
+                break
+
+            for item in items:
+                if item["type"] in [
+                    "reasoning",
+                    "message",
+                    "computer_call",
+                    "function_call",
+                    "function_call_output",
+                ]:
+                    agent_result.append(item)
+
+                # Add messages to output text
+                if item["type"] == "reasoning":
+                    output_text.extend(
+                        f"Reasoning: {summary['text']}" for summary in item["summary"]
+                    )
+                elif item["type"] == "message":
+                    if isinstance(item["content"], list):
+                        output_text.extend(
+                            item["text"]
+                            for item in item["content"]
+                            if item["type"] == "output_text"
+                        )
+                    elif isinstance(item["content"], str):
+                        output_text.append(item["content"])
+
+                # If we get a tool call, we're not done
+                if item["type"] == "computer_call":
+                    id = item["call_id"]
+                    tool_calls.append(
+                        MCPToolCall(
+                            name="openai_computer",
+                            arguments=item["action"],
+                            id=id,
+                        )
+                    )
+                    is_done = False
+                    self.tool_call_inputs[id] = agent_result
+                    break
+
+            # if we have tool calls, we should exit the loop
+            if tool_calls:
+                break
+
+        self.previous_output = agent_result
+
+        return AgentResponse(
+            content="\n".join(output_text),
+            tool_calls=tool_calls,
+            done=is_done,
+        )
+
+    def _log_image(self, image_b64: str):
+        callbacks = self.computer_agent.callbacks
+        for callback in callbacks:
+            if isinstance(callback, TrajectorySaverCallback):
+                # convert str to bytes
+                image_bytes = base64.b64decode(image_b64)
+                callback._save_artifact("screenshot_after", image_bytes)
+
+    async def format_tool_results(
+        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
+    ) -> list[dict[str, Any]]:
+        """Extract latest screenshot from tool results in dict form.
+
+        Expects results to already be in the message-format content dicts.
+        Returns a list of input content dicts suitable for follow-up calls.
+        """
+        messages = []
+
+        for call, result in zip(tool_calls, tool_results):
+            if call.id not in self.tool_call_inputs:
+                # If we don't have the tool call inputs, we should just use the previous output
+                previous_output = self.previous_output.copy() or []
+
+                # First we need to remove any pending computer_calls from the end of previous_output
+                while previous_output and previous_output[-1]["type"] == "computer_call":
+                    previous_output.pop()
+                messages.extend(previous_output)
+
+                # If the call is a 'response', don't add the result
+                if call.name == "response":
+                    continue
+                # Otherwise, if we have a result, we should add it to the messages
+                content = [
+                    (
+                        {"type": "input_text", "text": content.text}
+                        if isinstance(content, types.TextContent)
+                        else (
+                            {
+                                "type": "input_image",
+                                "image_url": f"data:image/png;base64,{content.data}",
+                            }
+                            if isinstance(content, types.ImageContent)
+                            else {"type": "input_text", "text": ""}
+                        )
+                    )
+                    for content in result.content
+                ]
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": content,
+                    }
+                )
+
+                continue
+
+            # Add the assistant's computer call
+            messages.extend(self.tool_call_inputs[call.id])
+
+            if result.isError:
+                error_text = "".join(
+                    [
+                        content.text
+                        for content in result.content
+                        if isinstance(content, types.TextContent)
+                    ]
+                )
+
+                # Replace computer call with failed tool call
+                messages.pop()
+                messages.extend(
+                    make_failed_tool_call_items(
+                        tool_name=call.name,
+                        tool_kwargs=call.arguments or {},
+                        error_message=error_text,
+                        call_id=call.id,
+                    )
+                )
+            else:
+                # Get the latest screenshot
+                screenshots = [
+                    content.data
+                    for content in result.content
+                    if isinstance(content, types.ImageContent)
+                ]
+
+                # Add the resulting screenshot
+                if screenshots:
+                    self._log_image(screenshots[0])
+                    self.last_screenshot_b64 = screenshots[0]
+                    messages.append(
+                        {
+                            "type": "computer_call_output",
+                            "call_id": call.id,
+                            "output": {
+                                "type": "input_image",
+                                "image_url": f"data:image/png;base64,{screenshots[0]}",
+                            },
+                        }
+                    )
+                else:
+                    # Otherwise, replace computer call with failed tool call
+                    messages.pop()
+                    messages.extend(
+                        make_failed_tool_call_items(
+                            tool_name=call.name,
+                            tool_kwargs=call.arguments or {},
+                            error_message="No screenshots returned.",
+                            call_id=call.id,
+                        )
+                    )
+
+        return messages
+
+
+__all__ = [
+    "MCPComputerAgent",
+]
@@ -0,0 +1,297 @@
+"""HUD ComputerAgent wrapper and Fake AsyncOpenAI client.
+
+Provides FakeAsyncOpenAI that adapts our ComputerAgent to the OpenAI Responses
+interface needed by HUD's OperatorAgent. It implements only `responses.create`
+and returns an OpenAI Response object with `id` and `output` fields, where `output` is a list of
+OpenAI-like response blocks. We intentionally only support a single-step call
+by consuming the first yielded result from `ComputerAgent.run()`.
+"""
+
+import time
+import traceback
+import uuid
+from typing import Any, Dict, List, Optional
+
+from cua_agent.agent import ComputerAgent as BaseComputerAgent
+from cua_agent.callbacks import PromptInstructionsCallback
+from hud.agents import OperatorAgent
+from hud.tools.computer.settings import computer_settings
+
+# OpenAI Responses typed models (required)
+from openai.types.responses import (
+    Response,
+    ResponseComputerToolCall,
+    ResponseInputParam,
+    ResponseOutputItem,
+    ResponseOutputMessage,
+    ResponseOutputText,
+    ResponseReasoningItem,
+    ResponseUsage,
+)
+from PIL import Image
+
+
+def _map_agent_output_to_openai_blocks(
+    output_items: List[Dict[str, Any]],
+) -> List[ResponseOutputItem]:
+    """Map our agent output items to OpenAI ResponseOutputItem typed models.
+
+    Only a subset is supported: computer_call, assistant message (text), and reasoning.
+    Unknown types are ignored.
+    """
+    blocks: List[ResponseOutputItem] = []
+    for item in output_items or []:
+        t = item.get("type")
+        if t == "computer_call":
+            comp = ResponseComputerToolCall.model_validate(
+                {
+                    "id": item.get("id") or f"cu_{uuid.uuid4().hex}",
+                    "type": "computer_call",
+                    "call_id": item["call_id"],
+                    "action": item["action"],
+                    "pending_safety_checks": item.get("pending_safety_checks", []),
+                    "status": "completed",
+                }
+            )
+            blocks.append(comp)
+            # we will exit early here as the responses api only supports a single step
+            break
+        elif t == "message" and item.get("role") == "assistant":
+            content_blocks: List[ResponseOutputText] = []
+            for c in item.get("content", []) or []:
+                content_blocks.append(
+                    ResponseOutputText.model_validate(
+                        {
+                            "type": "output_text",
+                            "text": c["text"],
+                            "annotations": [],
+                        }
+                    )
+                )
+            if content_blocks:
+                msg = ResponseOutputMessage.model_validate(
+                    {
+                        "id": item.get("id") or f"msg_{uuid.uuid4()}",
+                        "type": "message",
+                        "role": "assistant",
+                        "status": "completed",
+                        "content": [ct.model_dump() for ct in content_blocks],
+                    }
+                )
+                blocks.append(msg)
+        elif t == "reasoning":
+            reasoning = ResponseReasoningItem.model_validate(
+                {
+                    "id": item.get("id") or f"rsn_{uuid.uuid4()}",
+                    "type": "reasoning",
+                    "summary": item["summary"],
+                }
+            )
+            blocks.append(reasoning)
+        # Unhandled types are ignored
+    return blocks
+
+
+def _to_plain_dict_list(items: Any) -> List[Dict[str, Any]]:
+    out: List[Dict[str, Any]] = []
+    for it in list(items):
+        if hasattr(it, "model_dump"):
+            out.append(it.model_dump())  # type: ignore[attr-defined]
+        elif isinstance(it, dict):
+            out.append(it)
+        else:
+            # Strict: rely on default __dict__ if present
+            out.append(dict(it))  # may raise if not mapping
+    return out
+
+
+class FakeAsyncOpenAI:
+    """Minimal fake OpenAI client with only `responses.create` implemented.
+
+    It uses a provided `ComputerAgent` instance to produce a single-step
+    response compatible with HUD's OperatorAgent loop.
+    """
+
+    def __init__(self, computer_agent: BaseComputerAgent) -> None:
+        self._agent = computer_agent
+        self.responses = self._Responses(self)
+
+    class _Responses:
+        def __init__(self, parent: "FakeAsyncOpenAI") -> None:
+            # Caches for cross-call context when using previous_response_id
+            self.blocks_cache: Dict[str, ResponseInputParam | ResponseOutputItem] = {}
+            self.context_cache: Dict[str, List[str]] = {}
+            self.agent = parent._agent
+
+        async def create(
+            self,
+            *,
+            model: str,
+            input: ResponseInputParam,
+            tools: Optional[List[Dict[str, Any]]] = None,
+            instructions: Optional[str] = None,
+            previous_response_id: Optional[str] = None,
+            max_retries: int = 5,
+            **_: Any,
+        ) -> Any:
+            for attempt in range(max_retries):
+                # Prepend cached blocks from previous_response_id to input
+                full_input = input
+                if previous_response_id is not None:
+                    prev_block_ids = self.context_cache[previous_response_id]
+                    prev_blocks = [self.blocks_cache[b_id] for b_id in prev_block_ids]
+                    full_input = _to_plain_dict_list(prev_blocks + input)
+
+                # Pre-pend instructions message
+                effective_input = full_input
+                if instructions:
+                    effective_input = [
+                        {
+                            "role": "user",
+                            "content": instructions,
+                        }
+                    ] + full_input
+
+                # Run a single iteration of the ComputerAgent
+                agent_result: Optional[Dict[str, Any]] = None
+                async for result in self.agent.run(effective_input):  # type: ignore[arg-type]
+                    agent_result = result
+                    break
+                assert agent_result is not None, "Agent failed to produce result"
+
+                output = _map_agent_output_to_openai_blocks(agent_result["output"])
+                usage = agent_result["usage"]
+
+                # Cache conversation context using the last response id
+                block_ids: List[str] = []
+                blocks_to_cache = full_input + output
+                for b in blocks_to_cache:
+                    bid = getattr(b, "id", None) or f"tmp-{hash(repr(b))}"
+                    self.blocks_cache[bid] = b  # type: ignore[assignment]
+                    block_ids.append(bid)
+                response_id = agent_result.get("id") or f"fake-{int(time.time()*1000)}"
+                self.context_cache[response_id] = block_ids
+
+                try:
+                    return Response.model_validate(
+                        {
+                            "id": response_id,
+                            "created_at": time.time(),
+                            "object": "response",
+                            "model": model,
+                            "output": output,
+                            "parallel_tool_calls": False,
+                            "tool_choice": "auto",
+                            "tools": [],
+                            "previous_response_id": previous_response_id,
+                            "usage": ResponseUsage.model_validate(
+                                {
+                                    "input_tokens": usage.get("input_tokens", 0),
+                                    "output_tokens": usage.get("output_tokens", 0),
+                                    "total_tokens": usage.get("total_tokens", 0),
+                                    "input_tokens_details": usage.get(
+                                        "input_tokens_details", {"cached_tokens": 0}
+                                    ),
+                                    "output_tokens_details": usage.get(
+                                        "output_tokens_details", {"reasoning_tokens": 0}
+                                    ),
+                                }
+                            ),
+                        }
+                    )
+                except Exception as e:
+                    print(
+                        f"Error while validating agent response (attempt {attempt + 1}/{max_retries}): ",
+                        e,
+                    )
+                    if attempt == max_retries - 1:
+                        print(traceback.format_exc())
+                        raise e
+
+
+# ---------------------------------------------------------------------------
+# Proxy OperatorAgent (moved from __init__.py)
+# ---------------------------------------------------------------------------
+
+
+class ProxyOperatorAgent(OperatorAgent):
+    """OperatorAgent that proxies model calls through our ComputerAgent.
+
+    Accepts the same config keys we pass via hud.run_dataset `agent_config`:
+    - model: str | None
+    - allowed_tools: list[str] | None
+    Additional kwargs are forwarded to OperatorAgent (if any are supported).
+    """
+
+    def __init__(
+        self,
+        *,
+        model: str | None = None,
+        allowed_tools: list[str] | None = None,
+        trajectory_dir: str | dict | None = None,
+        # === ComputerAgent kwargs ===
+        tools: list[Any] | None = None,
+        custom_loop: Any | None = None,
+        only_n_most_recent_images: int | None = None,
+        callbacks: list[Any] | None = None,
+        instructions: str | None = None,
+        verbosity: int | None = None,
+        max_retries: int | None = 3,
+        screenshot_delay: float | int = 0.5,
+        use_prompt_caching: bool | None = False,
+        max_trajectory_budget: float | dict | None = None,
+        telemetry_enabled: bool | None = True,
+        **kwargs: Any,
+    ) -> None:
+        model = model or "computer-use-preview"
+        allowed_tools = allowed_tools or ["openai_computer"]
+
+        computer_shim = {
+            "screenshot": lambda: Image.new(
+                "RGB",
+                (computer_settings.OPENAI_COMPUTER_WIDTH, computer_settings.OPENAI_COMPUTER_HEIGHT),
+            ),
+            "environment": "linux",
+            "dimensions": (
+                computer_settings.OPENAI_COMPUTER_WIDTH,
+                computer_settings.OPENAI_COMPUTER_HEIGHT,
+            ),
+        }
+        # Build tools ensuring the computer_shim is included
+        agent_tools: list[Any] = [computer_shim]
+        if tools:
+            agent_tools.extend(tools)
+
+        # Build callbacks, injecting prompt instructions if provided
+        agent_callbacks = list(callbacks or [])
+        if instructions:
+            agent_callbacks.append(PromptInstructionsCallback(instructions))
+
+        computer_agent = BaseComputerAgent(
+            model=model,
+            tools=agent_tools,
+            custom_loop=custom_loop,
+            only_n_most_recent_images=only_n_most_recent_images,
+            callbacks=agent_callbacks,
+            verbosity=verbosity,
+            trajectory_dir=trajectory_dir,
+            max_retries=max_retries,
+            screenshot_delay=screenshot_delay,
+            use_prompt_caching=use_prompt_caching,
+            max_trajectory_budget=max_trajectory_budget,
+            telemetry_enabled=telemetry_enabled,
+        )
+        model_client = FakeAsyncOpenAI(computer_agent)
+
+        super().__init__(
+            model_client=model_client,  # type: ignore[arg-type]
+            model=model,
+            allowed_tools=allowed_tools,
+            **kwargs,
+        )
+
+
+__all__ = [
+    "FakeAsyncOpenAI",
+    "ProxyOperatorAgent",
+]