chore: import upstream snapshot with attribution

libs/python/agent/cua_agent/callbacks/__init__.py

@@ -0,0 +1,26 @@
+"""
+Callback system for ComputerAgent preprocessing and postprocessing hooks.
+"""
+
+from .base import AsyncCallbackHandler
+from .budget_manager import BudgetManagerCallback
+from .image_retention import ImageRetentionCallback
+from .logging import LoggingCallback
+from .operator_validator import OperatorNormalizerCallback
+from .otel import OtelCallback, OtelErrorCallback
+from .prompt_instructions import PromptInstructionsCallback
+from .telemetry import TelemetryCallback
+from .trajectory_saver import TrajectorySaverCallback
+
+__all__ = [
+    "AsyncCallbackHandler",
+    "ImageRetentionCallback",
+    "LoggingCallback",
+    "TrajectorySaverCallback",
+    "BudgetManagerCallback",
+    "TelemetryCallback",
+    "OtelCallback",
+    "OtelErrorCallback",
+    "OperatorNormalizerCallback",
+    "PromptInstructionsCallback",
+]

libs/python/agent/cua_agent/callbacks/base.py

@@ -0,0 +1,167 @@
+"""
+Base callback handler interface for ComputerAgent preprocessing and postprocessing hooks.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Union
+
+
+class AsyncCallbackHandler(ABC):
+    """
+    Base class for async callback handlers that can preprocess messages before
+    the agent loop and postprocess output after the agent loop.
+    """
+
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Called at the start of an agent run loop."""
+        pass
+
+    async def on_run_end(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> None:
+        """Called at the end of an agent run loop."""
+        pass
+
+    async def on_run_continue(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> bool:
+        """Called during agent run loop to determine if execution should continue.
+
+        Args:
+            kwargs: Run arguments
+            old_items: Original messages
+            new_items: New messages generated during run
+
+        Returns:
+            True to continue execution, False to stop
+        """
+        return True
+
+    async def on_llm_start(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Called before messages are sent to the agent loop.
+
+        Args:
+            messages: List of message dictionaries to preprocess
+
+        Returns:
+            List of preprocessed message dictionaries
+        """
+        return messages
+
+    async def on_llm_end(self, output: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Called after the agent loop returns output.
+
+        Args:
+            output: List of output message dictionaries to postprocess
+
+        Returns:
+            List of postprocessed output dictionaries
+        """
+        return output
+
+    async def on_computer_call_start(self, item: Dict[str, Any]) -> None:
+        """
+        Called when a computer call is about to start.
+
+        Args:
+            item: The computer call item dictionary
+        """
+        pass
+
+    async def on_computer_call_end(
+        self, item: Dict[str, Any], result: List[Dict[str, Any]]
+    ) -> None:
+        """
+        Called when a computer call has completed.
+
+        Args:
+            item: The computer call item dictionary
+            result: The result of the computer call
+        """
+        pass
+
+    async def on_function_call_start(self, item: Dict[str, Any]) -> None:
+        """
+        Called when a function call is about to start.
+
+        Args:
+            item: The function call item dictionary
+        """
+        pass
+
+    async def on_function_call_end(
+        self, item: Dict[str, Any], result: List[Dict[str, Any]]
+    ) -> None:
+        """
+        Called when a function call has completed.
+
+        Args:
+            item: The function call item dictionary
+            result: The result of the function call
+        """
+        pass
+
+    async def on_text(self, item: Dict[str, Any]) -> None:
+        """
+        Called when a text message is encountered.
+
+        Args:
+            item: The message item dictionary
+        """
+        pass
+
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        """
+        Called when an API call is about to start.
+
+        Args:
+            kwargs: The kwargs being passed to the API call
+        """
+        pass
+
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """
+        Called when an API call has completed.
+
+        Args:
+            kwargs: The kwargs that were passed to the API call
+            result: The result of the API call
+        """
+        pass
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """
+        Called when usage information is received.
+
+        Args:
+            usage: The usage information
+        """
+        pass
+
+    async def on_screenshot(self, screenshot: Union[str, bytes], name: str = "screenshot") -> None:
+        """
+        Called when a screenshot is taken.
+
+        Args:
+            screenshot: The screenshot image
+            name: The name of the screenshot
+        """
+        pass
+
+    async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
+        """
+        Called when responses are received.
+
+        Args:
+            kwargs: The kwargs being passed to the agent loop
+            responses: The responses received
+        """
+        pass

libs/python/agent/cua_agent/callbacks/budget_manager.py

@@ -0,0 +1,56 @@
+from typing import Any, Dict, List
+
+from .base import AsyncCallbackHandler
+
+
+class BudgetExceededError(Exception):
+    """Exception raised when budget is exceeded."""
+
+    pass
+
+
+class BudgetManagerCallback(AsyncCallbackHandler):
+    """Budget manager callback that tracks usage costs and can stop execution when budget is exceeded."""
+
+    def __init__(
+        self, max_budget: float, reset_after_each_run: bool = True, raise_error: bool = False
+    ):
+        """
+        Initialize BudgetManagerCallback.
+
+        Args:
+            max_budget: Maximum budget allowed
+            reset_after_each_run: Whether to reset budget after each run
+            raise_error: Whether to raise an error when budget is exceeded
+        """
+        self.max_budget = max_budget
+        self.reset_after_each_run = reset_after_each_run
+        self.raise_error = raise_error
+        self.total_cost = 0.0
+
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Reset budget if configured to do so."""
+        if self.reset_after_each_run:
+            self.total_cost = 0.0
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Track usage costs."""
+        if "response_cost" in usage:
+            self.total_cost += usage["response_cost"]
+
+    async def on_run_continue(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> bool:
+        """Check if budget allows continuation."""
+        if self.total_cost >= self.max_budget:
+            if self.raise_error:
+                raise BudgetExceededError(
+                    f"Budget exceeded: ${self.total_cost} >= ${self.max_budget}"
+                )
+            else:
+                print(f"Budget exceeded: ${self.total_cost} >= ${self.max_budget}")
+            return False
+        return True

libs/python/agent/cua_agent/callbacks/image_retention.py

@@ -0,0 +1,95 @@
+"""
+Image retention callback handler that limits the number of recent images in message history.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from .base import AsyncCallbackHandler
+
+
+class ImageRetentionCallback(AsyncCallbackHandler):
+    """
+    Callback handler that applies image retention policy to limit the number
+    of recent images in message history to prevent context window overflow.
+    """
+
+    def __init__(self, only_n_most_recent_images: Optional[int] = None):
+        """
+        Initialize the image retention callback.
+
+        Args:
+            only_n_most_recent_images: If set, only keep the N most recent images in message history
+        """
+        self.only_n_most_recent_images = only_n_most_recent_images
+
+    async def on_llm_start(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Apply image retention policy to messages before sending to agent loop.
+
+        Args:
+            messages: List of message dictionaries
+
+        Returns:
+            List of messages with image retention policy applied
+        """
+        if self.only_n_most_recent_images is None:
+            return messages
+
+        return self._apply_image_retention(messages)
+
+    def _apply_image_retention(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Apply image retention policy to keep only the N most recent images.
+
+        Removes computer_call_output items with image_url and their corresponding computer_call items,
+        keeping only the most recent N image pairs based on only_n_most_recent_images setting.
+
+        Args:
+            messages: List of message dictionaries
+
+        Returns:
+            Filtered list of messages with image retention applied
+        """
+        if self.only_n_most_recent_images is None:
+            return messages
+
+        # Gather indices of all computer_call_output messages that contain an image_url
+        output_indices: List[int] = []
+        for idx, msg in enumerate(messages):
+            if msg.get("type") == "computer_call_output":
+                out = msg.get("output")
+                if isinstance(out, dict) and ("image_url" in out):
+                    output_indices.append(idx)
+
+        # Nothing to trim
+        if len(output_indices) <= self.only_n_most_recent_images:
+            return messages
+
+        # Determine which outputs to keep (most recent N)
+        keep_output_indices = set(output_indices[-self.only_n_most_recent_images :])
+
+        # Build set of indices to remove in one pass
+        to_remove: set[int] = set()
+
+        for idx in output_indices:
+            if idx in keep_output_indices:
+                continue  # keep this screenshot and its context
+
+            to_remove.add(idx)  # remove the computer_call_output itself
+
+            # Remove the immediately preceding computer_call with matching call_id (if present)
+            call_id = messages[idx].get("call_id")
+            prev_idx = idx - 1
+            if (
+                prev_idx >= 0
+                and messages[prev_idx].get("type") == "computer_call"
+                and messages[prev_idx].get("call_id") == call_id
+            ):
+                to_remove.add(prev_idx)
+                # Check a single reasoning immediately before that computer_call
+                r_idx = prev_idx - 1
+                if r_idx >= 0 and messages[r_idx].get("type") == "reasoning":
+                    to_remove.add(r_idx)
+
+        # Construct filtered list
+        filtered = [m for i, m in enumerate(messages) if i not in to_remove]
+        return filtered

libs/python/agent/cua_agent/callbacks/logging.py

@@ -0,0 +1,260 @@
+"""
+Logging callback for ComputerAgent that provides configurable logging of agent lifecycle events.
+"""
+
+import json
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+from .base import AsyncCallbackHandler
+
+
+def sanitize_image_urls(data: Any) -> Any:
+    """
+    Recursively search for 'image_url' keys and set their values to '[omitted]'.
+
+    Args:
+        data: Any data structure (dict, list, or primitive type)
+
+    Returns:
+        A deep copy of the data with all 'image_url' values replaced with '[omitted]'
+    """
+    if isinstance(data, dict):
+        # Create a copy of the dictionary
+        sanitized = {}
+        for key, value in data.items():
+            if key == "image_url":
+                sanitized[key] = "[omitted]"
+            else:
+                # Recursively sanitize the value
+                sanitized[key] = sanitize_image_urls(value)
+        return sanitized
+
+    elif isinstance(data, list):
+        # Recursively sanitize each item in the list
+        return [sanitize_image_urls(item) for item in data]
+
+    else:
+        # For primitive types (str, int, bool, None, etc.), return as-is
+        return data
+
+
+class LoggingCallback(AsyncCallbackHandler):
+    """
+    Callback handler that logs agent lifecycle events with configurable verbosity.
+
+    Logging levels:
+    - DEBUG: All events including API calls, message preprocessing, and detailed outputs
+    - INFO: Major lifecycle events (start/end, messages, outputs)
+    - WARNING: Only warnings and errors
+    - ERROR: Only errors
+    """
+
+    def __init__(self, logger: Optional[logging.Logger] = None, level: int = logging.INFO):
+        """
+        Initialize the logging callback.
+
+        Args:
+            logger: Logger instance to use. If None, creates a logger named 'agent.ComputerAgent'
+            level: Logging level (logging.DEBUG, logging.INFO, etc.)
+        """
+        self.logger = logger or logging.getLogger("agent.ComputerAgent")
+        self.level = level
+
+        # Set up logger if it doesn't have handlers
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+            self.logger.setLevel(level)
+
+    def _update_usage(self, usage: Dict[str, Any]) -> None:
+        """Update total usage statistics."""
+
+        def add_dicts(target: Dict[str, Any], source: Dict[str, Any]) -> None:
+            for key, value in source.items():
+                if isinstance(value, dict):
+                    if key not in target:
+                        target[key] = {}
+                    add_dicts(target[key], value)
+                else:
+                    if key not in target:
+                        target[key] = 0
+                    target[key] += value
+
+        add_dicts(self.total_usage, usage)
+
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Called before the run starts."""
+        self.total_usage = {}
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Called when usage information is received."""
+        self._update_usage(usage)
+
+    async def on_run_end(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> None:
+        """Called after the run ends."""
+
+        def format_dict(d, indent=0):
+            lines = []
+            prefix = f" - {' ' * indent}"
+            for key, value in d.items():
+                if isinstance(value, dict):
+                    lines.append(f"{prefix}{key}:")
+                    lines.extend(format_dict(value, indent + 1))
+                elif isinstance(value, float):
+                    lines.append(f"{prefix}{key}: ${value:.4f}")
+                else:
+                    lines.append(f"{prefix}{key}: {value}")
+            return lines
+
+        formatted_output = "\n".join(format_dict(self.total_usage))
+        self.logger.info(f"Total usage:\n{formatted_output}")
+
+    async def on_llm_start(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Called before LLM processing starts."""
+        if self.logger.isEnabledFor(logging.INFO):
+            self.logger.info(f"LLM processing started with {len(messages)} messages")
+        if self.logger.isEnabledFor(logging.DEBUG):
+            sanitized_messages = [sanitize_image_urls(msg) for msg in messages]
+            self.logger.debug(f"LLM input messages: {json.dumps(sanitized_messages, indent=2)}")
+        return messages
+
+    async def on_llm_end(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Called after LLM processing ends."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            sanitized_messages = [sanitize_image_urls(msg) for msg in messages]
+            self.logger.debug(f"LLM output: {json.dumps(sanitized_messages, indent=2)}")
+        return messages
+
+    async def on_computer_call_start(self, item: Dict[str, Any]) -> None:
+        """Called when a computer call starts."""
+        action = item.get("action", {})
+        action_type = action.get("type", "unknown")
+        action_args = {k: v for k, v in action.items() if k != "type"}
+
+        # INFO level logging for the action
+        self.logger.info(f"Computer: {action_type}({action_args})")
+
+        # DEBUG level logging for full details
+        if self.logger.isEnabledFor(logging.DEBUG):
+            self.logger.debug(f"Computer call started: {json.dumps(action, indent=2)}")
+
+    async def on_computer_call_end(self, item: Dict[str, Any], result: Any) -> None:
+        """Called when a computer call ends."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            action = item.get("action", "unknown")
+            self.logger.debug(f"Computer call completed: {json.dumps(action, indent=2)}")
+            if result:
+                sanitized_result = sanitize_image_urls(result)
+                self.logger.debug(f"Computer call result: {json.dumps(sanitized_result, indent=2)}")
+
+    async def on_function_call_start(self, item: Dict[str, Any]) -> None:
+        """Called when a function call starts."""
+        name = item.get("name", "unknown")
+        arguments = item.get("arguments", "{}")
+
+        # INFO level logging for the function call
+        self.logger.info(f"Function: {name}({arguments})")
+
+        # DEBUG level logging for full details
+        if self.logger.isEnabledFor(logging.DEBUG):
+            self.logger.debug(f"Function call started: {name}")
+
+    async def on_function_call_end(self, item: Dict[str, Any], result: Any) -> None:
+        """Called when a function call ends."""
+        # INFO level logging for function output (similar to function_call_output)
+        if result:
+            # Handle both list and direct result formats
+            if isinstance(result, list) and len(result) > 0:
+                output = (
+                    result[0].get("output", str(result))
+                    if isinstance(result[0], dict)
+                    else str(result[0])
+                )
+            else:
+                output = str(result)
+
+            # Truncate long outputs
+            if len(output) > 100:
+                output = output[:100] + "..."
+
+            self.logger.info(f"Output: {output}")
+
+        # DEBUG level logging for full details
+        if self.logger.isEnabledFor(logging.DEBUG):
+            name = item.get("name", "unknown")
+            self.logger.debug(f"Function call completed: {name}")
+            if result:
+                self.logger.debug(f"Function call result: {json.dumps(result, indent=2)}")
+
+    async def on_text(self, item: Dict[str, Any]) -> None:
+        """Called when a text message is encountered."""
+        # Get the role to determine if it's Agent or User
+        role = item.get("role", "unknown")
+        content_items = item.get("content", [])
+
+        # Process content items to build display text
+        text_parts = []
+        for content_item in content_items:
+            content_type = content_item.get("type", "output_text")
+            if content_type == "output_text":
+                text_content = content_item.get("text", "")
+                if not text_content.strip():
+                    text_parts.append("[empty]")
+                else:
+                    # Truncate long text and add ellipsis
+                    if len(text_content) > 2048:
+                        text_parts.append(text_content[:2048] + "...")
+                    else:
+                        text_parts.append(text_content)
+            else:
+                # Non-text content, show as [type]
+                text_parts.append(f"[{content_type}]")
+
+        # Join all text parts
+        display_text = "".join(text_parts) if text_parts else "[empty]"
+
+        # Log with appropriate level and format
+        if role == "assistant":
+            self.logger.info(f"Agent: {display_text}")
+        elif role == "user":
+            self.logger.info(f"User: {display_text}")
+        else:
+            # Fallback for unknown roles, use debug level
+            if self.logger.isEnabledFor(logging.DEBUG):
+                self.logger.debug(f"Text message ({role}): {display_text}")
+
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        """Called when an API call is about to start."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            model = kwargs.get("model", "unknown")
+            self.logger.debug(f"API call starting for model: {model}")
+            # Log sanitized messages if present
+            if "messages" in kwargs:
+                sanitized_messages = sanitize_image_urls(kwargs["messages"])
+                self.logger.debug(f"API call messages: {json.dumps(sanitized_messages, indent=2)}")
+            elif "input" in kwargs:
+                sanitized_input = sanitize_image_urls(kwargs["input"])
+                self.logger.debug(f"API call input: {json.dumps(sanitized_input, indent=2)}")
+
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """Called when an API call has completed."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            model = kwargs.get("model", "unknown")
+            self.logger.debug(f"API call completed for model: {model}")
+            self.logger.debug(
+                f"API call result: {json.dumps(sanitize_image_urls(result), indent=2)}"
+            )
+
+    async def on_screenshot(self, item: Union[str, bytes], name: str = "screenshot") -> None:
+        """Called when a screenshot is taken."""
+        if self.logger.isEnabledFor(logging.DEBUG):
+            image_size = len(item) / 1024
+            self.logger.debug(f"Screenshot captured: {name} {image_size:.2f} KB")

libs/python/agent/cua_agent/callbacks/operator_validator.py

@@ -0,0 +1,140 @@
+"""
+OperatorValidatorCallback
+
+Ensures agent output actions conform to expected schemas by fixing common issues:
+- click: add default button='left' if missing
+- keypress: wrap keys string into a list
+- etc.
+
+This runs in on_llm_end, which receives the output array (AgentMessage[] as dicts).
+The purpose is to avoid spending another LLM call to fix broken computer call syntax when possible.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+from .base import AsyncCallbackHandler
+
+
+class OperatorNormalizerCallback(AsyncCallbackHandler):
+    """Normalizes common computer call hallucinations / errors in computer call syntax."""
+
+    async def on_llm_end(self, output: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        # Mutate in-place as requested, but still return the list for chaining
+        for item in output or []:
+            if item.get("type") != "computer_call":
+                continue
+            action = item.get("action")
+            if not isinstance(action, dict):
+                continue
+
+            # rename mouse click actions to "click"
+            for mouse_btn in ["left", "right", "wheel", "back", "forward"]:
+                if action.get("type", "") == f"{mouse_btn}_click":
+                    action["type"] = "click"
+                    action["button"] = mouse_btn
+            # rename hotkey actions to "keypress"
+            for alias in ["hotkey", "key", "press", "key_press"]:
+                if action.get("type", "") == alias:
+                    action["type"] = "keypress"
+            # assume click actions
+            if "button" in action and "type" not in action:
+                action["type"] = "click"
+            if "click" in action and "type" not in action:
+                action["type"] = "click"
+            if ("scroll_x" in action or "scroll_y" in action) and "type" not in action:
+                action["type"] = "scroll"
+            if "text" in action and "type" not in action:
+                action["type"] = "type"
+
+            action_type = action.get("type")
+
+            def _keep_keys(action: Dict[str, Any], keys_to_keep: List[str]):
+                """Keep only the provided keys on action; delete everything else.
+                Always ensures required 'type' is present if listed in keys_to_keep.
+                """
+                for key in list(action.keys()):
+                    if key not in keys_to_keep:
+                        del action[key]
+
+            # rename "coordinate" to "x", "y"
+            if "coordinate" in action:
+                action["x"] = action["coordinate"][0]
+                action["y"] = action["coordinate"][1]
+                del action["coordinate"]
+            if action_type == "click":
+                # convert "click" to "button"
+                if "button" not in action and "click" in action:
+                    action["button"] = action["click"]
+                    del action["click"]
+                # default button to "left"
+                action["button"] = action.get("button", "left")
+            # add default scroll x, y if missing
+            if action_type == "scroll":
+                action["scroll_x"] = action.get("scroll_x", 0)
+                action["scroll_y"] = action.get("scroll_y", 0)
+            # ensure keys arg is a list (normalize aliases first)
+            if action_type == "keypress":
+                keys = action.get("keys")
+                for keys_alias in ["keypress", "key", "press", "key_press", "text"]:
+                    if keys_alias in action:
+                        action["keys"] = action[keys_alias]
+                        del action[keys_alias]
+                keys = action.get("keys")
+                if isinstance(keys, str):
+                    action["keys"] = keys.replace("-", "+").split("+") if len(keys) > 1 else [keys]
+            required_keys_by_type = {
+                # OpenAI actions
+                "click": ["type", "button", "x", "y"],
+                "double_click": ["type", "x", "y"],
+                "drag": ["type", "path"],
+                "keypress": ["type", "keys"],
+                "move": ["type", "x", "y"],
+                "screenshot": ["type"],
+                "scroll": ["type", "scroll_x", "scroll_y", "x", "y"],
+                "type": ["type", "text"],
+                "wait": ["type"],
+                # Anthropic actions
+                "left_mouse_down": ["type", "x", "y"],
+                "left_mouse_up": ["type", "x", "y"],
+                "triple_click": ["type", "button", "x", "y"],
+            }
+            keep = required_keys_by_type.get(action_type or "")
+            if keep:
+                _keep_keys(action, keep)
+
+        # # Second pass: if an assistant message is immediately followed by a computer_call,
+        # # replace the assistant message itself with a reasoning message with summary text.
+        # if isinstance(output, list):
+        #     for i, item in enumerate(output):
+        #         # AssistantMessage shape: { type: 'message', role: 'assistant', content: OutputContent[] }
+        #         if item.get("type") == "message" and item.get("role") == "assistant":
+        #             next_idx = i + 1
+        #             if next_idx >= len(output):
+        #                 continue
+        #             next_item = output[next_idx]
+        #             if not isinstance(next_item, dict):
+        #                 continue
+        #             if next_item.get("type") != "computer_call":
+        #                 continue
+        #             contents = item.get("content") or []
+        #             # Extract text from OutputContent[]
+        #             text_parts: List[str] = []
+        #             if isinstance(contents, list):
+        #                 for c in contents:
+        #                     if isinstance(c, dict) and c.get("type") == "output_text" and isinstance(c.get("text"), str):
+        #                         text_parts.append(c["text"])
+        #             text_content = "\n".join(text_parts).strip()
+        #             # Replace assistant message with reasoning message
+        #             output[i] = {
+        #                 "type": "reasoning",
+        #                 "summary": [
+        #                     {
+        #                         "type": "summary_text",
+        #                         "text": text_content,
+        #                     }
+        #                 ],
+        #             }
+
+        return output

libs/python/agent/cua_agent/callbacks/otel.py

@@ -0,0 +1,212 @@
+"""
+OpenTelemetry callback handler for Computer-Use Agent (cua-agent).
+
+Instruments agent operations for the Four Golden Signals:
+- Latency: Operation duration
+- Traffic: Operation counts
+- Errors: Error counts
+- Saturation: Concurrent operations
+"""
+
+import time
+from typing import Any, Dict, List, Optional
+
+from .base import AsyncCallbackHandler
+
+# Import OTEL functions - these are available when cua-core[telemetry] is installed
+try:
+    from cua_core.telemetry import (
+        create_span,
+        is_otel_enabled,
+        record_error,
+        record_operation,
+        record_tokens,
+        track_concurrent,
+    )
+
+    OTEL_AVAILABLE = True
+except ImportError:
+    OTEL_AVAILABLE = False
+
+    def is_otel_enabled() -> bool:
+        return False
+
+
+class OtelCallback(AsyncCallbackHandler):
+    """
+    OpenTelemetry callback handler for instrumentation.
+
+    Tracks:
+    - Agent session lifecycle (start/end)
+    - Agent run lifecycle (start/end with duration)
+    - Individual steps (with duration)
+    - Computer actions (with duration)
+    - Token usage
+    - Errors
+    """
+
+    def __init__(self, agent: Any):
+        """
+        Initialize OTEL callback.
+
+        Args:
+            agent: The ComputerAgent instance
+        """
+        self.agent = agent
+        self.model = getattr(agent, "model", "unknown")
+
+        # Timing state
+        self.run_start_time: Optional[float] = None
+        self.step_start_time: Optional[float] = None
+        self.step_count = 0
+
+        # Span management
+        self._session_span: Optional[Any] = None
+        self._run_span: Optional[Any] = None
+
+        # Track concurrent sessions
+        self._concurrent_tracker: Optional[Any] = None
+
+    def _get_agent_type(self) -> str:
+        """Get the agent loop type name."""
+        if hasattr(self.agent, "agent_loop") and self.agent.agent_loop is not None:
+            return type(self.agent.agent_loop).__name__
+        return "unknown"
+
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Called at the start of an agent run loop."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        self.run_start_time = time.perf_counter()
+        self.step_start_time = self.run_start_time
+        self.step_count = 0
+
+    async def on_run_end(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> None:
+        """Called at the end of an agent run loop."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        if self.run_start_time is not None:
+            duration = time.perf_counter() - self.run_start_time
+
+            # Record run metrics
+            record_operation(
+                operation="agent.run",
+                duration_seconds=duration,
+                status="success",
+                model=self.model,
+                steps=self.step_count,
+            )
+
+        self.run_start_time = None
+
+    async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
+        """Called when responses are received (each step)."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        self.step_count += 1
+        current_time = time.perf_counter()
+
+        # Calculate step duration if we have a start time
+        if self.step_start_time is not None:
+            step_duration = current_time - self.step_start_time
+            record_operation(
+                operation="agent.step",
+                duration_seconds=step_duration,
+                status="success",
+                model=self.model,
+                step_number=self.step_count,
+            )
+
+        # Start timing next step
+        self.step_start_time = current_time
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Called when usage information is received."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+
+        if prompt_tokens > 0 or completion_tokens > 0:
+            record_tokens(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                model=self.model,
+            )
+
+    async def on_computer_call_start(self, item: Dict[str, Any]) -> None:
+        """Called when a computer call is about to start."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+    async def on_computer_call_end(
+        self, item: Dict[str, Any], result: List[Dict[str, Any]]
+    ) -> None:
+        """Called when a computer call has completed."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        action = item.get("action", {})
+        action_type = action.get("type", "unknown")
+
+        # Record computer action metric
+        # Note: We don't have precise timing here, so we record with 0 duration
+        # The actual timing should be done in the computer module
+        record_operation(
+            operation=f"computer.action.{action_type}",
+            duration_seconds=0,  # Timing handled elsewhere
+            status="success",
+            model=self.model,
+        )
+
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        """Called when an LLM API call is about to start."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """Called when an LLM API call has completed."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+
+class OtelErrorCallback(AsyncCallbackHandler):
+    """
+    Callback that captures errors and sends them to OTEL.
+
+    Should be added early in the callback chain to catch all errors.
+    """
+
+    def __init__(self, agent: Any):
+        """
+        Initialize error callback.
+
+        Args:
+            agent: The ComputerAgent instance
+        """
+        self.agent = agent
+        self.model = getattr(agent, "model", "unknown")
+
+    async def on_error(self, error: Exception, context: Dict[str, Any]) -> None:
+        """Called when an error occurs during agent execution."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        error_type = type(error).__name__
+        operation = context.get("operation", "unknown")
+
+        # Record error metric
+        record_error(
+            error_type=error_type,
+            operation=operation,
+            model=self.model,
+        )

libs/python/agent/cua_agent/callbacks/pii_anonymization.py

@@ -0,0 +1,99 @@
+"""
+PII anonymization callback handler using Microsoft Presidio for text and image redaction.
+"""
+
+import base64
+import io
+import logging
+from typing import Any, Dict, List, Optional, Tuple
+
+from .base import AsyncCallbackHandler
+
+try:
+    # TODO: Add Presidio dependencies
+    from PIL import Image
+
+    PRESIDIO_AVAILABLE = True
+except ImportError:
+    PRESIDIO_AVAILABLE = False
+
+logger = logging.getLogger(__name__)
+
+
+class PIIAnonymizationCallback(AsyncCallbackHandler):
+    """
+    Callback handler that anonymizes PII in text and images using Microsoft Presidio.
+
+    This handler:
+    1. Anonymizes PII in messages before sending to the agent loop
+    2. Deanonymizes PII in tool calls and message outputs after the agent loop
+    3. Redacts PII from images in computer_call_output messages
+    """
+
+    def __init__(
+        self,
+        # TODO: Any extra kwargs if needed
+    ):
+        """
+        Initialize the PII anonymization callback.
+
+        Args:
+            anonymize_text: Whether to anonymize text content
+            anonymize_images: Whether to redact images
+            entities_to_anonymize: List of entity types to anonymize (None for all)
+            anonymization_operator: Presidio operator to use ("replace", "mask", "redact", etc.)
+            image_redaction_color: RGB color for image redaction
+        """
+        if not PRESIDIO_AVAILABLE:
+            raise ImportError(
+                "Presidio is not available. Install with: "
+                "pip install cua-agent[pii-anonymization]"
+            )
+
+        # TODO: Implement __init__
+
+    async def on_llm_start(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Anonymize PII in messages before sending to agent loop.
+
+        Args:
+            messages: List of message dictionaries
+
+        Returns:
+            List of messages with PII anonymized
+        """
+        anonymized_messages = []
+        for msg in messages:
+            anonymized_msg = await self._anonymize_message(msg)
+            anonymized_messages.append(anonymized_msg)
+
+        return anonymized_messages
+
+    async def on_llm_end(self, output: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Deanonymize PII in tool calls and message outputs after agent loop.
+
+        Args:
+            output: List of output dictionaries
+
+        Returns:
+            List of output with PII deanonymized for tool calls
+        """
+        deanonymized_output = []
+        for item in output:
+            # Only deanonymize tool calls and computer_call messages
+            if item.get("type") in ["computer_call", "computer_call_output"]:
+                deanonymized_item = await self._deanonymize_item(item)
+                deanonymized_output.append(deanonymized_item)
+            else:
+                deanonymized_output.append(item)
+
+        return deanonymized_output
+
+    async def _anonymize_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
+        # TODO: Implement _anonymize_message
+        return message
+
+    async def _deanonymize_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
+        # TODO: Implement _deanonymize_item
+        return item

libs/python/agent/cua_agent/callbacks/prompt_instructions.py

@@ -0,0 +1,47 @@
+"""
+Prompt instructions callback.
+
+This callback allows simple prompt engineering by pre-pending a user
+instructions message to the start of the conversation before each LLM call.
+
+Usage:
+
+    from cua_agent.callbacks import PromptInstructionsCallback
+    agent = ComputerAgent(
+        model="openai/computer-use-preview",
+        callbacks=[PromptInstructionsCallback("Follow these rules...")]
+    )
+
+"""
+
+from typing import Any, Dict, List, Optional
+
+from .base import AsyncCallbackHandler
+
+
+class PromptInstructionsCallback(AsyncCallbackHandler):
+    """
+    Prepend a user instructions message to the message list.
+
+    This is a minimal, non-invasive way to guide the agent's behavior without
+    modifying agent loops or tools. It works with any provider/loop since it
+    only alters the messages array before sending to the model.
+    """
+
+    def __init__(self, instructions: Optional[str]) -> None:
+        self.instructions = instructions
+
+    async def on_llm_start(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        # Pre-pend instructions message
+        if not self.instructions:
+            return messages
+
+        # Ensure we don't duplicate if already present at the front
+        if messages and isinstance(messages[0], dict):
+            first = messages[0]
+            if first.get("role") == "user" and first.get("content") == self.instructions:
+                return messages
+
+        return [
+            {"role": "user", "content": self.instructions},
+        ] + messages

libs/python/agent/cua_agent/callbacks/telemetry.py

@@ -0,0 +1,247 @@
+"""
+Telemetry callback handler for Computer-Use Agent (cua-agent)
+"""
+
+import platform
+import time
+import uuid
+from typing import Any, Dict, List, Optional, Union
+
+from cua_core.telemetry import (
+    is_telemetry_enabled,
+    record_event,
+)
+
+from .base import AsyncCallbackHandler
+
+SYSTEM_INFO = {
+    "os": platform.system().lower(),
+    "os_version": platform.release(),
+    "python_version": platform.python_version(),
+}
+
+
+class TelemetryCallback(AsyncCallbackHandler):
+    """
+    Telemetry callback handler for Computer-Use Agent (cua-agent)
+
+    Tracks agent usage, performance metrics, and optionally trajectory data.
+    """
+
+    def __init__(self, agent, log_trajectory: bool = False):
+        """
+        Initialize telemetry callback.
+
+        Args:
+            agent: The ComputerAgent instance
+            log_trajectory: Whether to log full trajectory items (opt-in)
+        """
+        self.agent = agent
+        self.log_trajectory = log_trajectory
+
+        # Generate session/run IDs
+        self.session_id = str(uuid.uuid4())
+        self.run_id = None
+
+        # Track timing and metrics
+        self.run_start_time = None
+        self.step_count = 0
+        self.step_start_time = None
+        self.total_usage = {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0,
+            "response_cost": 0.0,
+        }
+
+        # Record agent initialization
+        if is_telemetry_enabled():
+            self._record_agent_initialization()
+
+    def _record_agent_initialization(self) -> None:
+        """Record agent type/model and session initialization."""
+        # Get the agent loop type (class name)
+        agent_type = "unknown"
+        if hasattr(self.agent, "agent_loop") and self.agent.agent_loop is not None:
+            agent_type = type(self.agent.agent_loop).__name__
+
+        agent_info = {
+            "session_id": self.session_id,
+            "agent_type": agent_type,
+            "model": getattr(self.agent, "model", "unknown"),
+            **SYSTEM_INFO,
+        }
+
+        # Include VM name if available
+        vm_name = self._get_vm_name()
+        if vm_name:
+            agent_info["vm_name"] = vm_name
+
+        record_event("agent_session_start", agent_info)
+
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Called at the start of an agent run loop."""
+        if not is_telemetry_enabled():
+            return
+
+        self.run_id = str(uuid.uuid4())
+        self.run_start_time = time.time()
+        self.step_count = 0
+
+        # Calculate input context size
+        input_context_size = self._calculate_context_size(old_items)
+
+        run_data = {
+            "session_id": self.session_id,
+            "run_id": self.run_id,
+            "start_time": self.run_start_time,
+            "input_context_size": input_context_size,
+            "num_existing_messages": len(old_items),
+        }
+
+        # Include VM name if available
+        vm_name = self._get_vm_name()
+        if vm_name:
+            run_data["vm_name"] = vm_name
+
+        # Log trajectory if opted in
+        if self.log_trajectory:
+            trajectory = self._extract_trajectory(old_items)
+            if trajectory:
+                run_data["uploaded_trajectory"] = trajectory
+
+        record_event("agent_run_start", run_data)
+
+    async def on_run_end(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> None:
+        """Called at the end of an agent run loop."""
+        if not is_telemetry_enabled() or not self.run_start_time:
+            return
+
+        run_duration = time.time() - self.run_start_time
+
+        run_data = {
+            "session_id": self.session_id,
+            "run_id": self.run_id,
+            "end_time": time.time(),
+            "duration_seconds": run_duration,
+            "num_steps": self.step_count,
+            "total_usage": self.total_usage.copy(),
+        }
+
+        # Include VM name if available
+        vm_name = self._get_vm_name()
+        if vm_name:
+            run_data["vm_name"] = vm_name
+
+        # Log trajectory if opted in
+        if self.log_trajectory:
+            trajectory = self._extract_trajectory(new_items)
+            if trajectory:
+                run_data["uploaded_trajectory"] = trajectory
+
+        record_event("agent_run_end", run_data)
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Called when usage information is received."""
+        if not is_telemetry_enabled():
+            return
+
+        # Accumulate usage stats
+        self.total_usage["prompt_tokens"] += usage.get("prompt_tokens", 0)
+        self.total_usage["completion_tokens"] += usage.get("completion_tokens", 0)
+        self.total_usage["total_tokens"] += usage.get("total_tokens", 0)
+        self.total_usage["response_cost"] += usage.get("response_cost", 0.0)
+
+        # Record individual usage event
+        usage_data = {
+            "session_id": self.session_id,
+            "run_id": self.run_id,
+            "step": self.step_count,
+            **usage,
+        }
+
+        record_event("agent_usage", usage_data)
+
+    async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
+        """Called when responses are received."""
+        if not is_telemetry_enabled():
+            return
+
+        self.step_count += 1
+        step_duration = None
+
+        if self.step_start_time:
+            step_duration = time.time() - self.step_start_time
+
+        self.step_start_time = time.time()
+
+        step_data = {
+            "session_id": self.session_id,
+            "run_id": self.run_id,
+            "step": self.step_count,
+            "timestamp": self.step_start_time,
+        }
+
+        if step_duration is not None:
+            step_data["duration_seconds"] = step_duration
+
+        record_event("agent_step", step_data)
+
+    def _get_vm_name(self) -> Optional[str]:
+        """Extract VM name from agent's computer handler if available."""
+        try:
+            if hasattr(self.agent, "computer_handler") and self.agent.computer_handler:
+                handler = self.agent.computer_handler
+                # Check if it's a cuaComputerHandler with a cua_computer
+                if hasattr(handler, "cua_computer"):
+                    computer = handler.cua_computer
+                    if hasattr(computer, "config") and hasattr(computer.config, "name"):
+                        return computer.config.name
+        except Exception:
+            pass
+        return None
+
+    def _calculate_context_size(self, items: List[Dict[str, Any]]) -> int:
+        """Calculate approximate context size in tokens/characters."""
+        total_size = 0
+
+        for item in items:
+            if item.get("type") == "message" and "content" in item:
+                content = item["content"]
+                if isinstance(content, str):
+                    total_size += len(content)
+                elif isinstance(content, list):
+                    for part in content:
+                        if isinstance(part, dict) and "text" in part:
+                            total_size += len(part["text"])
+            elif "content" in item and isinstance(item["content"], str):
+                total_size += len(item["content"])
+
+        return total_size
+
+    def _extract_trajectory(self, items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Extract trajectory items that should be logged."""
+        trajectory = []
+
+        for item in items:
+            # Include user messages, assistant messages, reasoning, computer calls, and computer outputs
+            if (
+                item.get("role") == "user"  # User inputs
+                or (
+                    item.get("type") == "message" and item.get("role") == "assistant"
+                )  # Model outputs
+                or item.get("type") == "reasoning"  # Reasoning traces
+                or item.get("type") == "computer_call"  # Computer actions
+                or item.get("type") == "computer_call_output"  # Computer outputs
+            ):
+                # Create a copy of the item with timestamp
+                trajectory_item = item.copy()
+                trajectory_item["logged_at"] = time.time()
+                trajectory.append(trajectory_item)
+
+        return trajectory

libs/python/agent/cua_agent/callbacks/trajectory_saver.py

@@ -0,0 +1,660 @@
+"""
+Trajectory saving callback handler for ComputerAgent.
+"""
+
+import base64
+import io
+import json
+import os
+import uuid
+from copy import deepcopy
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+try:
+    from typing import override
+except ImportError:
+    from typing_extensions import override
+
+from PIL import Image, ImageDraw
+
+from .base import AsyncCallbackHandler
+
+
+def sanitize_image_urls(data: Any) -> Any:
+    """
+    Recursively search for 'image_url' keys and set their values to '[omitted]'.
+
+    Args:
+        data: Any data structure (dict, list, or primitive type)
+
+    Returns:
+        A deep copy of the data with all 'image_url' values replaced with '[omitted]'
+    """
+    if isinstance(data, dict):
+        # Create a copy of the dictionary
+        sanitized = {}
+        for key, value in data.items():
+            if key == "image_url":
+                sanitized[key] = "[omitted]"
+            else:
+                # Recursively sanitize the value
+                sanitized[key] = sanitize_image_urls(value)
+        return sanitized
+
+    elif isinstance(data, list):
+        # Recursively sanitize each item in the list
+        return [sanitize_image_urls(item) for item in data]
+
+    else:
+        # For primitive types (str, int, bool, None, etc.), return as-is
+        return data
+
+
+def extract_computer_call_outputs(
+    items: List[Dict[str, Any]], screenshot_dir: Optional[Path]
+) -> List[Dict[str, Any]]:
+    """
+    Save any base64-encoded screenshots from computer_call_output or function_call_output
+    entries to files and replace their image_url with the saved file path when a call_id is present.
+
+    Only operates if screenshot_dir is provided and exists; otherwise returns items unchanged.
+
+    Args:
+        items: List of message/result dicts potentially containing computer_call_output
+               or function_call_output entries
+        screenshot_dir: Directory to write screenshots into
+
+    Returns:
+        A new list with updated image_url fields when applicable.
+    """
+    if not items:
+        return items
+    if not screenshot_dir or not screenshot_dir.exists():
+        return items
+
+    updated: List[Dict[str, Any]] = []
+    for item in items:
+        # work on a shallow copy; deep copy nested 'output' if we modify it
+        msg = dict(item)
+        try:
+            if msg.get("type") == "computer_call_output":
+                call_id = msg.get("call_id")
+                output = msg.get("output", {})
+                image_url = output.get("image_url")
+                if call_id and isinstance(image_url, str) and image_url.startswith("data:"):
+                    # derive extension from MIME type e.g. data:image/png;base64,
+                    try:
+                        ext = image_url.split(";", 1)[0].split("/")[-1]
+                        if not ext:
+                            ext = "png"
+                    except Exception:
+                        ext = "png"
+                    out_path = screenshot_dir / f"{call_id}.{ext}"
+                    # write file if it doesn't exist
+                    if not out_path.exists():
+                        try:
+                            b64_payload = image_url.split(",", 1)[1]
+                            img_bytes = base64.b64decode(b64_payload)
+                            out_path.parent.mkdir(parents=True, exist_ok=True)
+                            with open(out_path, "wb") as f:
+                                f.write(img_bytes)
+                        except Exception:
+                            # if anything fails, skip modifying this message
+                            pass
+                    # update image_url to file path
+                    new_output = dict(output)
+                    new_output["image_url"] = str(out_path)
+                    msg["output"] = new_output
+
+            elif msg.get("type") == "function_call_output":
+                # Handle function_call_output from GPT 5.4 / BrowserTool
+                call_id = msg.get("call_id")
+                output = msg.get("output", "")
+
+                # Parse output if it's a string
+                if isinstance(output, str):
+                    try:
+                        output_dict = json.loads(output)
+                    except (json.JSONDecodeError, TypeError):
+                        output_dict = None
+                else:
+                    output_dict = output
+
+                if isinstance(output_dict, dict) and call_id:
+                    image_data = None
+                    image_key = None
+
+                    # Format 1: {"type": "input_image", "image_url": "data:image/png;base64,..."}
+                    if output_dict.get("type") == "input_image":
+                        image_url = output_dict.get("image_url", "")
+                        if isinstance(image_url, str) and image_url.startswith("data:"):
+                            image_data = image_url.split(",", 1)[1] if "," in image_url else None
+                            image_key = "image_url"
+
+                    # Format 2: {"success": True, "screenshot": "base64data"}
+                    elif output_dict.get("screenshot"):
+                        image_data = output_dict.get("screenshot")
+                        image_key = "screenshot"
+
+                    if image_data and image_key:
+                        out_path = screenshot_dir / f"{call_id}.png"
+                        if not out_path.exists():
+                            try:
+                                img_bytes = base64.b64decode(image_data)
+                                out_path.parent.mkdir(parents=True, exist_ok=True)
+                                with open(out_path, "wb") as f:
+                                    f.write(img_bytes)
+                            except Exception:
+                                pass
+                        # Update output to reference file path
+                        new_output_dict = dict(output_dict)
+                        new_output_dict[image_key] = str(out_path)
+                        msg["output"] = json.dumps(new_output_dict)
+
+            elif msg.get("role") == "user":
+                # Handle user messages with input_image content (GPT-5.4 sibling screenshot messages)
+                # These accompany function_call_output for computer calls
+                content = msg.get("content", [])
+                if isinstance(content, list):
+                    new_content = []
+                    content_modified = False
+                    for content_item in content:
+                        if (
+                            isinstance(content_item, dict)
+                            and content_item.get("type") == "input_image"
+                        ):
+                            image_url = content_item.get("image_url", "")
+                            if isinstance(image_url, str) and image_url.startswith("data:"):
+                                # Generate a unique ID for this screenshot
+                                screenshot_id = str(uuid.uuid4())[:8]
+                                try:
+                                    ext = image_url.split(";", 1)[0].split("/")[-1]
+                                    if not ext:
+                                        ext = "png"
+                                except Exception:
+                                    ext = "png"
+                                out_path = screenshot_dir / f"user_screenshot_{screenshot_id}.{ext}"
+                                if not out_path.exists():
+                                    try:
+                                        b64_payload = image_url.split(",", 1)[1]
+                                        img_bytes = base64.b64decode(b64_payload)
+                                        out_path.parent.mkdir(parents=True, exist_ok=True)
+                                        with open(out_path, "wb") as f:
+                                            f.write(img_bytes)
+                                    except Exception:
+                                        new_content.append(content_item)
+                                        continue
+                                # Update image_url to file path
+                                new_item = dict(content_item)
+                                new_item["image_url"] = str(out_path)
+                                new_content.append(new_item)
+                                content_modified = True
+                            else:
+                                new_content.append(content_item)
+                        else:
+                            new_content.append(content_item)
+                    if content_modified:
+                        msg["content"] = new_content
+
+        except Exception:
+            # do not block on malformed entries; keep original
+            pass
+        updated.append(msg)
+    return updated
+
+
+class TrajectorySaverCallback(AsyncCallbackHandler):
+    """
+    Callback handler that saves agent trajectories to disk.
+
+    Saves each run as a separate trajectory with unique ID, and each turn
+    within the trajectory gets its own folder with screenshots and responses.
+    """
+
+    def __init__(
+        self, trajectory_dir: str, reset_on_run: bool = True, screenshot_dir: Optional[str] = None
+    ):
+        """
+        Initialize trajectory saver.
+
+        Args:
+            trajectory_dir: Base directory to save trajectories
+            reset_on_run: If True, reset trajectory_id/turn/artifact on each run.
+                         If False, continue using existing trajectory_id if set.
+        """
+        self.trajectory_dir = Path(trajectory_dir)
+        self.trajectory_id: Optional[str] = None
+        self.current_turn: int = 0
+        self.current_artifact: int = 0
+        self.model: Optional[str] = None
+        self.total_usage: Dict[str, Any] = {}
+        self.reset_on_run = reset_on_run
+        # Optional directory to store extracted screenshots from metadata/new_items
+        self.screenshot_dir: Optional[Path] = Path(screenshot_dir) if screenshot_dir else None
+
+        # Ensure trajectory directory exists
+        self.trajectory_dir.mkdir(parents=True, exist_ok=True)
+
+        # Ensure screenshot directory exists if specified
+        if self.screenshot_dir:
+            self.screenshot_dir.mkdir(parents=True, exist_ok=True)
+
+    def _get_turn_dir(self) -> Path:
+        """Get the directory for the current turn."""
+        if not self.trajectory_id:
+            raise ValueError("Trajectory not initialized - call _on_run_start first")
+
+        # format: trajectory_id/turn_000
+        turn_dir = self.trajectory_dir / self.trajectory_id / f"turn_{self.current_turn:03d}"
+        turn_dir.mkdir(parents=True, exist_ok=True)
+        return turn_dir
+
+    def _save_artifact(self, name: str, artifact: Union[str, bytes, Dict[str, Any]]) -> None:
+        """Save an artifact to the current turn directory."""
+        turn_dir = self._get_turn_dir()
+        if isinstance(artifact, bytes):
+            # format: turn_000/0000_name.png
+            artifact_filename = f"{self.current_artifact:04d}_{name}"
+            artifact_path = turn_dir / f"{artifact_filename}.png"
+            with open(artifact_path, "wb") as f:
+                f.write(artifact)
+        else:
+            # format: turn_000/0000_name.json
+            artifact_filename = f"{self.current_artifact:04d}_{name}"
+            artifact_path = turn_dir / f"{artifact_filename}.json"
+            # add created_at
+            if isinstance(artifact, dict):
+                artifact = artifact.copy()
+                artifact["created_at"] = str(uuid.uuid1().time)
+            with open(artifact_path, "w") as f:
+                json.dump(sanitize_image_urls(artifact), f, indent=2)
+        self.current_artifact += 1
+
+    def _update_usage(self, usage: Dict[str, Any]) -> None:
+        """Update total usage statistics."""
+
+        def add_dicts(target: Dict[str, Any], source: Dict[str, Any]) -> None:
+            for key, value in source.items():
+                if isinstance(value, dict):
+                    if key not in target:
+                        target[key] = {}
+                    add_dicts(target[key], value)
+                else:
+                    if key not in target:
+                        target[key] = 0
+                    target[key] += value
+
+        add_dicts(self.total_usage, usage)
+
+    @override
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Initialize trajectory tracking for a new run."""
+        model = kwargs.get("model", "unknown")
+
+        # Only reset trajectory state if reset_on_run is True or no trajectory exists
+        if self.reset_on_run or not self.trajectory_id:
+            model_name_short = model.split("+")[-1].split("/")[-1].lower()[:16]
+            if "+" in model:
+                model_name_short = model.split("+")[0].lower()[:4] + "_" + model_name_short
+            # strip non-alphanumeric characters from model_name_short
+            model_name_short = "".join(c for c in model_name_short if c.isalnum() or c == "_")
+
+            # id format: yyyy-mm-dd_model_hhmmss_uuid[:4]
+            now = datetime.now()
+            self.trajectory_id = f"{now.strftime('%Y-%m-%d')}_{model_name_short}_{now.strftime('%H%M%S')}_{str(uuid.uuid4())[:4]}"
+            self.current_turn = 0
+            self.current_artifact = 0
+            self.model = model
+            self.total_usage = {}
+
+            # Create trajectory directory
+            trajectory_path = self.trajectory_dir / self.trajectory_id
+            trajectory_path.mkdir(parents=True, exist_ok=True)
+
+            # Save trajectory metadata (optionally extract screenshots to screenshot_dir)
+            kwargs_to_save = kwargs.copy()
+            try:
+                if "messages" in kwargs_to_save:
+                    kwargs_to_save["messages"] = extract_computer_call_outputs(
+                        kwargs_to_save["messages"], self.screenshot_dir
+                    )
+            except Exception:
+                # If extraction fails, fall back to original messages
+                pass
+            metadata = {
+                "trajectory_id": self.trajectory_id,
+                "created_at": str(uuid.uuid1().time),
+                "status": "running",
+                "kwargs": kwargs_to_save,
+            }
+
+            with open(trajectory_path / "metadata.json", "w") as f:
+                json.dump(metadata, f, indent=2)
+        else:
+            # Continue with existing trajectory - just update model if needed
+            self.model = model
+
+    @override
+    async def on_run_end(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> None:
+        """Finalize run tracking by updating metadata with completion status, usage, and new items."""
+        if not self.trajectory_id:
+            return
+
+        # Update metadata with completion status, total usage, and new items
+        trajectory_path = self.trajectory_dir / self.trajectory_id
+        metadata_path = trajectory_path / "metadata.json"
+
+        # Read existing metadata
+        if metadata_path.exists():
+            with open(metadata_path, "r") as f:
+                metadata = json.load(f)
+        else:
+            metadata = {}
+
+        # Update metadata with completion info
+        # Optionally extract screenshots from new_items before persisting
+        new_items_to_save = new_items
+        try:
+            new_items_to_save = extract_computer_call_outputs(new_items, self.screenshot_dir)
+        except Exception:
+            pass
+
+        metadata.update(
+            {
+                "status": "completed",
+                "completed_at": str(uuid.uuid1().time),
+                "total_usage": self.total_usage,
+                "new_items": new_items_to_save,
+                "total_turns": self.current_turn,
+            }
+        )
+
+        # Save updated metadata
+        with open(metadata_path, "w") as f:
+            json.dump(metadata, f, indent=2)
+
+    @override
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        if not self.trajectory_id:
+            return
+
+        self._save_artifact("api_start", {"kwargs": kwargs})
+
+    @override
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """Save API call result."""
+        if not self.trajectory_id:
+            return
+
+        self._save_artifact("api_result", {"kwargs": kwargs, "result": result})
+
+    @override
+    async def on_screenshot(self, screenshot: Union[str, bytes], name: str = "screenshot") -> None:
+        """Save a screenshot."""
+        if isinstance(screenshot, str):
+            screenshot = base64.b64decode(screenshot)
+        self._save_artifact(name, screenshot)
+
+    @override
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Called when usage information is received."""
+        self._update_usage(usage)
+
+    @override
+    async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
+        """Save responses to the current turn directory and update usage statistics."""
+        if not self.trajectory_id:
+            return
+
+        # Save responses
+        turn_dir = self._get_turn_dir()
+        response_data = {
+            "timestamp": str(uuid.uuid1().time),
+            "model": self.model,
+            "kwargs": kwargs,
+            "response": responses,
+        }
+
+        self._save_artifact("agent_response", response_data)
+
+        # Increment turn counter
+        self.current_turn += 1
+
+    def _draw_crosshair_on_image(self, image_bytes: bytes, x: int, y: int) -> bytes:
+        """
+        Draw a red dot and crosshair at the specified coordinates on the image.
+
+        Args:
+            image_bytes: The original image as bytes
+            x: X coordinate for the crosshair
+            y: Y coordinate for the crosshair
+
+        Returns:
+            Modified image as bytes with red dot and crosshair
+        """
+        # Open the image
+        image = Image.open(io.BytesIO(image_bytes))
+        draw = ImageDraw.Draw(image)
+
+        # Draw crosshair lines (red, 2px thick)
+        crosshair_size = 20
+        line_width = 2
+        color = "red"
+
+        # Horizontal line
+        draw.line([(x - crosshair_size, y), (x + crosshair_size, y)], fill=color, width=line_width)
+        # Vertical line
+        draw.line([(x, y - crosshair_size), (x, y + crosshair_size)], fill=color, width=line_width)
+
+        # Draw center dot (filled circle)
+        dot_radius = 3
+        draw.ellipse(
+            [(x - dot_radius, y - dot_radius), (x + dot_radius, y + dot_radius)], fill=color
+        )
+
+        # Convert back to bytes
+        output = io.BytesIO()
+        image.save(output, format="PNG")
+        return output.getvalue()
+
+    @override
+    async def on_computer_call_end(
+        self, item: Dict[str, Any], result: List[Dict[str, Any]]
+    ) -> None:
+        """
+        Called when a computer call has completed.
+        Saves screenshots and computer call output.
+        """
+        if not self.trajectory_id:
+            return
+
+        self._save_artifact("computer_call_result", {"item": item, "result": result})
+
+        # Check if action has x/y coordinates and there's a screenshot in the result
+        action = item.get("action", {})
+        if "x" in action and "y" in action:
+            # Look for screenshot in the result
+            for result_item in result:
+                if (
+                    result_item.get("type") == "computer_call_output"
+                    and result_item.get("output", {}).get("type") == "input_image"
+                ):
+
+                    image_url = result_item["output"]["image_url"]
+
+                    # Extract base64 image data
+                    if image_url.startswith("data:image/"):
+                        # Format: data:image/png;base64,<base64_data>
+                        base64_data = image_url.split(",", 1)[1]
+                    else:
+                        # Assume it's just base64 data
+                        base64_data = image_url
+
+                    try:
+                        # Decode the image
+                        image_bytes = base64.b64decode(base64_data)
+
+                        # Draw crosshair at the action coordinates
+                        annotated_image = self._draw_crosshair_on_image(
+                            image_bytes, int(action["x"]), int(action["y"])
+                        )
+
+                        # Save as screenshot_action
+                        self._save_artifact("screenshot_action", annotated_image)
+
+                    except Exception as e:
+                        # If annotation fails, just log and continue
+                        print(f"Failed to annotate screenshot: {e}")
+
+                    break  # Only process the first screenshot found
+
+        # Increment turn counter
+        self.current_turn += 1
+
+    @override
+    async def on_function_call_end(
+        self, item: Dict[str, Any], result: List[Dict[str, Any]]
+    ) -> None:
+        """
+        Called when a function call has completed.
+        Saves screenshots and function call output for GPT 5.4 / BrowserTool.
+        """
+        if not self.trajectory_id:
+            return
+
+        self._save_artifact("function_call_result", {"item": item, "result": result})
+
+        # Extract coordinates from function call arguments if present
+        x_coord, y_coord = None, None
+        try:
+            arguments = item.get("arguments", "{}")
+            if isinstance(arguments, str):
+                args_dict = json.loads(arguments)
+            else:
+                args_dict = arguments
+
+            # Check for coordinate array format (BrowserTool style)
+            coord = args_dict.get("coordinate")
+            if coord and isinstance(coord, list) and len(coord) >= 2:
+                x_coord, y_coord = coord[0], coord[1]
+            # Check for x/y format (computer_use style)
+            elif "x" in args_dict and "y" in args_dict:
+                x_coord, y_coord = args_dict.get("x"), args_dict.get("y")
+        except (json.JSONDecodeError, TypeError):
+            pass
+
+        # Look for screenshot in the result
+        screenshot_found = False
+        for result_item in result:
+            if screenshot_found:
+                break
+
+            if result_item.get("type") == "function_call_output":
+                output = result_item.get("output", "")
+
+                # Parse output if it's a string
+                if isinstance(output, str):
+                    try:
+                        output_dict = json.loads(output)
+                    except (json.JSONDecodeError, TypeError):
+                        # Try to evaluate as Python literal (for stringified dicts)
+                        try:
+                            import ast
+
+                            output_dict = ast.literal_eval(output)
+                        except (ValueError, SyntaxError):
+                            continue
+                else:
+                    output_dict = output
+
+                if not isinstance(output_dict, dict):
+                    continue
+
+                # Extract screenshot from various formats
+                image_data = None
+
+                # Format 1: {"type": "input_image", "image_url": "data:image/png;base64,..."}
+                if output_dict.get("type") == "input_image":
+                    image_url = output_dict.get("image_url", "")
+                    if image_url.startswith("data:image/"):
+                        image_data = image_url.split(",", 1)[1]
+                    elif image_url:
+                        image_data = image_url
+
+                # Format 2: {"success": True, "screenshot": "base64data"}
+                elif output_dict.get("screenshot"):
+                    image_data = output_dict.get("screenshot")
+
+                if image_data:
+                    try:
+                        # Decode the image
+                        image_bytes = base64.b64decode(image_data)
+
+                        # If we have coordinates, draw crosshair annotation
+                        if (
+                            x_coord is not None
+                            and y_coord is not None
+                            and x_coord != 0
+                            and y_coord != 0
+                        ):
+                            annotated_image = self._draw_crosshair_on_image(
+                                image_bytes, int(x_coord), int(y_coord)
+                            )
+                            self._save_artifact("screenshot_action", annotated_image)
+                        else:
+                            # Save plain screenshot without crosshair
+                            self._save_artifact("screenshot", image_bytes)
+
+                        screenshot_found = True
+
+                    except Exception as e:
+                        # If processing fails, just log and continue
+                        print(f"Failed to process screenshot from function call: {e}")
+
+            # Handle sibling user messages with input_image content (GPT-5.4 computer calls)
+            # These accompany function_call_output and contain the actual screenshot
+            elif result_item.get("role") == "user":
+                content = result_item.get("content", [])
+                if isinstance(content, list):
+                    for content_item in content:
+                        if (
+                            isinstance(content_item, dict)
+                            and content_item.get("type") == "input_image"
+                        ):
+                            image_url = content_item.get("image_url", "")
+                            if isinstance(image_url, str) and image_url.startswith("data:"):
+                                try:
+                                    b64_payload = image_url.split(",", 1)[1]
+                                    image_bytes = base64.b64decode(b64_payload)
+
+                                    # If we have coordinates, draw crosshair annotation
+                                    if (
+                                        x_coord is not None
+                                        and y_coord is not None
+                                        and x_coord != 0
+                                        and y_coord != 0
+                                    ):
+                                        annotated_image = self._draw_crosshair_on_image(
+                                            image_bytes, int(x_coord), int(y_coord)
+                                        )
+                                        self._save_artifact("screenshot_action", annotated_image)
+                                    else:
+                                        # Save plain screenshot without crosshair
+                                        self._save_artifact("screenshot", image_bytes)
+
+                                    screenshot_found = True
+                                    break
+
+                                except Exception as e:
+                                    # If processing fails, just log and continue
+                                    print(f"Failed to process screenshot from user message: {e}")
+
+        # Increment turn counter
+        self.current_turn += 1