chore: import upstream snapshot with attribution

libs/python/agent/.bumpversion.cfg

@@ -0,0 +1,10 @@
+[bumpversion]
+current_version = 0.8.4
+commit = True
+tag = True
+tag_name = agent-v{new_version}
+message = Bump cua-agent to v{new_version}
+
+[bumpversion:file:pyproject.toml]
+search = version = "{current_version}"
+replace = version = "{new_version}"

libs/python/agent/README.md

@@ -0,0 +1,5 @@
+# Cua Agent
+
+Computer-Use framework with liteLLM integration for running agentic workflows on macOS, Windows, and Linux sandboxes.
+
+**[Documentation](https://cua.ai/docs/cua/reference/agent-sdk)** - Installation, guides, and configuration.

libs/python/agent/benchmarks/.gitignore

@@ -0,0 +1,3 @@
+output/
+interactive_output/
+*_results.md

libs/python/agent/benchmarks/README.md

@@ -0,0 +1,76 @@
+# Computer Agent Benchmarks
+
+This directory contains benchmarks designed to test agent providers in the Computer Agent SDK against reference agent implementations.
+
+## Overview
+
+The benchmark system evaluates models on GUI grounding tasks, specifically click prediction accuracy. It supports both:
+
+- **Computer Agent SDK providers** (using model strings like `"huggingface-local/HelloKKMe/GTA1-7B"`)
+- **Reference agent implementations** (custom model classes implementing the `ModelProtocol`)
+
+## Available Benchmarks
+
+### 1. ScreenSpot-v2 (`ss-v2.py`)
+
+- **Dataset**: ScreenSpot-v2 (click-only GUI grounding)
+- **Format**: Standard resolution screenshots
+- **Task**: Predict click coordinates given an instruction and image
+- **Metrics**: Accuracy, Error Rate, Timing, VRAM usage
+
+### 2. ScreenSpot-Pro (`ss-pro.py`)
+
+- **Dataset**: ScreenSpot-Pro (high-resolution click-only GUI grounding)
+- **Format**: High-resolution screenshots
+- **Task**: Predict click coordinates given an instruction and image
+- **Metrics**: Accuracy, Error Rate, Timing, VRAM usage
+
+### 3. Interactive Testing (`interactive.py`)
+
+- **Real-time testing**: Take screenshots and visualize model predictions
+- **Commands**:
+  - Type instruction → test all models on last screenshot
+  - `screenshot` → take screenshot
+  - `models` → list available models
+  - `quit`/`exit` → exit tool
+- **Output**: Visual predictions with crosshairs for each model
+
+## Running Benchmarks
+
+### 1. Configure Models
+
+Edit `utils.py` to specify which models you want to test in `get_available_models()`.
+
+### 2. Run Benchmark
+
+```bash
+# ScreenSpot-v2 benchmark
+python ss-v2.py --samples 50
+
+# ScreenSpot-Pro benchmark
+python ss-pro.py --samples 50
+
+# Interactive testing
+python interactive.py
+```
+
+## Output
+
+### Console Output
+
+```
+Model Results:
+  Accuracy: 85.50% (171/200)
+  Avg Time: 1.23s (0.89s - 2.45s)
+  VRAM Usage: 4.5GB (max) / 3.4GB (avg)
+```
+
+### Generated Files
+
+- **Markdown Report**: `*_results.md` with detailed results tables
+- **Visualizations**: `output/` directory with prediction visualizations
+- **Interactive Output**: `interactive_output/` for interactive session results
+
+## Contributing
+
+To add a new reference model, follow the instructions in [contrib.md](contrib.md).

libs/python/agent/benchmarks/contrib.md

@@ -0,0 +1,166 @@
+# Contributing Reference Agent Implementations
+
+This guide explains how to add your own reference agent implementations to the benchmark system.
+
+## Adding Reference Agent Implementations
+
+### 1. Implement the ModelProtocol
+
+Create a new file in `models/` directory implementing the `ModelProtocol`:
+
+```python
+from models.base import ModelProtocol
+from typing import Optional, Tuple
+from PIL import Image
+
+class YourModelName(ModelProtocol):
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self._model = None
+
+    @property
+    def model_name(self) -> str:
+        return self.model_path
+
+    async def load_model(self) -> None:
+        """Load the model into memory."""
+        # Your model loading logic here
+        pass
+
+    async def unload_model(self) -> None:
+        """Unload the model from memory."""
+        # Your model cleanup logic here
+        pass
+
+    async def predict_click(self, image: Image.Image, instruction: str) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates for the given image and instruction.
+
+        Args:
+            image: PIL Image to analyze
+            instruction: Text instruction describing what to click
+
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        # Your prediction logic here
+        return (x, y)  # Return predicted coordinates
+```
+
+### 2. Register Your Model
+
+Add your model to the `get_available_models()` function in `utils.py`:
+
+```python
+def get_available_models() -> List[Union[str, ModelProtocol]]:
+    models = [
+        # Computer Agent SDK providers
+        "huggingface-local/HelloKKMe/GTA1-7B",
+
+        # Reference implementations
+        GTA1Model("HelloKKMe/GTA1-7B"),
+        YourModelName("path/to/your/model"),  # Add your model here
+    ]
+    return models
+```
+
+### 3. Test Your Implementation
+
+Before submitting, test your model with the interactive tool:
+
+```bash
+python interactive.py
+```
+
+This will help you verify that your model loads correctly and produces reasonable predictions.
+
+## Example: Adding a New Model
+
+Here's a complete example of adding a hypothetical "MyVisionModel":
+
+1. **Create `models/my_vision_model.py`:**
+
+```python
+import torch
+from transformers import AutoModel, AutoProcessor
+from models.base import ModelProtocol
+from typing import Optional, Tuple
+from PIL import Image
+
+class MyVisionModel(ModelProtocol):
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.model = None
+        self.processor = None
+
+    @property
+    def model_name(self) -> str:
+        return f"MyVisionModel({self.model_path})"
+
+    async def load_model(self) -> None:
+        """Load the model and processor."""
+        self.processor = AutoProcessor.from_pretrained(self.model_path)
+        self.model = AutoModel.from_pretrained(
+            self.model_path,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+
+    async def unload_model(self) -> None:
+        """Clean up model resources."""
+        del self.model
+        del self.processor
+        self.model = None
+        self.processor = None
+        torch.cuda.empty_cache()
+
+    async def predict_click(self, image: Image.Image, instruction: str) -> Optional[Tuple[int, int]]:
+        """Predict click coordinates."""
+        try:
+            # Preprocess inputs
+            inputs = self.processor(
+                text=instruction,
+                images=image,
+                return_tensors="pt"
+            )
+
+            # Run inference
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+
+            # Extract coordinates (model-specific logic)
+            x, y = self._extract_coordinates(outputs)
+            return (int(x), int(y))
+
+        except Exception as e:
+            print(f"Prediction failed: {e}")
+            return None
+
+    def _extract_coordinates(self, outputs):
+        """Extract x, y coordinates from model outputs."""
+        # Your model-specific coordinate extraction logic
+        pass
+```
+
+2. **Update `models/__init__.py`:**
+
+```python
+from .gta1 import GTA1Model
+from .my_vision_model import MyVisionModel
+
+__all__ = ["GTA1Model", "MyVisionModel"]
+```
+
+3. **Update `utils.py`:**
+
+```python
+from models import GTA1Model, MyVisionModel
+
+def get_available_models() -> List[Union[str, ModelProtocol]]:
+    models = [
+        "huggingface-local/HelloKKMe/GTA1-7B",
+        GTA1Model("HelloKKMe/GTA1-7B"),
+        MyVisionModel("my-org/my-vision-model"),  # Add here
+    ]
+    return models
+```

libs/python/agent/benchmarks/interactive.py

@@ -0,0 +1,201 @@
+#!/usr/bin/env python3
+"""
+Interactive Click Prediction Tool
+
+Takes screenshots and allows testing multiple models interactively.
+Models are loaded/unloaded one at a time to avoid memory issues.
+"""
+
+import asyncio
+import os
+from datetime import datetime
+from typing import Any, Dict, List
+
+from utils import (
+    ModelWrapper,
+    get_available_models,
+    save_prediction_visualization,
+    take_screenshot,
+)
+
+
+async def predict_with_all_models(image, instruction: str, models) -> List[Dict[str, Any]]:
+    """
+    Predict click coordinates with all models sequentially.
+
+    Args:
+        image: PIL Image to analyze
+        instruction: Instruction text
+        models: List of model instances
+
+    Returns:
+        List of prediction results
+    """
+    predictions = []
+
+    for model in models:
+        model_wrapper = ModelWrapper(model)
+        print(f"\n🔄 Loading {model_wrapper.model_name}...")
+
+        try:
+            # Load model
+            await model_wrapper.load_model()
+
+            # Predict
+            coords = await model_wrapper.predict_click(image, instruction)
+
+            predictions.append(
+                {"model_name": model_wrapper.model_name, "coords": coords, "error": None}
+            )
+
+            if coords:
+                print(f"✅ {model_wrapper.model_name}: ({coords[0]}, {coords[1]})")
+            else:
+                print(f"❌ {model_wrapper.model_name}: No prediction")
+
+        except Exception as e:
+            print(f"❌ {model_wrapper.model_name}: ERROR - {str(e)}")
+            predictions.append(
+                {"model_name": model_wrapper.model_name, "coords": None, "error": str(e)}
+            )
+
+        finally:
+            # Always unload model to free memory
+            try:
+                await model_wrapper.unload_model()
+                print(f"🗑️  Unloaded {model_wrapper.model_name}")
+            except Exception as e:
+                print(f"⚠️  Error unloading {model_wrapper.model_name}: {e}")
+
+    return predictions
+
+
+def print_header():
+    """Print the interactive tool header."""
+    print("=" * 60)
+    print("🖱️  Interactive Click Prediction Tool")
+    print("=" * 60)
+    print("Commands:")
+    print("  • Type an instruction to test models on last screenshot")
+    print("  • 'screenshot' - Take a new screenshot")
+    print("  • 'models' - List available models")
+    print("  • 'quit' or 'exit' - Exit the tool")
+    print("=" * 60)
+    print("💡 Tip: Take a screenshot first, then send instructions to test models!")
+
+
+def print_models(models):
+    """Print available models."""
+    print("\n📋 Available Models:")
+    for i, model in enumerate(models, 1):
+        if isinstance(model, str):
+            print(f"  {i}. {model}")
+        else:
+            print(f"  {i}. models.{model.__class__.__name__}")
+
+
+async def main():
+    """
+    Main interactive loop.
+    """
+    print_header()
+
+    # Get available models
+    models = get_available_models()
+    print_models(models)
+
+    # Create output directory for visualizations
+    output_dir = "interactive_output"
+    os.makedirs(output_dir, exist_ok=True)
+
+    session_count = 0
+    last_screenshot = None
+    screenshot_timestamp = None
+
+    while True:
+        try:
+            # Get user input
+            print(f"\n{'='*40}")
+            user_input = input("🎯 Enter instruction (or command): ").strip()
+
+            if not user_input:
+                continue
+
+            # Handle commands
+            if user_input.lower() in ["quit", "exit", "q"]:
+                print("👋 Goodbye!")
+                break
+
+            elif user_input.lower() == "models":
+                print_models(models)
+                continue
+
+            elif user_input.lower() == "screenshot":
+                print("📸 Taking screenshot...")
+                try:
+                    last_screenshot = take_screenshot()
+                    screenshot_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                    screenshot_path = os.path.join(
+                        output_dir, f"screenshot_{screenshot_timestamp}.png"
+                    )
+                    last_screenshot.save(screenshot_path)
+                    print(f"✅ Screenshot captured and saved to: {screenshot_path}")
+                    print(f"📝 Ready for instructions! Screenshot size: {last_screenshot.size}")
+                except Exception as e:
+                    print(f"❌ Error taking screenshot: {e}")
+                continue
+
+            # Handle instruction input
+            if last_screenshot is None:
+                print(
+                    "⚠️  No screenshot available! Please take a screenshot first using 'screenshot' command."
+                )
+                continue
+
+            session_count += 1
+            print(f"\n🎯 Session {session_count}: '{user_input}'")
+            print(f"📷 Using screenshot from: {screenshot_timestamp}")
+
+            # Predict with all models using last screenshot
+            print(f"\n🤖 Testing {len(models)} models on screenshot...")
+            predictions = await predict_with_all_models(last_screenshot, user_input, models)
+
+            # Display results summary
+            print("\n📊 Results Summary:")
+            print("-" * 50)
+            for pred in predictions:
+                if pred["coords"]:
+                    print(f"✅ {pred['model_name']}: ({pred['coords'][0]}, {pred['coords'][1]})")
+                elif pred["error"]:
+                    print(f"❌ {pred['model_name']}: ERROR - {pred['error']}")
+                else:
+                    print(f"❌ {pred['model_name']}: No prediction")
+
+            # Save visualization
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            vis_filename = f"session_{session_count:03d}_{timestamp}.png"
+            vis_path = os.path.join(output_dir, vis_filename)
+
+            try:
+                save_prediction_visualization(last_screenshot, user_input, predictions, vis_path)
+                print(f"\n💾 Visualization saved to: {vis_path}")
+            except Exception as e:
+                print(f"⚠️  Error saving visualization: {e}")
+
+            print(f"\n✨ Session {session_count} completed!")
+
+        except KeyboardInterrupt:
+            print("\n\n👋 Interrupted by user. Goodbye!")
+            break
+        except Exception as e:
+            print(f"\n❌ Unexpected error: {e}")
+            print("Continuing...")
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("\n👋 Goodbye!")
+    except Exception as e:
+        print(f"❌ Fatal error: {e}")

libs/python/agent/benchmarks/models/__init__.py

@@ -0,0 +1,3 @@
+from .base import ModelProtocol
+
+__all__ = ["ModelProtocol"]

libs/python/agent/benchmarks/models/base.py

@@ -0,0 +1,39 @@
+"""
+Base protocol for benchmark models.
+"""
+
+from typing import Optional, Protocol, Tuple
+
+from PIL import Image
+
+
+class ModelProtocol(Protocol):
+    """Protocol for benchmark models that can predict click coordinates."""
+
+    @property
+    def model_name(self) -> str:
+        """Return the name of the model."""
+        ...
+
+    async def load_model(self) -> None:
+        """Load the model into memory."""
+        ...
+
+    async def unload_model(self) -> None:
+        """Unload the model from memory."""
+        ...
+
+    async def predict_click(
+        self, image: Image.Image, instruction: str
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates for the given image and instruction.
+
+        Args:
+            image: PIL Image to analyze
+            instruction: Text instruction describing what to click
+
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        ...

libs/python/agent/benchmarks/models/gta1.py

@@ -0,0 +1,158 @@
+"""
+GTA1 model implementation for benchmarking.
+"""
+
+import gc
+import re
+from typing import Optional, Tuple
+
+import torch
+from PIL import Image
+from qwen_vl_utils import process_vision_info, smart_resize
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+
+from .base import ModelProtocol
+
+
+class GTA1Model:
+    """Ground truth GTA1 model implementation."""
+
+    def __init__(self, model_path: str = "HelloKKMe/GTA1-7B"):
+        self.model_path = model_path
+        self.model = None
+        self.processor = None
+        self.max_new_tokens = 32
+
+        self.system_prompt = """
+You are an expert UI element locator. Given a GUI image and a user's element description, provide the coordinates of the specified element as a single (x,y) point. The image resolution is height {height} and width {width}. For elements with area, return the center point.
+
+Output the coordinate pair exactly:
+(x,y)
+""".strip()
+
+    @property
+    def model_name(self) -> str:
+        """Return the name of the model."""
+        return f"GTA1-{self.model_path.split('/')[-1]}"
+
+    async def load_model(self) -> None:
+        """Load the model into memory."""
+        if self.model is None:
+            print(f"Loading GTA1 model: {self.model_path}")
+            self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                self.model_path, torch_dtype=torch.bfloat16, device_map="auto"
+            )
+            self.processor = AutoProcessor.from_pretrained(
+                self.model_path, min_pixels=3136, max_pixels=4096 * 2160
+            )
+            print("GTA1 model loaded successfully")
+
+    async def unload_model(self) -> None:
+        """Unload the model from memory."""
+        if self.model is not None:
+            print("Unloading GTA1 model from GPU...")
+            del self.model
+            del self.processor
+            self.model = None
+            self.processor = None
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            print("GTA1 model unloaded")
+
+    def _extract_coordinates(self, raw_string: str) -> Tuple[int, int]:
+        """Extract coordinates from model output."""
+        try:
+            matches = re.findall(r"\((-?\d*\.?\d+),\s*(-?\d*\.?\d+)\)", raw_string)
+            return tuple(map(int, map(float, matches[0])))  # type: ignore
+        except:
+            return (0, 0)
+
+    async def predict_click(
+        self, image: Image.Image, instruction: str
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates for the given image and instruction.
+
+        Args:
+            image: PIL Image to analyze
+            instruction: Text instruction describing what to click
+
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        if self.model is None or self.processor is None:
+            await self.load_model()
+
+        assert self.processor is not None
+        assert self.model is not None
+
+        try:
+            width, height = image.width, image.height
+
+            # Resize image according to processor requirements
+            resized_height, resized_width = smart_resize(
+                image.height,
+                image.width,
+                factor=self.processor.image_processor.patch_size
+                * self.processor.image_processor.merge_size,
+                min_pixels=self.processor.image_processor.min_pixels,
+                max_pixels=self.processor.image_processor.max_pixels,
+            )
+            resized_image = image.resize((resized_width, resized_height))
+            scale_x, scale_y = width / resized_width, height / resized_height
+
+            # Prepare messages
+            system_message = {
+                "role": "system",
+                "content": self.system_prompt.format(height=resized_height, width=resized_width),
+            }
+
+            user_message = {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": resized_image},
+                    {"type": "text", "text": instruction},
+                ],
+            }
+
+            # Process inputs
+            image_inputs, video_inputs = process_vision_info([system_message, user_message])  # type: ignore
+            text = self.processor.apply_chat_template(
+                [system_message, user_message], tokenize=False, add_generation_prompt=True
+            )
+            inputs = self.processor(
+                text=[text],
+                images=image_inputs,
+                videos=video_inputs,
+                padding=True,
+                return_tensors="pt",
+            )
+            inputs = inputs.to(self.model.device)
+
+            # Generate prediction
+            output_ids = self.model.generate(
+                **inputs,
+                max_new_tokens=self.max_new_tokens,
+                do_sample=False,
+                temperature=1.0,
+                use_cache=True,
+            )
+            generated_ids = [
+                output_ids[len(input_ids) :]
+                for input_ids, output_ids in zip(inputs.input_ids, output_ids)
+            ]
+            output_text = self.processor.batch_decode(
+                generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+            )[0]
+
+            # Extract and rescale coordinates
+            pred_x, pred_y = self._extract_coordinates(output_text)
+            pred_x = int(pred_x * scale_x)
+            pred_y = int(pred_y * scale_y)
+
+            return (pred_x, pred_y)
+
+        except Exception as e:
+            print(f"Error in GTA1 prediction: {e}")
+            return None

libs/python/agent/benchmarks/ss-pro.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+"""
+ScreenSpot-Pro Benchmark Script
+
+Evaluates models on the ScreenSpot-Pro dataset for click prediction accuracy.
+Supports both ComputerAgent model strings and custom model classes.
+"""
+
+import argparse
+import asyncio
+import random
+import statistics
+import time
+from typing import Optional
+
+from datasets import load_dataset
+from tqdm import tqdm
+from utils import (
+    ModelWrapper,
+    get_available_models,
+    get_gpu_memory,
+    is_click_in_bbox,
+    save_results_to_markdown,
+    save_visualizations,
+)
+
+
+async def evaluate_model(
+    model_wrapper: ModelWrapper, dataset, max_samples: Optional[int] = None
+) -> dict:
+    """
+    Evaluate a model on the ScreenSpot-Pro dataset.
+
+    Args:
+        model_wrapper: ModelWrapper instance
+        dataset: ScreenSpot-Pro dataset (list of samples)
+        max_samples: Maximum number of samples to evaluate (None for all)
+
+    Returns:
+        Dictionary with evaluation results
+    """
+    print(f"\nEvaluating model: {model_wrapper.model_name}")
+
+    # Load model
+    await model_wrapper.load_model()
+
+    total_samples = len(dataset)
+    if max_samples is not None:
+        total_samples = min(max_samples, total_samples)
+
+    correct_predictions = 0
+    error_predictions = 0
+    results = []
+
+    for i in tqdm(range(total_samples), desc=f"Evaluating {model_wrapper.model_name}"):
+        sample = dataset[i]
+
+        # Extract sample data
+        image = sample["image"]
+        instruction = sample["instruction"]
+        bbox = sample["bbox"]  # [x1, y1, x2, y2]
+        sample_id = sample["img_filename"]
+
+        # Predict click coordinates with timing
+        start_time = time.time()
+        click_coords = await model_wrapper.predict_click(image, instruction)
+        prediction_time = time.time() - start_time
+
+        # Check if prediction is correct
+        is_correct = is_click_in_bbox(click_coords, bbox)
+
+        if is_correct:
+            correct_predictions += 1
+
+        results.append(
+            {
+                "id": sample_id,
+                "instruction": instruction,
+                "bbox": bbox,
+                "predicted_coords": click_coords,
+                "is_correct": is_correct,
+                "failed": False,
+                "prediction_time": prediction_time,
+            }
+        )
+
+    # Unload model
+    await model_wrapper.unload_model()
+
+    # Calculate metrics
+    accuracy = correct_predictions / total_samples if total_samples > 0 else 0.0
+    error_rate = error_predictions / total_samples if total_samples > 0 else 0.0
+
+    # Calculate timing statistics
+    successful_times = [r["prediction_time"] for r in results if not r["failed"]]
+    avg_prediction_time = sum(successful_times) / len(successful_times) if successful_times else 0.0
+    median_prediction_time = statistics.median(successful_times) if successful_times else 0.0
+    min_prediction_time = min(successful_times) if successful_times else 0.0
+    max_prediction_time = max(successful_times) if successful_times else 0.0
+
+    # Get VRAM statistics
+    vram_stats = model_wrapper.get_vram_stats()
+
+    return {
+        "model_name": model_wrapper.model_name,
+        "total_samples": total_samples,
+        "correct_predictions": correct_predictions,
+        "failed_predictions": error_predictions,
+        "accuracy": accuracy,
+        "failure_rate": error_rate,
+        "avg_prediction_time": avg_prediction_time,
+        "median_prediction_time": median_prediction_time,
+        "min_prediction_time": min_prediction_time,
+        "max_prediction_time": max_prediction_time,
+        "vram_max_mb": vram_stats["max_mb"],
+        "vram_avg_mb": vram_stats["avg_mb"],
+        "results": results,
+    }
+
+
+async def main():
+    """
+    Main function to run the benchmark.
+    """
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="ScreenSpot-Pro Benchmark Script")
+    parser.add_argument(
+        "--samples", type=int, default=300, help="Number of samples to evaluate (default: 300)"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, help="Random seed for shuffling (default: 42)"
+    )
+    args = parser.parse_args()
+
+    # Set random seed
+    random.seed(args.seed)
+
+    # Load dataset
+    print("Loading ScreenSpot-Pro dataset...")
+    ds = load_dataset("lmms-lab/ScreenSpot-Pro")
+    dataset = ds["train"]  # type: ignore
+    # Convert to list to support indexing
+    dataset_list = list(dataset)
+    print(f"Dataset loaded: {len(dataset_list)} samples")
+
+    # Shuffle dataset with seed
+    random.shuffle(dataset_list)
+    print(f"Dataset shuffled with seed {args.seed}")
+
+    # Get available models
+    models = get_available_models()
+
+    # Evaluation settings
+    max_samples = args.samples  # Use command line argument
+
+    # Run evaluations
+    all_results = []
+
+    for model in models:
+        model_wrapper = ModelWrapper(model)
+        result = await evaluate_model(model_wrapper, dataset_list, max_samples)
+        all_results.append(result)
+
+        # Print summary
+        print(f"\n{result['model_name']} Results:")
+        print(f"  Accuracy: {result['accuracy']*100:.2f}%")
+        print(f"  Correct: {result['correct_predictions']}/{result['total_samples']}")
+        print(f"  Errors: {result['failed_predictions']}")
+        print(f"  Error Rate: {result['failure_rate']*100:.2f}%")
+        print(f"  Avg Time: {result['avg_prediction_time']:.2f}s")
+        print(f"  Median Time: {result['median_prediction_time']:.2f}s")
+        print(
+            f"  Time Range: {result['min_prediction_time']:.2f}s - {result['max_prediction_time']:.2f}s"
+        )
+        print(f"  VRAM Max: {result['vram_max_mb']:.1f}MB")
+        print(f"  VRAM Avg: {result['vram_avg_mb']:.1f}MB")
+
+        # Print GPU memory info
+        gpu_memory = get_gpu_memory()
+        if gpu_memory and gpu_memory[0] > 0:
+            print(f"  GPU Free Memory: {gpu_memory[0]:.1f}MB")
+
+    # Save results
+    if all_results:
+        save_results_to_markdown(all_results)
+        save_visualizations(all_results, dataset_list)
+        print("\nBenchmark completed successfully!")
+    else:
+        print("\nNo successful evaluations completed.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

libs/python/agent/benchmarks/ss-v2.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""
+ScreenSpot-v2 Benchmark Script
+
+Evaluates models on the ScreenSpot-v2 dataset for click prediction accuracy.
+Supports both ComputerAgent model strings and custom model classes.
+"""
+
+import argparse
+import asyncio
+import random
+import statistics
+import time
+from typing import Optional
+
+from datasets import load_dataset
+from tqdm import tqdm
+from utils import (
+    ModelWrapper,
+    get_available_models,
+    get_gpu_memory,
+    is_click_in_bbox,
+    save_results_to_markdown,
+    save_visualizations,
+)
+
+
+async def evaluate_model(
+    model_wrapper: ModelWrapper, samples, max_samples: Optional[int] = None
+) -> dict:
+    """
+    Evaluate a model on any iterable of samples.
+
+    Args:
+        model_wrapper: ModelWrapper instance
+        samples: Iterable of dicts with keys: image, bbox, instruction
+        max_samples: Maximum number of samples to evaluate (None for all)
+
+    Returns:
+        Dictionary with evaluation results
+    """
+    print(f"\nEvaluating model: {model_wrapper.model_name}")
+
+    # Load model
+    await model_wrapper.load_model()
+
+    # Convert to list if needed and limit samples
+    if hasattr(samples, "__len__"):
+        total_samples = len(samples)
+        if max_samples is not None:
+            total_samples = min(max_samples, total_samples)
+        sample_list = list(samples)[:total_samples]
+    else:
+        # For iterators, take max_samples or all
+        sample_list = list(samples)
+        if max_samples is not None:
+            sample_list = sample_list[:max_samples]
+        total_samples = len(sample_list)
+
+    correct_predictions = 0
+    error_predictions = 0
+    results = []
+
+    for i, sample in enumerate(tqdm(sample_list, desc=f"Evaluating {model_wrapper.model_name}")):
+        # Extract required data (only these 3 keys matter)
+        image = sample["image"]
+        instruction = sample["instruction"]
+        bbox = sample["bbox"]  # [x1, y1, x2, y2]
+
+        # Predict click coordinates with timing
+        start_time = time.time()
+        click_coords = await model_wrapper.predict_click(image, instruction)
+        prediction_time = time.time() - start_time
+
+        # Check if prediction is correct
+        is_correct = is_click_in_bbox(click_coords, bbox)
+
+        if is_correct:
+            correct_predictions += 1
+
+        results.append(
+            {
+                "sample_idx": i,
+                "instruction": instruction,
+                "bbox": bbox,
+                "predicted_coords": click_coords,
+                "is_correct": is_correct,
+                "failed": False,
+                "prediction_time": prediction_time,
+            }
+        )
+
+    # Unload model
+    await model_wrapper.unload_model()
+
+    # Calculate metrics
+    accuracy = correct_predictions / total_samples if total_samples > 0 else 0.0
+    error_rate = error_predictions / total_samples if total_samples > 0 else 0.0
+
+    # Calculate timing statistics
+    successful_times = [r["prediction_time"] for r in results if not r["failed"]]
+    avg_prediction_time = sum(successful_times) / len(successful_times) if successful_times else 0.0
+    median_prediction_time = statistics.median(successful_times) if successful_times else 0.0
+    min_prediction_time = min(successful_times) if successful_times else 0.0
+    max_prediction_time = max(successful_times) if successful_times else 0.0
+
+    # Get VRAM statistics
+    vram_stats = model_wrapper.get_vram_stats()
+
+    return {
+        "model_name": model_wrapper.model_name,
+        "total_samples": total_samples,
+        "correct_predictions": correct_predictions,
+        "failed_predictions": error_predictions,
+        "accuracy": accuracy,
+        "failure_rate": error_rate,
+        "avg_prediction_time": avg_prediction_time,
+        "median_prediction_time": median_prediction_time,
+        "min_prediction_time": min_prediction_time,
+        "max_prediction_time": max_prediction_time,
+        "vram_max_mb": vram_stats["max_mb"],
+        "vram_avg_mb": vram_stats["avg_mb"],
+        "results": results,
+    }
+
+
+async def main():
+    """
+    Main function to run the benchmark.
+    """
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="ScreenSpot-v2 Benchmark Script")
+    parser.add_argument(
+        "--samples", type=int, default=500, help="Number of samples to evaluate (default: 500)"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, help="Random seed for shuffling (default: 42)"
+    )
+    args = parser.parse_args()
+
+    # Set random seed
+    random.seed(args.seed)
+
+    # Load dataset
+    print("Loading ScreenSpot-v2 dataset...")
+    ds = load_dataset("lmms-lab/ScreenSpot-v2")
+    dataset = ds["train"]  # type: ignore
+    # Convert to simple list of dicts with only required keys
+    samples = []
+    for item in dataset:
+        # Convert dataset item to dict if needed
+        item_dict = dict(item) if hasattr(item, "keys") else item
+
+        # Convert ScreenSpot-v2 bbox format [x, y, w, h] to [x1, y1, x2, y2]
+        bbox_xywh = item_dict["bbox"]  # type: ignore
+        x, y, w, h = bbox_xywh
+        bbox_xyxy = [x, y, x + w, y + h]
+
+        samples.append(
+            {
+                "image": item_dict["image"],  # type: ignore
+                "instruction": item_dict["instruction"],  # type: ignore
+                "bbox": bbox_xyxy,
+            }
+        )
+    print(f"Dataset loaded: {len(samples)} samples")
+
+    # Shuffle samples with seed
+    random.shuffle(samples)
+    print(f"Samples shuffled with seed {args.seed}")
+
+    # Get available models
+    models = get_available_models()
+
+    # Evaluation settings
+    max_samples = args.samples  # Use command line argument
+
+    # Run evaluations
+    all_results = []
+
+    for model in models:
+        model_wrapper = ModelWrapper(model)
+        result = await evaluate_model(model_wrapper, samples, max_samples)
+        all_results.append(result)
+
+        # Print summary
+        print(f"\n{result['model_name']} Results:")
+        print(f"  Accuracy: {result['accuracy']*100:.2f}%")
+        print(f"  Correct: {result['correct_predictions']}/{result['total_samples']}")
+        print(f"  Errors: {result['failed_predictions']}")
+        print(f"  Error Rate: {result['failure_rate']*100:.2f}%")
+        print(f"  Avg Time: {result['avg_prediction_time']:.2f}s")
+        print(f"  Median Time: {result['median_prediction_time']:.2f}s")
+        print(
+            f"  Time Range: {result['min_prediction_time']:.2f}s - {result['max_prediction_time']:.2f}s"
+        )
+        print(f"  VRAM Max: {result['vram_max_mb']:.1f}MB")
+        print(f"  VRAM Avg: {result['vram_avg_mb']:.1f}MB")
+
+        # Print GPU memory info
+        gpu_memory = get_gpu_memory()
+        if gpu_memory and gpu_memory[0] > 0:
+            print(f"  GPU Free Memory: {gpu_memory[0]:.1f}MB")
+
+    # Save results
+    if all_results:
+        save_results_to_markdown(
+            all_results, "screenspot_v2_results.md", title="ScreenSpot-v2 Benchmark Results"
+        )
+        save_visualizations(all_results, samples)
+        print("\nBenchmark completed successfully!")
+    else:
+        print("\nNo successful evaluations completed.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

libs/python/agent/benchmarks/utils.py

@@ -0,0 +1,444 @@
+#!/usr/bin/env python3
+"""
+Shared utilities for ScreenSpot-Pro benchmarking and interactive testing.
+"""
+
+import dotenv
+
+dotenv.load_dotenv()
+
+import asyncio
+import base64
+import gc
+import os
+import statistics
+import subprocess as sp
+import sys
+from datetime import datetime
+from io import BytesIO
+from typing import List, Optional, Tuple, Union
+
+import torch
+from PIL import Image, ImageDraw
+from tqdm import tqdm
+
+# Add parent directory to path for imports
+sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
+from cua_agent.agent import ComputerAgent
+from models.base import ModelProtocol
+
+
+def get_gpu_memory() -> List[int]:
+    """
+    Get GPU memory usage using nvidia-smi.
+
+    Returns:
+        List of free memory values in MB for each GPU
+    """
+    try:
+        command = "nvidia-smi --query-gpu=memory.free --format=csv"
+        memory_free_info = sp.check_output(command.split()).decode("ascii").split("\n")[:-1][1:]
+        memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
+        return memory_free_values
+    except (sp.CalledProcessError, FileNotFoundError, IndexError):
+        # Fallback to torch if nvidia-smi is not available
+        if torch.cuda.is_available():
+            device = torch.cuda.current_device()
+            total = torch.cuda.get_device_properties(device).total_memory / 1024 / 1024
+            reserved = torch.cuda.memory_reserved(device) / 1024 / 1024
+            return [int(total - reserved)]
+        return [0]
+
+
+def get_vram_usage() -> dict:
+    """
+    Get current VRAM usage statistics.
+
+    Returns:
+        Dictionary with VRAM usage info (in MB)
+    """
+    if torch.cuda.is_available():
+        device = torch.cuda.current_device()
+        allocated = torch.cuda.memory_allocated(device) / 1024 / 1024  # Convert to MB
+        reserved = torch.cuda.memory_reserved(device) / 1024 / 1024  # Convert to MB
+        total = torch.cuda.get_device_properties(device).total_memory / 1024 / 1024
+        return {
+            "allocated_mb": allocated,
+            "reserved_mb": reserved,
+            "total_mb": total,
+            "free_mb": total - reserved,
+        }
+    else:
+        return {"allocated_mb": 0.0, "reserved_mb": 0.0, "total_mb": 0.0, "free_mb": 0.0}
+
+
+def get_available_models() -> List[Union[str, ModelProtocol]]:
+    """
+    Get list of available models for testing.
+
+    Returns:
+        List of model strings and model classes
+    """
+    local_provider = "huggingface-local/"  # Options: huggingface-local/ or mlx/
+
+    # from models.gta1 import GTA1Model
+
+    models = [
+        # === ComputerAgent model strings ===
+        "openai/computer-use-preview",
+        "anthropic/claude-opus-4-20250514",
+        # f"{local_provider}HelloKKMe/GTA1-7B",
+        # f"{local_provider}HelloKKMe/GTA1-32B",
+        "openai/computer-use-preview+openai/gpt-4o-mini",
+        "anthropic/claude-opus-4-20250514+openai/gpt-4o-mini",
+        # === Reference model classes ===
+        # GTA1Model("HelloKKMe/GTA1-7B"),
+        # GTA1Model("HelloKKMe/GTA1-32B"),
+    ]
+
+    return models
+
+
+def is_click_in_bbox(click_coords: Optional[Tuple[int, int]], bbox: List[int]) -> bool:
+    """
+    Check if click coordinates are within the bounding box.
+
+    Args:
+        click_coords: (x, y) coordinates or None
+        bbox: [x1, y1, x2, y2] bounding box
+
+    Returns:
+        True if click is within bbox, False otherwise
+    """
+    if click_coords is None:
+        return False
+
+    x, y = click_coords
+    x1, y1, x2, y2 = bbox
+
+    return x1 <= x <= x2 and y1 <= y <= y2
+
+
+def image_to_base64(image: Image.Image) -> str:
+    """
+    Convert PIL Image to base64 string.
+
+    Args:
+        image: PIL Image
+
+    Returns:
+        Base64 encoded image string
+    """
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
+
+
+class ModelWrapper:
+    """
+    Wrapper to provide unified interface for both ComputerAgent and custom models.
+    """
+
+    def __init__(self, model: Union[str, ModelProtocol]):
+        self.model = model
+        self.is_computer_agent = isinstance(model, str)
+        self.agent: Optional[ComputerAgent] = None
+        self.vram_usage_history: List[float] = []  # Track VRAM usage over time
+
+        if self.is_computer_agent:
+            self.model_name = str(model)
+        else:
+            self.model_name = (
+                f"{model.__class__.__name__}('{getattr(model, 'model_name', 'unknown')}')"
+            )
+
+    async def load_model(self) -> None:
+        """Load the model."""
+        if self.is_computer_agent:
+            self.agent = ComputerAgent(model=str(self.model))
+        else:
+            await self.model.load_model()  # type: ignore
+
+        # Record initial VRAM usage after loading
+        vram_info = get_vram_usage()
+        self.vram_usage_history.append(vram_info["allocated_mb"])
+
+    async def unload_model(self) -> None:
+        """Unload the model."""
+        if not self.is_computer_agent:
+            await self.model.unload_model()  # type: ignore
+        else:
+            del self.agent
+            self.agent = None
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        # Record VRAM usage after unloading
+        vram_info = get_vram_usage()
+        self.vram_usage_history.append(vram_info["allocated_mb"])
+
+    def get_vram_stats(self) -> dict:
+        """Get VRAM usage statistics for this model."""
+        if not self.vram_usage_history:
+            return {"max_mb": 0.0, "avg_mb": 0.0}
+
+        return {
+            "max_mb": max(self.vram_usage_history),
+            "avg_mb": sum(self.vram_usage_history) / len(self.vram_usage_history),
+        }
+
+    async def predict_click(
+        self, image: Image.Image, instruction: str
+    ) -> Optional[Tuple[int, int]]:
+        """Predict click coordinates."""
+        # Record VRAM usage before prediction
+        vram_info = get_vram_usage()
+        self.vram_usage_history.append(vram_info["allocated_mb"])
+
+        if self.is_computer_agent:
+            if self.agent is None:
+                await self.load_model()
+
+            if self.agent is not None:
+                image_b64 = image_to_base64(image)
+                result = await self.agent.predict_click(
+                    instruction=instruction, image_b64=image_b64
+                )
+
+                # Record VRAM usage after prediction
+                vram_info = get_vram_usage()
+                self.vram_usage_history.append(vram_info["allocated_mb"])
+
+                return result
+            return None
+        else:
+            result = await self.model.predict_click(image, instruction)  # type: ignore
+
+            # Record VRAM usage after prediction
+            vram_info = get_vram_usage()
+            self.vram_usage_history.append(vram_info["allocated_mb"])
+
+            return result
+
+
+def save_results_to_markdown(
+    all_results: List[dict],
+    output_file: str = "screenspot_pro_results.md",
+    title: str = "ScreenSpot-Pro Benchmark Results",
+) -> None:
+    """
+    Save evaluation results to a markdown table.
+
+    Args:
+        all_results: List of evaluation results for each model
+        output_file: Output markdown file path
+    """
+    with open(output_file, "w", encoding="utf-8") as f:
+        f.write(f"# {title}\n\n")
+        f.write(f"**Evaluation Date:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
+
+        # Summary table
+        f.write("## Summary\n\n")
+        f.write(
+            "| Model | Total Samples | Correct | Errors | Accuracy | Error Rate | Avg Time (s) | Median Time (s) | Time Range (s) | VRAM Max (GB) | VRAM Avg (GB) |\n"
+        )
+        f.write(
+            "|-------|---------------|---------|--------|----------|------------|--------------|-----------------|----------------|---------------|---------------|\n"
+        )
+
+        for result in all_results:
+            model_name = result["model_name"]
+            total = result["total_samples"]
+            correct = result["correct_predictions"]
+            errors = result["failed_predictions"]
+            accuracy = result["accuracy"] * 100
+            error_rate = result["failure_rate"] * 100
+            avg_time = result.get("avg_prediction_time", 0.0)
+            median_time = result.get("median_prediction_time", 0.0)
+            min_time = result.get("min_prediction_time", 0.0)
+            max_time = result.get("max_prediction_time", 0.0)
+            time_range = f"{min_time:.2f} - {max_time:.2f}"
+            vram_max = result.get("vram_max_mb", 0.0) / 1024
+            vram_avg = result.get("vram_avg_mb", 0.0) / 1024
+
+            f.write(
+                f"| {model_name} | {total} | {correct} | {errors} | {accuracy:.2f}% | {error_rate:.2f}% | {avg_time:.2f} | {median_time:.2f} | {time_range} | {vram_max:.1f} | {vram_avg:.1f} |\n"
+            )
+
+        # Detailed results for each model
+        for result in all_results:
+            f.write(f"\n## {result['model_name']} - Detailed Results\n\n")
+            f.write(
+                "| Sample Index | Instruction | BBox | Predicted | Correct | Error | Time (s) |\n"
+            )
+            f.write("|-----------|-------------|------|-----------|---------|-------|----------|\n")
+
+            for sample_result in result["results"][:10]:  # Show first 10 samples
+                sample_idx = sample_result["sample_idx"]
+                instruction = (
+                    sample_result["instruction"][:50] + "..."
+                    if len(sample_result["instruction"]) > 50
+                    else sample_result["instruction"]
+                )
+                bbox = str(sample_result["bbox"])
+                predicted = (
+                    str(sample_result["predicted_coords"])
+                    if sample_result["predicted_coords"]
+                    else "None"
+                )
+                correct = "PASS" if sample_result["is_correct"] else "FAIL"
+                error = "YES" if sample_result["failed"] else "NO"
+                pred_time = sample_result.get("prediction_time", 0.0)
+
+                f.write(
+                    f"| {sample_idx} | {instruction} | {bbox} | {predicted} | {correct} | {error} | {pred_time:.2f} |\n"
+                )
+
+            if len(result["results"]) > 10:
+                f.write(f"\n*Showing first 10 of {len(result['results'])} samples*\n")
+
+    print(f"\nResults saved to: {output_file}")
+
+
+def save_visualizations(all_results: List[dict], samples, output_dir: str = "output") -> None:
+    """
+    Save visualizations of predicted coordinates vs bboxes to an output folder.
+
+    Args:
+        all_results: List of evaluation results for each model
+        samples: List of sample dicts with image, bbox, instruction keys
+        output_dir: Output directory path
+    """
+    os.makedirs(output_dir, exist_ok=True)
+
+    for result in all_results:
+        model_name = result["model_name"].replace("/", "_").replace("\\", "_")
+        model_dir = os.path.join(output_dir, model_name)
+        os.makedirs(model_dir, exist_ok=True)
+
+        print(f"Saving visualizations for {result['model_name']}...")
+
+        # Save first 10 samples for visualization
+        for i, sample_result in enumerate(
+            tqdm(result["results"][:10], desc=f"Saving {model_name} visualizations")
+        ):
+            # Get sample data using index
+            sample_idx = sample_result["sample_idx"]
+
+            if sample_idx < len(samples):
+                sample = samples[sample_idx]
+                image = sample["image"].copy()  # Make a copy to avoid modifying original
+            else:
+                print(f"Warning: Could not find sample at index {sample_idx}")
+                continue
+
+            bbox = sample_result["bbox"]
+            predicted_coords = sample_result["predicted_coords"]
+            is_correct = sample_result["is_correct"]
+
+            # Draw on image
+            draw = ImageDraw.Draw(image)
+
+            # Draw bounding box (ground truth) in green
+            x1, y1, x2, y2 = bbox
+            draw.rectangle([x1, y1, x2, y2], outline="green", width=3)
+            draw.text((x1, y1 - 20), "Ground Truth", fill="green")
+
+            # Draw predicted click in red or blue
+            if predicted_coords is not None:
+                px, py = predicted_coords
+                color = "blue" if is_correct else "red"
+                # Draw crosshair
+                crosshair_size = 15
+                draw.line(
+                    [(px - crosshair_size, py), (px + crosshair_size, py)], fill=color, width=3
+                )
+                draw.line(
+                    [(px, py - crosshair_size), (px, py + crosshair_size)], fill=color, width=3
+                )
+                draw.text((px + 10, py - 20), f"Predicted ({px},{py})", fill=color)
+
+            # Add status text
+            status = "CORRECT" if is_correct else "INCORRECT"
+            status_color = "blue" if is_correct else "red"
+            draw.text((10, 10), f"Status: {status}", fill=status_color)
+            draw.text(
+                (10, 30), f"Instruction: {sample_result['instruction'][:50]}...", fill="black"
+            )
+
+            # Save image
+            filename = f"sample_{i+1:02d}_idx{sample_idx}_{status.lower()}.png"
+            filepath = os.path.join(model_dir, filename)
+            image.save(filepath)
+
+        print(f"Visualizations saved to: {model_dir}")
+
+
+def save_prediction_visualization(
+    image: Image.Image,
+    instruction: str,
+    predictions: List[dict],
+    output_file: str = "interactive_prediction.png",
+) -> None:
+    """
+    Save visualization of multiple model predictions on a single image.
+
+    Args:
+        image: PIL Image to visualize
+        instruction: Instruction text
+        predictions: List of prediction dicts with keys: model_name, coords, error
+        output_file: Output file path
+    """
+    # Create a copy of the image
+    vis_image = image.copy()
+    draw = ImageDraw.Draw(vis_image)
+
+    # Colors for different models
+    colors = ["red", "blue", "orange", "purple", "brown", "pink", "gray", "olive"]
+
+    # Draw predictions
+    for i, pred in enumerate(predictions):
+        color = colors[i % len(colors)]
+        model_name = pred["model_name"]
+        coords = pred.get("coords")
+        error = pred.get("error")
+
+        if coords is not None:
+            px, py = coords
+            # Draw crosshair
+            crosshair_size = 20
+            draw.line([(px - crosshair_size, py), (px + crosshair_size, py)], fill=color, width=4)
+            draw.line([(px, py - crosshair_size), (px, py + crosshair_size)], fill=color, width=4)
+            # Draw model name
+            draw.text((px + 15, py + 15), f"{model_name}: ({px},{py})", fill=color)
+        else:
+            # Draw error text
+            draw.text((10, 50 + i * 20), f"{model_name}: ERROR - {error}", fill=color)
+
+    # Add instruction at the top
+    draw.text((10, 10), f"Instruction: {instruction}", fill="black")
+
+    # Save image
+    vis_image.save(output_file)
+    print(f"Prediction visualization saved to: {output_file}")
+
+
+def take_screenshot() -> Image.Image:
+    """
+    Take a screenshot of the current screen.
+
+    Returns:
+        PIL Image of the screenshot
+    """
+    try:
+        from PIL import ImageGrab
+
+        screenshot = ImageGrab.grab()
+        return screenshot
+    except ImportError:
+        print("PIL/Pillow not installed. Please install it with: pip install pillow")
+        raise
+    except Exception as e:
+        print(f"Error taking screenshot: {e}")
+        raise

libs/python/agent/cua_agent/__init__.py

@@ -0,0 +1,49 @@
+"""
+agent - Decorator-based Computer Use Agent with liteLLM integration
+"""
+
+import logging
+import sys
+
+# Import loops to register them
+from . import loops
+from .agent import ComputerAgent
+from .decorators import register_agent
+from .types import AgentResponse, Messages
+
+__all__ = ["register_agent", "ComputerAgent", "Messages", "AgentResponse"]
+
+__version__ = "0.4.0"
+
+logger = logging.getLogger(__name__)
+
+# Initialize telemetry when the package is imported
+try:
+    # Import from core telemetry for basic functions
+    from cua_core.telemetry import (
+        is_telemetry_enabled,
+        record_event,
+    )
+
+    # Check if telemetry is enabled
+    if is_telemetry_enabled():
+        logger.info("Telemetry is enabled")
+
+        # Record package initialization
+        record_event(
+            "module_init",
+            {
+                "module": "agent",
+                "version": __version__,
+                "python_version": sys.version,
+            },
+        )
+
+    else:
+        logger.info("Telemetry is disabled")
+except ImportError as e:
+    # Telemetry not available
+    logger.warning(f"Telemetry not available: {e}")
+except Exception as e:
+    # Other issues with telemetry
+    logger.warning(f"Error initializing telemetry: {e}")

libs/python/agent/cua_agent/__main__.py

@@ -0,0 +1,22 @@
+"""
+Entry point for running agent CLI module.
+
+Usage:
+    python -m agent.cli <model_string>
+"""
+
+import asyncio
+import sys
+
+from .cli import main
+
+if __name__ == "__main__":
+    # Check if 'cli' is specified as the module
+    if len(sys.argv) > 1 and sys.argv[1] == "cli":
+        # Remove 'cli' from arguments and run CLI
+        sys.argv.pop(1)
+        asyncio.run(main())
+    else:
+        print("Usage: python -m agent.cli <model_string>")
+        print("Example: python -m agent.cli openai/computer-use-preview")
+        sys.exit(1)

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

@@ -0,0 +1,19 @@
+"""
+Adapters package for agent - Custom LLM adapters for LiteLLM
+"""
+
+from .azure_ml_adapter import AzureMLAdapter
+from .cua_adapter import CUAAdapter
+from .huggingfacelocal_adapter import HuggingFaceLocalAdapter
+from .human_adapter import HumanAdapter
+from .mlxvlm_adapter import MLXVLMAdapter
+from .yutori_adapter import YutoriAdapter
+
+__all__ = [
+    "AzureMLAdapter",
+    "HuggingFaceLocalAdapter",
+    "HumanAdapter",
+    "MLXVLMAdapter",
+    "CUAAdapter",
+    "YutoriAdapter",
+]

libs/python/agent/cua_agent/adapters/azure_ml_adapter.py

@@ -0,0 +1,283 @@
+"""
+Azure ML Custom Provider Adapter for LiteLLM.
+
+This adapter provides direct OpenAI-compatible API access to Azure ML endpoints
+without message transformation, specifically for models like Fara-7B that require
+exact OpenAI message formatting.
+"""
+
+import json
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
+
+import httpx
+from litellm import acompletion, completion
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+
+
+class AzureMLAdapter(CustomLLM):
+    """
+    Azure ML Adapter for OpenAI-compatible endpoints.
+
+    Makes direct HTTP calls to Azure ML foundry inference endpoints
+    using the OpenAI-compatible API format without transforming messages.
+
+    Usage:
+        model = "azure_ml/Fara-7B"
+        api_base = "https://foundry-inference-xxx.centralus.inference.ml.azure.com"
+        api_key = "your-api-key"
+
+        response = litellm.completion(
+            model=model,
+            messages=[...],
+            api_base=api_base,
+            api_key=api_key
+        )
+    """
+
+    def __init__(self, **kwargs):
+        """Initialize the adapter."""
+        super().__init__()
+        self._client: Optional[httpx.Client] = None
+        self._async_client: Optional[httpx.AsyncClient] = None
+
+    def _get_client(self) -> httpx.Client:
+        """Get or create sync HTTP client."""
+        if self._client is None:
+            self._client = httpx.Client(timeout=600.0)
+        return self._client
+
+    def _get_async_client(self) -> httpx.AsyncClient:
+        """Get or create async HTTP client."""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(timeout=600.0)
+        return self._async_client
+
+    def _prepare_request(self, **kwargs) -> tuple[str, dict, dict]:
+        """
+        Prepare the HTTP request without transforming messages.
+
+        Applies Azure ML workaround: double-encodes function arguments to work around
+        Azure ML's bug where it parses arguments before validation.
+
+        Returns:
+            Tuple of (url, headers, json_data)
+        """
+        # Extract required params
+        api_base = kwargs.get("api_base")
+        api_key = kwargs.get("api_key")
+        model = kwargs.get("model", "").replace("azure_ml/", "")
+        messages = kwargs.get("messages", [])
+
+        if not api_base:
+            raise ValueError("api_base is required for azure_ml provider")
+        if not api_key:
+            raise ValueError("api_key is required for azure_ml provider")
+
+        # Build OpenAI-compatible endpoint URL
+        base_url = api_base.rstrip("/")
+        url = f"{base_url}/chat/completions"
+
+        # Prepare headers
+        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+
+        # WORKAROUND for Azure ML bug:
+        # Azure ML incorrectly parses the arguments field before validation,
+        # causing it to reject valid JSON strings. We double-encode arguments
+        # so that after Azure ML's parse, they remain as strings.
+        messages_copy = []
+        for message in messages:
+            msg_copy = message.copy()
+
+            # Check if message has tool_calls that need double-encoding
+            if "tool_calls" in msg_copy:
+                tool_calls_copy = []
+                for tool_call in msg_copy["tool_calls"]:
+                    tc_copy = tool_call.copy()
+
+                    if "function" in tc_copy and "arguments" in tc_copy["function"]:
+                        func_copy = tc_copy["function"].copy()
+                        arguments = func_copy["arguments"]
+
+                        # If arguments is already a string, double-encode it
+                        if isinstance(arguments, str):
+                            func_copy["arguments"] = json.dumps(arguments)
+
+                        tc_copy["function"] = func_copy
+
+                    tool_calls_copy.append(tc_copy)
+
+                msg_copy["tool_calls"] = tool_calls_copy
+
+            messages_copy.append(msg_copy)
+
+        # Prepare request body with double-encoded messages
+        json_data = {"model": model, "messages": messages_copy}
+
+        # Add optional parameters if provided
+        optional_params = [
+            "temperature",
+            "top_p",
+            "n",
+            "stream",
+            "stop",
+            "max_tokens",
+            "presence_penalty",
+            "frequency_penalty",
+            "logit_bias",
+            "user",
+            "response_format",
+            "seed",
+            "tools",
+            "tool_choice",
+        ]
+
+        for param in optional_params:
+            if param in kwargs and kwargs[param] is not None:
+                json_data[param] = kwargs[param]
+
+        return url, headers, json_data
+
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        """
+        Synchronous completion method.
+
+        Makes a direct HTTP POST to Azure ML's OpenAI-compatible endpoint.
+        """
+        url, headers, json_data = self._prepare_request(**kwargs)
+
+        client = self._get_client()
+        response = client.post(url, headers=headers, json=json_data)
+        response.raise_for_status()
+
+        # Parse response
+        response_json = response.json()
+
+        # Return using litellm's completion with the actual response
+        return completion(
+            model=f"azure_ml/{kwargs.get('model', '')}",
+            mock_response=response_json["choices"][0]["message"]["content"],
+            messages=kwargs.get("messages", []),
+        )
+
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        """
+        Asynchronous completion method.
+
+        Makes a direct async HTTP POST to Azure ML's OpenAI-compatible endpoint.
+        """
+        url, headers, json_data = self._prepare_request(**kwargs)
+
+        client = self._get_async_client()
+        response = await client.post(url, headers=headers, json=json_data)
+        response.raise_for_status()
+
+        # Parse response
+        response_json = response.json()
+
+        # Return using litellm's acompletion with the actual response
+        return await acompletion(
+            model=f"azure_ml/{kwargs.get('model', '')}",
+            mock_response=response_json["choices"][0]["message"]["content"],
+            messages=kwargs.get("messages", []),
+        )
+
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        """
+        Synchronous streaming method.
+
+        Makes a streaming HTTP POST to Azure ML's OpenAI-compatible endpoint.
+        """
+        url, headers, json_data = self._prepare_request(**kwargs)
+        json_data["stream"] = True
+
+        client = self._get_client()
+
+        with client.stream("POST", url, headers=headers, json=json_data) as response:
+            response.raise_for_status()
+
+            for line in response.iter_lines():
+                if line.startswith("data: "):
+                    data = line[6:]  # Remove "data: " prefix
+                    if data == "[DONE]":
+                        break
+
+                    try:
+                        chunk_json = json.loads(data)
+                        delta = chunk_json["choices"][0].get("delta", {})
+                        content = delta.get("content", "")
+                        finish_reason = chunk_json["choices"][0].get("finish_reason")
+
+                        generic_streaming_chunk: GenericStreamingChunk = {
+                            "finish_reason": finish_reason,
+                            "index": 0,
+                            "is_finished": finish_reason is not None,
+                            "text": content,
+                            "tool_use": None,
+                            "usage": chunk_json.get(
+                                "usage",
+                                {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+                            ),
+                        }
+
+                        yield generic_streaming_chunk
+                    except json.JSONDecodeError:
+                        continue
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        """
+        Asynchronous streaming method.
+
+        Makes an async streaming HTTP POST to Azure ML's OpenAI-compatible endpoint.
+        """
+        url, headers, json_data = self._prepare_request(**kwargs)
+        json_data["stream"] = True
+
+        client = self._get_async_client()
+
+        async with client.stream("POST", url, headers=headers, json=json_data) as response:
+            response.raise_for_status()
+
+            async for line in response.aiter_lines():
+                if line.startswith("data: "):
+                    data = line[6:]  # Remove "data: " prefix
+                    if data == "[DONE]":
+                        break
+
+                    try:
+                        chunk_json = json.loads(data)
+                        delta = chunk_json["choices"][0].get("delta", {})
+                        content = delta.get("content", "")
+                        finish_reason = chunk_json["choices"][0].get("finish_reason")
+
+                        generic_streaming_chunk: GenericStreamingChunk = {
+                            "finish_reason": finish_reason,
+                            "index": 0,
+                            "is_finished": finish_reason is not None,
+                            "text": content,
+                            "tool_use": None,
+                            "usage": chunk_json.get(
+                                "usage",
+                                {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+                            ),
+                        }
+
+                        yield generic_streaming_chunk
+                    except json.JSONDecodeError:
+                        continue
+
+    def __del__(self):
+        """Cleanup HTTP clients."""
+        if self._client is not None:
+            self._client.close()
+        if self._async_client is not None:
+            import asyncio
+
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_running():
+                    loop.create_task(self._async_client.aclose())
+                else:
+                    loop.run_until_complete(self._async_client.aclose())
+            except Exception:
+                pass

libs/python/agent/cua_agent/adapters/cua_adapter.py

@@ -0,0 +1,250 @@
+import os
+from typing import Any, AsyncIterator, Iterator
+
+from cua_core.http import cua_version_headers
+from litellm import acompletion, completion
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+
+
+class CUAAdapter(CustomLLM):
+    def __init__(self, base_url: str | None = None, api_key: str | None = None, **_: Any):
+        super().__init__()
+        self.base_url = base_url or os.environ.get("CUA_BASE_URL") or "https://inference.cua.ai/v1"
+        self.api_key = (
+            api_key or os.environ.get("CUA_INFERENCE_API_KEY") or os.environ.get("CUA_API_KEY")
+        )
+
+    def _normalize_model(self, model: str) -> str:
+        """Strip known prefixes to get the base model name."""
+        known_prefixes = ("cua/", "anthropic/", "gemini/", "google/", "openai/")
+        result = model
+        for prefix in known_prefixes:
+            if result.startswith(prefix):
+                result = result[len(prefix) :]
+        return result
+
+    def _resolve_route(self, model: str, api_base: str) -> tuple[str, str]:
+        """Return (prefixed_model, api_base) for the CUA inference API."""
+        if "anthropic/" in model:
+            return f"anthropic/{self._normalize_model(model)}", api_base.removesuffix("/v1")
+        elif "gemini/" in model or "google/" in model:
+            return f"gemini/{self._normalize_model(model)}", api_base + "/gemini"
+        else:
+            return f"openai/{self._normalize_model(model)}", api_base
+
+    def _resolve_api_key(self, kwargs: dict | None = None) -> str:
+        """Resolve the CUA API key, raising a clear error if missing.
+
+        Checks kwargs (from ComputerAgent api_key param) then falls back
+        to self.api_key (from CUA_API_KEY / CUA_INFERENCE_API_KEY env vars).
+
+        This validation must run before the inner litellm call because that
+        call uses an anthropic/ or openai/ model prefix, which would cause
+        litellm to fall back to ANTHROPIC_API_KEY from env — sending the
+        wrong key to the CUA inference endpoint.
+        """
+        resolved = (kwargs.get("api_key") if kwargs else None) or self.api_key
+        if not resolved:
+            raise ValueError(
+                "No CUA API key provided for cua/ model inference. "
+                "Please either set the CUA_API_KEY environment variable "
+                "or pass api_key to ComputerAgent()."
+            )
+        return resolved
+
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        model, api_base = self._resolve_route(
+            kwargs.get("model", ""), kwargs.get("api_base") or self.base_url
+        )
+
+        api_key = self._resolve_api_key(kwargs)
+
+        # Ensure the CUA inference API always receives Bearer auth;
+        # merge caller headers first, then force Authorization so it cannot be overridden.
+        extra_headers = {}
+        if "extra_headers" in kwargs:
+            extra_headers.update(kwargs.pop("extra_headers"))
+        extra_headers["Authorization"] = f"Bearer {api_key}"
+
+        params = {
+            "model": model,
+            "messages": kwargs.get("messages", []),
+            "api_base": api_base,
+            "api_key": api_key,
+            "extra_headers": extra_headers,
+            "stream": False,
+        }
+
+        # Forward tools if provided
+        if "tools" in kwargs:
+            params["tools"] = kwargs["tools"]
+
+        if "optional_params" in kwargs:
+            protected_keys = {"api_key", "extra_headers", "model", "api_base", "stream"}
+            filtered = {
+                k: v for k, v in kwargs["optional_params"].items() if k not in protected_keys
+            }
+            params.update(filtered)
+            del kwargs["optional_params"]
+
+        if "headers" in kwargs:
+            params["headers"] = kwargs["headers"]
+            del kwargs["headers"]
+
+        # Always include CUA version headers
+        version_hdrs = cua_version_headers()
+        if version_hdrs:
+            params["headers"] = {**version_hdrs, **params.get("headers", {})}
+
+        # Print dropped parameters
+        original_keys = set(kwargs.keys())
+        used_keys = set(params.keys())  # Only these are extracted from kwargs
+        ignored_keys = {
+            "litellm_params",
+            "client",
+            "print_verbose",
+            "acompletion",
+            "timeout",
+            "logging_obj",
+            "encoding",
+            "custom_prompt_dict",
+            "model_response",
+            "logger_fn",
+        }
+        dropped_keys = original_keys - used_keys - ignored_keys
+        if dropped_keys:
+            dropped_keyvals = {k: kwargs[k] for k in dropped_keys}
+            # print(f"CUAAdapter.completion: Dropped parameters: {dropped_keyvals}")
+
+        return completion(**params)  # type: ignore
+
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        model, api_base = self._resolve_route(
+            kwargs.get("model", ""), kwargs.get("api_base") or self.base_url
+        )
+
+        api_key = self._resolve_api_key(kwargs)
+
+        # Ensure the CUA inference API always receives Bearer auth;
+        # merge caller headers first, then force Authorization so it cannot be overridden.
+        extra_headers = {}
+        if "extra_headers" in kwargs:
+            extra_headers.update(kwargs.pop("extra_headers"))
+        extra_headers["Authorization"] = f"Bearer {api_key}"
+
+        params = {
+            "model": model,
+            "messages": kwargs.get("messages", []),
+            "api_base": api_base,
+            "api_key": api_key,
+            "extra_headers": extra_headers,
+            "stream": False,
+        }
+
+        # Forward tools if provided
+        if "tools" in kwargs:
+            params["tools"] = kwargs["tools"]
+
+        if "optional_params" in kwargs:
+            protected_keys = {"api_key", "extra_headers", "model", "api_base", "stream"}
+            filtered = {
+                k: v for k, v in kwargs["optional_params"].items() if k not in protected_keys
+            }
+            params.update(filtered)
+            del kwargs["optional_params"]
+
+        if "headers" in kwargs:
+            params["headers"] = kwargs["headers"]
+            del kwargs["headers"]
+
+        # Always include CUA version headers
+        version_hdrs = cua_version_headers()
+        if version_hdrs:
+            params["headers"] = {**version_hdrs, **params.get("headers", {})}
+
+        # Print dropped parameters
+        original_keys = set(kwargs.keys())
+        used_keys = set(params.keys())  # Only these are extracted from kwargs
+        ignored_keys = {
+            "litellm_params",
+            "client",
+            "print_verbose",
+            "acompletion",
+            "timeout",
+            "logging_obj",
+            "encoding",
+            "custom_prompt_dict",
+            "model_response",
+            "logger_fn",
+        }
+        dropped_keys = original_keys - used_keys - ignored_keys
+        if dropped_keys:
+            dropped_keyvals = {k: kwargs[k] for k in dropped_keys}
+            # print(f"CUAAdapter.acompletion: Dropped parameters: {dropped_keyvals}")
+
+        response = await acompletion(**params)  # type: ignore
+
+        return response
+
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        params = dict(kwargs)
+        model, api_base = self._resolve_route(
+            params.get("model", ""), params.get("api_base") or self.base_url
+        )
+        api_key = self._resolve_api_key(kwargs)
+
+        # Ensure the CUA inference API always receives Bearer auth;
+        # merge caller headers first, then force Authorization so it cannot be overridden.
+        extra_headers = {}
+        if "extra_headers" in params:
+            extra_headers.update(params.pop("extra_headers"))
+        extra_headers["Authorization"] = f"Bearer {api_key}"
+
+        params.update(
+            {
+                "model": model,
+                "api_base": api_base,
+                "api_key": api_key,
+                "extra_headers": extra_headers,
+                "stream": True,
+            }
+        )
+        # Always include CUA version headers
+        version_hdrs = cua_version_headers()
+        if version_hdrs:
+            params["headers"] = {**version_hdrs, **params.get("headers", {})}
+        # Yield chunks directly from LiteLLM's streaming generator
+        for chunk in completion(**params):  # type: ignore
+            yield chunk  # type: ignore
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        params = dict(kwargs)
+        model, api_base = self._resolve_route(
+            params.get("model", ""), params.get("api_base") or self.base_url
+        )
+        api_key = self._resolve_api_key(kwargs)
+
+        # Ensure the CUA inference API always receives Bearer auth;
+        # merge caller headers first, then force Authorization so it cannot be overridden.
+        extra_headers = {}
+        if "extra_headers" in params:
+            extra_headers.update(params.pop("extra_headers"))
+        extra_headers["Authorization"] = f"Bearer {api_key}"
+
+        params.update(
+            {
+                "model": model,
+                "api_base": api_base,
+                "api_key": api_key,
+                "extra_headers": extra_headers,
+                "stream": True,
+            }
+        )
+        # Always include CUA version headers
+        version_hdrs = cua_version_headers()
+        if version_hdrs:
+            params["headers"] = {**version_hdrs, **params.get("headers", {})}
+        stream = await acompletion(**params)  # type: ignore
+        async for chunk in stream:  # type: ignore
+            yield chunk  # type: ignore

libs/python/agent/cua_agent/adapters/huggingfacelocal_adapter.py

@@ -0,0 +1,186 @@
+import asyncio
+import functools
+import warnings
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
+
+from litellm import acompletion, completion
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+
+# Try to import HuggingFace dependencies
+try:
+    import torch
+    from transformers import AutoModelForImageTextToText, AutoProcessor
+
+    HF_AVAILABLE = True
+except ImportError:
+    HF_AVAILABLE = False
+
+from .models import load_model as load_model_handler
+
+
+class HuggingFaceLocalAdapter(CustomLLM):
+    """HuggingFace Local Adapter for running vision-language models locally."""
+
+    def __init__(self, device: str = "auto", trust_remote_code: bool = False, **kwargs):
+        """Initialize the adapter.
+
+        Args:
+            device: Device to load model on ("auto", "cuda", "cpu", etc.)
+            trust_remote_code: Whether to trust remote code
+            **kwargs: Additional arguments
+        """
+        super().__init__()
+        self.device = device
+        self.trust_remote_code = trust_remote_code
+        # Cache for model handlers keyed by model_name
+        self._handlers: Dict[str, Any] = {}
+        self._executor = ThreadPoolExecutor(max_workers=1)  # Single thread pool
+
+    def _get_handler(self, model_name: str):
+        """Get or create a model handler for the given model name."""
+        if model_name not in self._handlers:
+            self._handlers[model_name] = load_model_handler(
+                model_name=model_name, device=self.device, trust_remote_code=self.trust_remote_code
+            )
+        return self._handlers[model_name]
+
+    def _convert_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Convert OpenAI format messages to HuggingFace format.
+
+        Args:
+            messages: Messages in OpenAI format
+
+        Returns:
+            Messages in HuggingFace format
+        """
+        converted_messages = []
+
+        for message in messages:
+            converted_message = {"role": message["role"], "content": []}
+
+            content = message.get("content", [])
+            if isinstance(content, str):
+                # Simple text content
+                converted_message["content"].append({"type": "text", "text": content})
+            elif isinstance(content, list):
+                # Multi-modal content
+                for item in content:
+                    if item.get("type") == "text":
+                        converted_message["content"].append(
+                            {"type": "text", "text": item.get("text", "")}
+                        )
+                    elif item.get("type") == "image_url":
+                        # Convert image_url format to image format
+                        image_url = item.get("image_url", {}).get("url", "")
+                        converted_message["content"].append({"type": "image", "image": image_url})
+
+            converted_messages.append(converted_message)
+
+        return converted_messages
+
+    def _generate(self, **kwargs) -> str:
+        """Generate response using the local HuggingFace model.
+
+        Args:
+            **kwargs: Keyword arguments containing messages and model info
+
+        Returns:
+            Generated text response
+        """
+        if not HF_AVAILABLE:
+            raise ImportError(
+                "HuggingFace transformers dependencies not found. "
+                'Please install with: pip install "cua-agent[uitars-hf]"'
+            )
+
+        # Extract messages and model from kwargs
+        messages = kwargs.get("messages", [])
+        model_name = kwargs.get("model", "ByteDance-Seed/UI-TARS-1.5-7B")
+        max_new_tokens = kwargs.get("max_tokens", 128)
+
+        # Warn about ignored kwargs
+        ignored_kwargs = set(kwargs.keys()) - {"messages", "model", "max_tokens"}
+        if ignored_kwargs:
+            warnings.warn(f"Ignoring unsupported kwargs: {ignored_kwargs}")
+
+        # Convert messages to HuggingFace format
+        hf_messages = self._convert_messages(messages)
+
+        # Delegate to model handler
+        handler = self._get_handler(model_name)
+        generated_text = handler.generate(hf_messages, max_new_tokens=max_new_tokens)
+        return generated_text
+
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        """Synchronous completion method.
+
+        Returns:
+            ModelResponse with generated text
+        """
+        generated_text = self._generate(**kwargs)
+
+        return completion(
+            model=f"huggingface-local/{kwargs['model']}",
+            mock_response=generated_text,
+        )
+
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        """Asynchronous completion method.
+
+        Returns:
+            ModelResponse with generated text
+        """
+        # Run _generate in thread pool to avoid blocking
+        loop = asyncio.get_event_loop()
+        generated_text = await loop.run_in_executor(
+            self._executor, functools.partial(self._generate, **kwargs)
+        )
+
+        return await acompletion(
+            model=f"huggingface-local/{kwargs['model']}",
+            mock_response=generated_text,
+        )
+
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        """Synchronous streaming method.
+
+        Returns:
+            Iterator of GenericStreamingChunk
+        """
+        generated_text = self._generate(**kwargs)
+
+        generic_streaming_chunk: GenericStreamingChunk = {
+            "finish_reason": "stop",
+            "index": 0,
+            "is_finished": True,
+            "text": generated_text,
+            "tool_use": None,
+            "usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+        }
+
+        yield generic_streaming_chunk
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        """Asynchronous streaming method.
+
+        Returns:
+            AsyncIterator of GenericStreamingChunk
+        """
+        # Run _generate in thread pool to avoid blocking
+        loop = asyncio.get_event_loop()
+        generated_text = await loop.run_in_executor(
+            self._executor, functools.partial(self._generate, **kwargs)
+        )
+
+        generic_streaming_chunk: GenericStreamingChunk = {
+            "finish_reason": "stop",
+            "index": 0,
+            "is_finished": True,
+            "text": generated_text,
+            "tool_use": None,
+            "usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+        }
+
+        yield generic_streaming_chunk

libs/python/agent/cua_agent/adapters/human_adapter.py

@@ -0,0 +1,350 @@
+import asyncio
+import os
+from typing import Any, AsyncIterator, Dict, Iterator, List
+
+import requests
+from litellm import acompletion, completion
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+
+
+class HumanAdapter(CustomLLM):
+    """Human Adapter for human-in-the-loop completions.
+
+    This adapter sends completion requests to a human completion server
+    where humans can review and respond to AI requests.
+    """
+
+    def __init__(self, base_url: str | None = None, timeout: float = 300.0, **kwargs):
+        """Initialize the human adapter.
+
+        Args:
+            base_url: Base URL for the human completion server.
+                     Defaults to HUMAN_BASE_URL environment variable or http://localhost:8002
+            timeout: Timeout in seconds for waiting for human response
+            **kwargs: Additional arguments
+        """
+        super().__init__()
+        self.base_url = base_url or os.getenv("HUMAN_BASE_URL", "http://localhost:8002")
+        self.timeout = timeout
+
+        # Ensure base_url doesn't end with slash
+        self.base_url = self.base_url.rstrip("/")
+
+    def _queue_completion(self, messages: List[Dict[str, Any]], model: str) -> str:
+        """Queue a completion request and return the call ID.
+
+        Args:
+            messages: Messages in OpenAI format
+            model: Model name
+
+        Returns:
+            Call ID for tracking the request
+
+        Raises:
+            Exception: If queueing fails
+        """
+        try:
+            response = requests.post(
+                f"{self.base_url}/queue", json={"messages": messages, "model": model}, timeout=10
+            )
+            response.raise_for_status()
+            return response.json()["id"]
+        except requests.RequestException as e:
+            raise Exception(f"Failed to queue completion request: {e}")
+
+    def _wait_for_completion(self, call_id: str) -> Dict[str, Any]:
+        """Wait for human to complete the call.
+
+        Args:
+            call_id: ID of the queued completion call
+
+        Returns:
+            Dict containing response and/or tool_calls
+
+        Raises:
+            TimeoutError: If timeout is exceeded
+            Exception: If completion fails
+        """
+        import time
+
+        start_time = time.time()
+
+        while True:
+            try:
+                # Check status
+                status_response = requests.get(f"{self.base_url}/status/{call_id}")
+                status_response.raise_for_status()
+                status_data = status_response.json()
+
+                if status_data["status"] == "completed":
+                    result = {}
+                    if "response" in status_data and status_data["response"]:
+                        result["response"] = status_data["response"]
+                    if "tool_calls" in status_data and status_data["tool_calls"]:
+                        result["tool_calls"] = status_data["tool_calls"]
+                    return result
+                elif status_data["status"] == "failed":
+                    error_msg = status_data.get("error", "Unknown error")
+                    raise Exception(f"Completion failed: {error_msg}")
+
+                # Check timeout
+                if time.time() - start_time > self.timeout:
+                    raise TimeoutError(
+                        f"Timeout waiting for human response after {self.timeout} seconds"
+                    )
+
+                # Wait before checking again
+                time.sleep(1.0)
+
+            except requests.RequestException as e:
+                if time.time() - start_time > self.timeout:
+                    raise TimeoutError(f"Timeout waiting for human response: {e}")
+                # Continue trying if we haven't timed out
+                time.sleep(1.0)
+
+    async def _async_wait_for_completion(self, call_id: str) -> Dict[str, Any]:
+        """Async version of wait_for_completion.
+
+        Args:
+            call_id: ID of the queued completion call
+
+        Returns:
+            Dict containing response and/or tool_calls
+
+        Raises:
+            TimeoutError: If timeout is exceeded
+            Exception: If completion fails
+        """
+        import time
+
+        import aiohttp
+
+        start_time = time.time()
+
+        async with aiohttp.ClientSession() as session:
+            while True:
+                try:
+                    # Check status
+                    async with session.get(f"{self.base_url}/status/{call_id}") as response:
+                        response.raise_for_status()
+                        status_data = await response.json()
+
+                    if status_data["status"] == "completed":
+                        result = {}
+                        if "response" in status_data and status_data["response"]:
+                            result["response"] = status_data["response"]
+                        if "tool_calls" in status_data and status_data["tool_calls"]:
+                            result["tool_calls"] = status_data["tool_calls"]
+                        return result
+                    elif status_data["status"] == "failed":
+                        error_msg = status_data.get("error", "Unknown error")
+                        raise Exception(f"Completion failed: {error_msg}")
+
+                    # Check timeout
+                    if time.time() - start_time > self.timeout:
+                        raise TimeoutError(
+                            f"Timeout waiting for human response after {self.timeout} seconds"
+                        )
+
+                    # Wait before checking again
+                    await asyncio.sleep(1.0)
+
+                except Exception as e:
+                    if time.time() - start_time > self.timeout:
+                        raise TimeoutError(f"Timeout waiting for human response: {e}")
+                    # Continue trying if we haven't timed out
+                    await asyncio.sleep(1.0)
+
+    def _generate_response(self, messages: List[Dict[str, Any]], model: str) -> Dict[str, Any]:
+        """Generate a human response for the given messages.
+
+        Args:
+            messages: Messages in OpenAI format
+            model: Model name
+
+        Returns:
+            Dict containing response and/or tool_calls
+        """
+        # Queue the completion request
+        call_id = self._queue_completion(messages, model)
+
+        # Wait for human response
+        response = self._wait_for_completion(call_id)
+
+        return response
+
+    async def _async_generate_response(
+        self, messages: List[Dict[str, Any]], model: str
+    ) -> Dict[str, Any]:
+        """Async version of _generate_response.
+
+        Args:
+            messages: Messages in OpenAI format
+            model: Model name
+
+        Returns:
+            Dict containing response and/or tool_calls
+        """
+        # Queue the completion request (sync operation)
+        call_id = self._queue_completion(messages, model)
+
+        # Wait for human response (async)
+        response = await self._async_wait_for_completion(call_id)
+
+        return response
+
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        """Synchronous completion method.
+
+        Returns:
+            ModelResponse with human-generated text or tool calls
+        """
+        messages = kwargs.get("messages", [])
+        model = kwargs.get("model", "human")
+
+        # Generate human response
+        human_response_data = self._generate_response(messages, model)
+
+        # Create ModelResponse with proper structure
+        import time
+        import uuid
+
+        from litellm.types.utils import Choices, Message, ModelResponse
+
+        # Create message content based on response type
+        if "tool_calls" in human_response_data and human_response_data["tool_calls"]:
+            # Tool calls response
+            message = Message(
+                role="assistant",
+                content=human_response_data.get("response", ""),
+                tool_calls=human_response_data["tool_calls"],
+            )
+        else:
+            # Text response
+            message = Message(role="assistant", content=human_response_data.get("response", ""))
+
+        choice = Choices(finish_reason="stop", index=0, message=message)
+
+        result = ModelResponse(
+            id=f"human-{uuid.uuid4()}",
+            choices=[choice],
+            created=int(time.time()),
+            model=f"human/{model}",
+            object="chat.completion",
+        )
+
+        return result
+
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        """Asynchronous completion method.
+
+        Returns:
+            ModelResponse with human-generated text or tool calls
+        """
+        messages = kwargs.get("messages", [])
+        model = kwargs.get("model", "human")
+
+        # Generate human response
+        human_response_data = await self._async_generate_response(messages, model)
+
+        # Create ModelResponse with proper structure
+        import time
+        import uuid
+
+        from litellm.types.utils import Choices, Message, ModelResponse
+
+        # Create message content based on response type
+        if "tool_calls" in human_response_data and human_response_data["tool_calls"]:
+            # Tool calls response
+            message = Message(
+                role="assistant",
+                content=human_response_data.get("response", ""),
+                tool_calls=human_response_data["tool_calls"],
+            )
+        else:
+            # Text response
+            message = Message(role="assistant", content=human_response_data.get("response", ""))
+
+        choice = Choices(finish_reason="stop", index=0, message=message)
+
+        result = ModelResponse(
+            id=f"human-{uuid.uuid4()}",
+            choices=[choice],
+            created=int(time.time()),
+            model=f"human/{model}",
+            object="chat.completion",
+        )
+
+        return result
+
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        """Synchronous streaming method.
+
+        Yields:
+            Streaming chunks with human-generated text or tool calls
+        """
+        messages = kwargs.get("messages", [])
+        model = kwargs.get("model", "human")
+
+        # Generate human response
+        human_response_data = self._generate_response(messages, model)
+
+        import time
+
+        # Handle tool calls vs text response
+        if "tool_calls" in human_response_data and human_response_data["tool_calls"]:
+            # Stream tool calls as a single chunk
+            generic_chunk: GenericStreamingChunk = {
+                "finish_reason": "tool_calls",
+                "index": 0,
+                "is_finished": True,
+                "text": human_response_data.get("response", ""),
+                "tool_use": human_response_data["tool_calls"],
+                "usage": {"completion_tokens": 1, "prompt_tokens": 0, "total_tokens": 1},
+            }
+            yield generic_chunk
+        else:
+            # Stream text response
+            response_text = human_response_data.get("response", "")
+            generic_chunk: GenericStreamingChunk = {
+                "finish_reason": "stop",
+                "index": 0,
+                "is_finished": True,
+                "text": response_text,
+                "tool_use": None,
+                "usage": {
+                    "completion_tokens": len(response_text.split()),
+                    "prompt_tokens": 0,
+                    "total_tokens": len(response_text.split()),
+                },
+            }
+            yield generic_chunk
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        """Asynchronous streaming method.
+
+        Yields:
+            Streaming chunks with human-generated text or tool calls
+        """
+        messages = kwargs.get("messages", [])
+        model = kwargs.get("model", "human")
+
+        # Generate human response
+        human_response = await self._async_generate_response(messages, model)
+
+        # Return as single streaming chunk
+        generic_streaming_chunk: GenericStreamingChunk = {
+            "finish_reason": "stop",
+            "index": 0,
+            "is_finished": True,
+            "text": human_response,
+            "tool_use": None,
+            "usage": {
+                "completion_tokens": len(human_response.split()),
+                "prompt_tokens": 0,
+                "total_tokens": len(human_response.split()),
+            },
+        }
+
+        yield generic_streaming_chunk

libs/python/agent/cua_agent/adapters/mlxvlm_adapter.py

@@ -0,0 +1,370 @@
+import asyncio
+import base64
+import functools
+import io
+import math
+import re
+import warnings
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, cast
+
+from litellm import acompletion, completion
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+from PIL import Image
+
+# Try to import MLX dependencies
+try:
+    import mlx.core as mx
+    from mlx_vlm import generate, load
+    from mlx_vlm.prompt_utils import apply_chat_template
+    from mlx_vlm.utils import load_config
+    from transformers.tokenization_utils import PreTrainedTokenizer
+
+    MLX_AVAILABLE = True
+except ImportError:
+    MLX_AVAILABLE = False
+
+# Constants for smart_resize
+IMAGE_FACTOR = 28
+MIN_PIXELS = 100 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+
+def round_by_factor(number: float, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: float, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: float, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+class MLXVLMAdapter(CustomLLM):
+    """MLX VLM Adapter for running vision-language models locally using MLX."""
+
+    def __init__(self, **kwargs):
+        """Initialize the adapter.
+
+        Args:
+            **kwargs: Additional arguments
+        """
+        super().__init__()
+
+        self.models = {}  # Cache for loaded models
+        self.processors = {}  # Cache for loaded processors
+        self.configs = {}  # Cache for loaded configs
+        self._executor = ThreadPoolExecutor(max_workers=1)  # Single thread pool
+
+    def _load_model_and_processor(self, model_name: str):
+        """Load model and processor if not already cached.
+
+        Args:
+            model_name: Name of the model to load
+
+        Returns:
+            Tuple of (model, processor, config)
+        """
+        if not MLX_AVAILABLE:
+            raise ImportError("MLX VLM dependencies not available. Please install mlx-vlm.")
+
+        if model_name not in self.models:
+            # Load model and processor
+            model_obj, processor = load(
+                model_name, processor_kwargs={"min_pixels": MIN_PIXELS, "max_pixels": MAX_PIXELS}
+            )
+            config = load_config(model_name)
+
+            # Cache them
+            self.models[model_name] = model_obj
+            self.processors[model_name] = processor
+            self.configs[model_name] = config
+
+        return self.models[model_name], self.processors[model_name], self.configs[model_name]
+
+    def _process_coordinates(
+        self, text: str, original_size: Tuple[int, int], model_size: Tuple[int, int]
+    ) -> str:
+        """Process coordinates in box tokens based on image resizing using smart_resize approach.
+
+        Args:
+            text: Text containing box tokens
+            original_size: Original image size (width, height)
+            model_size: Model processed image size (width, height)
+
+        Returns:
+            Text with processed coordinates
+        """
+        # Find all box tokens
+        box_pattern = r"<\|box_start\|>\((\d+),\s*(\d+)\)<\|box_end\|>"
+
+        def process_coords(match):
+            model_x, model_y = int(match.group(1)), int(match.group(2))
+            # Scale coordinates from model space to original image space
+            # Both original_size and model_size are in (width, height) format
+            new_x = int(model_x * original_size[0] / model_size[0])  # Width
+            new_y = int(model_y * original_size[1] / model_size[1])  # Height
+            return f"<|box_start|>({new_x},{new_y})<|box_end|>"
+
+        return re.sub(box_pattern, process_coords, text)
+
+    def _convert_messages(self, messages: List[Dict[str, Any]]) -> Tuple[
+        List[Dict[str, Any]],
+        List[Image.Image],
+        Dict[int, Tuple[int, int]],
+        Dict[int, Tuple[int, int]],
+    ]:
+        """Convert OpenAI format messages to MLX VLM format and extract images.
+
+        Args:
+            messages: Messages in OpenAI format
+
+        Returns:
+            Tuple of (processed_messages, images, original_sizes, model_sizes)
+        """
+        processed_messages = []
+        images = []
+        original_sizes = {}  # Track original sizes of images for coordinate mapping
+        model_sizes = {}  # Track model processed sizes
+        image_index = 0
+
+        for message in messages:
+            processed_message = {"role": message["role"], "content": []}
+
+            content = message.get("content", [])
+            if isinstance(content, str):
+                # Simple text content
+                processed_message["content"] = content
+            elif isinstance(content, list):
+                # Multi-modal content
+                processed_content = []
+                for item in content:
+                    if item.get("type") == "text":
+                        processed_content.append({"type": "text", "text": item.get("text", "")})
+                    elif item.get("type") == "image_url":
+                        image_url = item.get("image_url", {}).get("url", "")
+                        pil_image = None
+
+                        if image_url.startswith("data:image/"):
+                            # Extract base64 data
+                            base64_data = image_url.split(",")[1]
+                            # Convert base64 to PIL Image
+                            image_data = base64.b64decode(base64_data)
+                            pil_image = Image.open(io.BytesIO(image_data))
+                        else:
+                            # Handle file path or URL
+                            pil_image = Image.open(image_url)
+
+                        # Store original image size for coordinate mapping
+                        original_size = pil_image.size
+                        original_sizes[image_index] = original_size
+
+                        # Use smart_resize to determine model size
+                        # Note: smart_resize expects (height, width) but PIL gives (width, height)
+                        height, width = original_size[1], original_size[0]
+                        new_height, new_width = smart_resize(height, width)
+                        # Store model size in (width, height) format for consistent coordinate processing
+                        model_sizes[image_index] = (new_width, new_height)
+
+                        # Resize the image using the calculated dimensions from smart_resize
+                        resized_image = pil_image.resize((new_width, new_height))
+                        images.append(resized_image)
+
+                        # Add image placeholder to content
+                        processed_content.append({"type": "image"})
+
+                        image_index += 1
+
+                processed_message["content"] = processed_content
+
+            processed_messages.append(processed_message)
+
+        return processed_messages, images, original_sizes, model_sizes
+
+    def _generate(self, **kwargs) -> str:
+        """Generate response using the local MLX VLM model.
+
+        Args:
+            **kwargs: Keyword arguments containing messages and model info
+
+        Returns:
+            Generated text response
+        """
+        messages = kwargs.get("messages", [])
+        model_name = kwargs.get("model", "mlx-community/UI-TARS-1.5-7B-4bit")
+        max_tokens = kwargs.get("max_tokens", 128)
+
+        # Warn about ignored kwargs
+        ignored_kwargs = set(kwargs.keys()) - {"messages", "model", "max_tokens"}
+        if ignored_kwargs:
+            warnings.warn(f"Ignoring unsupported kwargs: {ignored_kwargs}")
+
+        # Load model and processor
+        model, processor, config = self._load_model_and_processor(model_name)
+
+        # Convert messages and extract images
+        processed_messages, images, original_sizes, model_sizes = self._convert_messages(messages)
+
+        # Process user text input with box coordinates after image processing
+        # Swap original_size and model_size arguments for inverse transformation
+        for msg_idx, msg in enumerate(processed_messages):
+            if msg.get("role") == "user" and isinstance(msg.get("content"), str):
+                content = msg.get("content", "")
+                if (
+                    "<|box_start|>" in content
+                    and original_sizes
+                    and model_sizes
+                    and 0 in original_sizes
+                    and 0 in model_sizes
+                ):
+                    orig_size = original_sizes[0]
+                    model_size = model_sizes[0]
+                    # Swap arguments to perform inverse transformation for user input
+                    processed_messages[msg_idx]["content"] = self._process_coordinates(
+                        content, model_size, orig_size
+                    )
+
+        try:
+            # Format prompt according to model requirements using the processor directly
+            prompt = processor.apply_chat_template(
+                processed_messages, tokenize=False, add_generation_prompt=True, return_tensors="pt"
+            )
+            tokenizer = cast(PreTrainedTokenizer, processor)
+
+            # Generate response
+            text_content, usage = generate(
+                model,
+                tokenizer,
+                str(prompt),
+                images,  # type: ignore
+                verbose=False,
+                max_tokens=max_tokens,
+            )
+
+        except Exception as e:
+            raise RuntimeError(f"Error generating response: {str(e)}") from e
+
+        # Process coordinates in the response back to original image space
+        if original_sizes and model_sizes and 0 in original_sizes and 0 in model_sizes:
+            # Get original image size and model size (using the first image)
+            orig_size = original_sizes[0]
+            model_size = model_sizes[0]
+
+            # Check if output contains box tokens that need processing
+            if "<|box_start|>" in text_content:
+                # Process coordinates from model space back to original image space
+                text_content = self._process_coordinates(text_content, orig_size, model_size)
+
+        return text_content
+
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        """Synchronous completion method.
+
+        Returns:
+            ModelResponse with generated text
+        """
+        generated_text = self._generate(**kwargs)
+
+        result = completion(
+            model=f"mlx/{kwargs.get('model', 'mlx-community/UI-TARS-1.5-7B-4bit')}",
+            mock_response=generated_text,
+        )
+        return cast(ModelResponse, result)
+
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        """Asynchronous completion method.
+
+        Returns:
+            ModelResponse with generated text
+        """
+        # Run _generate in thread pool to avoid blocking
+        loop = asyncio.get_event_loop()
+        generated_text = await loop.run_in_executor(
+            self._executor, functools.partial(self._generate, **kwargs)
+        )
+
+        result = await acompletion(
+            model=f"mlx/{kwargs.get('model', 'mlx-community/UI-TARS-1.5-7B-4bit')}",
+            mock_response=generated_text,
+        )
+        return cast(ModelResponse, result)
+
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        """Synchronous streaming method.
+
+        Returns:
+            Iterator of GenericStreamingChunk
+        """
+        generated_text = self._generate(**kwargs)
+
+        generic_streaming_chunk: GenericStreamingChunk = {
+            "finish_reason": "stop",
+            "index": 0,
+            "is_finished": True,
+            "text": generated_text,
+            "tool_use": None,
+            "usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+        }
+
+        yield generic_streaming_chunk
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        """Asynchronous streaming method.
+
+        Returns:
+            AsyncIterator of GenericStreamingChunk
+        """
+        # Run _generate in thread pool to avoid blocking
+        loop = asyncio.get_event_loop()
+        generated_text = await loop.run_in_executor(
+            self._executor, functools.partial(self._generate, **kwargs)
+        )
+
+        generic_streaming_chunk: GenericStreamingChunk = {
+            "finish_reason": "stop",
+            "index": 0,
+            "is_finished": True,
+            "text": generated_text,
+            "tool_use": None,
+            "usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
+        }
+
+        yield generic_streaming_chunk

libs/python/agent/cua_agent/adapters/models/__init__.py

@@ -0,0 +1,41 @@
+from typing import Optional
+
+try:
+    from transformers import AutoConfig
+
+    HF_AVAILABLE = True
+except ImportError:
+    HF_AVAILABLE = False
+
+from .generic import GenericHFModel
+from .internvl import InternVLModel
+from .opencua import OpenCUAModel
+from .qwen2_5_vl import Qwen2_5_VLModel
+
+
+def load_model(model_name: str, device: str = "auto", trust_remote_code: bool = False):
+    """Factory function to load and return the right model handler instance.
+
+    - If the underlying transformers config class matches OpenCUA, return OpenCUAModel
+    - Otherwise, return GenericHFModel
+    """
+    if not HF_AVAILABLE:
+        raise ImportError(
+            'HuggingFace transformers dependencies not found. Install with: pip install "cua-agent[uitars-hf]"'
+        )
+    cfg = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code)
+    cls = cfg.__class__.__name__
+    print(f"cls: {cls}")
+    if "OpenCUA" in cls:
+        return OpenCUAModel(
+            model_name=model_name, device=device, trust_remote_code=trust_remote_code
+        )
+    elif "Qwen2_5_VL" in cls:
+        return Qwen2_5_VLModel(
+            model_name=model_name, device=device, trust_remote_code=trust_remote_code
+        )
+    elif "InternVL" in cls:
+        return InternVLModel(
+            model_name=model_name, device=device, trust_remote_code=trust_remote_code
+        )
+    return GenericHFModel(model_name=model_name, device=device, trust_remote_code=trust_remote_code)

libs/python/agent/cua_agent/adapters/models/generic.py

@@ -0,0 +1,78 @@
+from typing import Any, Dict, List, Optional
+
+# Hugging Face imports are local to avoid hard dependency at module import
+try:
+    import torch  # type: ignore
+    from transformers import AutoModel, AutoProcessor  # type: ignore
+
+    HF_AVAILABLE = True
+except Exception:
+    HF_AVAILABLE = False
+
+
+class GenericHFModel:
+    """Generic Hugging Face vision-language model handler.
+    Loads an AutoModelForImageTextToText and AutoProcessor and generates text.
+    """
+
+    def __init__(
+        self, model_name: str, device: str = "auto", trust_remote_code: bool = False
+    ) -> None:
+        if not HF_AVAILABLE:
+            raise ImportError(
+                'HuggingFace transformers dependencies not found. Install with: pip install "cua-agent[uitars-hf]"'
+            )
+        self.model_name = model_name
+        self.device = device
+        self.model = None
+        self.processor = None
+        self.trust_remote_code = trust_remote_code
+        self._load()
+
+    def _load(self) -> None:
+        # Load model
+        self.model = AutoModel.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.float16,
+            device_map=self.device,
+            attn_implementation="sdpa",
+            trust_remote_code=self.trust_remote_code,
+        )
+        # Load processor
+        self.processor = AutoProcessor.from_pretrained(
+            self.model_name,
+            min_pixels=3136,
+            max_pixels=4096 * 2160,
+            device_map=self.device,
+            trust_remote_code=self.trust_remote_code,
+        )
+
+    def generate(self, messages: List[Dict[str, Any]], max_new_tokens: int = 128) -> str:
+        """Generate text for the given HF-format messages.
+        messages: [{ role, content: [{type:'text'|'image', text|image}] }]
+        """
+        assert self.model is not None and self.processor is not None
+        # Apply chat template and tokenize
+        inputs = self.processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        )
+        # Move inputs to the same device as model
+        inputs = inputs.to(self.model.device)
+        # Generate
+        with torch.no_grad():
+            generated_ids = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
+        # Trim prompt tokens from output
+        generated_ids_trimmed = [
+            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        # Decode
+        output_text = self.processor.batch_decode(
+            generated_ids_trimmed,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )
+        return output_text[0] if output_text else ""

libs/python/agent/cua_agent/adapters/models/internvl.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+# Hugging Face imports are local to avoid hard dependency at module import
+try:
+    import base64  # type: ignore
+    from io import BytesIO  # type: ignore
+
+    # Attempt to import InternVL's model dependencies
+    import einops as _  # type: ignore
+    import requests  # type: ignore
+    import timm as _  # type: ignore
+    import torch  # type: ignore
+    import torchvision.transforms as T  # type: ignore
+    from PIL import Image  # type: ignore
+    from torchvision.transforms.functional import InterpolationMode  # type: ignore
+    from transformers import AutoModel, AutoTokenizer  # type: ignore
+
+    HF_AVAILABLE = True
+except Exception:
+    HF_AVAILABLE = False
+
+
+class InternVLModel:
+    """Generic Hugging Face vision-language model handler.
+    Uses InternVL's native `model.chat()` interface with `AutoTokenizer`.
+    Provides preprocessing to support multi-turn conversations with multiple images.
+    """
+
+    def __init__(
+        self, model_name: str, device: str = "auto", trust_remote_code: bool = False
+    ) -> None:
+        if not HF_AVAILABLE:
+            raise ImportError(
+                'InternVL dependencies not found. Install with: pip install "cua-agent[internvl-hf]"'
+            )
+        self.model_name = model_name
+        self.device = device
+        self.model = None
+        self.tokenizer = None
+        self.trust_remote_code = trust_remote_code
+        self._load()
+
+    def _load(self) -> None:
+        # Load model
+        self.model = AutoModel.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            use_flash_attn=True,
+            device_map=self.device,
+            trust_remote_code=self.trust_remote_code,
+        ).eval()
+        # Load tokenizer (InternVL requires trust_remote_code=True and often use_fast=False)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.model_name,
+            trust_remote_code=self.trust_remote_code,
+            use_fast=False,
+        )
+
+    # ---- Image preprocessing utilities adapted from InternVL docs ----
+    IMAGENET_MEAN = (0.485, 0.456, 0.406)
+    IMAGENET_STD = (0.229, 0.224, 0.225)
+
+    def _build_transform(self, input_size: int) -> T.Compose:
+        MEAN, STD = self.IMAGENET_MEAN, self.IMAGENET_STD
+        transform = T.Compose(
+            [
+                T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
+                T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+                T.ToTensor(),
+                T.Normalize(mean=MEAN, std=STD),
+            ]
+        )
+        return transform
+
+    def _find_closest_aspect_ratio(
+        self,
+        aspect_ratio: float,
+        target_ratios: List[tuple],
+        width: int,
+        height: int,
+        image_size: int,
+    ):
+        best_ratio_diff = float("inf")
+        best_ratio = (1, 1)
+        area = width * height
+        for ratio in target_ratios:
+            target_aspect_ratio = ratio[0] / ratio[1]
+            ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+            if ratio_diff < best_ratio_diff:
+                best_ratio_diff = ratio_diff
+                best_ratio = ratio
+            elif ratio_diff == best_ratio_diff:
+                if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                    best_ratio = ratio
+        return best_ratio
+
+    def _dynamic_preprocess(
+        self,
+        image: Image.Image,
+        min_num: int = 1,
+        max_num: int = 12,
+        image_size: int = 448,
+        use_thumbnail: bool = True,
+    ) -> List[Image.Image]:
+        orig_width, orig_height = image.size
+        aspect_ratio = orig_width / orig_height
+
+        target_ratios = set(
+            (i, j)
+            for n in range(min_num, max_num + 1)
+            for i in range(1, n + 1)
+            for j in range(1, n + 1)
+            if i * j <= max_num and i * j >= min_num
+        )
+        target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+        target_aspect_ratio = self._find_closest_aspect_ratio(
+            aspect_ratio, target_ratios, orig_width, orig_height, image_size
+        )
+
+        target_width = image_size * target_aspect_ratio[0]
+        target_height = image_size * target_aspect_ratio[1]
+        blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+        resized_img = image.resize((target_width, target_height))
+        processed_images: List[Image.Image] = []
+        for i in range(blocks):
+            box = (
+                (i % (target_width // image_size)) * image_size,
+                (i // (target_width // image_size)) * image_size,
+                ((i % (target_width // image_size)) + 1) * image_size,
+                ((i // (target_width // image_size)) + 1) * image_size,
+            )
+            split_img = resized_img.crop(box)
+            processed_images.append(split_img)
+        assert len(processed_images) == blocks
+        if use_thumbnail and len(processed_images) != 1:
+            thumbnail_img = image.resize((image_size, image_size))
+            processed_images.append(thumbnail_img)
+        return processed_images
+
+    def _load_image_from_source(self, src: str) -> Image.Image:
+        """Load PIL image from various sources: data URL, http(s), or local path."""
+        if src.startswith("data:image/"):
+            # data URL base64
+            header, b64data = src.split(",", 1)
+            img_bytes = base64.b64decode(b64data)
+            return Image.open(BytesIO(img_bytes)).convert("RGB")
+        if src.startswith("http://") or src.startswith("https://"):
+            resp = requests.get(src, timeout=10)
+            resp.raise_for_status()
+            return Image.open(BytesIO(resp.content)).convert("RGB")
+        # Assume local file path
+        return Image.open(src).convert("RGB")
+
+    def _images_to_pixel_values(
+        self, images: List[Image.Image], input_size: int = 448, max_num: int = 12
+    ):
+        transform = self._build_transform(input_size=input_size)
+        pixel_values_list = []
+        num_patches_list: List[int] = []
+        for img in images:
+            tiles = self._dynamic_preprocess(
+                img, image_size=input_size, use_thumbnail=True, max_num=max_num
+            )
+            pv = [transform(tile) for tile in tiles]
+            pv = torch.stack(pv)
+            num_patches_list.append(pv.shape[0])
+            pixel_values_list.append(pv)
+        if not pixel_values_list:
+            return None, []
+        pixel_values = torch.cat(pixel_values_list)
+        return pixel_values, num_patches_list
+
+    def generate(self, messages: List[Dict[str, Any]], max_new_tokens: int = 128) -> str:
+        """Generate text for the given HF-format messages.
+        messages: [{ role, content: [{type:'text'|'image', text|image}] }]
+
+        This implementation constructs InternVL-compatible inputs and uses
+        `model.chat(tokenizer, pixel_values, question, history=...)` to avoid
+        relying on AutoProcessor (which fails for some tokenizers).
+        """
+        assert self.model is not None and self.tokenizer is not None
+
+        # Build textual context and collect images and the final question
+        context_lines: List[str] = []
+        all_images: List[Image.Image] = []
+        last_user_text_parts: List[str] = []
+
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", [])
+            if isinstance(content, str):
+                content_items = [{"type": "text", "text": content}]
+            else:
+                content_items = content
+
+            if role == "user":
+                # Collect text and images
+                parts_text: List[str] = []
+                for item in content_items:
+                    if item.get("type") == "text":
+                        t = item.get("text", "")
+                        if t:
+                            parts_text.append(t)
+                    elif item.get("type") == "image":
+                        url = item.get("image", "")
+                        if url:
+                            try:
+                                all_images.append(self._load_image_from_source(url))
+                            except Exception:
+                                # Ignore failed image loads but keep going
+                                pass
+                text = "\n".join(parts_text).strip()
+                if text:
+                    context_lines.append(f"User: {text}")
+                # Track last user text separately for question
+                last_user_text_parts = parts_text or last_user_text_parts
+            elif role == "assistant":
+                # Only keep text content for history
+                parts_text = [
+                    item.get("text", "") for item in content_items if item.get("type") == "text"
+                ]
+                text = "\n".join(parts_text).strip()
+                if text:
+                    context_lines.append(f"Assistant: {text}")
+
+        # Prepare pixel values for all collected images (across turns)
+        pixel_values = None
+        num_patches_list: List[int] = []
+        if all_images:
+            pixel_values, num_patches_list = self._images_to_pixel_values(
+                all_images, input_size=448, max_num=12
+            )
+            if pixel_values is not None:
+                # Convert dtype/device as in docs
+                pixel_values = pixel_values.to(torch.bfloat16)
+                # Chat API expects tensors on CUDA when model is on CUDA
+                try:
+                    pixel_values = pixel_values.to(self.model.device)
+                except Exception:
+                    pass
+
+        # Build question with any prior context and numbered image placeholders
+        if all_images:
+            # Separate images layout: Image-1: <image> ... then question text
+            prefix_lines = [f"Image-{i+1}: <image>" for i in range(len(all_images))]
+            prefix = "\n".join(prefix_lines) + "\n"
+        else:
+            prefix = ""
+
+        last_user_text = "\n".join(last_user_text_parts).strip()
+        # Combine prior text-only turns as context to emulate multi-turn
+        context_text = "\n".join(context_lines[:-1]) if len(context_lines) > 1 else ""
+        base_question = last_user_text if last_user_text else "Describe the image(s) in detail."
+        if context_text:
+            question = (context_text + "\n" + prefix + base_question).strip()
+        else:
+            question = (prefix + base_question).strip()
+
+        # Generation config
+        generation_config = dict(max_new_tokens=max_new_tokens, do_sample=False)
+
+        # Call InternVL chat
+        try:
+            if pixel_values is None:
+                # Pure-text conversation (embed prior turns in question)
+                response = self.model.chat(self.tokenizer, None, question, generation_config)
+            else:
+                # Multi-image: pass num_patches_list if >1 image
+                if len(num_patches_list) > 1:
+                    response = self.model.chat(
+                        self.tokenizer,
+                        pixel_values,
+                        question,
+                        generation_config,
+                        num_patches_list=num_patches_list,
+                    )
+                else:
+                    response = self.model.chat(
+                        self.tokenizer, pixel_values, question, generation_config
+                    )
+        except Exception as e:
+            # Fallback: return empty string to avoid crashing the adapter
+            return ""
+
+        return response or ""

libs/python/agent/cua_agent/adapters/models/opencua.py

@@ -0,0 +1,115 @@
+import base64
+import re
+from io import BytesIO
+from typing import Any, Dict, List
+
+try:
+    import blobfile as _  # assert blobfile is installed
+    import torch  # type: ignore
+    from PIL import Image  # type: ignore
+    from transformers import (  # type: ignore
+        AutoImageProcessor,
+        AutoModel,
+        AutoTokenizer,
+    )
+
+    OPENCUA_AVAILABLE = True
+except Exception:
+    OPENCUA_AVAILABLE = False
+
+
+class OpenCUAModel:
+    """OpenCUA model handler using AutoTokenizer, AutoModel and AutoImageProcessor."""
+
+    def __init__(
+        self, model_name: str, device: str = "auto", trust_remote_code: bool = False
+    ) -> None:
+        if not OPENCUA_AVAILABLE:
+            raise ImportError(
+                'OpenCUA requirements not found. Install with: pip install "cua-agent[opencua-hf]"'
+            )
+        self.model_name = model_name
+        self.device = device
+        self.model = None
+        self.tokenizer = None
+        self.image_processor = None
+        self.trust_remote_code = trust_remote_code
+        self._load()
+
+    def _load(self) -> None:
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.model_name, trust_remote_code=self.trust_remote_code
+        )
+        self.model = AutoModel.from_pretrained(
+            self.model_name,
+            torch_dtype="auto",
+            device_map=self.device,
+            trust_remote_code=self.trust_remote_code,
+            attn_implementation="sdpa",
+        )
+        self.image_processor = AutoImageProcessor.from_pretrained(
+            self.model_name, trust_remote_code=self.trust_remote_code
+        )
+
+    @staticmethod
+    def _extract_last_image_b64(messages: List[Dict[str, Any]]) -> str:
+        # Expect HF-format messages with content items type: "image" with data URL
+        for msg in reversed(messages):
+            for item in reversed(msg.get("content", [])):
+                if isinstance(item, dict) and item.get("type") == "image":
+                    url = item.get("image", "")
+                    if isinstance(url, str) and url.startswith("data:image/"):
+                        return url.split(",", 1)[1]
+        return ""
+
+    def generate(self, messages: List[Dict[str, Any]], max_new_tokens: int = 512) -> str:
+        assert (
+            self.model is not None
+            and self.tokenizer is not None
+            and self.image_processor is not None
+        )
+
+        # Tokenize text side using chat template
+        input_ids = self.tokenizer.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True
+        )
+        input_ids = torch.tensor([input_ids]).to(self.model.device)
+
+        # Prepare image inputs from last data URL image
+        image_b64 = self._extract_last_image_b64(messages)
+        pixel_values = None
+        grid_thws = None
+        if image_b64:
+            image = Image.open(BytesIO(base64.b64decode(image_b64))).convert("RGB")
+            image_info = self.image_processor.preprocess(images=[image])
+            pixel_values = torch.tensor(image_info["pixel_values"]).to(
+                dtype=torch.bfloat16, device=self.model.device
+            )
+            grid_thws = (
+                torch.tensor(image_info["image_grid_thw"])
+                if "image_grid_thw" in image_info
+                else None
+            )
+
+        gen_kwargs: Dict[str, Any] = {
+            "max_new_tokens": max_new_tokens,
+            "temperature": 0,
+        }
+        if pixel_values is not None:
+            gen_kwargs["pixel_values"] = pixel_values
+        if grid_thws is not None:
+            gen_kwargs["grid_thws"] = grid_thws
+
+        with torch.no_grad():
+            generated_ids = self.model.generate(
+                input_ids,
+                **gen_kwargs,
+            )
+
+        # Remove prompt tokens
+        prompt_len = input_ids.shape[1]
+        generated_ids = generated_ids[:, prompt_len:]
+        output_text = self.tokenizer.batch_decode(
+            generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        return output_text

libs/python/agent/cua_agent/adapters/models/qwen2_5_vl.py

@@ -0,0 +1,78 @@
+from typing import Any, Dict, List, Optional
+
+# Hugging Face imports are local to avoid hard dependency at module import
+try:
+    import torch  # type: ignore
+    from transformers import AutoModelForImageTextToText, AutoProcessor  # type: ignore
+
+    HF_AVAILABLE = True
+except Exception:
+    HF_AVAILABLE = False
+
+
+class Qwen2_5_VLModel:
+    """Qwen2.5-VL Hugging Face vision-language model handler.
+    Loads an AutoModelForImageTextToText and AutoProcessor and generates text.
+    """
+
+    def __init__(
+        self, model_name: str, device: str = "auto", trust_remote_code: bool = False
+    ) -> None:
+        if not HF_AVAILABLE:
+            raise ImportError(
+                'HuggingFace transformers dependencies not found. Install with: pip install "cua-agent[uitars-hf]"'
+            )
+        self.model_name = model_name
+        self.device = device
+        self.model = None
+        self.processor = None
+        self.trust_remote_code = trust_remote_code
+        self._load()
+
+    def _load(self) -> None:
+        # Load model
+        self.model = AutoModelForImageTextToText.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.bfloat16,
+            device_map=self.device,
+            attn_implementation="sdpa",
+            trust_remote_code=self.trust_remote_code,
+        )
+        # Load processor
+        self.processor = AutoProcessor.from_pretrained(
+            self.model_name,
+            min_pixels=3136,
+            max_pixels=4096 * 2160,
+            device_map=self.device,
+            trust_remote_code=self.trust_remote_code,
+        )
+
+    def generate(self, messages: List[Dict[str, Any]], max_new_tokens: int = 128) -> str:
+        """Generate text for the given HF-format messages.
+        messages: [{ role, content: [{type:'text'|'image', text|image}] }]
+        """
+        assert self.model is not None and self.processor is not None
+        # Apply chat template and tokenize
+        inputs = self.processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        )
+        # Move inputs to the same device as model
+        inputs = inputs.to(self.model.device)
+        # Generate
+        with torch.no_grad():
+            generated_ids = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
+        # Trim prompt tokens from output
+        generated_ids_trimmed = [
+            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        # Decode
+        output_text = self.processor.batch_decode(
+            generated_ids_trimmed,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )
+        return output_text[0] if output_text else ""

libs/python/agent/cua_agent/adapters/yutori_adapter.py

@@ -0,0 +1,99 @@
+"""
+Yutori adapter for litellm - routes yutori/ prefixed models to the Yutori API.
+"""
+
+import os
+from typing import Any, AsyncIterator, Iterator
+
+from litellm import acompletion, completion
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.utils import GenericStreamingChunk, ModelResponse
+
+YUTORI_API_BASE = "https://api.yutori.com/v1"
+
+
+class YutoriAdapter(CustomLLM):
+    def __init__(self, base_url: str | None = None, api_key: str | None = None, **_: Any):
+        super().__init__()
+        self.base_url = base_url or os.environ.get("YUTORI_API_BASE") or YUTORI_API_BASE
+        self.api_key = api_key or os.environ.get("YUTORI_API_KEY")
+
+    def _normalize_model(self, model: str) -> str:
+        """Strip the yutori/ prefix to get the bare model name."""
+        if model.startswith("yutori/"):
+            return model[len("yutori/") :]
+        return model
+
+    def _resolve_api_key(self, kwargs: dict | None = None) -> str:
+        """Resolve the Yutori API key, raising a clear error if missing."""
+        resolved = (kwargs.get("api_key") if kwargs else None) or self.api_key
+        if not resolved:
+            raise ValueError(
+                "No Yutori API key provided. "
+                "Please either set the YUTORI_API_KEY environment variable "
+                "or pass api_key to ComputerAgent()."
+            )
+        return resolved
+
+    def _build_params(self, kwargs: dict) -> dict:
+        """Build parameters for the inner litellm call."""
+        model = self._normalize_model(kwargs.get("model", ""))
+        api_key = self._resolve_api_key(kwargs)
+
+        extra_headers = {}
+        if "extra_headers" in kwargs:
+            extra_headers.update(kwargs.pop("extra_headers"))
+        extra_headers["Authorization"] = f"Bearer {api_key}"
+
+        params = {
+            "model": f"openai/{model}",
+            "messages": kwargs.get("messages", []),
+            "api_base": self.base_url,
+            "api_key": api_key,
+            "extra_headers": extra_headers,
+            "stream": False,
+        }
+
+        # Forward tools if provided
+        if "tools" in kwargs:
+            params["tools"] = kwargs["tools"]
+        if "tool_choice" in kwargs:
+            params["tool_choice"] = kwargs["tool_choice"]
+
+        # Forward optional generation params
+        for key in (
+            "temperature",
+            "top_p",
+            "max_completion_tokens",
+            "max_tokens",
+            "response_format",
+        ):
+            if key in kwargs:
+                params[key] = kwargs[key]
+
+        if "optional_params" in kwargs:
+            protected_keys = {"api_key", "extra_headers", "model", "api_base", "stream"}
+            filtered = {
+                k: v for k, v in kwargs["optional_params"].items() if k not in protected_keys
+            }
+            params.update(filtered)
+
+        if "headers" in kwargs:
+            params["headers"] = kwargs["headers"]
+
+        return params
+
+    def completion(self, *args, **kwargs) -> ModelResponse:
+        params = self._build_params(kwargs)
+        return completion(**params)  # type: ignore
+
+    async def acompletion(self, *args, **kwargs) -> ModelResponse:
+        params = self._build_params(kwargs)
+        response = await acompletion(**params)  # type: ignore
+        return response
+
+    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
+        raise NotImplementedError("Yutori n1 does not support streaming.")
+
+    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
+        raise NotImplementedError("Yutori n1 does not support streaming.")

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

libs/python/agent/cua_agent/cli.py

@@ -0,0 +1,515 @@
+"""
+CLI chat interface for agent - Computer Use Agent
+
+Usage:
+    python -m agent.cli <model_string>
+
+Examples:
+    python -m agent.cli openai/computer-use-preview
+    python -m agent.cli anthropic/claude-sonnet-4-5-20250929
+    python -m agent.cli omniparser+anthropic/claude-sonnet-4-5-20250929
+"""
+
+try:
+    import argparse
+    import asyncio
+    import base64
+    import json
+    import os
+    import platform
+    import sys
+    import time
+    from pathlib import Path
+    from typing import Any, Dict, List
+
+    import dotenv
+
+    try:
+        from PIL import Image, ImageDraw
+
+        PIL_AVAILABLE = True
+    except Exception:
+        PIL_AVAILABLE = False
+    from yaspin import yaspin
+except ImportError:
+    if __name__ == "__main__":
+        raise ImportError(
+            "CLI dependencies not found. " 'Please install with: pip install "cua-agent[cli]"'
+        )
+
+# Load environment variables
+dotenv.load_dotenv()
+
+
+# Color codes for terminal output
+class Colors:
+    RESET = "\033[0m"
+    BOLD = "\033[1m"
+    DIM = "\033[2m"
+
+    # Text colors
+    RED = "\033[31m"
+    GREEN = "\033[32m"
+    YELLOW = "\033[33m"
+    BLUE = "\033[34m"
+    MAGENTA = "\033[35m"
+    CYAN = "\033[36m"
+    WHITE = "\033[37m"
+    GRAY = "\033[90m"
+
+    # Background colors
+    BG_RED = "\033[41m"
+    BG_GREEN = "\033[42m"
+    BG_YELLOW = "\033[43m"
+    BG_BLUE = "\033[44m"
+
+
+def print_colored(
+    text: str,
+    color: str = "",
+    bold: bool = False,
+    dim: bool = False,
+    end: str = "\n",
+    right: str = "",
+):
+    """Print colored text to terminal with optional right-aligned text."""
+    prefix = ""
+    if bold:
+        prefix += Colors.BOLD
+    if dim:
+        prefix += Colors.DIM
+    if color:
+        prefix += color
+
+    if right:
+        # Get terminal width (default to 80 if unable to determine)
+        try:
+            import shutil
+
+            terminal_width = shutil.get_terminal_size().columns
+        except:
+            terminal_width = 80
+
+        # Add right margin
+        terminal_width -= 1
+
+        # Calculate padding needed
+        # Account for ANSI escape codes not taking visual space
+        visible_left_len = len(text)
+        visible_right_len = len(right)
+        padding = terminal_width - visible_left_len - visible_right_len
+
+        if padding > 0:
+            output = f"{prefix}{text}{' ' * padding}{right}{Colors.RESET}"
+        else:
+            # If not enough space, just put a single space between
+            output = f"{prefix}{text} {right}{Colors.RESET}"
+    else:
+        output = f"{prefix}{text}{Colors.RESET}"
+
+    print(output, end=end)
+
+
+def print_action(action_type: str, details: Dict[str, Any], total_cost: float):
+    """Print computer action with nice formatting."""
+    # Format action details
+    args_str = ""
+    if action_type == "click" and "x" in details and "y" in details:
+        args_str = f"_{details.get('button', 'left')}({details['x']}, {details['y']})"
+    elif action_type == "type" and "text" in details:
+        text = details["text"]
+        if len(text) > 50:
+            text = text[:47] + "..."
+        args_str = f'("{text}")'
+    elif action_type == "key" and "text" in details:
+        args_str = f"('{details['text']}')"
+    elif action_type == "scroll" and "x" in details and "y" in details:
+        args_str = f"({details['x']}, {details['y']})"
+
+    if total_cost > 0:
+        print_colored(f"🛠️  {action_type}{args_str}", dim=True, right=f"💸 ${total_cost:.2f}")
+    else:
+        print_colored(f"🛠️  {action_type}{args_str}", dim=True)
+
+
+def print_welcome(model: str, agent_loop: str, container_name: str):
+    """Print welcome message."""
+    print_colored(f"Connected to {container_name} ({model}, {agent_loop})")
+    print_colored("Type 'exit' to quit.", dim=True)
+
+
+async def ainput(prompt: str = ""):
+    return await asyncio.to_thread(input, prompt)
+
+
+async def chat_loop(
+    agent, model: str, container_name: str, initial_prompt: str = "", show_usage: bool = True
+):
+    """Main chat loop with the agent."""
+    print_welcome(model, agent.agent_config_info.agent_class.__name__, container_name)
+
+    history = []
+
+    if initial_prompt:
+        history.append({"role": "user", "content": initial_prompt})
+
+    total_cost = 0
+
+    while True:
+        if len(history) == 0 or history[-1].get("role") != "user":
+            # Get user input with prompt
+            print_colored("> ", end="")
+            user_input = await ainput()
+
+            if user_input.lower() in ["exit", "quit", "q"]:
+                print_colored("\n👋 Goodbye!")
+                break
+
+            if not user_input:
+                continue
+
+            # Add user message to history
+            history.append({"role": "user", "content": user_input})
+
+        # Stream responses from the agent with spinner
+        with yaspin(text="Thinking...", spinner="line", attrs=["dark"]) as spinner:
+            spinner.hide()
+
+            async for result in agent.run(history):
+                # Add agent responses to history
+                history.extend(result.get("output", []))
+
+                if show_usage:
+                    total_cost += result.get("usage", {}).get("response_cost", 0)
+
+                # Process and display the output
+                for item in result.get("output", []):
+                    if item.get("type") == "message" and item.get("role") == "assistant":
+                        # Display agent text response
+                        content = item.get("content", [])
+                        for content_part in content:
+                            if content_part.get("text"):
+                                text = content_part.get("text", "").strip()
+                                if text:
+                                    spinner.hide()
+                                    print_colored(text)
+
+                    elif item.get("type") == "computer_call":
+                        # Display computer action
+                        action = item.get("action", {})
+                        action_type = action.get("type", "")
+                        if action_type:
+                            spinner.hide()
+                            print_action(action_type, action, total_cost)
+                            spinner.text = f"Performing {action_type}..."
+                            spinner.show()
+
+                    elif item.get("type") == "function_call":
+                        # Display function call
+                        function_name = item.get("name", "")
+                        spinner.hide()
+                        print_colored(f"🔧 Calling function: {function_name}", dim=True)
+                        spinner.text = f"Calling {function_name}..."
+                        spinner.show()
+
+                    elif item.get("type") == "function_call_output":
+                        # Display function output (dimmed)
+                        output = item.get("output", "")
+                        if output and len(output.strip()) > 0:
+                            spinner.hide()
+                            print_colored(f"📤 {output}", dim=True)
+
+            spinner.hide()
+            if show_usage and total_cost > 0:
+                print_colored(f"Total cost: ${total_cost:.2f}", dim=True)
+
+
+async def main():
+    """Main CLI function."""
+    parser = argparse.ArgumentParser(
+        description="Cua Agent CLI - Interactive computer use assistant",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python -m agent.cli openai/computer-use-preview
+  python -m agent.cli anthropic/claude-sonnet-4-5-20250929
+  python -m agent.cli omniparser+anthropic/claude-sonnet-4-5-20250929
+  python -m agent.cli huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B
+        """,
+    )
+
+    parser.add_argument(
+        "model",
+        help="Model string (e.g., 'openai/computer-use-preview', 'anthropic/claude-sonnet-4-5-20250929')",
+    )
+
+    parser.add_argument(
+        "--provider",
+        choices=["cloud", "lume", "winsandbox", "docker"],
+        default="cloud",
+        help="Computer provider to use: cloud (default), lume, winsandbox, or docker",
+    )
+
+    parser.add_argument(
+        "--images",
+        type=int,
+        default=3,
+        help="Number of recent images to keep in context (default: 3)",
+    )
+
+    parser.add_argument("--trajectory", action="store_true", help="Save trajectory for debugging")
+
+    parser.add_argument("--budget", type=float, help="Maximum budget for the session (in dollars)")
+
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
+
+    parser.add_argument(
+        "-p",
+        "--prompt",
+        type=str,
+        help="Initial prompt to send to the agent. Leave blank for interactive mode.",
+    )
+
+    parser.add_argument(
+        "--prompt-file",
+        type=Path,
+        help="Path to a UTF-8 text file whose contents will be used as the initial prompt. If provided, overrides --prompt.",
+    )
+
+    parser.add_argument(
+        "--predict-click",
+        dest="predict_click",
+        type=str,
+        help="Instruction for click prediction. If set, runs predict_click, draws crosshair on a fresh screenshot, saves and opens it.",
+    )
+
+    parser.add_argument("-c", "--cache", action="store_true", help="Tell the API to enable caching")
+
+    parser.add_argument(
+        "-u", "--usage", action="store_true", help="Show total cost of the agent runs"
+    )
+
+    parser.add_argument(
+        "-r",
+        "--max-retries",
+        type=int,
+        default=3,
+        help="Maximum number of retries for the LLM API calls",
+    )
+
+    # Provider override credentials
+    parser.add_argument(
+        "--api-key",
+        dest="api_key",
+        type=str,
+        help="API key override for the model provider (passed to ComputerAgent)",
+    )
+    parser.add_argument(
+        "--api-base",
+        dest="api_base",
+        type=str,
+        help="API base URL override for the model provider (passed to ComputerAgent)",
+    )
+
+    args = parser.parse_args()
+
+    # Check for required environment variables
+    container_name = os.getenv("CUA_CONTAINER_NAME")
+    cua_api_key = os.getenv("CUA_API_KEY")
+
+    # Prompt for missing environment variables (container name always required)
+    if not container_name:
+        if args.provider == "cloud":
+            print_colored("CUA_CONTAINER_NAME not set.", dim=True)
+            print_colored("You can get a Cua container at https://cua.ai/", dim=True)
+            container_name = input("Enter your Cua container name: ").strip()
+            if not container_name:
+                print_colored("❌ Container name is required.")
+                sys.exit(1)
+        else:
+            container_name = "cli-sandbox"
+
+    # Only require API key for cloud provider
+    if args.provider == "cloud" and not cua_api_key:
+        print_colored("CUA_API_KEY not set.", dim=True)
+        cua_api_key = input("Enter your Cua API key: ").strip()
+        if not cua_api_key:
+            print_colored("❌ API key is required for cloud provider.")
+            sys.exit(1)
+
+    # Check for provider-specific API keys based on model
+    provider_api_keys = {
+        "openai/": "OPENAI_API_KEY",
+        "anthropic/": "ANTHROPIC_API_KEY",
+    }
+
+    # Find matching provider and check for API key
+    for prefix, env_var in provider_api_keys.items():
+        if prefix in args.model:
+            if not os.getenv(env_var):
+                print_colored(f"{env_var} not set.", dim=True)
+                api_key = input(f"Enter your {env_var.replace('_', ' ').title()}: ").strip()
+                if not api_key:
+                    print_colored(f"❌ {env_var.replace('_', ' ').title()} is required.")
+                    sys.exit(1)
+                # Set the environment variable for the session
+                os.environ[env_var] = api_key
+            break
+
+    # Import here to avoid import errors if dependencies are missing
+    try:
+        from computer import Computer
+        from cua_agent import ComputerAgent
+    except ImportError as e:
+        print_colored(f"❌ Import error: {e}", Colors.RED, bold=True)
+        print_colored("Make sure agent and computer libraries are installed.", Colors.YELLOW)
+        sys.exit(1)
+
+    # Resolve provider -> os_type, provider_type, api key requirement
+    provider_map = {
+        "cloud": ("linux", "cloud", True),
+        "lume": ("macos", "lume", False),
+        "winsandbox": ("windows", "winsandbox", False),
+        "docker": ("linux", "docker", False),
+    }
+    os_type, provider_type, needs_api_key = provider_map[args.provider]
+
+    computer_kwargs = {
+        "os_type": os_type,
+        "provider_type": provider_type,
+        "name": container_name,
+    }
+    if needs_api_key:
+        computer_kwargs["api_key"] = cua_api_key  # type: ignore
+
+    # Create computer instance
+    async with Computer(**computer_kwargs) as computer:  # type: ignore
+
+        # Create agent
+        agent_kwargs = {
+            "model": args.model,
+            "tools": [computer],
+            "trust_remote_code": True,  # needed for some local models (e.g., InternVL, OpenCUA)
+            "verbosity": 20 if args.verbose else 30,  # DEBUG vs WARNING
+            "max_retries": args.max_retries,
+        }
+
+        # Thread API credentials to agent if provided
+        if args.api_key:
+            agent_kwargs["api_key"] = args.api_key
+        if args.api_base:
+            agent_kwargs["api_base"] = args.api_base
+
+        if args.images > 0:
+            agent_kwargs["only_n_most_recent_images"] = args.images
+
+        if args.trajectory:
+            agent_kwargs["trajectory_dir"] = "trajectories"
+
+        if args.budget:
+            agent_kwargs["max_trajectory_budget"] = {
+                "max_budget": args.budget,
+                "raise_error": True,
+                "reset_after_each_run": False,
+            }
+
+        if args.cache:
+            agent_kwargs["use_prompt_caching"] = True
+
+        agent = ComputerAgent(**agent_kwargs)
+
+        # If predict-click mode is requested, run once and exit
+        if args.predict_click:
+            if not PIL_AVAILABLE:
+                print_colored(
+                    "❌ Pillow (PIL) is required for --predict-click visualization. Install with: pip install pillow",
+                    Colors.RED,
+                    bold=True,
+                )
+                sys.exit(1)
+
+            instruction = args.predict_click
+            print_colored(f"Predicting click for: '{instruction}'", Colors.CYAN)
+
+            # Take a fresh screenshot FIRST
+            try:
+                img_bytes = await computer.interface.screenshot()
+            except Exception as e:
+                print_colored(f"❌ Failed to take screenshot: {e}", Colors.RED, bold=True)
+                sys.exit(1)
+
+            # Encode screenshot to base64 for predict_click
+            try:
+                image_b64 = base64.b64encode(img_bytes).decode("utf-8")
+            except Exception as e:
+                print_colored(f"❌ Failed to encode screenshot: {e}", Colors.RED, bold=True)
+                sys.exit(1)
+
+            try:
+                coords = await agent.predict_click(instruction, image_b64=image_b64)
+            except Exception as e:
+                print_colored(f"❌ predict_click failed: {e}", Colors.RED, bold=True)
+                sys.exit(1)
+
+            if not coords:
+                print_colored("⚠️  No coordinates returned.", Colors.YELLOW)
+                sys.exit(2)
+
+            x, y = coords
+            print_colored(f"✅ Predicted coordinates: ({x}, {y})", Colors.GREEN)
+
+            try:
+                from io import BytesIO
+
+                with Image.open(BytesIO(img_bytes)) as img:
+                    img = img.convert("RGB")
+                    draw = ImageDraw.Draw(img)
+                    # Draw crosshair
+                    size = 12
+                    color = (255, 0, 0)
+                    draw.line([(x - size, y), (x + size, y)], fill=color, width=3)
+                    draw.line([(x, y - size), (x, y + size)], fill=color, width=3)
+                    # Optional small circle
+                    r = 6
+                    draw.ellipse([(x - r, y - r), (x + r, y + r)], outline=color, width=2)
+
+                    out_path = Path.cwd() / f"predict_click_{int(time.time())}.png"
+                    img.save(out_path)
+                    print_colored(f"🖼️  Saved to {out_path}")
+
+                    # Open the image with default viewer
+                    try:
+                        system = platform.system().lower()
+                        if system == "windows":
+                            os.startfile(str(out_path))  # type: ignore[attr-defined]
+                        elif system == "darwin":
+                            os.system(f'open "{out_path}"')
+                        else:
+                            os.system(f'xdg-open "{out_path}"')
+                    except Exception:
+                        pass
+            except Exception as e:
+                print_colored(f"❌ Failed to render/save screenshot: {e}", Colors.RED, bold=True)
+                sys.exit(1)
+
+            # Done
+            sys.exit(0)
+
+        # Resolve initial prompt from --prompt-file or --prompt
+        initial_prompt = args.prompt or ""
+        if args.prompt_file:
+            try:
+                initial_prompt = args.prompt_file.read_text(encoding="utf-8")
+            except Exception as e:
+                print_colored(f"❌ Failed to read --prompt-file: {e}", Colors.RED, bold=True)
+                sys.exit(1)
+
+        # Start chat loop (default interactive mode)
+        await chat_loop(agent, args.model, container_name, initial_prompt, args.usage)
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except (KeyboardInterrupt, EOFError) as _:
+        print_colored("\n\n👋 Goodbye!")

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

@@ -0,0 +1,61 @@
+"""
+Computer handler factory and interface definitions.
+
+This module provides a factory function to create computer handlers from different
+computer interface types, supporting both the ComputerHandler protocol and the
+Computer library interface.
+"""
+
+try:
+    from computer import Computer as cuaComputer
+
+    from .cua import cuaComputerHandler
+except ImportError:
+    cuaComputer = None  # type: ignore[assignment,misc]
+    cuaComputerHandler = None  # type: ignore[assignment]
+
+try:
+    from cua_sandbox import Sandbox as cuaSandbox
+except ImportError:
+    cuaSandbox = None  # type: ignore[assignment,misc]
+
+from .base import AsyncComputerHandler
+from .custom import CustomComputerHandler
+from .sandbox import SandboxComputerHandler
+
+
+def is_agent_computer(computer):
+    """Check if the given computer is a ComputerHandler or Cua Computer."""
+    return (
+        isinstance(computer, AsyncComputerHandler)
+        or (cuaComputer is not None and isinstance(computer, cuaComputer))
+        or (cuaSandbox is not None and isinstance(computer, cuaSandbox))
+        or (isinstance(computer, dict))
+    )  # and "screenshot" in computer)
+
+
+async def make_computer_handler(computer):
+    """
+    Create a computer handler from a computer interface.
+
+    Args:
+        computer: Either a ComputerHandler instance, Computer instance,
+                  Sandbox instance, or dict of functions
+
+    Returns:
+        ComputerHandler: A computer handler instance
+
+    Raises:
+        ValueError: If the computer type is not supported
+    """
+    if isinstance(computer, AsyncComputerHandler):
+        return computer
+    if cuaComputer is not None and isinstance(computer, cuaComputer):
+        computer_handler = cuaComputerHandler(computer)
+        await computer_handler._initialize()
+        return computer_handler
+    if cuaSandbox is not None and isinstance(computer, cuaSandbox):
+        return SandboxComputerHandler(computer)
+    if isinstance(computer, dict):
+        return CustomComputerHandler(computer)
+    raise ValueError(f"Unsupported computer type: {type(computer)}")

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

@@ -0,0 +1,83 @@
+"""
+Base computer interface protocol for agent interactions.
+"""
+
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Protocol,
+    Union,
+    runtime_checkable,
+)
+
+
+@runtime_checkable
+class AsyncComputerHandler(Protocol):
+    """Protocol defining the interface for computer interactions."""
+
+    # ==== Computer-Use-Preview Action Space ====
+
+    async def get_environment(self) -> Literal["windows", "mac", "linux", "browser"]:
+        """Get the current environment type."""
+        ...
+
+    async def get_dimensions(self) -> tuple[int, int]:
+        """Get screen dimensions as (width, height)."""
+        ...
+
+    async def screenshot(self, text: Optional[str] = None) -> str:
+        """Take a screenshot and return as base64 string.
+
+        Args:
+            text: Optional descriptive text (for compatibility with GPT-4o models, ignored)
+        """
+        ...
+
+    async def click(self, x: int, y: int, button: str = "left") -> None:
+        """Click at coordinates with specified button."""
+        ...
+
+    async def double_click(self, x: int, y: int) -> None:
+        """Double click at coordinates."""
+        ...
+
+    async def scroll(self, x: int, y: int, scroll_x: int, scroll_y: int) -> None:
+        """Scroll at coordinates with specified scroll amounts."""
+        ...
+
+    async def type(self, text: str) -> None:
+        """Type text."""
+        ...
+
+    async def wait(self, ms: int = 1000) -> None:
+        """Wait for specified milliseconds."""
+        ...
+
+    async def move(self, x: int, y: int) -> None:
+        """Move cursor to coordinates."""
+        ...
+
+    async def keypress(self, keys: Union[List[str], str]) -> None:
+        """Press key combination."""
+        ...
+
+    async def drag(self, path: List[Dict[str, int]]) -> None:
+        """Drag along specified path."""
+        ...
+
+    async def get_current_url(self) -> str:
+        """Get current URL (for browser environments)."""
+        ...
+
+    # ==== Anthropic Action Space ====
+
+    async def left_mouse_down(self, x: Optional[int] = None, y: Optional[int] = None) -> None:
+        """Left mouse down at coordinates."""
+        ...
+
+    async def left_mouse_up(self, x: Optional[int] = None, y: Optional[int] = None) -> None:
+        """Left mouse up at coordinates."""
+        ...

libs/python/agent/cua_agent/computers/cua.py

@@ -0,0 +1,184 @@
+"""
+Computer handler implementation for OpenAI computer-use-preview protocol.
+"""
+
+import base64
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from computer import Computer
+
+from .base import AsyncComputerHandler
+
+
+class cuaComputerHandler(AsyncComputerHandler):
+    """Computer handler that implements the Computer protocol using the computer interface."""
+
+    def __init__(self, cua_computer: Computer):
+        """Initialize with a computer interface (from tool schema)."""
+        self.cua_computer = cua_computer
+        self.interface = None
+
+    async def _initialize(self):
+        if hasattr(self.cua_computer, "_initialized") and not self.cua_computer._initialized:
+            await self.cua_computer.run()
+        self.interface = self.cua_computer.interface
+
+    # ==== Computer-Use-Preview Action Space ====
+
+    async def get_environment(self) -> Literal["windows", "mac", "linux", "browser"]:
+        """Get the current environment type."""
+        # TODO: detect actual environment
+        return "linux"
+
+    async def get_dimensions(self) -> tuple[int, int]:
+        """Get screen dimensions as (width, height)."""
+        assert self.interface is not None
+        screen_size = await self.interface.get_screen_size()
+        return screen_size["width"], screen_size["height"]
+
+    async def screenshot(self, text: Optional[str] = None) -> str:
+        """Take a screenshot and return as base64 string.
+
+        Args:
+            text: Optional descriptive text (for compatibility with GPT-4o models, ignored)
+        """
+        assert self.interface is not None
+        screenshot_bytes = await self.interface.screenshot()
+        return base64.b64encode(screenshot_bytes).decode("utf-8")
+
+    async def click(self, x: int, y: int, button: str = "left") -> None:
+        """Click at coordinates with specified button."""
+        assert self.interface is not None
+        if button == "left":
+            await self.interface.left_click(x, y)
+        elif button == "right":
+            await self.interface.right_click(x, y)
+        else:
+            # Default to left click for unknown buttons
+            await self.interface.left_click(x, y)
+
+    async def double_click(self, x: int, y: int) -> None:
+        """Double click at coordinates."""
+        assert self.interface is not None
+        await self.interface.double_click(x, y)
+
+    async def right_click(self, x: int, y: int) -> None:
+        """Right click at coordinates."""
+        assert self.interface is not None
+        await self.interface.right_click(x, y)
+
+    async def scroll(self, x: int, y: int, scroll_x: int, scroll_y: int) -> None:
+        """Scroll at coordinates with specified scroll amounts."""
+        assert self.interface is not None
+        await self.interface.move_cursor(x, y)
+        await self.interface.scroll(scroll_x, scroll_y)
+
+    async def type(self, text: str) -> None:
+        """Type text."""
+        assert self.interface is not None
+        await self.interface.type_text(text)
+
+    async def wait(self, ms: int = 1000) -> None:
+        """Wait for specified milliseconds."""
+        assert self.interface is not None
+        import asyncio
+
+        await asyncio.sleep(ms / 1000.0)
+
+    async def move(self, x: int, y: int) -> None:
+        """Move cursor to coordinates."""
+        assert self.interface is not None
+        await self.interface.move_cursor(x, y)
+
+    async def keypress(self, keys: Union[List[str], str]) -> None:
+        """Press key combination."""
+        assert self.interface is not None
+        if isinstance(keys, str):
+            keys = keys.replace("-", "+").split("+")
+        if len(keys) == 1:
+            await self.interface.press_key(keys[0])
+        else:
+            # Handle key combinations
+            await self.interface.hotkey(*keys)
+
+    async def drag(
+        self,
+        path: Optional[List[Dict[str, int]]] = None,
+        start_x: Optional[int] = None,
+        start_y: Optional[int] = None,
+        end_x: Optional[int] = None,
+        end_y: Optional[int] = None,
+    ) -> None:
+        """Drag along specified path or from start to end coordinates.
+
+        Supports two formats:
+        - path: List of {x, y} points to drag through
+        - start_x, start_y, end_x, end_y: Simple drag from start to end
+        """
+        assert self.interface is not None
+
+        # If start/end coordinates provided, convert to path format
+        if start_x is not None and start_y is not None and end_x is not None and end_y is not None:
+            path = [{"x": start_x, "y": start_y}, {"x": end_x, "y": end_y}]
+
+        if not path:
+            return
+
+        # Start drag from first point
+        start = path[0]
+        await self.interface.mouse_down(start["x"], start["y"])
+
+        # Move through path
+        for point in path[1:]:
+            await self.interface.move_cursor(point["x"], point["y"])
+
+        # End drag at last point
+        end = path[-1]
+        await self.interface.mouse_up(end["x"], end["y"])
+
+    async def terminate(self, status: str = "success") -> Dict[str, Any]:
+        """Terminate the current task and report its completion status.
+
+        Args:
+            status: Status of the task ("success" or "failure")
+
+        Returns:
+            Dict with terminated flag and status
+        """
+        return {"success": True, "status": status, "terminated": True}
+
+    async def get_current_url(self) -> str:
+        """Get current URL (for browser environments)."""
+        # This would need to be implemented based on the specific browser interface
+        # For now, return empty string
+        return ""
+
+    # ==== Anthropic Computer Action Space ====
+    async def left_mouse_down(self, x: Optional[int] = None, y: Optional[int] = None) -> None:
+        """Left mouse down at coordinates."""
+        assert self.interface is not None
+        await self.interface.mouse_down(x, y, button="left")
+
+    async def left_mouse_up(self, x: Optional[int] = None, y: Optional[int] = None) -> None:
+        """Left mouse up at coordinates."""
+        assert self.interface is not None
+        await self.interface.mouse_up(x, y, button="left")
+
+    # ==== Browser Control Methods (via Playwright) ====
+    async def playwright_exec(
+        self, command: str, params: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Execute a Playwright browser command.
+
+        Supports: visit_url, click, type, scroll, web_search, screenshot,
+                  get_current_url, go_back, go_forward
+
+        Args:
+            command: The browser command to execute
+            params: Command parameters
+
+        Returns:
+            Dict containing the command result
+        """
+        assert self.interface is not None
+        return await self.interface.playwright_exec(command, params or {})

libs/python/agent/cua_agent/computers/custom.py

@@ -0,0 +1,216 @@
+"""
+Custom computer handler implementation that accepts a dictionary of functions.
+"""
+
+import base64
+import io
+from typing import Any, Callable, Dict, List, Literal, Optional, Union
+
+from PIL import Image
+
+from .base import AsyncComputerHandler
+
+
+class CustomComputerHandler(AsyncComputerHandler):
+    """Computer handler that implements the Computer protocol using a dictionary of custom functions."""
+
+    def __init__(self, functions: Dict[str, Callable]):
+        """
+        Initialize with a dictionary of functions.
+
+        Args:
+            functions: Dictionary where keys are method names and values are callable functions.
+                      Only 'screenshot' is required, all others are optional.
+
+        Raises:
+            ValueError: If required 'screenshot' function is not provided.
+        """
+        if "screenshot" not in functions:
+            raise ValueError("'screenshot' function is required in functions dictionary")
+
+        self.functions = functions
+        self._last_screenshot_size: Optional[tuple[int, int]] = None
+
+    async def _call_function(self, func, *args, **kwargs):
+        """
+        Call a function, handling both async and sync functions.
+
+        Args:
+            func: The function to call
+            *args: Positional arguments to pass to the function
+            **kwargs: Keyword arguments to pass to the function
+
+        Returns:
+            The result of the function call
+        """
+        import asyncio
+        import inspect
+
+        if callable(func):
+            if inspect.iscoroutinefunction(func):
+                return await func(*args, **kwargs)
+            else:
+                return func(*args, **kwargs)
+        else:
+            return func
+
+    async def _get_value(self, attribute: str):
+        """
+        Get value for an attribute, checking both 'get_{attribute}' and '{attribute}' keys.
+
+        Args:
+            attribute: The attribute name to look for
+
+        Returns:
+            The value from the functions dict, called if callable, returned directly if not
+        """
+        # Check for 'get_{attribute}' first
+        get_key = f"get_{attribute}"
+        if get_key in self.functions:
+            return await self._call_function(self.functions[get_key])
+
+        # Check for '{attribute}'
+        if attribute in self.functions:
+            return await self._call_function(self.functions[attribute])
+
+        return None
+
+    def _to_b64_str(self, img: Union[bytes, Image.Image, str]) -> str:
+        """
+        Convert image to base64 string.
+
+        Args:
+            img: Image as bytes, PIL Image, or base64 string
+
+        Returns:
+            str: Base64 encoded image string
+        """
+        if isinstance(img, str):
+            # Already a base64 string
+            return img
+        elif isinstance(img, bytes):
+            # Raw bytes
+            return base64.b64encode(img).decode("utf-8")
+        elif isinstance(img, Image.Image):
+            # PIL Image
+            buffer = io.BytesIO()
+            img.save(buffer, format="PNG")
+            return base64.b64encode(buffer.getvalue()).decode("utf-8")
+        else:
+            raise ValueError(f"Unsupported image type: {type(img)}")
+
+    # ==== Computer-Use-Preview Action Space ====
+
+    async def get_environment(self) -> Literal["windows", "mac", "linux", "browser"]:
+        """Get the current environment type."""
+        result = await self._get_value("environment")
+        if result is None:
+            return "linux"
+        assert result in ["windows", "mac", "linux", "browser"]
+        return result  # type: ignore
+
+    async def get_dimensions(self) -> tuple[int, int]:
+        """Get screen dimensions as (width, height)."""
+        result = await self._get_value("dimensions")
+        if result is not None:
+            return result  # type: ignore
+
+        # Fallback: use last screenshot size if available
+        if not self._last_screenshot_size:
+            await self.screenshot()
+        assert self._last_screenshot_size is not None, "Failed to get screenshot size"
+
+        return self._last_screenshot_size
+
+    async def screenshot(self, text: Optional[str] = None) -> str:
+        """Take a screenshot and return as base64 string.
+
+        Args:
+            text: Optional descriptive text (for compatibility with GPT-4o models, ignored)
+        """
+        result = await self._call_function(self.functions["screenshot"])
+        b64_str = self._to_b64_str(result)  # type: ignore
+
+        # Try to extract dimensions for fallback use
+        try:
+            if isinstance(result, Image.Image):
+                self._last_screenshot_size = result.size
+            elif isinstance(result, bytes):
+                # Try to decode bytes to get dimensions
+                img = Image.open(io.BytesIO(result))
+                self._last_screenshot_size = img.size
+        except Exception:
+            # If we can't get dimensions, that's okay
+            pass
+
+        return b64_str
+
+    async def click(self, x: int, y: int, button: str = "left") -> None:
+        """Click at coordinates with specified button."""
+        if "click" in self.functions:
+            await self._call_function(self.functions["click"], x, y, button)
+        # No-op if not implemented
+
+    async def double_click(self, x: int, y: int) -> None:
+        """Double click at coordinates."""
+        if "double_click" in self.functions:
+            await self._call_function(self.functions["double_click"], x, y)
+        # No-op if not implemented
+
+    async def scroll(self, x: int, y: int, scroll_x: int, scroll_y: int) -> None:
+        """Scroll at coordinates with specified scroll amounts."""
+        if "scroll" in self.functions:
+            await self._call_function(self.functions["scroll"], x, y, scroll_x, scroll_y)
+        # No-op if not implemented
+
+    async def type(self, text: str) -> None:
+        """Type text."""
+        if "type" in self.functions:
+            await self._call_function(self.functions["type"], text)
+        # No-op if not implemented
+
+    async def wait(self, ms: int = 1000) -> None:
+        """Wait for specified milliseconds."""
+        if "wait" in self.functions:
+            await self._call_function(self.functions["wait"], ms)
+        else:
+            # Default implementation
+            import asyncio
+
+            await asyncio.sleep(ms / 1000.0)
+
+    async def move(self, x: int, y: int) -> None:
+        """Move cursor to coordinates."""
+        if "move" in self.functions:
+            await self._call_function(self.functions["move"], x, y)
+        # No-op if not implemented
+
+    async def keypress(self, keys: Union[List[str], str]) -> None:
+        """Press key combination."""
+        if "keypress" in self.functions:
+            await self._call_function(self.functions["keypress"], keys)
+        # No-op if not implemented
+
+    async def drag(self, path: List[Dict[str, int]]) -> None:
+        """Drag along specified path."""
+        if "drag" in self.functions:
+            await self._call_function(self.functions["drag"], path)
+        # No-op if not implemented
+
+    async def get_current_url(self) -> str:
+        """Get current URL (for browser environments)."""
+        if "get_current_url" in self.functions:
+            return await self._get_value("current_url")  # type: ignore
+        return ""  # Default fallback
+
+    async def left_mouse_down(self, x: Optional[int] = None, y: Optional[int] = None) -> None:
+        """Left mouse down at coordinates."""
+        if "left_mouse_down" in self.functions:
+            await self._call_function(self.functions["left_mouse_down"], x, y)
+        # No-op if not implemented
+
+    async def left_mouse_up(self, x: Optional[int] = None, y: Optional[int] = None) -> None:
+        """Left mouse up at coordinates."""
+        if "left_mouse_up" in self.functions:
+            await self._call_function(self.functions["left_mouse_up"], x, y)
+        # No-op if not implemented

libs/python/agent/cua_agent/computers/sandbox.py

@@ -0,0 +1,125 @@
+"""
+Computer handler implementation wrapping a cua_sandbox.Sandbox instance.
+"""
+
+import asyncio
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from .base import AsyncComputerHandler
+
+
+class SandboxComputerHandler(AsyncComputerHandler):
+    """Computer handler that adapts a cua_sandbox.Sandbox to the AsyncComputerHandler protocol."""
+
+    def __init__(self, sandbox: Any):
+        self._sandbox = sandbox
+
+    # ==== Computer-Use-Preview Action Space ====
+
+    async def get_environment(self) -> Literal["windows", "mac", "linux", "browser"]:
+        return await self._sandbox.get_environment()
+
+    async def get_dimensions(self) -> tuple[int, int]:
+        return await self._sandbox.get_dimensions()
+
+    async def screenshot(self, text: Optional[str] = None) -> str:
+        return await self._sandbox.screenshot_base64()
+
+    async def click(self, x: int, y: int, button: str = "left") -> None:
+        if button == "right":
+            await self._sandbox.mouse.right_click(x, y)
+        else:
+            await self._sandbox.mouse.click(x, y, button=button)
+
+    async def double_click(self, x: int, y: int) -> None:
+        await self._sandbox.mouse.double_click(x, y)
+
+    async def right_click(self, x: int, y: int) -> None:
+        await self._sandbox.mouse.right_click(x, y)
+
+    async def scroll(self, x: int, y: int, scroll_x: int, scroll_y: int) -> None:
+        await self._sandbox.mouse.scroll(x, y, scroll_x=scroll_x, scroll_y=scroll_y)
+
+    async def type(self, text: str) -> None:
+        await self._sandbox.keyboard.type(text)
+
+    async def wait(self, ms: int = 1000) -> None:
+        await asyncio.sleep(ms / 1000.0)
+
+    async def move(self, x: int, y: int) -> None:
+        await self._sandbox.mouse.move(x, y)
+
+    # Maps Anthropic/X11 key names → pynput Key attribute names used by computer-server
+    _KEY_NAME_MAP = {
+        "return": "enter",
+        "backspace": "backspace",
+        "delete": "delete",
+        "del": "delete",
+        "escape": "esc",
+        "esc": "esc",
+        "tab": "tab",
+        "space": "space",
+        " ": "space",
+        "ctrl": "ctrl",
+        "control": "ctrl",
+        "shift": "shift",
+        "alt": "alt",
+        "super": "cmd",
+        "meta": "cmd",
+        "cmd": "cmd",
+        "command": "cmd",
+        "win": "cmd",
+        "up": "up",
+        "down": "down",
+        "left": "left",
+        "right": "right",
+        "home": "home",
+        "end": "end",
+        "pageup": "page_up",
+        "page_up": "page_up",
+        "pgup": "page_up",
+        "pagedown": "page_down",
+        "page_down": "page_down",
+        "pgdn": "page_down",
+        "insert": "insert",
+        "ins": "insert",
+        "caps_lock": "caps_lock",
+        **{f"f{i}": f"f{i}" for i in range(1, 21)},
+    }
+
+    async def keypress(self, keys: Union[List[str], str]) -> None:
+        if isinstance(keys, str):
+            keys = [keys]
+        normalized = []
+        for k in keys:
+            mapped = self._KEY_NAME_MAP.get(k.lower(), k.lower() if len(k) > 1 else k)
+            normalized.append(mapped)
+        await self._sandbox.keyboard.keypress(normalized)
+
+    async def drag(
+        self,
+        path: Optional[List[Dict[str, int]]] = None,
+        start_x: Optional[int] = None,
+        start_y: Optional[int] = None,
+        end_x: Optional[int] = None,
+        end_y: Optional[int] = None,
+    ) -> None:
+        if start_x is not None and start_y is not None and end_x is not None and end_y is not None:
+            path = [{"x": start_x, "y": start_y}, {"x": end_x, "y": end_y}]
+        if not path:
+            return
+        await self._sandbox.mouse.drag(path)
+
+    async def get_current_url(self) -> str:
+        return ""
+
+    async def terminate(self, status: str = "success") -> Dict[str, Any]:
+        return {"success": True, "status": status, "terminated": True}
+
+    # ==== Anthropic Action Space ====
+
+    async def left_mouse_down(self, x: Optional[int] = None, y: Optional[int] = None) -> None:
+        await self._sandbox.mouse.mouse_down(x, y, button="left")
+
+    async def left_mouse_up(self, x: Optional[int] = None, y: Optional[int] = None) -> None:
+        await self._sandbox.mouse.mouse_up(x, y, button="left")

libs/python/agent/cua_agent/decorators.py

@@ -0,0 +1,93 @@
+"""
+Decorators for agent - agent_loop decorator
+"""
+
+from typing import List, Optional
+
+from .types import AgentConfigInfo
+
+# Global registry
+_agent_configs: List[AgentConfigInfo] = []
+
+
+def register_agent(models: str, priority: int = 0, tool_type: Optional[str] = None):
+    """
+    Decorator to register an AsyncAgentConfig class.
+
+    Args:
+        models: Regex pattern to match supported models
+        priority: Priority for agent selection (higher = more priority)
+        tool_type: Required tool type for this model ("browser" | "mobile" | None).
+                   Specialized models (like FARA) declare their required tool type,
+                   and ComputerAgent will auto-wrap tools accordingly.
+                   General models (like Claude) leave this as None for full flexibility.
+    """
+
+    def decorator(agent_class: type):
+        # Validate that the class implements AsyncAgentConfig protocol
+        if not hasattr(agent_class, "predict_step"):
+            raise ValueError(
+                f"Agent class {agent_class.__name__} must implement predict_step method"
+            )
+        if not hasattr(agent_class, "predict_click"):
+            raise ValueError(
+                f"Agent class {agent_class.__name__} must implement predict_click method"
+            )
+        if not hasattr(agent_class, "get_capabilities"):
+            raise ValueError(
+                f"Agent class {agent_class.__name__} must implement get_capabilities method"
+            )
+
+        # Register the agent config
+        config_info = AgentConfigInfo(
+            agent_class=agent_class,
+            models_regex=models,
+            priority=priority,
+            tool_type=tool_type,
+        )
+        _agent_configs.append(config_info)
+
+        # Sort by priority (highest first)
+        _agent_configs.sort(key=lambda x: x.priority, reverse=True)
+
+        return agent_class
+
+    return decorator
+
+
+def get_agent_configs() -> List[AgentConfigInfo]:
+    """Get all registered agent configs"""
+    return _agent_configs.copy()
+
+
+def _strip_cua_prefix(model: str) -> str:
+    """Strip the ``cua/<provider>/`` routing prefix so the bare model name
+    can be matched against registered agent patterns.
+
+    Examples:
+        cua/google/gemini-3-flash-preview  ->  gemini-3-flash-preview
+        cua/anthropic/claude-sonnet-4-6    ->  claude-sonnet-4-6
+        gemini-3-flash-preview             ->  gemini-3-flash-preview  (unchanged)
+    """
+    parts = model.split("/")
+    if parts[0] == "cua" and len(parts) >= 3:
+        return "/".join(parts[2:])
+    return model
+
+
+def find_agent_config(model: str) -> Optional[AgentConfigInfo]:
+    """Find the best matching agent config for a model.
+
+    For each registered config (checked in priority order), tries the
+    original model string first and then the bare model name with the
+    ``cua/<provider>/`` routing prefix stripped.  This ensures that
+    routed models (e.g. ``cua/google/gemini-3-flash-preview``) resolve
+    to the same agent loop as their bare counterparts.
+    """
+    stripped = _strip_cua_prefix(model)
+    for config_info in _agent_configs:
+        if config_info.matches_model(model):
+            return config_info
+        if stripped != model and config_info.matches_model(stripped):
+            return config_info
+    return None

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

@@ -0,0 +1,24 @@
+"""
+Human-in-the-Loop Completion Tool
+
+This package provides a human-in-the-loop completion system that allows
+AI agents to request human assistance for complex decisions or responses.
+
+Components:
+- server.py: FastAPI server with completion queue management
+- ui.py: Gradio UI for human interaction
+- __main__.py: Combined server and UI application
+
+Usage:
+    # Run the server and UI
+    python -m agent.human_tool
+
+    # Or run components separately
+    python -m agent.human_tool.server  # API server only
+    python -m agent.human_tool.ui      # UI only
+"""
+
+from .server import CompletionQueue, completion_queue
+from .ui import HumanCompletionUI, create_ui
+
+__all__ = ["CompletionQueue", "completion_queue", "HumanCompletionUI", "create_ui"]

libs/python/agent/cua_agent/human_tool/__main__.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+"""
+Human-in-the-Loop Completion Server and UI
+
+This module combines the FastAPI server for handling completion requests
+with a Gradio UI for human interaction.
+"""
+
+import gradio as gr
+from fastapi import FastAPI
+
+from .server import app as fastapi_app
+from .ui import create_ui
+
+# Create the Gradio demo
+gradio_demo = create_ui()
+
+# Mount Gradio on FastAPI
+CUSTOM_PATH = "/gradio"
+app = gr.mount_gradio_app(fastapi_app, gradio_demo, path=CUSTOM_PATH)
+
+
+# Add a redirect from root to Gradio UI
+@fastapi_app.get("/")
+async def redirect_to_ui():
+    """Redirect root to Gradio UI."""
+    return {
+        "message": "Human Completion Server is running",
+        "ui_url": "/gradio",
+        "api_docs": "/docs",
+    }
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    print("🚀 Starting Human-in-the-Loop Completion Server...")
+    print("📊 API Server: http://localhost:8002")
+    print("🎨 Gradio UI: http://localhost:8002/gradio")
+    print("📚 API Docs: http://localhost:8002/docs")
+
+    uvicorn.run(app, host="0.0.0.0", port=8002)

libs/python/agent/cua_agent/human_tool/server.py

@@ -0,0 +1,245 @@
+import asyncio
+import uuid
+from dataclasses import asdict, dataclass
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+
+class CompletionStatus(str, Enum):
+    PENDING = "pending"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+@dataclass
+class CompletionCall:
+    id: str
+    messages: List[Dict[str, Any]]
+    model: str
+    status: CompletionStatus
+    created_at: datetime
+    completed_at: Optional[datetime] = None
+    response: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    error: Optional[str] = None
+
+
+class ToolCall(BaseModel):
+    id: str
+    type: str = "function"
+    function: Dict[str, Any]
+
+
+class CompletionRequest(BaseModel):
+    messages: List[Dict[str, Any]]
+    model: str
+
+
+class CompletionResponse(BaseModel):
+    response: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+
+
+class CompletionQueue:
+    def __init__(self):
+        self._queue: Dict[str, CompletionCall] = {}
+        self._pending_order: List[str] = []
+        self._lock = asyncio.Lock()
+
+    async def add_completion(self, messages: List[Dict[str, Any]], model: str) -> str:
+        """Add a completion call to the queue."""
+        async with self._lock:
+            call_id = str(uuid.uuid4())
+            completion_call = CompletionCall(
+                id=call_id,
+                messages=messages,
+                model=model,
+                status=CompletionStatus.PENDING,
+                created_at=datetime.now(),
+            )
+            self._queue[call_id] = completion_call
+            self._pending_order.append(call_id)
+            return call_id
+
+    async def get_pending_calls(self) -> List[Dict[str, Any]]:
+        """Get all pending completion calls."""
+        async with self._lock:
+            pending_calls = []
+            for call_id in self._pending_order:
+                if (
+                    call_id in self._queue
+                    and self._queue[call_id].status == CompletionStatus.PENDING
+                ):
+                    call = self._queue[call_id]
+                    pending_calls.append(
+                        {
+                            "id": call.id,
+                            "model": call.model,
+                            "created_at": call.created_at.isoformat(),
+                            "messages": call.messages,
+                        }
+                    )
+            return pending_calls
+
+    async def get_call_status(self, call_id: str) -> Optional[Dict[str, Any]]:
+        """Get the status of a specific completion call."""
+        async with self._lock:
+            if call_id not in self._queue:
+                return None
+
+            call = self._queue[call_id]
+            result = {
+                "id": call.id,
+                "status": call.status.value,
+                "created_at": call.created_at.isoformat(),
+                "model": call.model,
+                "messages": call.messages,
+            }
+
+            if call.completed_at:
+                result["completed_at"] = call.completed_at.isoformat()
+            if call.response:
+                result["response"] = call.response
+            if call.tool_calls:
+                result["tool_calls"] = call.tool_calls
+            if call.error:
+                result["error"] = call.error
+
+            return result
+
+    async def complete_call(
+        self,
+        call_id: str,
+        response: Optional[str] = None,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ) -> bool:
+        """Mark a completion call as completed with a response or tool calls."""
+        async with self._lock:
+            if call_id not in self._queue:
+                return False
+
+            call = self._queue[call_id]
+            if call.status != CompletionStatus.PENDING:
+                return False
+
+            call.status = CompletionStatus.COMPLETED
+            call.completed_at = datetime.now()
+            call.response = response
+            call.tool_calls = tool_calls
+
+            # Remove from pending order
+            if call_id in self._pending_order:
+                self._pending_order.remove(call_id)
+
+            return True
+
+    async def fail_call(self, call_id: str, error: str) -> bool:
+        """Mark a completion call as failed with an error."""
+        async with self._lock:
+            if call_id not in self._queue:
+                return False
+
+            call = self._queue[call_id]
+            if call.status != CompletionStatus.PENDING:
+                return False
+
+            call.status = CompletionStatus.FAILED
+            call.completed_at = datetime.now()
+            call.error = error
+
+            # Remove from pending order
+            if call_id in self._pending_order:
+                self._pending_order.remove(call_id)
+
+            return True
+
+    async def wait_for_completion(self, call_id: str, timeout: float = 300.0) -> Optional[str]:
+        """Wait for a completion call to be completed and return the response."""
+        start_time = asyncio.get_event_loop().time()
+
+        while True:
+            status = await self.get_call_status(call_id)
+            if not status:
+                return None
+
+            if status["status"] == CompletionStatus.COMPLETED.value:
+                return status.get("response")
+            elif status["status"] == CompletionStatus.FAILED.value:
+                raise Exception(f"Completion failed: {status.get('error', 'Unknown error')}")
+
+            # Check timeout
+            if asyncio.get_event_loop().time() - start_time > timeout:
+                await self.fail_call(call_id, "Timeout waiting for human response")
+                raise TimeoutError("Timeout waiting for human response")
+
+            # Wait a bit before checking again
+            await asyncio.sleep(0.5)
+
+
+# Global queue instance
+completion_queue = CompletionQueue()
+
+# FastAPI app
+app = FastAPI(title="Human Completion Server", version="1.0.0")
+
+
+@app.post("/queue", response_model=Dict[str, str])
+async def queue_completion(request: CompletionRequest):
+    """Add a completion request to the queue."""
+    call_id = await completion_queue.add_completion(request.messages, request.model)
+    return {"id": call_id, "status": "queued"}
+
+
+@app.get("/pending")
+async def list_pending():
+    """List all pending completion calls."""
+    pending_calls = await completion_queue.get_pending_calls()
+    return {"pending_calls": pending_calls}
+
+
+@app.get("/status/{call_id}")
+async def get_status(call_id: str):
+    """Get the status of a specific completion call."""
+    status = await completion_queue.get_call_status(call_id)
+    if not status:
+        raise HTTPException(status_code=404, detail="Completion call not found")
+    return status
+
+
+@app.post("/complete/{call_id}")
+async def complete_call(call_id: str, response: CompletionResponse):
+    """Complete a call with a human response."""
+    success = await completion_queue.complete_call(
+        call_id, response=response.response, tool_calls=response.tool_calls
+    )
+    if success:
+        return {"status": "success", "message": "Call completed"}
+    else:
+        raise HTTPException(status_code=404, detail="Call not found or already completed")
+
+
+@app.post("/fail/{call_id}")
+async def fail_call(call_id: str, error: Dict[str, str]):
+    """Mark a call as failed."""
+    success = await completion_queue.fail_call(call_id, error.get("error", "Unknown error"))
+    if not success:
+        raise HTTPException(
+            status_code=404, detail="Completion call not found or already completed"
+        )
+    return {"status": "failed"}
+
+
+@app.get("/")
+async def root():
+    """Root endpoint."""
+    return {"message": "Human Completion Server is running"}
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8002)

libs/python/agent/cua_agent/human_tool/ui.py

@@ -0,0 +1,754 @@
+import base64
+import io
+import json
+import time
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import gradio as gr
+import requests
+from PIL import Image
+
+from .server import completion_queue
+
+
+class HumanCompletionUI:
+    def __init__(self, server_url: str = "http://localhost:8002"):
+        self.server_url = server_url
+        self.current_call_id: Optional[str] = None
+        self.refresh_interval = 2.0  # seconds
+        self.last_image = None  # Store the last image for display
+        # Track current interactive action controls
+        self.current_action_type: str = "click"
+        self.current_button: str = "left"
+        self.current_scroll_x: int = 0
+        self.current_scroll_y: int = -120
+
+    def format_messages_for_chatbot(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Format messages for display in gr.Chatbot with type='messages'."""
+        formatted = []
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            tool_calls = msg.get("tool_calls", [])
+
+            # Handle different content formats
+            if isinstance(content, list):
+                # Multi-modal content - can include text and images
+                formatted_content = []
+                for item in content:
+                    if item.get("type") == "text":
+                        text = item.get("text", "")
+                        if text.strip():  # Only add non-empty text
+                            formatted_content.append(text)
+                    elif item.get("type") == "image_url":
+                        image_url = item.get("image_url", {}).get("url", "")
+                        if image_url:
+                            # Check if it's a base64 image or URL
+                            if image_url.startswith("data:image"):
+                                # For base64 images, decode and create gr.Image
+                                try:
+                                    header, data = image_url.split(",", 1)
+                                    image_data = base64.b64decode(data)
+                                    image = Image.open(io.BytesIO(image_data))
+                                    formatted_content.append(gr.Image(value=image))
+                                except Exception as e:
+                                    print(f"Error loading image: {e}")
+                                    formatted_content.append(f"[Image loading error: {e}]")
+                            else:
+                                # For URL images, create gr.Image with URL
+                                formatted_content.append(gr.Image(value=image_url))
+
+                # Determine final content format
+                if len(formatted_content) == 1:
+                    content = formatted_content[0]
+                elif len(formatted_content) > 1:
+                    content = formatted_content
+                else:
+                    content = "[Empty content]"
+
+            # Ensure role is valid for Gradio Chatbot
+            if role not in ["user", "assistant"]:
+                role = "assistant" if role == "system" else "user"
+
+            # Invert roles for better display in human UI context
+            # (what the AI says becomes "user", what human should respond becomes "assistant")
+            if role == "user":
+                role = "assistant"
+            else:
+                role = "user"
+
+            # Add the main message if it has content
+            if content and str(content).strip():
+                formatted.append({"role": role, "content": content})
+
+            # Handle tool calls - create separate messages for each tool call
+            if tool_calls:
+                for tool_call in tool_calls:
+                    function_name = tool_call.get("function", {}).get("name", "unknown")
+                    arguments_str = tool_call.get("function", {}).get("arguments", "{}")
+
+                    try:
+                        # Parse arguments to format them nicely
+                        arguments = json.loads(arguments_str)
+                        formatted_args = json.dumps(arguments, indent=2)
+                    except json.JSONDecodeError:
+                        # If parsing fails, use the raw string
+                        formatted_args = arguments_str
+
+                    # Create a formatted message for the tool call
+                    tool_call_content = f"```json\n{formatted_args}\n```"
+
+                    formatted.append(
+                        {
+                            "role": role,
+                            "content": tool_call_content,
+                            "metadata": {"title": f"🛠️ Used {function_name}"},
+                        }
+                    )
+
+        return formatted
+
+    def get_pending_calls(self) -> List[Dict[str, Any]]:
+        """Get pending calls from the server."""
+        try:
+            response = requests.get(f"{self.server_url}/pending", timeout=5)
+            if response.status_code == 200:
+                return response.json().get("pending_calls", [])
+        except Exception as e:
+            print(f"Error fetching pending calls: {e}")
+        return []
+
+    def complete_call_with_response(self, call_id: str, response: str) -> bool:
+        """Complete a call with a text response."""
+        try:
+            response_data = {"response": response}
+            response_obj = requests.post(
+                f"{self.server_url}/complete/{call_id}", json=response_data, timeout=10
+            )
+            response_obj.raise_for_status()
+            return True
+        except requests.RequestException as e:
+            print(f"Error completing call: {e}")
+            return False
+
+    def complete_call_with_tool_calls(self, call_id: str, tool_calls: List[Dict[str, Any]]) -> bool:
+        """Complete a call with tool calls."""
+        try:
+            response_data = {"tool_calls": tool_calls}
+            response_obj = requests.post(
+                f"{self.server_url}/complete/{call_id}", json=response_data, timeout=10
+            )
+            response_obj.raise_for_status()
+            return True
+        except requests.RequestException as e:
+            print(f"Error completing call: {e}")
+            return False
+
+    def complete_call(
+        self,
+        call_id: str,
+        response: Optional[str] = None,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ) -> bool:
+        """Complete a call with either a response or tool calls."""
+        try:
+            response_data = {}
+            if response:
+                response_data["response"] = response
+            if tool_calls:
+                response_data["tool_calls"] = tool_calls
+
+            response_obj = requests.post(
+                f"{self.server_url}/complete/{call_id}", json=response_data, timeout=10
+            )
+            response_obj.raise_for_status()
+            return True
+        except requests.RequestException as e:
+            print(f"Error completing call: {e}")
+            return False
+
+    def get_last_image_from_messages(self, messages: List[Dict[str, Any]]) -> Optional[Any]:
+        """Extract the last image from the messages for display above conversation."""
+        last_image = None
+
+        for msg in reversed(messages):  # Start from the last message
+            content = msg.get("content", "")
+
+            if isinstance(content, list):
+                for item in reversed(content):  # Get the last image in the message
+                    if item.get("type") == "image_url":
+                        image_url = item.get("image_url", {}).get("url", "")
+                        if image_url:
+                            if image_url.startswith("data:image"):
+                                # For base64 images, create a gr.Image component
+                                try:
+                                    header, data = image_url.split(",", 1)
+                                    image_data = base64.b64decode(data)
+                                    image = Image.open(io.BytesIO(image_data))
+                                    return image
+                                except Exception as e:
+                                    print(f"Error loading image: {e}")
+                                    continue
+                            else:
+                                # For URL images, return the URL
+                                return image_url
+
+        return last_image
+
+    def refresh_pending_calls(self):
+        """Refresh the list of pending calls."""
+        pending_calls = self.get_pending_calls()
+
+        if not pending_calls:
+            return (
+                gr.update(choices=["latest"], value="latest"),  # dropdown
+                gr.update(value=None),  # image (no image)
+                gr.update(value=[]),  # chatbot (empty messages)
+                gr.update(interactive=False),  # submit button
+                gr.update(visible=False),  # click_actions_group hidden
+                gr.update(visible=False),  # actions_group hidden
+            )
+
+        # Sort pending calls by created_at to get oldest first
+        sorted_calls = sorted(pending_calls, key=lambda x: x.get("created_at", ""))
+
+        # Create choices for dropdown
+        choices = [("latest", "latest")]  # Add "latest" option first
+
+        for call in sorted_calls:
+            call_id = call["id"]
+            model = call.get("model", "unknown")
+            created_at = call.get("created_at", "")
+            # Format timestamp
+            try:
+                dt = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
+                time_str = dt.strftime("%H:%M:%S")
+            except:
+                time_str = created_at
+
+            choice_label = f"{call_id[:8]}... ({model}) - {time_str}"
+            choices.append((choice_label, call_id))
+
+        # Default to "latest" which shows the oldest pending conversation
+        selected_call_id = "latest"
+        if selected_call_id == "latest" and sorted_calls:
+            # Use the oldest call (first in sorted list)
+            selected_call = sorted_calls[0]
+            conversation = self.format_messages_for_chatbot(selected_call.get("messages", []))
+            self.current_call_id = selected_call["id"]
+            # Get the last image from messages
+            self.last_image = self.get_last_image_from_messages(selected_call.get("messages", []))
+        else:
+            conversation = []
+            self.current_call_id = None
+            self.last_image = None
+
+        return (
+            gr.update(choices=choices, value="latest"),
+            gr.update(value=self.last_image),
+            gr.update(value=conversation),
+            gr.update(interactive=bool(choices)),
+            gr.update(visible=True),  # click_actions_group visible when there is a call
+            gr.update(visible=True),  # actions_group visible when there is a call
+        )
+
+    def on_call_selected(self, selected_choice):
+        """Handle when a call is selected from the dropdown."""
+        if not selected_choice:
+            return (
+                gr.update(value=None),  # no image
+                gr.update(value=[]),  # empty chatbot
+                gr.update(interactive=False),
+                gr.update(visible=False),  # click_actions_group hidden
+                gr.update(visible=False),  # actions_group hidden
+            )
+
+        pending_calls = self.get_pending_calls()
+        if not pending_calls:
+            return (
+                gr.update(value=None),  # no image
+                gr.update(value=[]),  # empty chatbot
+                gr.update(interactive=False),
+                gr.update(visible=False),  # click_actions_group hidden
+                gr.update(visible=False),  # actions_group hidden
+            )
+
+        # Handle "latest" option
+        if selected_choice == "latest":
+            # Sort calls by created_at to get oldest first
+            sorted_calls = sorted(pending_calls, key=lambda x: x.get("created_at", ""))
+            selected_call = sorted_calls[0]  # Get the oldest call
+            call_id = selected_call["id"]
+        else:
+            # Extract call_id from the choice for specific calls
+            call_id = None
+            for call in pending_calls:
+                call_id_short = call["id"][:8]
+                if call_id_short in selected_choice:
+                    call_id = call["id"]
+                    break
+
+            if not call_id:
+                return (
+                    gr.update(value=None),  # no image
+                    gr.update(value=[]),  # empty chatbot
+                    gr.update(interactive=False),
+                )
+
+            # Find the selected call
+            selected_call = next((c for c in pending_calls if c["id"] == call_id), None)
+
+        if not selected_call:
+            return (
+                gr.update(value=None),  # no image
+                gr.update(value=[]),  # empty chatbot
+                gr.update(interactive=False),
+                gr.update(visible=False),  # click_actions_group hidden
+                gr.update(visible=False),  # actions_group hidden
+            )
+
+        conversation = self.format_messages_for_chatbot(selected_call.get("messages", []))
+        self.current_call_id = call_id
+        # Get the last image from messages
+        self.last_image = self.get_last_image_from_messages(selected_call.get("messages", []))
+
+        return (
+            gr.update(value=self.last_image),
+            gr.update(value=conversation),
+            gr.update(interactive=True),
+            gr.update(visible=True),  # click_actions_group visible
+            gr.update(visible=True),  # actions_group visible
+        )
+
+    def submit_response(self, response_text: str):
+        """Submit a text response to the current call."""
+        if not self.current_call_id:
+            return (
+                gr.update(value=response_text),  # keep response text
+                gr.update(value="❌ No call selected"),  # status
+            )
+
+        if not response_text.strip():
+            return (
+                gr.update(value=response_text),  # keep response text
+                gr.update(value="❌ Response cannot be empty"),  # status
+            )
+
+        success = self.complete_call_with_response(self.current_call_id, response_text)
+
+        if success:
+            status_msg = "✅ Response submitted successfully!"
+            return (
+                gr.update(value=""),  # clear response text
+                gr.update(value=status_msg),  # status
+            )
+        else:
+            return (
+                gr.update(value=response_text),  # keep response text
+                gr.update(value="❌ Failed to submit response"),  # status
+            )
+
+    def submit_action(self, action_type: str, **kwargs) -> str:
+        """Submit a computer action as a tool call."""
+        if not self.current_call_id:
+            return "❌ No call selected"
+
+        import uuid
+
+        # Create tool call structure
+        action_data = {"type": action_type, **kwargs}
+        tool_call = {
+            "id": f"call_{uuid.uuid4().hex[:24]}",
+            "type": "function",
+            "function": {"name": "computer", "arguments": json.dumps(action_data)},
+        }
+
+        success = self.complete_call_with_tool_calls(self.current_call_id, [tool_call])
+
+        if success:
+            return f"✅ {action_type.capitalize()} action submitted as tool call"
+        else:
+            return f"❌ Failed to submit {action_type} action"
+
+    def submit_click_action(
+        self, x: int, y: int, action_type: str = "click", button: str = "left"
+    ) -> str:
+        """Submit a coordinate-based action."""
+        if action_type == "click":
+            return self.submit_action(action_type, x=x, y=y, button=button)
+        else:
+            return self.submit_action(action_type, x=x, y=y)
+
+    def submit_type_action(self, text: str) -> str:
+        """Submit a type action."""
+        return self.submit_action("type", text=text)
+
+    def submit_hotkey_action(self, keys: str) -> str:
+        """Submit a hotkey action."""
+        return self.submit_action("keypress", keys=keys)
+
+    def submit_wait_action(self) -> str:
+        """Submit a wait action with no kwargs."""
+        return self.submit_action("wait")
+
+    def submit_description_click(
+        self, description: str, action_type: str = "click", button: str = "left"
+    ) -> str:
+        """Submit a description-based action."""
+        if action_type == "click":
+            return self.submit_action(action_type, element_description=description, button=button)
+        else:
+            return self.submit_action(action_type, element_description=description)
+
+    def wait_for_pending_calls(self, max_seconds: float = 10.0, check_interval: float = 0.2):
+        """Wait for pending calls to appear or until max_seconds elapsed.
+
+        This method loops and checks for pending calls at regular intervals,
+        returning as soon as a pending call is found or the maximum wait time is reached.
+
+        Args:
+            max_seconds: Maximum number of seconds to wait
+            check_interval: How often to check for pending calls (in seconds)
+        """
+        import time
+
+        start_time = time.time()
+
+        while time.time() - start_time < max_seconds:
+            # Check if there are any pending calls
+            pending_calls = self.get_pending_calls()
+            if pending_calls:
+                # Found pending calls, return immediately
+                return self.refresh_pending_calls()
+
+            # Wait before checking again
+            time.sleep(check_interval)
+
+        # Max wait time reached, return current state
+        return self.refresh_pending_calls()
+
+
+def create_ui():
+    """Create the Gradio interface."""
+    ui_handler = HumanCompletionUI()
+
+    with gr.Blocks(title="Human-in-the-Loop Agent Tool", fill_width=True) as demo:
+        gr.Markdown("# 🤖 Human-in-the-Loop Agent Tool")
+        gr.Markdown("Review AI conversation requests and provide human responses.")
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                with gr.Group():
+                    screenshot_image = gr.Image(
+                        label="Interactive Screenshot", interactive=False, height=600
+                    )
+
+                    # Action type selection for image clicks (wrapped for visibility control)
+                    with gr.Group(visible=False) as click_actions_group:
+                        with gr.Row():
+                            action_type_radio = gr.Dropdown(
+                                label="Interactive Action",
+                                choices=[
+                                    "click",
+                                    "double_click",
+                                    "move",
+                                    "left_mouse_up",
+                                    "left_mouse_down",
+                                    "scroll",
+                                ],
+                                value="click",
+                                scale=2,
+                            )
+                            action_button_radio = gr.Dropdown(
+                                label="Button",
+                                choices=["left", "right", "wheel", "back", "forward"],
+                                value="left",
+                                visible=True,
+                                scale=1,
+                            )
+                            scroll_x_input = gr.Number(
+                                label="scroll_x", value=0, visible=False, scale=1
+                            )
+                            scroll_y_input = gr.Number(
+                                label="scroll_y", value=-120, visible=False, scale=1
+                            )
+
+                    conversation_chatbot = gr.Chatbot(
+                        label="Conversation", height=500, buttons=["copy"]
+                    )
+
+            with gr.Column(scale=1):
+                with gr.Group():
+                    call_dropdown = gr.Dropdown(
+                        label="Select a pending conversation request",
+                        choices=["latest"],
+                        interactive=True,
+                        value="latest",
+                    )
+                    refresh_btn = gr.Button("🔄 Refresh", variant="secondary")
+                    status_display = gr.Textbox(
+                        label="Status", interactive=False, value="Ready to receive requests..."
+                    )
+
+                with gr.Group():
+                    response_text = gr.Textbox(
+                        label="Message", lines=3, placeholder="Enter your message here..."
+                    )
+                    submit_btn = gr.Button(
+                        "📤 Submit Message", variant="primary", interactive=False
+                    )
+
+                # Action Accordions (wrapped for visibility control)
+                with gr.Group(visible=False) as actions_group:
+                    with gr.Tabs():
+                        with gr.Tab("🖱️ Click Actions"):
+                            with gr.Group():
+                                description_text = gr.Textbox(
+                                    label="Element Description",
+                                    placeholder="e.g., 'Privacy and security option in left sidebar'",
+                                )
+                                with gr.Row():
+                                    description_action_type = gr.Dropdown(
+                                        label="Action",
+                                        choices=[
+                                            "click",
+                                            "double_click",
+                                            "move",
+                                            "left_mouse_up",
+                                            "left_mouse_down",
+                                        ],
+                                        value="click",
+                                    )
+                                    description_button = gr.Dropdown(
+                                        label="Button",
+                                        choices=["left", "right", "wheel", "back", "forward"],
+                                        value="left",
+                                    )
+                                description_submit_btn = gr.Button("Submit Click Action")
+
+                        with gr.Tab("📝 Type Action"):
+                            with gr.Group():
+                                type_text = gr.Textbox(
+                                    label="Text to Type", placeholder="Enter text to type..."
+                                )
+                                type_submit_btn = gr.Button("Submit Type")
+
+                        with gr.Tab("⌨️ Keypress Action"):
+                            with gr.Group():
+                                keypress_text = gr.Textbox(
+                                    label="Keys", placeholder="e.g., ctrl+c, alt+tab"
+                                )
+                                keypress_submit_btn = gr.Button("Submit Keypress")
+
+                        with gr.Tab("🧰 Misc Actions"):
+                            with gr.Group():
+                                misc_action_dropdown = gr.Dropdown(
+                                    label="Action", choices=["wait"], value="wait"
+                                )
+                                misc_submit_btn = gr.Button("Submit Action")
+
+        # Event handlers
+        refresh_btn.click(
+            fn=ui_handler.refresh_pending_calls,
+            outputs=[
+                call_dropdown,
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+        call_dropdown.change(
+            fn=ui_handler.on_call_selected,
+            inputs=[call_dropdown],
+            outputs=[
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+        def handle_image_click(evt: gr.SelectData):
+            if evt.index is not None:
+                x, y = evt.index
+                action_type = ui_handler.current_action_type or "click"
+                button = ui_handler.current_button or "left"
+                if action_type == "scroll":
+                    sx_i = int(ui_handler.current_scroll_x or 0)
+                    sy_i = int(ui_handler.current_scroll_y or 0)
+                    # Submit a scroll action with x,y position and scroll deltas
+                    result = ui_handler.submit_action(
+                        "scroll", x=x, y=y, scroll_x=sx_i, scroll_y=sy_i
+                    )
+                else:
+                    result = ui_handler.submit_click_action(x, y, action_type, button)
+                ui_handler.wait_for_pending_calls()
+                return result
+            return "No coordinates selected"
+
+        screenshot_image.select(fn=handle_image_click, outputs=[status_display]).then(
+            fn=ui_handler.wait_for_pending_calls,
+            outputs=[
+                call_dropdown,
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+        # Response submission
+        submit_btn.click(
+            fn=ui_handler.submit_response,
+            inputs=[response_text],
+            outputs=[response_text, status_display],
+        ).then(
+            fn=ui_handler.refresh_pending_calls,
+            outputs=[
+                call_dropdown,
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+        # Toggle visibility of controls based on action type
+        def toggle_action_controls(action_type):
+            # Button visible only for click
+            button_vis = gr.update(visible=(action_type == "click"))
+            # Scroll inputs visible only for scroll
+            scroll_x_vis = gr.update(visible=(action_type == "scroll"))
+            scroll_y_vis = gr.update(visible=(action_type == "scroll"))
+            # Update state
+            ui_handler.current_action_type = action_type or "click"
+            return button_vis, scroll_x_vis, scroll_y_vis
+
+        action_type_radio.change(
+            fn=toggle_action_controls,
+            inputs=[action_type_radio],
+            outputs=[action_button_radio, scroll_x_input, scroll_y_input],
+        )
+
+        # Keep other control values in ui_handler state
+        def on_button_change(val):
+            ui_handler.current_button = val or "left"
+
+        action_button_radio.change(fn=on_button_change, inputs=[action_button_radio])
+
+        def on_scroll_x_change(val):
+            try:
+                ui_handler.current_scroll_x = int(val) if val is not None else 0
+            except Exception:
+                ui_handler.current_scroll_x = 0
+
+        scroll_x_input.change(fn=on_scroll_x_change, inputs=[scroll_x_input])
+
+        def on_scroll_y_change(val):
+            try:
+                ui_handler.current_scroll_y = int(val) if val is not None else 0
+            except Exception:
+                ui_handler.current_scroll_y = 0
+
+        scroll_y_input.change(fn=on_scroll_y_change, inputs=[scroll_y_input])
+
+        type_submit_btn.click(
+            fn=ui_handler.submit_type_action, inputs=[type_text], outputs=[status_display]
+        ).then(
+            fn=ui_handler.wait_for_pending_calls,
+            outputs=[
+                call_dropdown,
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+        keypress_submit_btn.click(
+            fn=ui_handler.submit_hotkey_action, inputs=[keypress_text], outputs=[status_display]
+        ).then(
+            fn=ui_handler.wait_for_pending_calls,
+            outputs=[
+                call_dropdown,
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+        def handle_description_submit(description, action_type, button):
+            if description:
+                result = ui_handler.submit_description_click(description, action_type, button)
+                ui_handler.wait_for_pending_calls()
+                return result
+            return "Please enter a description"
+
+        description_submit_btn.click(
+            fn=handle_description_submit,
+            inputs=[description_text, description_action_type, description_button],
+            outputs=[status_display],
+        ).then(
+            fn=ui_handler.wait_for_pending_calls,
+            outputs=[
+                call_dropdown,
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+        # Misc action handler
+        def handle_misc_submit(selected_action):
+            if selected_action == "wait":
+                result = ui_handler.submit_wait_action()
+                ui_handler.wait_for_pending_calls()
+                return result
+            return f"Unsupported misc action: {selected_action}"
+
+        misc_submit_btn.click(
+            fn=handle_misc_submit, inputs=[misc_action_dropdown], outputs=[status_display]
+        ).then(
+            fn=ui_handler.wait_for_pending_calls,
+            outputs=[
+                call_dropdown,
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+        # Load initial data
+        demo.load(
+            fn=ui_handler.refresh_pending_calls,
+            outputs=[
+                call_dropdown,
+                screenshot_image,
+                conversation_chatbot,
+                submit_btn,
+                click_actions_group,
+                actions_group,
+            ],
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = create_ui()
+    demo.queue()
+    demo.launch(server_name="0.0.0.0", server_port=7860)

libs/python/agent/cua_agent/integrations/hud/__init__.py

@@ -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",
+]

libs/python/agent/cua_agent/integrations/hud/agent.py

@@ -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",
+]

libs/python/agent/cua_agent/integrations/hud/proxy.py

@@ -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",
+]

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

@@ -0,0 +1,50 @@
+"""
+Agent loops for agent
+"""
+
+# Import the loops to register them
+from . import (
+    anthropic,
+    composed_grounded,
+    fara,
+    gelato,
+    gemini,
+    generic_vlm,
+    glm45v,
+    gta1,
+    holo,
+    internvl,
+    moondream3,
+    omniparser,
+    openai,
+    opencua,
+    qwen3vl,
+    qwen35,
+    uiins,
+    uitars,
+    uitars2,
+    yutori,
+)
+
+__all__ = [
+    "anthropic",
+    "composed_grounded",
+    "gelato",
+    "gemini",
+    "generic_vlm",
+    "fara",
+    "glm45v",
+    "gta1",
+    "holo",
+    "internvl",
+    "moondream3",
+    "omniparser",
+    "openai",
+    "opencua",
+    "qwen3vl",
+    "qwen35",
+    "uiins",
+    "uitars",
+    "uitars2",
+    "yutori",
+]

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

@@ -0,0 +1,80 @@
+"""
+Base protocol for async agent configurations
+"""
+
+from abc import abstractmethod
+from typing import Any, Dict, List, Optional, Protocol, Tuple, Union
+
+from ..types import AgentCapability
+
+
+class AsyncAgentConfig(Protocol):
+    """Protocol defining the interface for async agent configurations."""
+
+    @abstractmethod
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **generation_config,
+    ) -> Dict[str, Any]:
+        """
+        Predict the next step based on input items.
+
+        Args:
+            messages: Input items following Responses format (message, function_call, computer_call)
+            model: Model name to use
+            tools: Optional list of tool schemas
+            max_retries: Maximum number of retries for failed API calls
+            stream: Whether to stream responses
+            computer_handler: Computer handler instance
+            _on_api_start: Callback for API start
+            _on_api_end: Callback for API end
+            _on_usage: Callback for usage tracking
+            _on_screenshot: Callback for screenshot events
+            **generation_config: Additional arguments to pass to the model provider
+                - api_key: Optional API key for the provider
+                - api_base: Optional API base URL for the provider
+
+        Returns:
+            Dictionary with "output" (output items) and "usage" array
+        """
+        ...
+
+    @abstractmethod
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **generation_config
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates based on image and instruction.
+
+        Args:
+            model: Model name to use
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+            **generation_config: Additional arguments to pass to the model provider
+                - api_key: Optional API key for the provider
+                - api_base: Optional API base URL for the provider
+
+        Returns:
+            None or tuple with (x, y) coordinates
+        """
+        ...
+
+    @abstractmethod
+    def get_capabilities(self) -> List[AgentCapability]:
+        """
+        Get list of capabilities supported by this agent config.
+
+        Returns:
+            List of capability strings (e.g., ["step", "click"])
+        """
+        ...

libs/python/agent/cua_agent/loops/composed_grounded.py

@@ -0,0 +1,316 @@
+"""
+Composed-grounded agent loop implementation that combines grounding and thinking models.
+Uses a two-stage approach: grounding model for element detection, thinking model for reasoning.
+"""
+
+import asyncio
+import base64
+import json
+import uuid
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from PIL import Image
+
+from ..agent import find_agent_config
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_completion_messages_to_responses_items,
+    convert_computer_calls_desc2xy,
+    convert_computer_calls_xy2desc,
+    convert_responses_items_to_completion_messages,
+    get_all_element_descriptions,
+)
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+
+GROUNDED_COMPUTER_TOOL_SCHEMA = {
+    "type": "function",
+    "function": {
+        "name": "computer",
+        "description": "Control a computer by taking screenshots and interacting with UI elements. This tool uses element descriptions to locate and interact with UI elements on the screen (e.g., 'red submit button', 'search text field', 'hamburger menu icon', 'close button in top right corner').",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "type": "string",
+                    "enum": [
+                        "screenshot",
+                        "click",
+                        "double_click",
+                        "drag",
+                        "type",
+                        "keypress",
+                        "scroll",
+                        "move",
+                        "wait",
+                        "get_current_url",
+                        "get_dimensions",
+                        "get_environment",
+                    ],
+                    "description": "The action to perform (required for all actions)",
+                },
+                "element_description": {
+                    "type": "string",
+                    "description": "Description of the element to interact with (required for click, double_click, move, scroll actions)",
+                },
+                "start_element_description": {
+                    "type": "string",
+                    "description": "Description of the element to start dragging from (required for drag action)",
+                },
+                "end_element_description": {
+                    "type": "string",
+                    "description": "Description of the element to drag to (required for drag action)",
+                },
+                "text": {
+                    "type": "string",
+                    "description": "The text to type (required for type action)",
+                },
+                "keys": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                    "description": "Key(s) to press (required for keypress action)",
+                },
+                "button": {
+                    "type": "string",
+                    "enum": ["left", "right", "wheel", "back", "forward"],
+                    "description": "The mouse button to use for click action (required for click and double_click action)",
+                },
+                "scroll_x": {
+                    "type": "integer",
+                    "description": "Horizontal scroll amount for scroll action (required for scroll action)",
+                },
+                "scroll_y": {
+                    "type": "integer",
+                    "description": "Vertical scroll amount for scroll action (required for scroll action)",
+                },
+            },
+            "required": ["action"],
+        },
+    },
+}
+
+
+def _prepare_tools_for_grounded(tool_schemas: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Prepare tools for grounded API format"""
+    grounded_tools = []
+
+    for schema in tool_schemas:
+        if schema["type"] == "computer":
+            grounded_tools.append(GROUNDED_COMPUTER_TOOL_SCHEMA)
+        else:
+            grounded_tools.append(schema)
+
+    return grounded_tools
+
+
+def get_last_computer_call_image(messages: List[Dict[str, Any]]) -> Optional[str]:
+    """Get the last computer call output image from messages."""
+    for message in reversed(messages):
+        if (
+            isinstance(message, dict)
+            and message.get("type") == "computer_call_output"
+            and isinstance(message.get("output"), dict)
+            and message["output"].get("type") == "input_image"
+        ):
+            image_url = message["output"].get("image_url", "")
+            if image_url.startswith("data:image/png;base64,"):
+                return image_url.split(",", 1)[1]
+    return None
+
+
+@register_agent(r".*\+.*", priority=1)
+class ComposedGroundedConfig(AsyncAgentConfig):
+    """
+    Composed-grounded agent configuration that uses both grounding and thinking models.
+
+    The model parameter should be in format: "grounding_model+thinking_model"
+    e.g., "huggingface-local/HelloKKMe/GTA1-7B+gemini/gemini-1.5-pro"
+    """
+
+    def __init__(self):
+        self.desc2xy: Dict[str, Tuple[float, float]] = {}
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Composed-grounded predict step implementation.
+
+        Process:
+        0. Store last computer call image, if none then take a screenshot
+        1. Convert computer calls from xy to descriptions
+        2. Convert responses items to completion messages
+        3. Call thinking model with litellm.acompletion
+        4. Convert completion messages to responses items
+        5. Get all element descriptions and populate desc2xy mapping
+        6. Convert computer calls from descriptions back to xy coordinates
+        7. Return output and usage
+        """
+        # Parse the composed model
+        if "+" not in model:
+            raise ValueError(
+                f"Composed model must be in format 'grounding_model+thinking_model', got: {model}"
+            )
+        grounding_model, thinking_model = model.split("+", 1)
+
+        pre_output_items = []
+
+        # Step 0: Store last computer call image, if none then take a screenshot
+        last_image_b64 = get_last_computer_call_image(messages)
+        if last_image_b64 is None:
+            # Take a screenshot
+            screenshot_b64 = await computer_handler.screenshot()  # type: ignore
+            if screenshot_b64:
+
+                call_id = uuid.uuid4().hex
+                pre_output_items += [
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "type": "output_text",
+                                "text": "Taking a screenshot to see the current computer screen.",
+                            }
+                        ],
+                    },
+                    {
+                        "action": {"type": "screenshot"},
+                        "call_id": call_id,
+                        "status": "completed",
+                        "type": "computer_call",
+                    },
+                    {
+                        "type": "computer_call_output",
+                        "call_id": call_id,
+                        "output": {
+                            "type": "input_image",
+                            "image_url": f"data:image/png;base64,{screenshot_b64}",
+                        },
+                    },
+                ]
+                last_image_b64 = screenshot_b64
+
+                # Call screenshot callback if provided
+                if _on_screenshot:
+                    await _on_screenshot(screenshot_b64)
+
+        tool_schemas = _prepare_tools_for_grounded(tools)  # type: ignore
+
+        # Step 1: Convert computer calls from xy to descriptions
+        input_messages = messages + pre_output_items
+        messages_with_descriptions = convert_computer_calls_xy2desc(input_messages, self.desc2xy)
+
+        # Step 2: Convert responses items to completion messages
+        completion_messages = convert_responses_items_to_completion_messages(
+            messages_with_descriptions, allow_images_in_tool_results=False
+        )
+
+        # Step 3: Call thinking model with litellm.acompletion
+        api_kwargs = {
+            "model": thinking_model,
+            "messages": completion_messages,
+            "tools": tool_schemas,
+            "max_retries": max_retries,
+            "stream": stream,
+            **kwargs,
+        }
+
+        if use_prompt_caching:
+            api_kwargs["use_prompt_caching"] = use_prompt_caching
+
+        # Call API start hook
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        # Make the completion call
+        response = await litellm.acompletion(**api_kwargs)
+
+        # Call API end hook
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        # Extract usage information
+        usage = {
+            **response.usage.model_dump(),  # type: ignore
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        # Step 4: Convert completion messages back to responses items format
+        response_dict = response.model_dump()  # type: ignore
+        choice_messages = [choice["message"] for choice in response_dict["choices"]]
+        thinking_output_items = []
+
+        for choice_message in choice_messages:
+            thinking_output_items.extend(
+                convert_completion_messages_to_responses_items([choice_message])
+            )
+
+        # Step 5: Get all element descriptions and populate desc2xy mapping
+        element_descriptions = get_all_element_descriptions(thinking_output_items)
+
+        if element_descriptions and last_image_b64:
+            # Use grounding model to predict coordinates for each description
+            grounding_agent_conf = find_agent_config(grounding_model)
+            if grounding_agent_conf:
+                grounding_agent = grounding_agent_conf.agent_class()
+
+                for desc in element_descriptions:
+                    for _ in range(3):  # try 3 times
+                        coords = await grounding_agent.predict_click(
+                            model=grounding_model, image_b64=last_image_b64, instruction=desc
+                        )
+                        if coords:
+                            self.desc2xy[desc] = coords
+                            break
+
+        # Step 6: Convert computer calls from descriptions back to xy coordinates
+        final_output_items = convert_computer_calls_desc2xy(thinking_output_items, self.desc2xy)
+
+        # Step 7: Return output and usage
+        return {"output": pre_output_items + final_output_items, "usage": usage}
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using the grounding model.
+
+        For composed models, uses only the grounding model part for click prediction.
+        """
+        # Parse the composed model to get grounding model
+        if "+" not in model:
+            raise ValueError(
+                f"Composed model must be in format 'grounding_model+thinking_model', got: {model}"
+            )
+        grounding_model, thinking_model = model.split("+", 1)
+
+        # Find and use the grounding agent
+        grounding_agent_conf = find_agent_config(grounding_model)
+        if grounding_agent_conf:
+            grounding_agent = grounding_agent_conf.agent_class()
+            return await grounding_agent.predict_click(
+                model=grounding_model, image_b64=image_b64, instruction=instruction, **kwargs
+            )
+
+        return None
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """Return the capabilities supported by this agent."""
+        return ["click", "step"]

libs/python/agent/cua_agent/loops/fara/__init__.py

@@ -0,0 +1,8 @@
+"""
+FARA-7B agent loop implementation.
+Original implementation from Microsoft: https://github.com/microsoft/Fara
+"""
+
+from .config import FaraVlmConfig
+
+__all__ = ("FaraVlmConfig",)

libs/python/agent/cua_agent/loops/fara/config.py

@@ -0,0 +1,661 @@
+"""FARA VLM agent configuration."""
+
+from __future__ import annotations
+
+import ast
+import json
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from litellm.responses.litellm_completion_transformation.transformation import (
+    LiteLLMCompletionResponsesConfig,
+)
+
+from ...decorators import register_agent
+from ...loops.base import AsyncAgentConfig
+from ...responses import (
+    make_click_item,
+    make_double_click_item,
+    make_drag_item,
+    make_keypress_item,
+    make_move_item,
+    make_output_text_item,
+    make_reasoning_item,
+    make_screenshot_item,
+    make_scroll_item,
+    make_type_item,
+    make_wait_item,
+)
+from ...types import AgentCapability
+from .helpers import (
+    _convert_responses_items_to_fara_messages,
+    build_nous_system,
+    parse_tool_call_from_text,
+)
+
+
+def _scale_fara_coordinates(
+    args: Dict[str, Any],
+    original_dims: Tuple[int, int],
+    resized_dims: Tuple[int, int],
+) -> Dict[str, Any]:
+    """
+    Scale FARA coordinates from resized image space to original viewport space.
+
+    FARA outputs pixel coordinates on the resized image (after smart_resize).
+    This scales them back to the original browser viewport, matching FARA's
+    convert_resized_coords_to_original() in fara_agent.py:
+        scale_x = og_w / rsz_w
+        return [coords[0] * scale_x, coords[1] * scale_y]
+
+    Args:
+        args: Action arguments containing "coordinate" key
+        original_dims: (width, height) of original browser viewport
+        resized_dims: (width, height) after smart_resize
+    """
+    coord = args.get("coordinate")
+    if not coord or not isinstance(coord, (list, tuple)) or len(coord) < 2:
+        return args
+
+    x, y = float(coord[0]), float(coord[1])
+    original_w, original_h = float(original_dims[0]), float(original_dims[1])
+    resized_w, resized_h = float(resized_dims[0]), float(resized_dims[1])
+
+    # Scale from resized to original: x_final = x * (original / resized)
+    scale_x = original_w / resized_w
+    scale_y = original_h / resized_h
+
+    x_scaled = max(0.0, min(original_w, x * scale_x))
+    y_scaled = max(0.0, min(original_h, y * scale_y))
+
+    return {**args, "coordinate": [round(x_scaled), round(y_scaled)]}
+
+
+def _fara_args_to_sdk_item(args: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """
+    Convert FARA model output args to SDK item using make_*_item helpers.
+
+    FARA format: {"action": "left_click", "coordinate": [100, 200]}
+    SDK format:  ResponseComputerToolCallParam with action={"type": "click", "x": 100, "y": 200}
+    """
+    action = args.get("action", "")
+    coordinate = args.get("coordinate", [0, 0])
+    x = coordinate[0] if len(coordinate) > 0 else 0
+    y = coordinate[1] if len(coordinate) > 1 else 0
+
+    # Click actions
+    if action in ("left_click", "click"):
+        return make_click_item(x=x, y=y, button="left")
+    if action == "right_click":
+        return make_click_item(x=x, y=y, button="right")
+    if action == "middle_click":
+        return make_click_item(x=x, y=y, button="wheel")
+    if action == "double_click":
+        return make_double_click_item(x=x, y=y)
+
+    # Type action
+    if action == "type":
+        return make_type_item(text=args.get("text", ""))
+
+    # Key action
+    if action in ("key", "keypress"):
+        keys = args.get("keys", [])
+        if isinstance(keys, str):
+            keys = keys.split("+")
+        return make_keypress_item(keys=keys)
+
+    # Move action
+    if action in ("mouse_move", "move"):
+        return make_move_item(x=x, y=y)
+
+    # Scroll action
+    if action == "scroll":
+        pixels = args.get("pixels") or 0  # Handle None explicitly
+        # FARA: positive = up, negative = down
+        scroll_y = -pixels  # SDK: positive = down
+        return make_scroll_item(x=x, y=y, scroll_x=0, scroll_y=scroll_y)
+
+    if action == "hscroll":
+        pixels = args.get("pixels") or 0  # Handle None explicitly
+        return make_scroll_item(x=x, y=y, scroll_x=pixels, scroll_y=0)
+
+    # Drag action
+    if action == "left_click_drag":
+        start_coord = args.get("start_coordinate", [0, 0])
+        end_coord = args.get("end_coordinate", [0, 0])
+        return make_drag_item(
+            path=[
+                {"x": start_coord[0], "y": start_coord[1]},
+                {"x": end_coord[0], "y": end_coord[1]},
+            ]
+        )
+
+    # Screenshot
+    if action == "screenshot":
+        return make_screenshot_item()
+
+    # Wait
+    if action == "wait":
+        return make_wait_item()
+
+    # Terminate - return None so no action is executed
+    # The caller checks for terminate action and adds an assistant message to stop the loop
+    if action == "terminate":
+        return None
+
+    # FARA browser-specific actions - create computer_call items directly
+    # agent.py uses getattr(computer, action_type) to call these methods
+    if action == "visit_url":
+        return {
+            "type": "computer_call",
+            "call_id": f"call_{id(args)}",
+            "action": {"type": "visit_url", "url": args.get("url", "")},
+            "pending_safety_checks": [],
+            "status": "completed",
+        }
+
+    if action == "web_search":
+        return {
+            "type": "computer_call",
+            "call_id": f"call_{id(args)}",
+            "action": {"type": "web_search", "query": args.get("query", "")},
+            "pending_safety_checks": [],
+            "status": "completed",
+        }
+
+    if action == "history_back":
+        return {
+            "type": "computer_call",
+            "call_id": f"call_{id(args)}",
+            "action": {"type": "history_back"},
+            "pending_safety_checks": [],
+            "status": "completed",
+        }
+
+    return None
+
+
+@register_agent(models=r"(?i).*fara-7b.*", tool_type="browser")
+class FaraVlmConfig(AsyncAgentConfig):
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        # Check if the last message is a terminate function_call_output
+        # If so, return a final assistant message to stop the loop
+        if messages:
+            last_msg = messages[-1]
+            if last_msg.get("type") in ("function_call_output", "computer_call_output"):
+                output_data = last_msg.get("output")
+
+                # Parse string if needed (could be JSON or Python dict literal)
+                if isinstance(output_data, str):
+                    try:
+                        output_data = json.loads(output_data)
+                    except:
+                        try:
+                            output_data = ast.literal_eval(output_data)
+                        except:
+                            pass
+
+                # Check if it's a terminate action output (contains "terminated": True)
+                if isinstance(output_data, dict) and output_data.get("terminated") is True:
+                    return {
+                        "output": [
+                            {
+                                "type": "message",
+                                "role": "assistant",
+                                "content": [{"type": "output_text", "text": "Task completed."}],
+                            }
+                        ],
+                        "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
+                    }
+
+        # Build messages using FARA's dedicated conversion layer
+        # This converts SDK format to FARA's native format (action + coordinate)
+        converted_msgs = _convert_responses_items_to_fara_messages(
+            messages, allow_images_in_tool_results=False
+        )
+
+        # Build function schemas from tools array
+        function_schemas = []
+        if tools:
+            from ...computers import is_agent_computer
+
+            for tool in tools:
+                tool_type = tool.get("type")
+
+                if tool_type == "computer":
+                    # For computer tools, use FARA_COMPUTER_TOOL schema
+                    computer = tool.get("computer")
+                    if computer and is_agent_computer(computer):
+                        function_schemas.append(FARA_COMPUTER_TOOL["function"])
+                elif tool_type == "function":
+                    # For function tools, use the provided function schema
+                    function_schema = tool.get("function")
+                    if function_schema:
+                        function_schemas.append(function_schema)
+
+        # If no tools provided or no computer tool found, use default FARA_COMPUTER_TOOL
+        if not function_schemas:
+            function_schemas = [FARA_COMPUTER_TOOL["function"]]
+
+        # Prepend Nous-generated system if available
+        nous_system = build_nous_system(function_schemas)
+        completion_messages = ([nous_system] if nous_system else []) + converted_msgs
+
+        # If there is no screenshot in the conversation, take one now and inject it.
+        def _has_any_image(msgs: List[Dict[str, Any]]) -> bool:
+            for m in msgs:
+                content = m.get("content")
+                if isinstance(content, list):
+                    for p in content:
+                        if isinstance(p, dict) and p.get("type") == "image_url":
+                            return True
+            return False
+
+        pre_output_items: List[Dict[str, Any]] = []
+        if not _has_any_image(completion_messages):
+            if computer_handler is None or not hasattr(computer_handler, "screenshot"):
+                raise RuntimeError(
+                    "No screenshots present and computer_handler.screenshot is not available."
+                )
+            screenshot_b64 = await computer_handler.screenshot()
+            if not screenshot_b64:
+                raise RuntimeError("Failed to capture screenshot from computer_handler.")
+
+            await _on_screenshot(screenshot_b64, "screenshot_before")
+
+            # Check if computer_handler has get_current_url method
+            screenshot_text = "Here is the next screenshot. Think about what to do next."
+            if hasattr(computer_handler, "get_current_url"):
+                try:
+                    current_url = await computer_handler.get_current_url()
+                    screenshot_text = f"Current URL: {current_url[:100]}\nHere is the next screenshot. Think about what to do next."
+                except Exception:
+                    # If get_current_url fails, fall back to default text
+                    pass
+
+            # Inject a user message with the screenshot so the model can see current context
+            screenshot_msg = {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{screenshot_b64}"},
+                    },
+                    {"type": "text", "text": screenshot_text},
+                ],
+            }
+            completion_messages.append(screenshot_msg)
+
+        # Smart-resize all screenshots and attach min/max pixel hints. Fail fast if deps missing.
+        # Track both original and resized dimensions for coordinate scaling.
+        last_original_w: Optional[int] = None
+        last_original_h: Optional[int] = None
+        last_rw: Optional[int] = None
+        last_rh: Optional[int] = None
+        MIN_PIXELS = 3136
+        MAX_PIXELS = 12845056
+        try:
+            import base64
+            import io
+
+            from PIL import Image  # type: ignore
+            from qwen_vl_utils import smart_resize  # type: ignore
+        except Exception:
+            raise ImportError(
+                "qwen-vl-utils not installed. Please install it with `pip install cua-agent[qwen]`."
+            )
+
+        for msg in completion_messages:
+            content = msg.get("content")
+            if not isinstance(content, list):
+                continue
+            for part in content:
+                if isinstance(part, dict) and part.get("type") == "image_url":
+                    url = ((part.get("image_url") or {}).get("url")) or ""
+                    # Expect data URL like data:image/png;base64,<b64>
+                    if url.startswith("data:") and "," in url:
+                        b64 = url.split(",", 1)[1]
+                        img_bytes = base64.b64decode(b64)
+                        im = Image.open(io.BytesIO(img_bytes))
+                        h, w = im.height, im.width
+                        rh, rw = smart_resize(
+                            h, w, factor=28, min_pixels=MIN_PIXELS, max_pixels=MAX_PIXELS
+                        )
+                        # Attach hints on this image block
+                        part["min_pixels"] = MIN_PIXELS
+                        part["max_pixels"] = MAX_PIXELS
+                        # Track both original and resized dimensions
+                        last_original_w, last_original_h = w, h
+                        last_rw, last_rh = rw, rh
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": completion_messages,
+            "max_retries": max_retries,
+            "stream": stream,
+            **{k: v for k, v in kwargs.items()},
+        }
+        if use_prompt_caching:
+            api_kwargs["use_prompt_caching"] = use_prompt_caching
+
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        response = await litellm.acompletion(**api_kwargs)
+
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        usage = {
+            **LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(  # type: ignore
+                response.usage
+            ).model_dump(),
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        # Extract response data
+        resp_dict = response.model_dump()  # type: ignore
+        choice = (resp_dict.get("choices") or [{}])[0]
+        message = choice.get("message") or {}
+        content_text = message.get("content") or ""
+        tool_calls_array = message.get("tool_calls") or []
+        reasoning_text = message.get("reasoning") or ""
+
+        output_items: List[Dict[str, Any]] = []
+        has_terminate = False
+
+        # Add reasoning if present (Ollama Cloud format)
+        if reasoning_text:
+            output_items.append(make_reasoning_item(reasoning_text))
+
+        # Extract thoughts (text before <tool_call> tag)
+        thoughts = ""
+        if "<tool_call>" in content_text:
+            thoughts = content_text.split("<tool_call>")[0].strip()
+
+        # Add thoughts as assistant message if present
+        if thoughts:
+            output_items.append(make_output_text_item(thoughts))
+
+        # Priority 1: Try to parse tool call from content text (OpenRouter format)
+        tool_call = parse_tool_call_from_text(content_text)
+
+        if tool_call and isinstance(tool_call, dict):
+            fn_name = tool_call.get("name") or "computer"
+            raw_args = tool_call.get("arguments") or {}
+
+            # Scale coordinates from resized image space to original viewport
+            if (
+                last_rw is None
+                or last_rh is None
+                or last_original_w is None
+                or last_original_h is None
+            ):
+                raise RuntimeError(
+                    "No screenshots found to derive dimensions for coordinate scaling."
+                )
+            args = _scale_fara_coordinates(
+                raw_args,
+                original_dims=(last_original_w, last_original_h),
+                resized_dims=(last_rw, last_rh),
+            )
+
+            # Convert FARA output to SDK format using make_*_item helpers
+            if fn_name in ("computer", "computer_use"):
+                item = _fara_args_to_sdk_item(args)
+                if item:
+                    output_items.append(item)
+                # Check for terminate (even if item is None)
+                if args.get("action") == "terminate":
+                    has_terminate = True
+
+        elif tool_calls_array:
+            # Priority 2: Use tool_calls field if present (Ollama Cloud format)
+            for tc in tool_calls_array:
+                function = tc.get("function", {})
+                fn_name = function.get("name", "computer")
+                args_str = function.get("arguments", "{}")
+
+                try:
+                    args = json.loads(args_str)
+
+                    # Scale coordinates from resized image space to original viewport
+                    if "coordinate" in args and last_rw is not None and last_rh is not None:
+                        if last_original_w is not None and last_original_h is not None:
+                            args = _scale_fara_coordinates(
+                                args,
+                                original_dims=(last_original_w, last_original_h),
+                                resized_dims=(last_rw, last_rh),
+                            )
+
+                    # Convert FARA output to SDK format
+                    if fn_name in ("computer", "computer_use"):
+                        item = _fara_args_to_sdk_item(args)
+                        if item:
+                            output_items.append(item)
+                        # Check for terminate (even if item is None)
+                        if args.get("action") == "terminate":
+                            has_terminate = True
+                except json.JSONDecodeError:
+                    pass
+
+        elif content_text:
+            # No tool calls found, return text response
+            output_items.append(make_output_text_item(content_text))
+
+        # If terminate detected, ensure LAST item is an assistant message to exit the loop
+        # The generic agent loop checks: while new_items[-1].get("role") != "assistant"
+        if has_terminate:
+            output_items.append(
+                {
+                    "type": "message",
+                    "role": "assistant",
+                    "content": [{"type": "output_text", "text": ""}],
+                }
+            )
+
+        # Prepend any pre_output_items (e.g., simulated screenshot-taking message)
+        return {"output": (pre_output_items + output_items), "usage": usage}
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        return ["step"]
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using Qwen3-VL via litellm.acompletion.
+
+        Only exposes a reduced tool schema with left_click to bias model to output a single click.
+        Returns (x, y) absolute pixels when screen dimensions can be obtained; otherwise normalized 0..1000 integers.
+        """
+        # Reduced tool
+        reduced_tool = {
+            "type": "function",
+            "function": {
+                **FARA_COMPUTER_TOOL["function"],
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "action": {"type": "string", "enum": ["left_click"]},
+                        "coordinate": {
+                            "description": "(x, y) in 0..1000 reference space",
+                            "type": "array",
+                            "items": {"type": ["number", "integer"]},
+                            "minItems": 2,
+                            "maxItems": 2,
+                        },
+                    },
+                    "required": ["action", "coordinate"],
+                },
+            },
+        }
+
+        # Build Nous system (lazy import inside helper already raises clear guidance if missing)
+        nous_system = build_nous_system([reduced_tool["function"]])
+
+        # Pre-process using smart_resize
+        min_pixels = 3136
+        max_pixels = 12845056
+        try:
+            # Lazy import to avoid hard dependency
+            import base64
+            import io
+
+            # If PIL is available, estimate size from image to derive smart bounds
+            from PIL import Image
+            from qwen_vl_utils import smart_resize  # type: ignore
+
+            img_bytes = base64.b64decode(image_b64)
+            im = Image.open(io.BytesIO(img_bytes))
+            h, w = im.height, im.width
+            rh, rw = smart_resize(h, w, factor=28, min_pixels=min_pixels, max_pixels=max_pixels)
+        except Exception:
+            raise ImportError(
+                "qwen-vl-utils not installed. Please install it with `pip install cua-agent[qwen]`."
+            )
+
+        messages = []
+        if nous_system:
+            messages.append(nous_system)
+        image_block: Dict[str, Any] = {
+            "type": "image_url",
+            "image_url": {"url": f"data:image/png;base64,{image_b64}"},
+            "min_pixels": min_pixels,
+            "max_pixels": max_pixels,
+        }
+        # Single user message with image and instruction, matching OpenAI-style content blocks
+        messages.append(
+            {
+                "role": "user",
+                "content": [
+                    image_block,
+                    {"type": "text", "text": instruction},
+                ],
+            }
+        )
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            **{k: v for k, v in kwargs.items()},
+        }
+        response = await litellm.acompletion(**api_kwargs)
+        resp = response.model_dump()  # type: ignore
+        choice = (resp.get("choices") or [{}])[0]
+        content_text = ((choice.get("message") or {}).get("content")) or ""
+        tool_call = parse_tool_call_from_text(content_text) or {}
+        args = tool_call.get("arguments") or {}
+        # Scale from resized image space to original viewport
+        args = _scale_fara_coordinates(
+            args,
+            original_dims=(w, h),
+            resized_dims=(rw, rh),
+        )
+        coord = args.get("coordinate")
+        if isinstance(coord, (list, tuple)) and len(coord) >= 2:
+            return int(coord[0]), int(coord[1])
+        return None
+
+
+# FARA-specific ComputerUse tool schema (OpenAI function tool format)
+# This schema is tailored for FARA-7B model and includes browser-specific actions
+# NOTE: Tool name MUST be "computer_use" to match what FARA-7B was trained on
+FARA_COMPUTER_TOOL: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "computer_use",
+        "description": (
+            "Use a mouse and keyboard to interact with a computer, and take screenshots.\n"
+            "* This is an interface to a desktop GUI. You do not have access to a terminal or applications menu. You must click on desktop icons to start applications.\n"
+            "* Some applications may take time to start or process actions, so you may need to wait and take successive screenshots to see the results of your actions. E.g. if you click on Firefox and a window doesn't open, try wait and taking another screenshot.\n"
+            "* The screen's resolution is 1000x1000.\n"
+            "* Whenever you intend to move the cursor to click on an element like an icon, you should consult a screenshot to determine the coordinates of the element before moving the cursor.\n"
+            "* If you tried clicking on a program or link but it failed to load, even after waiting, try adjusting your cursor position so that the tip of the cursor visually falls on the element that you want to click.\n"
+            "* Make sure to click any buttons, links, icons, etc with the cursor tip in the center of the element. Don't click boxes on their edges.\n"
+            "* Use terminate action when you have completed the task or cannot proceed further."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "description": "The action to perform.",
+                    "enum": [
+                        "key",
+                        "type",
+                        "mouse_move",
+                        "left_click",
+                        "left_click_drag",
+                        "right_click",
+                        "middle_click",
+                        "double_click",
+                        "triple_click",
+                        "scroll",
+                        "hscroll",
+                        "screenshot",
+                        "wait",
+                        "visit_url",
+                        "web_search",
+                        "history_back",
+                        "terminate",
+                    ],
+                    "type": "string",
+                },
+                "keys": {
+                    "description": "Required only by action=key.",
+                    "type": "array",
+                    "items": {"type": "string"},
+                },
+                "text": {
+                    "description": "Required only by action=type.",
+                    "type": "string",
+                },
+                "coordinate": {
+                    "description": "(x, y): Pixel coordinates from top-left.",
+                    "type": "array",
+                    "items": {"type": ["number", "integer"]},
+                    "minItems": 2,
+                    "maxItems": 2,
+                },
+                "pixels": {
+                    "description": "Scroll amount. Positive=up, negative=down. For scroll/hscroll.",
+                    "type": "number",
+                },
+                "time": {
+                    "description": "Seconds to wait (action=wait).",
+                    "type": "number",
+                },
+                "url": {
+                    "description": "The URL to visit. Required only by action=visit_url.",
+                    "type": "string",
+                },
+                "query": {
+                    "description": "The search query. Required only by action=web_search.",
+                    "type": "string",
+                },
+                "status": {
+                    "description": "Task completion status. Required only by action=terminate.",
+                    "type": "string",
+                    "enum": ["success", "failure"],
+                },
+            },
+            "required": ["action"],
+        },
+    },
+}

libs/python/agent/cua_agent/loops/fara/helpers.py

@@ -0,0 +1,817 @@
+# Source: https://github.com/QwenLM/Qwen-Agent/blob/main/qwen_agent/llm/fncall_prompts/nous_fncall_prompt.py
+
+import copy
+import json
+import os
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+from .schema import ContentItem, Message
+
+FN_CALL_TEMPLATE_QWEN = """# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{tool_descs}
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{{"name": <function-name>, "arguments": <args-json-object>}}
+</tool_call>"""
+
+FN_CALL_TEMPLATE = """You are a web automation agent that performs actions on websites to fulfill user requests by calling various tools.
+* You should stop execution at Critical Points. A Critical Point would be encountered in tasks like 'Checkout', 'Book', 'Purchase', 'Call', 'Email', 'Order', etc where a binding transaction/agreement would require the user's permission/personal or sensitive information (name, email, credit card, address, payment information, resume, etc) in order to complete a transaction (purchase, reservation, sign-up etc), or to communicate in a way that a human would be expected to do (call, email, apply to a job, etc).
+* Solve the task as far as you can up until a Critical Point:
+    - For example, if the task is to "call a restaurant to make a reservation", you should not actually make the call but should navigate to the restaurant's page and find the phone number.
+    - Similarly, if the task is to "order new size 12 running shoes" you should not actually place the order but should instead search for the right shoes that meet the criteria and add them to the cart.
+    - Some tasks, like answering questions, may not encounter a Critical Point at all.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{tool_descs}
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{{"name": <function-name>, "arguments": <args-json-object>}}
+</tool_call>"""
+
+
+SPECIAL_CODE_MODE = os.getenv("SPECIAL_CODE_MODE", "false").lower() == "true"
+CODE_TOOL_PATTERN = "code_interpreter"
+FN_CALL_TEMPLATE_WITH_CI = """# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{tool_descs}
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{{"name": <function-name>, "arguments": <args-json-object>}}
+</tool_call>
+For code parameters, use placeholders first, and then put the code within <code></code> XML tags, such as:
+<tool_call>
+{{"name": <function-name>, "arguments": {{"code": ""}}}}
+<code>
+Here is the code.
+</code>
+</tool_call>"""
+
+
+class NousFnCallPrompt:
+    def __init__(self, template_name: str = "default"):
+        """Initialize NousFnCallPrompt with a specific template.
+
+        Args:
+            template_name: Name of the template to use. Options:
+                          "default", "qwen", "with_ci"
+        """
+        self.template_name = template_name
+        self.template_map = {
+            "default": FN_CALL_TEMPLATE,
+            "qwen": FN_CALL_TEMPLATE_QWEN,
+            "with_ci": FN_CALL_TEMPLATE_WITH_CI,
+        }
+
+        if template_name not in self.template_map:
+            raise ValueError(
+                f"Unknown template_name: {template_name}. "
+                f"Available options: {list(self.template_map.keys())}"
+            )
+
+    def preprocess_fncall_messages(
+        self,
+        messages: List[Message],
+        functions: List[dict],
+        lang: Literal["en", "zh"],
+        parallel_function_calls: bool = True,
+        function_choice: Union[Literal["auto"], str] = "auto",
+    ) -> List[Message]:
+        del lang  # ignored
+        del parallel_function_calls  # ignored
+        if function_choice != "auto":
+            raise NotImplementedError
+
+        ori_messages = messages
+
+        # Change function_call responses to plaintext responses:
+        messages = []
+        for msg in copy.deepcopy(ori_messages):
+            role, content, reasoning_content = (
+                msg.role,
+                msg.content,
+                msg.reasoning_content,
+            )
+            if role in ("system", "user"):
+                messages.append(msg)
+            elif role == "assistant":
+                content = content or []
+                fn_call = msg.function_call
+                if fn_call:
+                    if (not SPECIAL_CODE_MODE) or (CODE_TOOL_PATTERN not in fn_call.name):
+                        fc = {
+                            "name": fn_call.name,
+                            "arguments": json.loads(fn_call.arguments),
+                        }
+                        fc = json.dumps(fc, ensure_ascii=False)
+                        fc = f"<tool_call>\n{fc}\n</tool_call>"
+                    else:
+                        para = json.loads(fn_call.arguments)
+                        code = para["code"]
+                        para["code"] = ""
+                        fc = {"name": fn_call.name, "arguments": para}
+                        fc = json.dumps(fc, ensure_ascii=False)
+                        fc = f"<tool_call>\n{fc}\n<code>\n{code}\n</code>\n</tool_call>"
+
+                    content.append(ContentItem(text=fc))
+                if messages[-1].role == "assistant":
+                    messages[-1].content.append(ContentItem(text="\n"))
+                    messages[-1].content.extend(content)
+                else:
+                    # TODO: Assuming there will only be one continuous reasoning_content here
+                    messages.append(
+                        Message(
+                            role=role,
+                            content=content,
+                            reasoning_content=reasoning_content,
+                        )
+                    )
+            elif role == "function":
+                assert isinstance(content, list)
+                assert len(content) == 1
+                assert content[0].text
+                fc = f"<tool_response>\n{content[0].text}\n</tool_response>"
+                content = [ContentItem(text=fc)]
+                if messages[-1].role == "user":
+                    messages[-1].content.append(ContentItem(text="\n"))
+                    messages[-1].content.extend(content)
+                else:
+                    messages.append(Message(role="user", content=content))
+            else:
+                raise TypeError
+
+        tool_descs = [{"type": "function", "function": f} for f in functions]
+        tool_names = [
+            function.get("name_for_model", function.get("name", "")) for function in functions
+        ]
+        tool_descs = "\n".join([json.dumps(f, ensure_ascii=False) for f in tool_descs])
+
+        # Select template based on configuration
+        if SPECIAL_CODE_MODE and any([CODE_TOOL_PATTERN in x for x in tool_names]):
+            selected_template = FN_CALL_TEMPLATE_WITH_CI
+        else:
+            selected_template = self.template_map[self.template_name]
+
+        tool_system = selected_template.format(tool_descs=tool_descs)
+        if messages[0].role == "system":
+            messages[0].content.append(ContentItem(text="\n\n" + tool_system))
+        else:
+            messages = [Message(role="system", content=[ContentItem(text=tool_system)])] + messages
+        return messages
+
+
+# Mainly for removing incomplete special tokens when streaming the output
+# This assumes that '<tool_call>\n{"name": "' is the special token for the NousFnCallPrompt
+def remove_incomplete_special_tokens(text: str) -> str:
+    if text in '<tool_call>\n{"name": "':
+        text = ""
+    return text
+
+
+def extract_fn(text: str):
+    fn_name, fn_args = "", ""
+    fn_name_s = '"name": "'
+    fn_name_e = '", "'
+    fn_args_s = '"arguments": '
+    i = text.find(fn_name_s)
+    k = text.find(fn_args_s)
+    if i > 0:
+        _text = text[i + len(fn_name_s) :]
+        j = _text.find(fn_name_e)
+        if j > -1:
+            fn_name = _text[:j]
+    if k > 0:
+        fn_args = text[k + len(fn_args_s) :]
+
+    if len(fn_args) > 5:
+        fn_args = fn_args[:-5]
+    else:
+        fn_args = ""
+    return fn_name, fn_args
+
+
+def build_nous_system(functions: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Use original FARA NousFnCallPrompt to generate a system message embedding tool schema."""
+    from .schema import ContentItem as NousContentItem
+    from .schema import Message as NousMessage
+
+    msgs = NousFnCallPrompt().preprocess_fncall_messages(
+        messages=[
+            NousMessage(
+                role="system", content=[NousContentItem(text="You are a helpful assistant.")]
+            )
+        ],
+        functions=functions,
+        lang="en",
+    )
+    sys = msgs[0].model_dump()
+    # Convert structured content to OpenAI-style content list
+    content = [{"type": "text", "text": c["text"]} for c in sys.get("content", [])]
+    return {"role": "system", "content": content}
+
+
+def fix_fara_tool_call_format(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Fix tool call format in conversation history for FARA compatibility.
+
+    The shared `convert_responses_items_to_completion_messages` function outputs:
+    - Tool name as "computer" (should be "computer_use")
+    - Action key as "type" (should be "action")
+
+    This function post-processes assistant messages to fix these issues.
+    """
+    import re
+
+    # Valid FARA action types
+    valid_actions = {
+        "left_click",
+        "right_click",
+        "middle_click",
+        "double_click",
+        "triple_click",
+        "click",
+        "type",
+        "key",
+        "scroll",
+        "hscroll",
+        "mouse_move",
+        "wait",
+        "visit_url",
+        "web_search",
+        "history_back",
+        "screenshot",
+        "terminate",
+    }
+
+    fixed_messages = []
+    for msg in messages:
+        if msg.get("role") != "assistant":
+            fixed_messages.append(msg)
+            continue
+
+        content = msg.get("content", "")
+        if not isinstance(content, str) or "<tool_call>" not in content:
+            fixed_messages.append(msg)
+            continue
+
+        # Find and fix all tool calls in the content
+        def fix_tool_call(match):
+            tool_call_content = match.group(1)
+            try:
+                tool_call = json.loads(tool_call_content)
+
+                # Fix tool name: "computer" -> "computer_use"
+                if tool_call.get("name") == "computer":
+                    tool_call["name"] = "computer_use"
+
+                # Fix arguments: "type" -> "action" and x/y -> coordinate
+                args = tool_call.get("arguments", {})
+                if isinstance(args, dict):
+                    # If "type" contains a valid action, rename to "action"
+                    if "type" in args and args["type"] in valid_actions:
+                        args["action"] = args.pop("type")
+
+                    # Convert internal x/y format back to FARA coordinate format
+                    if "x" in args and "y" in args and "coordinate" not in args:
+                        args["coordinate"] = [args.pop("x"), args.pop("y")]
+
+                    # Normalize action names: "click" -> "left_click"
+                    if args.get("action") == "click":
+                        args["action"] = "left_click"
+
+                    # Remove "button" field - FARA doesn't use it (action name implies button)
+                    args.pop("button", None)
+
+                    # If "action" is empty but we can infer from other keys
+                    if args.get("action") == "" and "coordinate" in args:
+                        args["action"] = "left_click"
+
+                    tool_call["arguments"] = args
+
+                return f"<tool_call>\n{json.dumps(tool_call)}\n</tool_call>"
+            except (json.JSONDecodeError, TypeError):
+                return match.group(0)  # Return original if parsing fails
+
+        # Match <tool_call>...</tool_call> or <tool_call>...</tool_call>
+        fixed_content = re.sub(
+            r"<tool_call>\s*(\{.*?\})\s*</tool_call>", fix_tool_call, content, flags=re.DOTALL
+        )
+
+        # Also handle malformed closing tags like <tool_call> used as closing
+        fixed_content = re.sub(
+            r"<tool_call>(\{.*?\})<tool_call>", fix_tool_call, fixed_content, flags=re.DOTALL
+        )
+
+        fixed_messages.append({**msg, "content": fixed_content})
+
+    return fixed_messages
+
+
+def parse_tool_call_from_text(text: str) -> Optional[Dict[str, Any]]:
+    """Extract JSON object within <tool_call>...</tool_call> from model text.
+
+    Accepts both </tool_call> and <tool_call> as closing tags for robustness.
+    Handles nested braces in JSON objects.
+    """
+    # Find the opening tag
+    start_idx = text.find("<tool_call>")
+    if start_idx == -1:
+        return None
+
+    # Find the start of JSON (first '{' after opening tag)
+    json_start = text.find("{", start_idx)
+    if json_start == -1:
+        return None
+
+    # Extract JSON by counting braces
+    brace_count = 0
+    json_end = json_start
+    for i in range(json_start, len(text)):
+        if text[i] == "{":
+            brace_count += 1
+        elif text[i] == "}":
+            brace_count -= 1
+            if brace_count == 0:
+                json_end = i + 1
+                break
+
+    if brace_count != 0:
+        return None
+
+    json_str = text[json_start:json_end]
+    try:
+        return json.loads(json_str)
+    except Exception:
+        return None
+
+
+async def unnormalize_coordinate(args: Dict[str, Any], dims: Tuple[int, int]) -> Dict[str, Any]:
+    """Coordinates appear in 0..1000 space, scale to actual screen size using dims if provided."""
+    coord = args.get("coordinate")
+    if not coord or not isinstance(coord, (list, tuple)) or len(coord) < 2:
+        return args
+    x, y = float(coord[0]), float(coord[1])
+    width, height = float(dims[0]), float(dims[1])
+    x_abs = max(0.0, min(width, (x / 1000.0) * width))
+    y_abs = max(0.0, min(height, (y / 1000.0) * height))
+    args = {**args, "coordinate": [round(x_abs), round(y_abs)]}
+    return args
+
+
+def convert_qwen_tool_args_to_computer_action(args: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """
+    Convert Qwen computer tool arguments to the Computer Calls action schema.
+
+    Qwen (example):
+        {"action": "left_click", "coordinate": [114, 68]}
+
+    Target (example):
+        {"action": "left_click", "x": 114, "y": 68}
+
+    Other mappings:
+    - right_click, middle_click, double_click (triple_click -> double_click)
+    - mouse_move -> { action: "move", x, y }
+    - key -> { action: "keypress", keys: [...] }
+    - type -> { action: "type", text }
+    - scroll/hscroll -> { action: "scroll", scroll_x, scroll_y, x, y }
+    - wait -> { action: "wait" }
+    - terminate/answer are not direct UI actions; return None for now
+    """
+    if not isinstance(args, dict):
+        return None
+
+    action = args.get("action")
+    if not isinstance(action, str):
+        return None
+
+    # Coordinates helper
+    coord = args.get("coordinate")
+    x = y = None
+    if isinstance(coord, (list, tuple)) and len(coord) >= 2:
+        try:
+            x = int(round(float(coord[0])))
+            y = int(round(float(coord[1])))
+        except Exception:
+            x = y = None
+
+    # Map actions
+    a = action.lower()
+    if a in {"left_click", "right_click", "middle_click", "double_click"}:
+        if x is None or y is None:
+            return None
+        return {"action": a, "x": x, "y": y}
+    if a == "triple_click":
+        # Approximate as double_click
+        if x is None or y is None:
+            return None
+        return {"action": "double_click", "x": x, "y": y}
+    if a == "mouse_move":
+        if x is None or y is None:
+            return None
+        return {"action": "move", "x": x, "y": y}
+    if a == "key":
+        keys = args.get("keys")
+        if isinstance(keys, list) and all(isinstance(k, str) for k in keys):
+            return {"action": "keypress", "keys": keys}
+        return None
+    if a == "type":
+        text = args.get("text")
+        if isinstance(text, str):
+            return {"action": "type", "text": text}
+        return None
+    if a in {"scroll", "hscroll"}:
+        pixels = args.get("pixels") or 0
+        try:
+            pixels_val = int(round(float(pixels)))
+        except Exception:
+            pixels_val = 0
+        scroll_x = pixels_val if a == "hscroll" else 0
+        scroll_y = pixels_val if a == "scroll" else 0
+        # Include cursor position if available (optional)
+        out: Dict[str, Any] = {"action": "scroll", "scroll_x": scroll_x, "scroll_y": scroll_y}
+        if x is not None and y is not None:
+            out.update({"x": x, "y": y})
+        return out
+    if a == "wait":
+        return {"action": "wait"}
+
+    # Non-UI or terminal actions: terminate/answer -> not mapped here
+    return None
+
+
+def convert_fara_args_to_browser_tool_format(args: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Convert FARA model output format to BrowserTool compatible format.
+
+    FARA model may output extra parameters that BrowserTool methods don't accept.
+    This function cleans up the arguments and maps them to the correct format.
+
+    Examples:
+        Input:  {"action": "click", "button": "left", "x": 378, "y": 144}
+        Output: {"action": "left_click", "coordinate": [378, 144]}
+
+        Input:  {"action": "visit_url", "url": "https://...", "text": "..."}
+        Output: {"action": "visit_url", "url": "https://..."}
+
+        Input:  {"action": "terminate", "url": "...", "text": "...", "status": "success"}
+        Output: {"action": "terminate", "status": "success"}
+    """
+    if not isinstance(args, dict):
+        return args
+
+    action = args.get("action", "")
+    if not isinstance(action, str):
+        return args
+
+    a = action.lower()
+    result: Dict[str, Any] = {"action": a}
+
+    # Handle coordinate-based actions
+    # Check for both coordinate array and separate x/y fields
+    coord = args.get("coordinate")
+    x = args.get("x")
+    y = args.get("y")
+
+    if coord and isinstance(coord, (list, tuple)) and len(coord) >= 2:
+        x, y = coord[0], coord[1]
+
+    # Click actions - normalize to left_click with coordinate
+    if a in {"click", "left_click"}:
+        if x is not None and y is not None:
+            result["action"] = "left_click"
+            result["coordinate"] = [x, y]
+        return result
+
+    if a in {"right_click", "middle_click", "double_click", "triple_click"}:
+        if x is not None and y is not None:
+            result["coordinate"] = [x, y]
+        return result
+
+    if a == "mouse_move":
+        if x is not None and y is not None:
+            result["coordinate"] = [x, y]
+        return result
+
+    if a == "left_click_drag":
+        if x is not None and y is not None:
+            result["coordinate"] = [x, y]
+        # Also handle start/end coordinates if present
+        start_coord = args.get("start_coordinate")
+        end_coord = args.get("end_coordinate")
+        if start_coord:
+            result["start_coordinate"] = start_coord
+        if end_coord:
+            result["end_coordinate"] = end_coord
+        return result
+
+    # Keyboard actions
+    if a == "key":
+        keys = args.get("keys")
+        if keys:
+            result["keys"] = keys
+        return result
+
+    if a == "type":
+        text = args.get("text")
+        if text:
+            result["text"] = text
+        # Include coordinate if typing at a specific location
+        if x is not None and y is not None:
+            result["coordinate"] = [x, y]
+        return result
+
+    # Scroll actions
+    if a in {"scroll", "hscroll"}:
+        pixels = args.get("pixels")
+        if pixels is not None:
+            result["pixels"] = pixels
+        if x is not None and y is not None:
+            result["coordinate"] = [x, y]
+        return result
+
+    # Browser-specific actions
+    if a == "visit_url":
+        url = args.get("url")
+        if url:
+            result["url"] = url
+        return result
+
+    if a == "web_search":
+        query = args.get("query")
+        if query:
+            result["query"] = query
+        return result
+
+    if a == "history_back":
+        return result
+
+    # Wait action
+    if a == "wait":
+        time_val = args.get("time")
+        if time_val is not None:
+            result["time"] = time_val
+        return result
+
+    # Screenshot action
+    if a == "screenshot":
+        return result
+
+    # Terminate action
+    if a == "terminate":
+        status = args.get("status", "success")
+        result["status"] = status
+        return result
+
+    # For any other action, return cleaned args (just action + known fields)
+    return result
+
+
+def _convert_responses_items_to_fara_messages(
+    messages: List[Dict[str, Any]],
+    allow_images_in_tool_results: bool = False,
+) -> List[Dict[str, Any]]:
+    """
+    Convert SDK responses_items format to FARA-compatible completion messages.
+
+    This is FARA's dedicated conversion layer (similar to Anthropic's pattern).
+    It handles the conversion from SDK's OpenAI-style format to FARA's native format:
+
+    SDK format:
+        {"type": "click", "x": 100, "y": 200, "button": "left"}
+
+    FARA format (in XML tool_call):
+        {"name": "computer_use", "arguments": {"action": "left_click", "coordinate": [100, 200]}}
+    """
+    completion_messages: List[Dict[str, Any]] = []
+
+    for message in messages:
+        msg_type = message.get("type")
+        role = message.get("role")
+
+        # Handle user messages
+        if role == "user" or msg_type == "user":
+            content = message.get("content", "")
+            if isinstance(content, list):
+                converted_content = []
+                for item in content:
+                    if isinstance(item, dict):
+                        item_type = item.get("type")
+                        if item_type == "input_image":
+                            image_url = item.get("image_url", "")
+                            if image_url and image_url != "[omitted]":
+                                converted_content.append(
+                                    {"type": "image_url", "image_url": {"url": image_url}}
+                                )
+                        elif item_type == "input_text":
+                            converted_content.append({"type": "text", "text": item.get("text", "")})
+                        elif item_type == "image_url":
+                            # Already in correct format
+                            converted_content.append(item)
+                        elif item_type == "text":
+                            converted_content.append(item)
+                        else:
+                            converted_content.append(item)
+                    else:
+                        converted_content.append({"type": "text", "text": str(item)})
+                completion_messages.append({"role": "user", "content": converted_content})
+            else:
+                completion_messages.append({"role": "user", "content": content})
+
+        # Handle assistant messages
+        elif role == "assistant" and msg_type == "message":
+            content = message.get("content", [])
+            if isinstance(content, str):
+                completion_messages.append({"role": "assistant", "content": content})
+            elif isinstance(content, list):
+                text_parts = []
+                for item in content:
+                    if isinstance(item, dict) and item.get("type") == "output_text":
+                        text_parts.append(item.get("text", ""))
+                completion_messages.append({"role": "assistant", "content": "\n".join(text_parts)})
+
+        # Handle reasoning
+        elif msg_type == "reasoning":
+            summary = message.get("summary", [])
+            reasoning_text = ""
+            if isinstance(summary, list) and summary:
+                for item in summary:
+                    if isinstance(item, dict) and item.get("type") == "summary_text":
+                        reasoning_text = item.get("text", "")
+                        break
+            if reasoning_text:
+                completion_messages.append({"role": "assistant", "content": reasoning_text})
+
+        # Handle computer_call - convert SDK format to FARA's XML tool_call format
+        elif msg_type == "computer_call":
+            action = message.get("action", {})
+            action_type = action.get("type")
+
+            # Convert SDK action to FARA format
+            fara_args = _sdk_action_to_fara_args(action)
+
+            # Build FARA's XML tool_call format
+            tool_call_json = json.dumps({"name": "computer_use", "arguments": fara_args})
+            tool_call_text = f"<tool_call>\n{tool_call_json}\n</tool_call>"
+
+            # Append to last assistant message or create new one
+            if completion_messages and completion_messages[-1].get("role") == "assistant":
+                prev_content = completion_messages[-1].get("content", "")
+                completion_messages[-1]["content"] = f"{prev_content}\n{tool_call_text}".strip()
+            else:
+                completion_messages.append({"role": "assistant", "content": tool_call_text})
+
+        # Handle computer_call_output - convert to FARA's tool_response format
+        elif msg_type == "computer_call_output":
+            output = message.get("output", {})
+
+            # Build response content
+            if isinstance(output, dict) and output.get("type") == "input_image":
+                image_url = output.get("image_url", "")
+                response_text = "<tool_response>\nAction executed successfully. Here is the next screenshot.\n</tool_response>"
+
+                # Add as user message with image
+                if allow_images_in_tool_results and image_url and image_url != "[omitted]":
+                    completion_messages.append(
+                        {
+                            "role": "user",
+                            "content": [
+                                {"type": "text", "text": response_text},
+                                {"type": "image_url", "image_url": {"url": image_url}},
+                            ],
+                        }
+                    )
+                else:
+                    completion_messages.append(
+                        {
+                            "role": "user",
+                            "content": [
+                                {"type": "text", "text": response_text},
+                            ],
+                        }
+                    )
+            elif isinstance(output, dict) and output.get("terminated"):
+                response_text = "<tool_response>\nTask terminated.\n</tool_response>"
+                completion_messages.append({"role": "user", "content": response_text})
+            else:
+                response_text = f"<tool_response>\n{json.dumps(output) if isinstance(output, dict) else str(output)}\n</tool_response>"
+                completion_messages.append({"role": "user", "content": response_text})
+
+        # Handle function_call (non-computer tools)
+        elif msg_type == "function_call":
+            fn_name = message.get("name", "")
+            fn_args = message.get("arguments", "{}")
+
+            tool_call_json = json.dumps(
+                {
+                    "name": fn_name,
+                    "arguments": json.loads(fn_args) if isinstance(fn_args, str) else fn_args,
+                }
+            )
+            tool_call_text = f"<tool_call>\n{tool_call_json}\n</tool_call>"
+
+            if completion_messages and completion_messages[-1].get("role") == "assistant":
+                prev_content = completion_messages[-1].get("content", "")
+                completion_messages[-1]["content"] = f"{prev_content}\n{tool_call_text}".strip()
+            else:
+                completion_messages.append({"role": "assistant", "content": tool_call_text})
+
+        # Handle function_call_output
+        elif msg_type == "function_call_output":
+            output = message.get("output", "")
+            response_text = f"<tool_response>\n{output}\n</tool_response>"
+            completion_messages.append({"role": "user", "content": response_text})
+
+    return completion_messages
+
+
+def _sdk_action_to_fara_args(action: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Convert SDK action format to FARA arguments format.
+
+    SDK format:  {"type": "click", "x": 100, "y": 200, "button": "left"}
+    FARA format: {"action": "left_click", "coordinate": [100, 200]}
+    """
+    action_type = action.get("type", "")
+
+    # Click actions
+    if action_type == "click":
+        button = action.get("button", "left")
+        action_name = {
+            "left": "left_click",
+            "right": "right_click",
+            "wheel": "middle_click",
+            "middle": "middle_click",
+        }.get(button, "left_click")
+        return {"action": action_name, "coordinate": [action.get("x", 0), action.get("y", 0)]}
+
+    if action_type == "double_click":
+        return {"action": "double_click", "coordinate": [action.get("x", 0), action.get("y", 0)]}
+
+    # Type action
+    if action_type == "type":
+        result = {"action": "type", "text": action.get("text", "")}
+        # Include coordinate if present (for click-then-type)
+        if "x" in action and "y" in action:
+            result["coordinate"] = [action.get("x", 0), action.get("y", 0)]
+        return result
+
+    # Keypress action
+    if action_type == "keypress":
+        keys = action.get("keys", [])
+        return {"action": "key", "keys": keys}
+
+    # Move action
+    if action_type in ("move", "mouse_move"):
+        return {"action": "mouse_move", "coordinate": [action.get("x", 0), action.get("y", 0)]}
+
+    # Scroll action
+    if action_type == "scroll":
+        scroll_x = action.get("scroll_x", 0)
+        scroll_y = action.get("scroll_y", 0)
+        # FARA uses pixels (positive = up/left, negative = down/right)
+        pixels = scroll_y if scroll_y != 0 else scroll_x
+        result = {"action": "scroll", "pixels": pixels}
+        if "x" in action and "y" in action:
+            result["coordinate"] = [action.get("x", 0), action.get("y", 0)]
+        return result
+
+    # Drag action
+    if action_type == "drag":
+        path = action.get("path", [])
+        if len(path) >= 2:
+            return {
+                "action": "left_click_drag",
+                "start_coordinate": [path[0].get("x", 0), path[0].get("y", 0)],
+                "end_coordinate": [path[-1].get("x", 0), path[-1].get("y", 0)],
+            }
+        return {"action": "left_click_drag"}
+
+    # Screenshot
+    if action_type == "screenshot":
+        return {"action": "screenshot"}
+
+    # Wait
+    if action_type == "wait":
+        return {"action": "wait"}
+
+    # Terminate
+    if action_type == "terminate":
+        return {"action": "terminate", "status": action.get("status", "success")}
+
+    # Fallback - return as-is with type renamed to action
+    return {"action": action_type, **{k: v for k, v in action.items() if k != "type"}}

libs/python/agent/cua_agent/loops/fara/schema.py

@@ -0,0 +1,143 @@
+# Source: https://github.com/QwenLM/Qwen-Agent/blob/main/qwen_agent/llm/schema.py
+
+from typing import List, Literal, Optional, Tuple, Union
+
+from pydantic import BaseModel, field_validator, model_validator
+
+
+class BaseModelCompatibleDict(BaseModel):
+    def __getitem__(self, item):
+        return getattr(self, item)
+
+    def __setitem__(self, key, value):
+        setattr(self, key, value)
+
+    def model_dump(self, **kwargs):
+        if "exclude_none" not in kwargs:
+            kwargs["exclude_none"] = True
+        return super().model_dump(**kwargs)
+
+    def model_dump_json(self, **kwargs):
+        if "exclude_none" not in kwargs:
+            kwargs["exclude_none"] = True
+        return super().model_dump_json(**kwargs)
+
+    def get(self, key, default=None):
+        try:
+            return getattr(self, key)
+        except AttributeError:
+            return default
+
+    def __str__(self):
+        return f"{self.model_dump()}"
+
+
+class FunctionCall(BaseModelCompatibleDict):
+    name: str
+    arguments: str
+
+    def __init__(self, name: str, arguments: str):
+        super().__init__(name=name, arguments=arguments)
+
+    def __repr__(self):
+        return f"FunctionCall({self.model_dump()})"
+
+
+class ContentItem(BaseModelCompatibleDict):
+    text: Optional[str] = None
+    image: Optional[str] = None
+    file: Optional[str] = None
+    audio: Optional[Union[str, dict]] = None
+    video: Optional[Union[str, list]] = None
+
+    def __init__(
+        self,
+        text: Optional[str] = None,
+        image: Optional[str] = None,
+        file: Optional[str] = None,
+        audio: Optional[Union[str, dict]] = None,
+        video: Optional[Union[str, list]] = None,
+    ):
+        super().__init__(text=text, image=image, file=file, audio=audio, video=video)
+
+    @model_validator(mode="after")
+    def check_exclusivity(self):
+        provided_fields = 0
+        if self.text is not None:
+            provided_fields += 1
+        if self.image:
+            provided_fields += 1
+        if self.file:
+            provided_fields += 1
+        if self.audio:
+            provided_fields += 1
+        if self.video:
+            provided_fields += 1
+
+        if provided_fields != 1:
+            raise ValueError(
+                "Exactly one of 'text', 'image', 'file', 'audio', or 'video' must be provided."
+            )
+        return self
+
+    def __repr__(self):
+        return f"ContentItem({self.model_dump()})"
+
+    def get_type_and_value(
+        self,
+    ) -> Tuple[Literal["text", "image", "file", "audio", "video"], str]:
+        ((t, v),) = self.model_dump().items()
+        assert t in ("text", "image", "file", "audio", "video")
+        return t, v
+
+    @property
+    def type(self) -> Literal["text", "image", "file", "audio", "video"]:
+        t, _ = self.get_type_and_value()
+        return t
+
+    @property
+    def value(self) -> str:
+        _, v = self.get_type_and_value()
+        return v
+
+
+class Message(BaseModelCompatibleDict):
+    role: str
+    content: Union[str, List[ContentItem]]
+    reasoning_content: Optional[Union[str, List[ContentItem]]] = None
+    name: Optional[str] = None
+    function_call: Optional[FunctionCall] = None
+    extra: Optional[dict] = None
+
+    def __init__(
+        self,
+        role: str,
+        content: Union[str, List[ContentItem]],
+        reasoning_content: Optional[Union[str, List[ContentItem]]] = None,
+        name: Optional[str] = None,
+        function_call: Optional[FunctionCall] = None,
+        extra: Optional[dict] = None,
+        **kwargs,
+    ):
+        if content is None:
+            content = ""
+        if reasoning_content is None:
+            reasoning_content = ""
+        super().__init__(
+            role=role,
+            content=content,
+            reasoning_content=reasoning_content,
+            name=name,
+            function_call=function_call,
+            extra=extra,
+        )
+
+    def __repr__(self):
+        return f"Message({self.model_dump()})"
+
+    @field_validator("role")
+    def role_checker(cls, value: str) -> str:
+        values = ["system", "user", "assistant", "function"]
+        if value not in values:
+            raise ValueError(f'{value} must be one of {",".join(values)}')
+        return value

libs/python/agent/cua_agent/loops/gelato.py

@@ -0,0 +1,183 @@
+"""
+Gelato agent loop implementation for click prediction using litellm.acompletion
+Model: https://huggingface.co/mlfoundations/Gelato-30B-A3B
+Code: https://github.com/mlfoundations/Gelato/tree/main
+"""
+
+import base64
+import math
+import re
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from PIL import Image
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..types import AgentCapability
+
+SYSTEM_PROMPT = """
+You are an expert UI element locator. Given a GUI image and a user's element description, provide the coordinates of the specified element as a single (x,y) point. For elements with area, return the center point.
+
+Output the coordinate pair exactly:
+(x,y)
+"""
+
+
+def extract_coordinates(raw_string):
+    """
+    Extract the coordinates from the raw string.
+    Args:
+        raw_string: str (e.g. "(100, 200)")
+    Returns:
+        x: float (e.g. 100.0)
+        y: float (e.g. 200.0)
+    """
+    try:
+        matches = re.findall(r"\((-?\d*\.?\d+),\s*(-?\d*\.?\d+)\)", raw_string)
+        return [tuple(map(int, match)) for match in matches][0]
+    except:
+        return 0, 0
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = 28,
+    min_pixels: int = 3136,
+    max_pixels: int = 8847360,
+) -> Tuple[int, int]:
+    """Smart resize function similar to qwen_vl_utils."""
+    # Calculate the total pixels
+    total_pixels = height * width
+
+    # If already within bounds, return original dimensions
+    if min_pixels <= total_pixels <= max_pixels:
+        # Round to nearest factor
+        new_height = (height // factor) * factor
+        new_width = (width // factor) * factor
+        return new_height, new_width
+
+    # Calculate scaling factor
+    if total_pixels > max_pixels:
+        scale = (max_pixels / total_pixels) ** 0.5
+    else:
+        scale = (min_pixels / total_pixels) ** 0.5
+
+    # Apply scaling
+    new_height = int(height * scale)
+    new_width = int(width * scale)
+
+    # Round to nearest factor
+    new_height = (new_height // factor) * factor
+    new_width = (new_width // factor) * factor
+
+    # Ensure minimum size
+    new_height = max(new_height, factor)
+    new_width = max(new_width, factor)
+
+    return new_height, new_width
+
+
+@register_agent(models=r".*Gelato.*")
+class GelatoConfig(AsyncAgentConfig):
+    """Gelato agent configuration implementing AsyncAgentConfig protocol for click prediction."""
+
+    def __init__(self):
+        self.current_model = None
+        self.last_screenshot_b64 = None
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        raise NotImplementedError()
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[float, float]]:
+        """
+        Predict click coordinates using UI-Ins model via litellm.acompletion.
+
+        Args:
+            model: The UI-Ins model name
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        # Decode base64 image
+        image_data = base64.b64decode(image_b64)
+        image = Image.open(BytesIO(image_data))
+        width, height = image.width, image.height
+
+        # Smart resize the image (similar to qwen_vl_utils)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=28,  # Default factor for Qwen models
+            min_pixels=3136,
+            max_pixels=4096 * 2160,
+        )
+        resized_image = image.resize((resized_width, resized_height))
+        scale_x, scale_y = width / resized_width, height / resized_height
+
+        # Convert resized image back to base64
+        buffered = BytesIO()
+        resized_image.save(buffered, format="PNG")
+        resized_image_b64 = base64.b64encode(buffered.getvalue()).decode()
+
+        # Prepare system and user messages
+        system_message = {
+            "role": "system",
+            "content": [{"type": "text", "text": SYSTEM_PROMPT.strip()}],
+        }
+
+        user_message = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/png;base64,{resized_image_b64}"},
+                },
+                {"type": "text", "text": instruction},
+            ],
+        }
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "messages": [system_message, user_message],
+            "max_tokens": 2056,
+            "temperature": 0.0,
+            **kwargs,
+        }
+
+        # Use liteLLM acompletion
+        response = await litellm.acompletion(**api_kwargs)
+
+        # Extract response text
+        output_text = response.choices[0].message.content  # type: ignore
+
+        # Extract and rescale coordinates
+        pred_x, pred_y = extract_coordinates(output_text)  # type: ignore
+        pred_x *= scale_x
+        pred_y *= scale_y
+
+        return (math.floor(pred_x), math.floor(pred_y))
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """Return the capabilities supported by this agent."""
+        return ["click"]

libs/python/agent/cua_agent/loops/generic_vlm.py

@@ -0,0 +1,601 @@
+"""
+Qwen3-VL agent loop implementation using litellm with function/tool calling.
+- Passes a ComputerUse tool schema to acompletion
+- Converts between Responses items and completion messages using helpers
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from litellm.responses.litellm_completion_transformation.transformation import (
+    LiteLLMCompletionResponsesConfig,
+)
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_completion_messages_to_responses_items,
+    convert_responses_items_to_completion_messages,
+    make_reasoning_item,
+)
+from ..types import AgentCapability
+
+# ComputerUse tool schema (OpenAI function tool format)
+QWEN3_COMPUTER_TOOL: Dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "computer",
+        "description": (
+            "Use a mouse and keyboard to interact with a computer, and take screenshots.\n"
+            "* This is an interface to a desktop GUI. You do not have access to a terminal or applications menu. You must click on desktop icons to start applications.\n"
+            "* Some applications may take time to start or process actions, so you may need to wait and take successive screenshots to see the results of your actions. E.g. if you click on Firefox and a window doesn't open, try wait and taking another screenshot.\n"
+            "* The screen's resolution is 1000x1000.\n"
+            "* Whenever you intend to move the cursor to click on an element like an icon, you should consult a screenshot to determine the coordinates of the element before moving the cursor.\n"
+            "* If you tried clicking on a program or link but it failed to load, even after waiting, try adjusting your cursor position so that the tip of the cursor visually falls on the element that you want to click.\n"
+            "* Make sure to click any buttons, links, icons, etc with the cursor tip in the center of the element. Don't click boxes on their edges."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "description": "The action to perform.",
+                    "enum": [
+                        "key",
+                        "type",
+                        "mouse_move",
+                        "left_click",
+                        "left_click_drag",
+                        "right_click",
+                        "middle_click",
+                        "double_click",
+                        "triple_click",
+                        "scroll",
+                        "hscroll",
+                        "screenshot",
+                        "wait",
+                        # "terminate",
+                        # "answer",
+                    ],
+                    "type": "string",
+                },
+                "keys": {
+                    "description": "Required only by action=key.",
+                    "type": "array",
+                    "items": {"type": "string"},
+                },
+                "text": {
+                    "description": "Required only by action=type and action=answer.",
+                    "type": "string",
+                },
+                "coordinate": {
+                    "description": "(x, y): Pixel coordinates from top-left.",
+                    "type": "array",
+                    "items": {"type": ["number", "integer"]},
+                    "minItems": 2,
+                    "maxItems": 2,
+                },
+                "pixels": {
+                    "description": "Scroll amount. Positive=up, negative=down. For scroll/hscroll.",
+                    "type": "number",
+                },
+                "time": {
+                    "description": "Seconds to wait (action=wait).",
+                    "type": "number",
+                },
+                # "status": {
+                #     "description": "Task status (action=terminate).",
+                #     "type": "string",
+                #     "enum": ["success", "failure"],
+                # },
+            },
+            "required": ["action"],
+        },
+    },
+}
+
+
+def _build_nous_system(functions: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Use qwen-agent NousFnCallPrompt to generate a system message embedding tool schema."""
+    try:
+        from qwen_agent.llm.fncall_prompts.nous_fncall_prompt import (
+            ContentItem as NousContentItem,
+        )
+        from qwen_agent.llm.fncall_prompts.nous_fncall_prompt import (
+            Message as NousMessage,
+        )
+        from qwen_agent.llm.fncall_prompts.nous_fncall_prompt import (
+            NousFnCallPrompt,
+        )
+    except ImportError:
+        raise ImportError(
+            "qwen-agent not installed. Please install it with `pip install cua-agent[qwen]`."
+        )
+    msgs = NousFnCallPrompt().preprocess_fncall_messages(
+        messages=[
+            NousMessage(
+                role="system", content=[NousContentItem(text="You are a helpful assistant.")]
+            )
+        ],
+        functions=functions,
+        lang="en",
+    )
+    sys = msgs[0].model_dump()
+    # Convert qwen-agent structured content to OpenAI-style content list
+    content = [{"type": "text", "text": c["text"]} for c in sys.get("content", [])]
+    return {"role": "system", "content": content}
+
+
+def _parse_tool_call_from_text(text: str) -> Optional[Dict[str, Any]]:
+    """Extract JSON object within <tool_call>...</tool_call> from model text."""
+    m = re.search(r"<tool_call>\s*(\{[\s\S]*?\})\s*</tool_call>", text)
+    if not m:
+        return None
+    try:
+        return json.loads(m.group(1))
+    except Exception:
+        return None
+
+
+async def _unnormalize_coordinate(args: Dict[str, Any], dims: Tuple[int, int]) -> Dict[str, Any]:
+    """Coordinates appear in 0..1000 space, scale to actual screen size using dims if provided."""
+    coord = args.get("coordinate")
+    if not coord or not isinstance(coord, (list, tuple)) or len(coord) < 2:
+        return args
+    x, y = float(coord[0]), float(coord[1])
+    width, height = float(dims[0]), float(dims[1])
+    x_abs = max(0.0, min(width, (x / 1000.0) * width))
+    y_abs = max(0.0, min(height, (y / 1000.0) * height))
+    args = {**args, "coordinate": [round(x_abs), round(y_abs)]}
+    return args
+
+
+def convert_qwen_tool_args_to_computer_action(args: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """
+    Convert Qwen computer tool arguments to the Computer Calls action schema.
+
+    Qwen (example):
+        {"action": "left_click", "coordinate": [114, 68]}
+
+    Target (example):
+        {"action": "left_click", "x": 114, "y": 68}
+
+    Other mappings:
+    - right_click, middle_click, double_click (triple_click -> double_click)
+    - mouse_move -> { action: "move", x, y }
+    - key -> { action: "keypress", keys: [...] }
+    - type -> { action: "type", text }
+    - scroll/hscroll -> { action: "scroll", scroll_x, scroll_y, x, y }
+    - wait -> { action: "wait" }
+    - terminate/answer are not direct UI actions; return None for now
+    """
+    if not isinstance(args, dict):
+        return None
+
+    action = args.get("action")
+    if not isinstance(action, str):
+        return None
+
+    # Coordinates helper
+    coord = args.get("coordinate")
+    x = y = None
+    if isinstance(coord, (list, tuple)) and len(coord) >= 2:
+        try:
+            x = int(round(float(coord[0])))
+            y = int(round(float(coord[1])))
+        except Exception:
+            x = y = None
+
+    # Map actions
+    a = action.lower()
+    if a in {"left_click", "right_click", "middle_click", "double_click"}:
+        if x is None or y is None:
+            return None
+        return {"action": a, "x": x, "y": y}
+    if a == "triple_click":
+        # Approximate as double_click
+        if x is None or y is None:
+            return None
+        return {"action": "double_click", "x": x, "y": y}
+    if a == "mouse_move":
+        if x is None or y is None:
+            return None
+        return {"action": "move", "x": x, "y": y}
+    if a == "key":
+        keys = args.get("keys")
+        if isinstance(keys, list) and all(isinstance(k, str) for k in keys):
+            return {"action": "keypress", "keys": keys}
+        return None
+    if a == "type":
+        text = args.get("text")
+        if isinstance(text, str):
+            return {"action": "type", "text": text}
+        return None
+    if a in {"scroll", "hscroll"}:
+        pixels = args.get("pixels") or 0
+        try:
+            pixels_val = int(round(float(pixels)))
+        except Exception:
+            pixels_val = 0
+        scroll_x = pixels_val if a == "hscroll" else 0
+        scroll_y = pixels_val if a == "scroll" else 0
+        # Include cursor position if available (optional)
+        out: Dict[str, Any] = {"action": "scroll", "scroll_x": scroll_x, "scroll_y": scroll_y}
+        if x is not None and y is not None:
+            out.update({"x": x, "y": y})
+        return out
+    if a == "wait":
+        return {"action": "wait"}
+
+    # Non-UI or terminal actions: terminate/answer -> not mapped here
+    return None
+
+
+@register_agent(models=r"(?i).*", priority=-100)
+class GenericVlmConfig(AsyncAgentConfig):
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        # Build messages using NousFnCallPrompt system with tool schema in text
+        # Start with converted conversation (images/text preserved)
+        converted_msgs = convert_responses_items_to_completion_messages(
+            messages,
+            allow_images_in_tool_results=False,
+        )
+
+        # Build function schemas from tools array
+        function_schemas = []
+        if tools:
+            from ..computers import is_agent_computer
+
+            for tool in tools:
+                tool_type = tool.get("type")
+
+                if tool_type == "computer":
+                    # For computer tools, use QWEN3_COMPUTER_TOOL schema
+                    computer = tool.get("computer")
+                    if computer and is_agent_computer(computer):
+                        function_schemas.append(QWEN3_COMPUTER_TOOL["function"])
+                elif tool_type == "function":
+                    # For function tools, use the provided function schema
+                    function_schema = tool.get("function")
+                    if function_schema:
+                        function_schemas.append(function_schema)
+
+        # If no tools provided or no computer tool found, use default QWEN3_COMPUTER_TOOL
+        if not function_schemas:
+            function_schemas = [QWEN3_COMPUTER_TOOL["function"]]
+
+        # Prepend Nous-generated system if available
+        nous_system = _build_nous_system(function_schemas)
+        completion_messages = ([nous_system] if nous_system else []) + converted_msgs
+
+        # If there is no screenshot in the conversation, take one now and inject it.
+        # Also record a pre_output_items assistant message to reflect action.
+        def _has_any_image(msgs: List[Dict[str, Any]]) -> bool:
+            for m in msgs:
+                content = m.get("content")
+                if isinstance(content, list):
+                    for p in content:
+                        if isinstance(p, dict) and p.get("type") == "image_url":
+                            return True
+            return False
+
+        def _has_screenshot_message(msgs: List[Dict[str, Any]]) -> bool:
+            """Check if messages already contain the 'Taking a screenshot' text."""
+            screenshot_text = "Taking a screenshot to see the current computer screen."
+            for m in msgs:
+                content = m.get("content")
+                if isinstance(content, str) and screenshot_text in content:
+                    return True
+                if isinstance(content, list):
+                    for p in content:
+                        if isinstance(p, dict) and p.get("type") == "text":
+                            if screenshot_text in (p.get("text") or ""):
+                                return True
+            return False
+
+        pre_output_items: List[Dict[str, Any]] = []
+        if not _has_any_image(completion_messages):
+            if computer_handler is None or not hasattr(computer_handler, "screenshot"):
+                raise RuntimeError(
+                    "No screenshots present and computer_handler.screenshot is not available."
+                )
+            screenshot_b64 = await computer_handler.screenshot()
+            if not screenshot_b64:
+                raise RuntimeError("Failed to capture screenshot from computer_handler.")
+            # Inject a user message with the screenshot so the model can see current context
+            completion_messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/png;base64,{screenshot_b64}"},
+                        },
+                        {"type": "text", "text": "Current screen"},
+                    ],
+                }
+            )
+            # Add assistant message to outputs to reflect the action, only if not already present
+            if not _has_screenshot_message(messages):
+                pre_output_items.append(
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": "Taking a screenshot to see the current computer screen.",
+                            }
+                        ],
+                    }
+                )
+
+        # Smart-resize all screenshots and attach min/max pixel hints. Fail fast if deps missing.
+        # Also record the last resized width/height to unnormalize coordinates later.
+        last_rw: Optional[int] = None
+        last_rh: Optional[int] = None
+        MIN_PIXELS = 3136
+        MAX_PIXELS = 12845056
+        try:
+            import base64
+            import io
+
+            from PIL import Image  # type: ignore
+            from qwen_vl_utils import smart_resize  # type: ignore
+        except Exception:
+            raise ImportError(
+                "qwen-vl-utils not installed. Please install it with `pip install cua-agent[qwen]`."
+            )
+
+        for msg in completion_messages:
+            content = msg.get("content")
+            if not isinstance(content, list):
+                continue
+            for part in content:
+                if isinstance(part, dict) and part.get("type") == "image_url":
+                    url = ((part.get("image_url") or {}).get("url")) or ""
+                    # Expect data URL like data:image/png;base64,<b64>
+                    if url.startswith("data:") and "," in url:
+                        b64 = url.split(",", 1)[1]
+                        img_bytes = base64.b64decode(b64)
+                        im = Image.open(io.BytesIO(img_bytes))
+                        h, w = im.height, im.width
+                        rh, rw = smart_resize(
+                            h, w, factor=32, min_pixels=MIN_PIXELS, max_pixels=MAX_PIXELS
+                        )
+                        # Attach hints on this image block
+                        part["min_pixels"] = MIN_PIXELS
+                        part["max_pixels"] = MAX_PIXELS
+                        last_rw, last_rh = rw, rh
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": completion_messages,
+            "max_retries": max_retries,
+            "stream": stream,
+            **{k: v for k, v in kwargs.items()},
+        }
+        if use_prompt_caching:
+            api_kwargs["use_prompt_caching"] = use_prompt_caching
+
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        response = await litellm.acompletion(**api_kwargs)
+
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        usage = {
+            **LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(  # type: ignore
+                response.usage
+            ).model_dump(),
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        # Extract response data
+        resp_dict = response.model_dump()  # type: ignore
+        choice = (resp_dict.get("choices") or [{}])[0]
+        message = choice.get("message") or {}
+        content_text = message.get("content") or ""
+        tool_calls_array = message.get("tool_calls") or []
+        reasoning_text = message.get("reasoning") or ""
+
+        output_items: List[Dict[str, Any]] = []
+
+        # Add reasoning if present (Ollama Cloud format)
+        if reasoning_text:
+            output_items.append(make_reasoning_item(reasoning_text))
+
+        # Priority 1: Try to parse tool call from content text (OpenRouter format)
+        tool_call = _parse_tool_call_from_text(content_text)
+
+        if tool_call and isinstance(tool_call, dict):
+            fn_name = tool_call.get("name") or "computer"
+            raw_args = tool_call.get("arguments") or {}
+            # Unnormalize coordinates to actual screen size using last resized dims
+            if last_rw is None or last_rh is None:
+                raise RuntimeError(
+                    "No screenshots found to derive dimensions for coordinate unnormalization."
+                )
+            args = await _unnormalize_coordinate(raw_args, (last_rw, last_rh))
+
+            # Build an OpenAI-style tool call so we can reuse the converter
+            fake_cm = {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "id": "call_0",
+                        "function": {
+                            "name": fn_name,
+                            "arguments": json.dumps(args),
+                        },
+                    }
+                ],
+            }
+            output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+        elif tool_calls_array:
+            # Priority 2: Use tool_calls field if present (Ollama Cloud format)
+            # Process and unnormalize coordinates in tool calls
+            processed_tool_calls = []
+            for tc in tool_calls_array:
+                function = tc.get("function", {})
+                fn_name = function.get("name", "computer")
+                args_str = function.get("arguments", "{}")
+
+                try:
+                    args = json.loads(args_str)
+
+                    # Unnormalize coordinates if present
+                    if "coordinate" in args and last_rw is not None and last_rh is not None:
+                        args = await _unnormalize_coordinate(args, (last_rw, last_rh))
+
+                    # Convert Qwen format to Computer Calls format if this is a computer tool
+                    if fn_name == "computer":
+                        converted_action = convert_qwen_tool_args_to_computer_action(args)
+                        if converted_action:
+                            args = converted_action
+
+                    processed_tool_calls.append(
+                        {
+                            "type": tc.get("type", "function"),
+                            "id": tc.get("id", "call_0"),
+                            "function": {
+                                "name": fn_name,
+                                "arguments": json.dumps(args),
+                            },
+                        }
+                    )
+                except json.JSONDecodeError:
+                    # Keep original if parsing fails
+                    processed_tool_calls.append(tc)
+
+            fake_cm = {
+                "role": "assistant",
+                "content": content_text if content_text else "",
+                "tool_calls": processed_tool_calls,
+            }
+            output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+        else:
+            # No tool calls found in either format, return text response
+            fake_cm = {"role": "assistant", "content": content_text}
+            output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+
+        # Prepend any pre_output_items (e.g., simulated screenshot-taking message)
+        return {"output": (pre_output_items + output_items), "usage": usage}
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        return ["step"]
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using Qwen3-VL via litellm.acompletion.
+
+        Only exposes a reduced tool schema with left_click to bias model to output a single click.
+        Returns (x, y) absolute pixels when screen dimensions can be obtained; otherwise normalized 0..1000 integers.
+        """
+        # Reduced tool
+        reduced_tool = {
+            "type": "function",
+            "function": {
+                **QWEN3_COMPUTER_TOOL["function"],
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "action": {"type": "string", "enum": ["left_click"]},
+                        "coordinate": {
+                            "description": "(x, y) in 0..1000 reference space",
+                            "type": "array",
+                            "items": {"type": ["number", "integer"]},
+                            "minItems": 2,
+                            "maxItems": 2,
+                        },
+                    },
+                    "required": ["action", "coordinate"],
+                },
+            },
+        }
+
+        # Build Nous system (lazy import inside helper already raises clear guidance if missing)
+        nous_system = _build_nous_system([reduced_tool["function"]])
+
+        # Pre-process using smart_resize
+        min_pixels = 3136
+        max_pixels = 12845056
+        try:
+            # Lazy import to avoid hard dependency
+            import base64
+            import io
+
+            # If PIL is available, estimate size from image to derive smart bounds
+            from PIL import Image
+            from qwen_vl_utils import smart_resize  # type: ignore
+
+            img_bytes = base64.b64decode(image_b64)
+            im = Image.open(io.BytesIO(img_bytes))
+            h, w = im.height, im.width
+            # Qwen notebook suggests factor=32 and a wide min/max range
+            rh, rw = smart_resize(h, w, factor=32, min_pixels=min_pixels, max_pixels=max_pixels)
+        except Exception:
+            raise ImportError(
+                "qwen-vl-utils not installed. Please install it with `pip install cua-agent[qwen]`."
+            )
+
+        messages = []
+        if nous_system:
+            messages.append(nous_system)
+        image_block: Dict[str, Any] = {
+            "type": "image_url",
+            "image_url": {"url": f"data:image/png;base64,{image_b64}"},
+            "min_pixels": min_pixels,
+            "max_pixels": max_pixels,
+        }
+        # Single user message with image and instruction, matching OpenAI-style content blocks
+        messages.append(
+            {
+                "role": "user",
+                "content": [
+                    image_block,
+                    {"type": "text", "text": instruction},
+                ],
+            }
+        )
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            **{k: v for k, v in kwargs.items()},
+        }
+        response = await litellm.acompletion(**api_kwargs)
+        resp = response.model_dump()  # type: ignore
+        choice = (resp.get("choices") or [{}])[0]
+        content_text = ((choice.get("message") or {}).get("content")) or ""
+        tool_call = _parse_tool_call_from_text(content_text) or {}
+        args = tool_call.get("arguments") or {}
+        args = await _unnormalize_coordinate(args, (rh, rw))
+        coord = args.get("coordinate")
+        if isinstance(coord, (list, tuple)) and len(coord) >= 2:
+            return int(coord[0]), int(coord[1])
+        return None

libs/python/agent/cua_agent/loops/glm45v.py

@@ -0,0 +1,907 @@
+"""
+GLM-4.5V agent loop implementation using liteLLM for GLM-4.5V model.
+Supports vision-language models for computer control with bounding box parsing.
+"""
+
+import asyncio
+import base64
+import json
+import re
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from litellm.responses.litellm_completion_transformation.transformation import (
+    LiteLLMCompletionResponsesConfig,
+)
+from litellm.types.utils import ModelResponse
+from PIL import Image
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_completion_messages_to_responses_items,
+    convert_responses_items_to_completion_messages,
+    make_click_item,
+    make_double_click_item,
+    make_drag_item,
+    make_input_image_item,
+    make_keypress_item,
+    make_output_text_item,
+    make_reasoning_item,
+    make_scroll_item,
+    make_type_item,
+    make_wait_item,
+)
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+
+# GLM-4.5V specific constants
+GLM_ACTION_SPACE = """
+### {left,right,middle}_click
+
+Call rule: `{left,right,middle}_click(start_box='[x,y]', element_info='')`
+{
+    'name': ['left_click', 'right_click', 'middle_click'],
+    'description': 'Perform a left/right/middle mouse click at the specified coordinates on the screen.',
+    'parameters': {
+        'type': 'object',
+        'properties': {
+            'start_box': {
+                'type': 'array',
+                'items': {
+                    'type': 'integer'
+                },
+                'description': 'Coordinates [x,y] where to perform the click, normalized to 0-999 range.'
+            },
+            'element_info': {
+                'type': 'string',
+                'description': 'Optional text description of the UI element being clicked.'
+            }
+        },
+        'required': ['start_box']
+    }
+}
+
+### hover
+
+Call rule: `hover(start_box='[x,y]', element_info='')`
+{
+    'name': 'hover',
+    'description': 'Move the mouse pointer to the specified coordinates without performing any click action.',
+    'parameters': {
+        'type': 'object',
+        'properties': {
+            'start_box': {
+                'type': 'array',
+                'items': {
+                    'type': 'integer'
+                },
+                'description': 'Coordinates [x,y] where to move the mouse pointer, normalized to 0-999 range.'
+            },
+            'element_info': {
+                'type': 'string',
+                'description': 'Optional text description of the UI element being hovered over.'
+            }
+        },
+        'required': ['start_box']
+    }
+}
+
+### left_double_click
+
+Call rule: `left_double_click(start_box='[x,y]', element_info='')`
+{
+    'name': 'left_double_click',
+    'description': 'Perform a left mouse double-click at the specified coordinates on the screen.',
+    'parameters': {
+        'type': 'object',
+        'properties': {
+            'start_box': {
+                'type': 'array',
+                'items': {
+                    'type': 'integer'
+                },
+                'description': 'Coordinates [x,y] where to perform the double-click, normalized to 0-999 range.'
+            },
+            'element_info': {
+                'type': 'string',
+                'description': 'Optional text description of the UI element being double-clicked.'
+            }
+        },
+        'required': ['start_box']
+    }
+}
+
+### left_drag
+
+Call rule: `left_drag(start_box='[x1,y1]', end_box='[x2,y2]', element_info='')`
+{
+    'name': 'left_drag',
+    'description': 'Drag the mouse from starting coordinates to ending coordinates while holding the left mouse button.',
+    'parameters': {
+        'type': 'object',
+        'properties': {
+            'start_box': {
+                'type': 'array',
+                'items': {
+                    'type': 'integer'
+                },
+                'description': 'Starting coordinates [x1,y1] for the drag operation, normalized to 0-999 range.'
+            },
+            'end_box': {
+                'type': 'array',
+                'items': {
+                    'type': 'integer'
+                },
+                'description': 'Ending coordinates [x2,y2] for the drag operation, normalized to 0-999 range.'
+            },
+            'element_info': {
+                'type': 'string',
+                'description': 'Optional text description of the UI element being dragged.'
+            }
+        },
+        'required': ['start_box', 'end_box']
+    }
+}
+
+### key
+
+Call rule: `key(keys='')`
+{
+    'name': 'key',
+    'description': 'Simulate pressing a single key or combination of keys on the keyboard.',
+    'parameters': {
+        'type': 'object',
+        'properties': {
+            'keys': {
+                'type': 'string',
+                'description': 'The key or key combination to press. Use '+' to separate keys in combinations (e.g., 'ctrl+c', 'alt+tab').'
+            }
+        },
+        'required': ['keys']
+    }
+}
+
+### type
+
+Call rule: `type(content='')`
+{
+    'name': 'type',
+    'description': 'Type text content into the currently focused text input field. This action only performs typing and does not handle field activation or clearing.',
+    'parameters': {
+        'type': 'object',
+        'properties': {
+            'content': {
+                'type': 'string',
+                'description': 'The text content to be typed into the active text field.'
+            }
+        },
+        'required': ['content']
+    }
+}
+
+### scroll
+
+Call rule: `scroll(start_box='[x,y]', direction='', step=5, element_info='')`
+{
+    'name': 'scroll',
+    'description': 'Scroll an element at the specified coordinates in the specified direction by a given number of wheel steps.',
+    'parameters': {
+        'type': 'object',
+        'properties': {
+            'start_box': {
+                'type': 'array',
+                'items': {
+                    'type': 'integer'
+                },
+                'description': 'Coordinates [x,y] of the element or area to scroll, normalized to 0-999 range.'
+            },
+            'direction': {
+                'type': 'string',
+                'enum': ['down', 'up'],
+                'description': 'The direction to scroll: 'down' or 'up'.'
+            },
+            'step': {
+                'type': 'integer',
+                'default': 5,
+                'description': 'Number of wheel steps to scroll, default is 5.'
+            },
+            'element_info': {
+                'type': 'string',
+                'description': 'Optional text description of the UI element being scrolled.'
+            }
+        },
+        'required': ['start_box', 'direction']
+    }
+}
+
+### WAIT
+
+Call rule: `WAIT()`
+{
+    'name': 'WAIT',
+    'description': 'Wait for 5 seconds before proceeding to the next action.',
+    'parameters': {
+        'type': 'object',
+        'properties': {},
+        'required': []
+    }
+}
+
+### DONE
+
+Call rule: `DONE()`
+{
+    'name': 'DONE',
+    'description': 'Indicate that the current task has been completed successfully and no further actions are needed.',
+    'parameters': {
+        'type': 'object',
+        'properties': {},
+        'required': []
+    }
+}
+
+### FAIL
+
+Call rule: `FAIL()`
+{
+    'name': 'FAIL',
+    'description': 'Indicate that the current task cannot be completed or is impossible to accomplish.',
+    'parameters': {
+        'type': 'object',
+        'properties': {},
+        'required': []
+    }
+}"""
+
+
+def encode_image_to_base64(image_path: str) -> str:
+    """Encode image file to base64 string with data URI."""
+    with open(image_path, "rb") as image_file:
+        encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
+        return f"data:image/png;base64,{encoded_string}"
+
+
+def parse_glm_response(response: str) -> Dict[str, Any]:
+    """
+    Parse GLM-4.5V response to extract action and memory.
+
+    The special tokens <|begin_of_box|> and <|end_of_box|> mark bounding boxes.
+    Coordinates are normalized values between 0 and 1000.
+    """
+    # Extract action from between special tokens
+    pattern = r"<\|begin_of_box\|>(.*?)<\|end_of_box\|>"
+    match = re.search(pattern, response)
+    if match:
+        action = match.group(1).strip()
+    else:
+        # Fallback: look for function call patterns
+        action_pattern = r"[\w_]+\([^)]*\)"
+        matches = re.findall(action_pattern, response)
+        action = matches[0] if matches else None
+
+    # Extract memory section
+    memory_pattern = r"Memory:(.*?)$"
+    memory_match = re.search(memory_pattern, response, re.DOTALL)
+    memory = memory_match.group(1).strip() if memory_match else "[]"
+
+    # Extract action text (everything before Memory:)
+    action_text_pattern = r"^(.*?)Memory:"
+    action_text_match = re.search(action_text_pattern, response, re.DOTALL)
+    action_text = action_text_match.group(1).strip() if action_text_match else response
+
+    # Clean up action text by removing special tokens
+    if action_text:
+        action_text = action_text.replace("<|begin_of_box|>", "").replace("<|end_of_box|>", "")
+
+    return {"action": action, "action_text": action_text, "memory": memory}
+
+
+def get_last_image_from_messages(messages: Messages) -> Optional[str]:
+    """Extract the last image from messages for processing."""
+    for message in reversed(messages):
+        if isinstance(message, dict):
+            if message.get("type") == "computer_call_output":
+                output = message.get("output", {})
+                if isinstance(output, dict) and output.get("type") == "input_image":
+                    image_url = output.get("image_url", "")
+                    if isinstance(image_url, str) and image_url.startswith("data:image/"):
+                        # Extract base64 part
+                        return image_url.split(",", 1)[1]
+            elif message.get("role") == "user":
+                content = message.get("content", [])
+                if isinstance(content, list):
+                    for item in reversed(content):
+                        if isinstance(item, dict) and item.get("type") == "image_url":
+                            image_url_obj = item.get("image_url", {})
+                            if isinstance(image_url_obj, dict):
+                                image_url = image_url_obj.get("url", "")
+                                if isinstance(image_url, str) and image_url.startswith(
+                                    "data:image/"
+                                ):
+                                    return image_url.split(",", 1)[1]
+    return None
+
+
+def convert_responses_items_to_glm45v_pc_prompt(
+    messages: Messages, task: str, memory: str = ""
+) -> List[Dict[str, Any]]:
+    """Convert responses items to GLM-4.5V PC prompt format with historical actions.
+
+    Args:
+        messages: List of message items from the conversation
+        task: The task description
+        memory: Current memory state
+
+    Returns:
+        List of content items for the prompt (text and image_url items)
+    """
+    action_space = GLM_ACTION_SPACE
+
+    # Template head
+    head_text = f"""You are a GUI Agent, and your primary task is to respond accurately to user requests or questions. In addition to directly answering the user's queries, you can also use tools or perform GUI operations directly until you fulfill the user's request or provide a correct answer. You should carefully read and understand the images and questions provided by the user, and engage in thinking and reflection when appropriate. The coordinates involved are all represented in thousandths (0-999).
+
+# Task:
+{task}
+
+# Task Platform
+Ubuntu
+
+# Action Space
+{action_space}
+
+# Historical Actions and Current Memory
+History:"""
+
+    # Template tail
+    tail_text = f"""
+Memory:
+{memory}
+# Output Format
+Plain text explanation with action(param='...')
+Memory:
+[{{"key": "value"}}, ...]
+
+# Some Additional Notes
+- I'll give you the most recent 4 history screenshots(shrunked to 50%*50%) along with the historical action steps.
+- You should put the key information you *have to remember* in a seperated memory part and I'll give it to you in the next round. The content in this part should be a dict list. If you no longer need some given information, you should remove it from the memory. Even if you don't need to remember anything, you should also output an empty list.
+- My computer's password is "password", feel free to use it when you need sudo rights.
+- For the thunderbird account "anonym-x2024@outlook.com", the password is "gTCI";=@y7|QJ0nDa_kN3Sb&>".
+
+Current Screenshot:
+"""
+
+    # Build history from messages
+    history = []
+    history_images = []
+
+    # Group messages into steps
+    current_step = []
+    step_num = 0
+
+    for message in messages:
+        msg_type = message.get("type")
+
+        if msg_type == "reasoning":
+            current_step.append(message)
+        elif msg_type == "message" and message.get("role") == "assistant":
+            current_step.append(message)
+        elif msg_type == "computer_call":
+            current_step.append(message)
+        elif msg_type == "computer_call_output":
+            current_step.append(message)
+            # End of step - process it
+            if current_step:
+                step_num += 1
+
+                # Extract bot thought from message content
+                bot_thought = ""
+                for item in current_step:
+                    if item.get("type") == "message" and item.get("role") == "assistant":
+                        content = item.get("content", [])
+                        for content_item in content:
+                            if content_item.get("type") == "output_text":
+                                bot_thought = content_item.get("text", "")
+                                break
+                        break
+
+                # Extract action from computer_call
+                action_text = ""
+                for item in current_step:
+                    if item.get("type") == "computer_call":
+                        action = item.get("action", {})
+                        action_type = action.get("type", "")
+
+                        if action_type == "click":
+                            x, y = action.get("x", 0), action.get("y", 0)
+                            # Convert to 0-999 range (assuming screen dimensions)
+                            # For now, use direct coordinates - this may need adjustment
+                            action_text = f"left_click(start_box='[{x},{y}]')"
+                        elif action_type == "double_click":
+                            x, y = action.get("x", 0), action.get("y", 0)
+                            action_text = f"left_double_click(start_box='[{x},{y}]')"
+                        elif action_type == "right_click":
+                            x, y = action.get("x", 0), action.get("y", 0)
+                            action_text = f"right_click(start_box='[{x},{y}]')"
+                        elif action_type == "drag":
+                            # Handle drag with path
+                            path = action.get("path", [])
+                            if len(path) >= 2:
+                                start = path[0]
+                                end = path[-1]
+                                action_text = f"left_drag(start_box='[{start.get('x', 0)},{start.get('y', 0)}]', end_box='[{end.get('x', 0)},{end.get('y', 0)}]')"
+                        elif action_type == "keypress":
+                            key = action.get("key", "")
+                            action_text = f"key(keys='{key}')"
+                        elif action_type == "type":
+                            text = action.get("text", "")
+                            action_text = f"type(content='{text}')"
+                        elif action_type == "scroll":
+                            x, y = action.get("x", 0), action.get("y", 0)
+                            direction = action.get("direction", "down")
+                            action_text = f"scroll(start_box='[{x},{y}]', direction='{direction}')"
+                        elif action_type == "wait":
+                            action_text = "WAIT()"
+                        break
+
+                # Extract screenshot from computer_call_output
+                screenshot_url = None
+                for item in current_step:
+                    if item.get("type") == "computer_call_output":
+                        output = item.get("output", {})
+                        if output.get("type") == "input_image":
+                            screenshot_url = output.get("image_url", "")
+                            break
+
+                # Store step info
+                step_info = {
+                    "step_num": step_num,
+                    "bot_thought": bot_thought,
+                    "action_text": action_text,
+                    "screenshot_url": screenshot_url,
+                }
+                history.append(step_info)
+
+                # Store screenshot for last 4 steps
+                if screenshot_url:
+                    history_images.append(screenshot_url)
+
+                current_step = []
+
+    # Build content array with head, history, and tail
+    content = []
+    current_text = head_text
+
+    total_history_steps = len(history)
+    history_image_count = min(4, len(history_images))  # Last 4 images
+
+    for step_idx, step_info in enumerate(history):
+        step_num = step_info["step_num"]
+        bot_thought = step_info["bot_thought"]
+        action_text = step_info["action_text"]
+
+        if step_idx < total_history_steps - history_image_count:
+            # For steps beyond the last 4, use text placeholder
+            current_text += f"\nstep {step_num}: Screenshot:(Omitted in context.) Thought: {bot_thought}\nAction: {action_text}"
+        else:
+            # For the last 4 steps, insert images
+            current_text += f"\nstep {step_num}: Screenshot:"
+            content.append({"type": "text", "text": current_text})
+
+            # Add image
+            img_idx = step_idx - (total_history_steps - history_image_count)
+            if img_idx < len(history_images):
+                content.append({"type": "image_url", "image_url": {"url": history_images[img_idx]}})
+
+            current_text = f" Thought: {bot_thought}\nAction: {action_text}"
+
+    # Add tail
+    current_text += tail_text
+    content.append({"type": "text", "text": current_text})
+
+    return content
+
+
+def model_dump(obj) -> Dict[str, Any]:
+    if isinstance(obj, dict):
+        return {k: model_dump(v) for k, v in obj.items()}
+    elif hasattr(obj, "model_dump"):
+        return obj.model_dump()
+    else:
+        return obj
+
+
+def convert_glm_completion_to_responses_items(
+    response: ModelResponse, image_width: int, image_height: int
+) -> List[Dict[str, Any]]:
+    """
+    Convert GLM-4.5V completion response to responses items format.
+
+    Args:
+        response: LiteLLM ModelResponse from GLM-4.5V
+        image_width: Original image width for coordinate scaling
+        image_height: Original image height for coordinate scaling
+
+    Returns:
+        List of response items in the proper format
+    """
+    import uuid
+
+    response_items = []
+
+    if not response.choices or not response.choices[0].message:
+        return response_items
+
+    message = response.choices[0].message
+    content = message.content or ""
+    reasoning_content = getattr(message, "reasoning_content", None)
+
+    # Add reasoning item if present
+    if reasoning_content:
+        reasoning_item = model_dump(make_reasoning_item(reasoning_content))
+        response_items.append(reasoning_item)
+
+    # Parse the content to extract action and text
+    parsed_response = parse_glm_response(content)
+    action = parsed_response.get("action", "")
+    action_text = parsed_response.get("action_text", "")
+
+    # Add message item with text content (excluding action and memory)
+    if action_text:
+        # Remove action from action_text if it's there
+        clean_text = action_text
+        if action and action in clean_text:
+            clean_text = clean_text.replace(action, "").strip()
+
+        # Remove memory section
+        memory_pattern = r"Memory:\s*\[.*?\]\s*$"
+        clean_text = re.sub(memory_pattern, "", clean_text, flags=re.DOTALL).strip()
+
+        if clean_text:
+            message_item = model_dump(make_output_text_item(clean_text))
+            response_items.append(message_item)
+
+    # Convert action to computer call if present
+    if action:
+        call_id = f"call_{uuid.uuid4().hex[:8]}"
+
+        # Parse different action types and create appropriate computer calls
+        if action.startswith("left_click"):
+            coord_match = re.search(r"start_box='?\[(\d+),\s*(\d+)\]'?", action)
+            if coord_match:
+                x, y = int(coord_match.group(1)), int(coord_match.group(2))
+                # Convert from 0-999 to actual pixel coordinates
+                actual_x = int((x / 999.0) * image_width)
+                actual_y = int((y / 999.0) * image_height)
+                computer_call = model_dump(make_click_item(actual_x, actual_y))
+                computer_call["call_id"] = call_id
+                computer_call["status"] = "completed"
+                response_items.append(computer_call)
+
+        elif action.startswith("right_click"):
+            coord_match = re.search(r"start_box='?\[(\d+),\s*(\d+)\]'?", action)
+            if coord_match:
+                x, y = int(coord_match.group(1)), int(coord_match.group(2))
+                actual_x = int((x / 999.0) * image_width)
+                actual_y = int((y / 999.0) * image_height)
+                computer_call = model_dump(make_click_item(actual_x, actual_y, button="right"))
+                computer_call["call_id"] = call_id
+                computer_call["status"] = "completed"
+                response_items.append(computer_call)
+
+        elif action.startswith("left_double_click"):
+            coord_match = re.search(r"start_box='?\[(\d+),\s*(\d+)\]'?", action)
+            if coord_match:
+                x, y = int(coord_match.group(1)), int(coord_match.group(2))
+                actual_x = int((x / 999.0) * image_width)
+                actual_y = int((y / 999.0) * image_height)
+                computer_call = model_dump(make_double_click_item(actual_x, actual_y))
+                computer_call["call_id"] = call_id
+                computer_call["status"] = "completed"
+                response_items.append(computer_call)
+
+        elif action.startswith("left_drag"):
+            start_match = re.search(r"start_box='?\[(\d+),\s*(\d+)\]'?", action)
+            end_match = re.search(r"end_box='?\[(\d+),\s*(\d+)\]'?", action)
+            if start_match and end_match:
+                x1, y1 = int(start_match.group(1)), int(start_match.group(2))
+                x2, y2 = int(end_match.group(1)), int(end_match.group(2))
+                actual_x1 = int((x1 / 999.0) * image_width)
+                actual_y1 = int((y1 / 999.0) * image_height)
+                actual_x2 = int((x2 / 999.0) * image_width)
+                actual_y2 = int((y2 / 999.0) * image_height)
+                # Create path for drag operation
+                drag_path = [{"x": actual_x1, "y": actual_y1}, {"x": actual_x2, "y": actual_y2}]
+                computer_call = model_dump(make_drag_item(drag_path))
+                computer_call["call_id"] = call_id
+                computer_call["status"] = "completed"
+                response_items.append(computer_call)
+
+        elif action.startswith("key"):
+            key_match = re.search(r"keys='([^']+)'", action)
+            if key_match:
+                keys = key_match.group(1)
+                # Split keys by '+' for key combinations, or use as single key
+                key_list = keys.split("+") if "+" in keys else [keys]
+                computer_call = model_dump(make_keypress_item(key_list))
+                computer_call["call_id"] = call_id
+                computer_call["status"] = "completed"
+                response_items.append(computer_call)
+
+        elif action.startswith("type"):
+            content_match = re.search(r"content='([^']*)'", action)
+            if content_match:
+                content = content_match.group(1)
+                computer_call = model_dump(make_type_item(content))
+                computer_call["call_id"] = call_id
+                computer_call["status"] = "completed"
+                response_items.append(computer_call)
+
+        elif action.startswith("scroll"):
+            coord_match = re.search(r"start_box='?\[(\d+),\s*(\d+)\]'?", action)
+            direction_match = re.search(r"direction='([^']+)'", action)
+            if coord_match and direction_match:
+                x, y = int(coord_match.group(1)), int(coord_match.group(2))
+                direction = direction_match.group(1)
+                actual_x = int((x / 999.0) * image_width)
+                actual_y = int((y / 999.0) * image_height)
+                # Convert direction to scroll amounts
+                scroll_x, scroll_y = 0, 0
+                if direction == "up":
+                    scroll_y = -5
+                elif direction == "down":
+                    scroll_y = 5
+                elif direction == "left":
+                    scroll_x = -5
+                elif direction == "right":
+                    scroll_x = 5
+                computer_call = model_dump(make_scroll_item(actual_x, actual_y, scroll_x, scroll_y))
+                computer_call["call_id"] = call_id
+                computer_call["status"] = "completed"
+                response_items.append(computer_call)
+
+        elif action == "WAIT()":
+            computer_call = model_dump(make_wait_item())
+            computer_call["call_id"] = call_id
+            computer_call["status"] = "completed"
+            response_items.append(computer_call)
+
+    return response_items
+
+
+@register_agent(models=r"(?i).*GLM-4\.5V.*")
+class Glm4vConfig(AsyncAgentConfig):
+    """GLM-4.5V agent configuration using liteLLM."""
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Predict the next step using GLM-4.5V model.
+
+        Args:
+            messages: Input messages following Responses format
+            model: Model name to use
+            tools: Optional list of tool schemas
+            max_retries: Maximum number of retries for API calls
+            stream: Whether to stream the response
+            computer_handler: Computer handler for taking screenshots
+            use_prompt_caching: Whether to use prompt caching
+            _on_api_start: Callback for API start
+            _on_api_end: Callback for API end
+            _on_usage: Callback for usage tracking
+            _on_screenshot: Callback for screenshot events
+
+        Returns:
+            Dict with "output" and "usage" keys
+        """
+        # Get the user instruction from the last user message
+        user_instruction = ""
+        for message in reversed(messages):
+            if isinstance(message, dict) and message.get("role") == "user":
+                content = message.get("content", "")
+                if isinstance(content, str):
+                    user_instruction = content
+                elif isinstance(content, list):
+                    for item in content:
+                        if isinstance(item, dict) and item.get("type") == "text":
+                            user_instruction = item.get("text", "")
+                            break
+                break
+
+        # Get the last image for processing
+        last_image_b64 = get_last_image_from_messages(messages)
+        if not last_image_b64 and computer_handler:
+            # Take a screenshot if no image available
+            screenshot_b64 = await computer_handler.screenshot()
+            if screenshot_b64:
+                last_image_b64 = screenshot_b64
+                if _on_screenshot:
+                    await _on_screenshot(screenshot_b64)
+
+        if not last_image_b64:
+            raise ValueError("No image available for GLM-4.5V processing")
+
+        # Convert responses items to GLM-4.5V PC prompt format with historical actions
+        prompt_content = convert_responses_items_to_glm45v_pc_prompt(
+            messages=messages,
+            task=user_instruction,
+            memory="[]",  # Initialize with empty memory for now
+        )
+
+        # Add the current screenshot to the end
+        prompt_content.append(
+            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{last_image_b64}"}}
+        )
+
+        # Prepare messages for liteLLM
+        litellm_messages = [
+            {"role": "system", "content": "You are a helpful GUI agent assistant."},
+            {"role": "user", "content": prompt_content},
+        ]
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "messages": litellm_messages,
+            # "max_tokens": 2048,
+            # "temperature": 0.001,
+            # "extra_body": {
+            #     "skip_special_tokens": False,
+            # }
+        }
+        api_kwargs.update({k: v for k, v in (kwargs or {}).items()})
+
+        # Add API callbacks
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        # Call liteLLM
+        response = await litellm.acompletion(**api_kwargs)
+
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        # Get image dimensions for coordinate scaling
+        image_width, image_height = 1920, 1080  # Default dimensions
+
+        # Try to get actual dimensions from the image
+        try:
+            image_data = base64.b64decode(last_image_b64)
+            image = Image.open(BytesIO(image_data))
+            image_width, image_height = image.size
+        except Exception:
+            pass  # Use default dimensions
+
+        # Convert GLM completion response to responses items
+        response_items = convert_glm_completion_to_responses_items(
+            response, image_width, image_height
+        )
+
+        # Extract usage information
+        response_usage = {
+            **LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(
+                response.usage
+            ).model_dump(),
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(response_usage)
+
+        # Create agent response
+        agent_response = {"output": response_items, "usage": response_usage}
+
+        return agent_response
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using GLM-4.5V model.
+
+        Args:
+            model: Model name to use
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+
+        Returns:
+            Tuple with (x, y) coordinates or None
+        """
+        try:
+            # Create a simple click instruction prompt
+            click_prompt = f"""You are a GUI agent. Look at the screenshot and identify where to click for: {instruction}
+
+Respond with a single click action in this format:
+left_click(start_box='[x,y]')
+
+Where x,y are coordinates normalized to 0-999 range."""
+
+            # Prepare messages for liteLLM
+            litellm_messages = [
+                {"role": "system", "content": "You are a helpful GUI agent assistant."},
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": click_prompt},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/png;base64,{image_b64}"},
+                        },
+                    ],
+                },
+            ]
+
+            # Prepare API call kwargs
+            api_kwargs = {
+                "model": model,
+                "messages": litellm_messages,
+                "max_tokens": 2056,
+                "temperature": 0.001,
+                "extra_body": {
+                    "skip_special_tokens": False,
+                },
+            }
+            api_kwargs.update({k: v for k, v in (kwargs or {}).items()})
+
+            # Call liteLLM
+            response = await litellm.acompletion(**api_kwargs)
+
+            # Extract response content
+            response_content = response.choices[0].message.content.strip()
+            print(response)
+
+            # Parse response for click coordinates
+            # Look for coordinates in the response, handling special tokens
+            coord_pattern = r"<\|begin_of_box\|>.*?left_click\(start_box='?\[(\d+),(\d+)\]'?\).*?<\|end_of_box\|>"
+            match = re.search(coord_pattern, response_content)
+
+            if not match:
+                # Fallback: look for coordinates without special tokens
+                coord_pattern = r"left_click\(start_box='?\[(\d+),(\d+)\]'?\)"
+                match = re.search(coord_pattern, response_content)
+
+            if match:
+                x, y = int(match.group(1)), int(match.group(2))
+
+                # Get actual image dimensions for scaling
+                try:
+                    image_data = base64.b64decode(image_b64)
+                    image = Image.open(BytesIO(image_data))
+                    image_width, image_height = image.size
+                except Exception:
+                    # Use default dimensions
+                    image_width, image_height = 1920, 1080
+
+                # Convert from 0-999 normalized coordinates to actual pixel coordinates
+                actual_x = int((x / 999.0) * image_width)
+                actual_y = int((y / 999.0) * image_height)
+
+                return (actual_x, actual_y)
+
+            return None
+
+        except Exception as e:
+            # Log error and return None
+            print(f"Error in predict_click: {e}")
+            return None
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """
+        Get list of capabilities supported by this agent config.
+
+        Returns:
+            List of capability strings
+        """
+        return ["step", "click"]

libs/python/agent/cua_agent/loops/gta1.py

@@ -0,0 +1,175 @@
+"""
+GTA1 agent loop implementation for click prediction using litellm.acompletion
+Paper: https://arxiv.org/pdf/2507.05791
+Code: https://github.com/Yan98/GTA1
+"""
+
+import asyncio
+import base64
+import json
+import math
+import re
+import uuid
+from io import BytesIO
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+
+import litellm
+from PIL import Image
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+
+SYSTEM_PROMPT = """
+You are an expert UI element locator. Given a GUI image and a user's element description, provide the coordinates of the specified element as a single (x,y) point. The image resolution is height {height} and width {width}. For elements with area, return the center point.
+
+Output the coordinate pair exactly:
+(x,y)
+""".strip()
+
+
+def extract_coordinates(raw_string: str) -> Tuple[float, float]:
+    """Extract coordinates from model output."""
+    try:
+        matches = re.findall(r"\((-?\d*\.?\d+),\s*(-?\d*\.?\d+)\)", raw_string)
+        return tuple(map(float, matches[0]))  # type: ignore
+    except:
+        return (0.0, 0.0)
+
+
+def smart_resize(
+    height: int, width: int, factor: int = 28, min_pixels: int = 3136, max_pixels: int = 8847360
+) -> Tuple[int, int]:
+    """Smart resize function similar to qwen_vl_utils."""
+    # Calculate the total pixels
+    total_pixels = height * width
+
+    # If already within bounds, return original dimensions
+    if min_pixels <= total_pixels <= max_pixels:
+        # Round to nearest factor
+        new_height = (height // factor) * factor
+        new_width = (width // factor) * factor
+        return new_height, new_width
+
+    # Calculate scaling factor
+    if total_pixels > max_pixels:
+        scale = (max_pixels / total_pixels) ** 0.5
+    else:
+        scale = (min_pixels / total_pixels) ** 0.5
+
+    # Apply scaling
+    new_height = int(height * scale)
+    new_width = int(width * scale)
+
+    # Round to nearest factor
+    new_height = (new_height // factor) * factor
+    new_width = (new_width // factor) * factor
+
+    # Ensure minimum size
+    new_height = max(new_height, factor)
+    new_width = max(new_width, factor)
+
+    return new_height, new_width
+
+
+@register_agent(models=r".*GTA1.*")
+class GTA1Config(AsyncAgentConfig):
+    """GTA1 agent configuration implementing AsyncAgentConfig protocol for click prediction."""
+
+    def __init__(self):
+        self.current_model = None
+        self.last_screenshot_b64 = None
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        raise NotImplementedError()
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[float, float]]:
+        """
+        Predict click coordinates using GTA1 model via litellm.acompletion.
+
+        Args:
+            model: The GTA1 model name
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        # Decode base64 image
+        image_data = base64.b64decode(image_b64)
+        image = Image.open(BytesIO(image_data))
+        width, height = image.width, image.height
+
+        # Smart resize the image (similar to qwen_vl_utils)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=28,  # Default factor for Qwen models
+            min_pixels=3136,
+            max_pixels=4096 * 2160,
+        )
+        resized_image = image.resize((resized_width, resized_height))
+        scale_x, scale_y = width / resized_width, height / resized_height
+
+        # Convert resized image back to base64
+        buffered = BytesIO()
+        resized_image.save(buffered, format="PNG")
+        resized_image_b64 = base64.b64encode(buffered.getvalue()).decode()
+
+        # Prepare system and user messages
+        system_message = {
+            "role": "system",
+            "content": SYSTEM_PROMPT.format(height=resized_height, width=resized_width),
+        }
+
+        user_message = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/png;base64,{resized_image_b64}"},
+                },
+                {"type": "text", "text": instruction},
+            ],
+        }
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "messages": [system_message, user_message],
+            "max_tokens": 2056,
+            "temperature": 0.0,
+            **kwargs,
+        }
+
+        # Use liteLLM acompletion
+        response = await litellm.acompletion(**api_kwargs)
+
+        # Extract response text
+        output_text = response.choices[0].message.content  # type: ignore
+
+        # Extract and rescale coordinates
+        pred_x, pred_y = extract_coordinates(output_text)  # type: ignore
+        pred_x *= scale_x
+        pred_y *= scale_y
+
+        return (math.floor(pred_x), math.floor(pred_y))
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """Return the capabilities supported by this agent."""
+        return ["click"]

libs/python/agent/cua_agent/loops/holo.py

@@ -0,0 +1,218 @@
+"""
+Holo 1.5 agent loop implementation for click prediction using litellm.acompletion.
+
+Implements the Holo1.5 grounding behavior:
+- Prompt asks for absolute pixel coordinates in JSON: {"action":"click_absolute","x":int,"y":int}
+- Optionally resizes the image using Qwen2-VL smart_resize parameters (via transformers AutoProcessor)
+- If resized, maps predicted coordinates back to the original screenshot resolution
+
+Note: We do NOT manually load the model; acompletions (via HuggingFaceLocalAdapter)
+will handle loading based on the provided model name.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from PIL import Image
+
+from ..decorators import register_agent
+from ..types import AgentCapability
+from .base import AsyncAgentConfig
+
+
+def _strip_hf_prefix(model: str) -> str:
+    """Strip provider prefixes like 'huggingface-local/' from model names for HF processor load."""
+    if "/" in model and model.lower().startswith("huggingface-local/"):
+        return model.split("/", 1)[1]
+    return model
+
+
+def _maybe_smart_resize(image: Image.Image, model: str) -> Tuple[Image.Image, Tuple[int, int]]:
+    """
+    Try to compute Qwen2-VL smart_resize output size using transformers AutoProcessor.
+
+    Returns (processed_image, (orig_w, orig_h)). If transformers or processor unavailable,
+    returns the original image and size without resizing.
+    """
+    orig_w, orig_h = image.size
+    try:
+        # Import lazily to avoid hard dependency if not installed
+        from transformers import AutoProcessor  # type: ignore
+        from transformers.models.qwen2_vl.image_processing_qwen2_vl import (  # type: ignore
+            smart_resize,
+        )
+
+        processor_name = _strip_hf_prefix(model)
+        processor = AutoProcessor.from_pretrained(processor_name)
+        image_processor = getattr(processor, "image_processor", None)
+        if image_processor is None:
+            return image, (orig_w, orig_h)
+
+        factor = getattr(image_processor, "patch_size", 14) * getattr(
+            image_processor, "merge_size", 1
+        )
+        min_pixels = getattr(image_processor, "min_pixels", 256 * 256)
+        max_pixels = getattr(image_processor, "max_pixels", 1536 * 1536)
+
+        resized_h, resized_w = smart_resize(
+            orig_h,
+            orig_w,
+            factor=factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+
+        if (resized_w, resized_h) == (orig_w, orig_h):
+            return image, (orig_w, orig_h)
+
+        processed = image.resize((resized_w, resized_h), resample=Image.Resampling.LANCZOS)
+        return processed, (orig_w, orig_h)
+    except Exception:
+        # If any failure (no transformers, processor load error), fall back to original
+        return image, (orig_w, orig_h)
+
+
+def _build_holo_prompt(instruction: str) -> str:
+    """Construct the Holo1.5 grounding prompt."""
+    # Keep it close to the cookbook while avoiding heavy schema generation
+    schema_hint = '{"action": "click_absolute", "x": <int>, "y": <int>}'
+    return (
+        "Localize an element on the GUI image according to the provided target and output a click position. "
+        f"You must output a valid JSON following the format: {schema_hint} "
+        f"Your target is: {instruction}"
+    )
+
+
+def _parse_click_json(output_text: str) -> Optional[Tuple[int, int]]:
+    """
+    Parse JSON from model output and extract x, y ints.
+    Tries to find the first JSON object substring if extra text is present.
+    """
+    try:
+        # Fast path: direct JSON
+        data = json.loads(output_text)
+    except Exception:
+        # Try to locate a JSON object within the text
+        start = output_text.find("{")
+        end = output_text.rfind("}")
+        if start == -1 or end == -1 or end <= start:
+            return None
+        try:
+            data = json.loads(output_text[start : end + 1])
+        except Exception:
+            return None
+
+    try:
+        x = int(data.get("x"))
+        y = int(data.get("y"))
+        return x, y
+    except Exception:
+        return None
+
+
+@register_agent(models=r"(?i).*(Holo1\.5|Hcompany/Holo1\.5).*")
+class HoloConfig(AsyncAgentConfig):
+    """Holo is a family of UI grounding models from H Company"""
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        # Holo models are only trained on UI localization tasks, not all-in-one agent
+        raise NotImplementedError()
+
+    async def predict_click(
+        self,
+        model: str,
+        image_b64: str,
+        instruction: str,
+        **kwargs,
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using Holo1.5 via litellm.acompletion.
+
+        - Optionally smart-resizes the image using Qwen2-VL rules if transformers are available
+        - Prompts for JSON with absolute pixel coordinates
+        - Parses x,y and maps back to original screenshot size if resized
+        """
+        try:
+            img_bytes = base64.b64decode(image_b64)
+            original_img = Image.open(BytesIO(img_bytes))
+        except Exception:
+            return None
+
+        # Optional preprocessing
+        processed_img, (orig_w, orig_h) = _maybe_smart_resize(original_img, model)
+
+        # If we resized, send the resized image; otherwise send original
+        img_to_send = processed_img
+        buf = BytesIO()
+        img_to_send.save(buf, format="PNG")
+        processed_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+
+        prompt = _build_holo_prompt(instruction)
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{processed_b64}"},
+                    },
+                    {"type": "text", "text": prompt},
+                ],
+            }
+        ]
+
+        api_kwargs = {
+            "model": model,
+            "messages": messages,
+            # Deterministic, small output
+            "max_tokens": kwargs.get("max_tokens", 256),
+            "temperature": kwargs.get("temperature", 0.0),
+        }
+
+        response = await litellm.acompletion(**api_kwargs)
+        output_text = (response.choices[0].message.content or "").strip()  # type: ignore
+
+        coords = _parse_click_json(output_text)
+        if coords is None:
+            return None
+
+        x, y = coords
+
+        # Map back to original size if we resized
+        proc_w, proc_h = img_to_send.size
+        if (proc_w, proc_h) != (orig_w, orig_h):
+            try:
+                sx = orig_w / float(proc_w)
+                sy = orig_h / float(proc_h)
+                x = int(round(x * sx))
+                y = int(round(y * sy))
+            except Exception:
+                # Fallback: clamp within original bounds
+                pass
+
+        # Clamp to original image bounds
+        x = max(0, min(orig_w - 1, x))
+        y = max(0, min(orig_h - 1, y))
+        return x, y
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        return ["click"]

libs/python/agent/cua_agent/loops/internvl.py

@@ -0,0 +1,180 @@
+"""
+InternVL agent loop implementation for click prediction using litellm.acompletion.
+
+Implements the ScreenSpot InternVL grounding baseline behavior:
+- Uses the exact grounding prompt format with <image> and <ref> tags
+- Expects coordinates in 0-1000 normalized range in formats [[x1,y1,x2,y2]] or [[x,y]]
+- Converts to pixel coordinates relative to the original screenshot size
+
+Note: We do NOT manually load the InternVL model; acompletions (via HuggingFaceLocalAdapter)
+will handle loading based on the provided model name.
+"""
+
+from __future__ import annotations
+
+import base64
+import math
+import re
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from PIL import Image
+
+from ..decorators import register_agent
+from ..types import AgentCapability
+from .composed_grounded import ComposedGroundedConfig
+
+# Regex patterns for extracting coordinates
+# Accept optional whitespace and optional decimal fractions
+_NUM = r"(\d+(?:\.\d+)?)"
+_POINT_PATTERN = re.compile(r"\[\[\s*" + _NUM + r"\s*,\s*" + _NUM + r"\s*\]\]")
+_BBOX_PATTERN = re.compile(
+    r"\[\[\s*" + _NUM + r"\s*,\s*" + _NUM + r"\s*,\s*" + _NUM + r"\s*,\s*" + _NUM + r"\s*\]\]"
+)
+
+
+def _extract_first_point(text: str) -> Optional[Tuple[float, float]]:
+    """Extract the first [[x,y]] as normalized (0-1000) floats."""
+    m = _POINT_PATTERN.search(text)
+    if not m:
+        return None
+    try:
+        x = float(m.group(1))
+        y = float(m.group(2))
+        return x, y
+    except Exception:
+        return None
+
+
+def _extract_last_bbox(text: str) -> Optional[Tuple[float, float, float, float]]:
+    """Extract the last [[x1,y1,x2,y2]] as normalized (0-1000) floats."""
+    matches = list(_BBOX_PATTERN.finditer(text))
+    if not matches:
+        return None
+    m = matches[-1]
+    try:
+        x1 = float(m.group(1))
+        y1 = float(m.group(2))
+        x2 = float(m.group(3))
+        y2 = float(m.group(4))
+        return x1, y1, x2, y2
+    except Exception:
+        return None
+
+
+def _scale_norm_to_pixels(x_norm: float, y_norm: float, width: int, height: int) -> Tuple[int, int]:
+    """Scale 0-1000 normalized coordinates to pixel coordinates for given image size."""
+    x_px = int(math.floor((x_norm / 1000.0) * width))
+    y_px = int(math.floor((y_norm / 1000.0) * height))
+    # Clamp to image bounds just in case
+    x_px = max(0, min(width - 1, x_px))
+    y_px = max(0, min(height - 1, y_px))
+    return x_px, y_px
+
+
+@register_agent(models=r"(?i).*InternVL.*")
+class InternVLConfig(ComposedGroundedConfig):
+    """InternVL agent configuration reusing ComposedGroundedConfig for steps and
+    overriding predict_click to implement ScreenSpot InternVL grounding baseline."""
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Fallback to a self-composed model"""
+        return await super().predict_step(
+            messages=messages,
+            model=f"{model}+{model}",
+            tools=tools,
+            max_retries=max_retries,
+            stream=stream,
+            computer_handler=computer_handler,
+            _on_api_start=_on_api_start,
+            _on_api_end=_on_api_end,
+            _on_usage=_on_usage,
+            _on_screenshot=_on_screenshot,
+            **kwargs,
+        )
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using InternVL via litellm.acompletion.
+
+        Behavior mirrors the ScreenSpot InternVL baseline:
+        - Prompt: "<image>\nPlease provide the bounding box coordinate of the UI element this user instruction describes: <ref>{instruction}</ref>. Answer in the format of [[x1, y1, x2, y2]]"
+        - Parse either [[x,y]] point or [[x1,y1,x2,y2]] bbox, using bbox center if point missing
+        - Coordinates are 0-1000 normalized; convert to pixel coordinates for the original screenshot
+        """
+        try:
+            # Decode image dimensions to scale the normalized outputs
+            img_bytes = base64.b64decode(image_b64)
+            image = Image.open(BytesIO(img_bytes))
+            width, height = image.size
+        except Exception:
+            # If decoding fails, proceed with a safe default size to avoid crash
+            width, height = 1920, 1080
+
+        # Build grounding prompt exactly like the baseline
+        grounding_prompt = (
+            f"Please provide the bounding box coordinate of the UI element this user instruction describes: <ref>{instruction}</ref>. "
+            f"Answer in the format of [[x1, y1, x2, y2]]"
+        )
+
+        # Prepare messages for LiteLLM
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{image_b64}"},
+                    },
+                    {"type": "text", "text": grounding_prompt},
+                ],
+            }
+        ]
+
+        # Call acompletion; HuggingFaceLocalAdapter/model handler will handle InternVL loading
+        api_kwargs = {
+            "model": model,
+            "messages": messages,
+            # Conservative generation params akin to baseline (deterministic)
+            "max_tokens": kwargs.get("max_tokens", 256),
+            "temperature": kwargs.get("temperature", 0.0),
+        }
+
+        response = await litellm.acompletion(**api_kwargs)
+        output_text = (response.choices[0].message.content or "").strip()  # type: ignore
+
+        # print(f"InternVL output: {output_text}")
+
+        # Try to parse a point first; if absent, parse bbox and take center
+        point = _extract_first_point(output_text)
+        if point is None:
+            bbox = _extract_last_bbox(output_text)
+            if bbox is None:
+                return None
+            x1, y1, x2, y2 = bbox
+            cx = (x1 + x2) / 2.0
+            cy = (y1 + y2) / 2.0
+            point = (cx, cy)
+
+        x_norm, y_norm = point
+        x_px, y_px = _scale_norm_to_pixels(x_norm, y_norm, width, height)
+        return (x_px, y_px)
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        return ["click", "step"]

libs/python/agent/cua_agent/loops/model_types.csv

@@ -0,0 +1,6 @@
+model,predict_step,predict_point
+anthropic,✅,✅
+openai,✅,✅
+uitars,✅,✅
+omniparser,❌,✅
+gta1,❌,✅

libs/python/agent/cua_agent/loops/moondream3.py

@@ -0,0 +1,493 @@
+"""
+Moondream3+ composed-grounded agent loop implementation.
+Grounding is handled by a local Moondream3 preview model via Transformers.
+Thinking is delegated to the trailing LLM in the composed model string: "moondream3+<thinking_model>".
+
+Differences from composed_grounded:
+- Provides a singleton Moondream3 client outside the class.
+- predict_click uses model.point(image, instruction, settings={"max_objects": 1}) and returns pixel coordinates.
+- If the last image was a screenshot (or we take one), run model.detect(image, "all form ui") to get bboxes, then
+  run model.caption on each cropped bbox to label it. Overlay labels on the screenshot and emit via _on_screenshot.
+- Add a user message listing all detected form UI names so the thinker can reference them.
+- If the thinking model doesn't support vision, filter out image content before calling litellm.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import uuid
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from PIL import Image, ImageDraw, ImageFont
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_completion_messages_to_responses_items,
+    convert_computer_calls_desc2xy,
+    convert_computer_calls_xy2desc,
+    convert_responses_items_to_completion_messages,
+    get_all_element_descriptions,
+)
+from ..types import AgentCapability
+
+_MOONDREAM_SINGLETON = None
+
+
+def get_moondream_model() -> Any:
+    """Get a singleton instance of the Moondream3 preview model."""
+    global _MOONDREAM_SINGLETON
+    if _MOONDREAM_SINGLETON is None:
+        try:
+            import torch
+            from transformers import AutoModelForCausalLM
+
+            _MOONDREAM_SINGLETON = AutoModelForCausalLM.from_pretrained(
+                "moondream/moondream3-preview",
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+                device_map="cuda",
+            )
+        except ImportError as e:
+            raise RuntimeError(
+                "moondream3 requires torch and transformers. Install with: pip install cua-agent[moondream3]"
+            ) from e
+    return _MOONDREAM_SINGLETON
+
+
+def _decode_image_b64(image_b64: str) -> Image.Image:
+    data = base64.b64decode(image_b64)
+    return Image.open(io.BytesIO(data)).convert("RGB")
+
+
+def _image_to_b64(img: Image.Image) -> str:
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+
+def _supports_vision(model: str) -> bool:
+    """Heuristic vision support detection for thinking model."""
+    m = model.lower()
+    vision_markers = [
+        "gpt-4o",
+        "gpt-4.1",
+        "o1",
+        "o3",
+        "claude-3",
+        "claude-3.5",
+        "sonnet",
+        "haiku",
+        "opus",
+        "gemini-1.5",
+        "llava",
+    ]
+    return any(v in m for v in vision_markers)
+
+
+def _filter_images_from_completion_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    filtered: List[Dict[str, Any]] = []
+    for msg in messages:
+        msg_copy = {**msg}
+        content = msg_copy.get("content")
+        if isinstance(content, list):
+            msg_copy["content"] = [c for c in content if c.get("type") != "image_url"]
+        filtered.append(msg_copy)
+    return filtered
+
+
+def _annotate_detect_and_label_ui(base_img: Image.Image, model_md) -> Tuple[str, List[str]]:
+    """Detect UI elements with Moondream, caption each, draw labels with backgrounds.
+
+    Args:
+        base_img: PIL image of the screenshot (RGB or RGBA). Will be copied/converted internally.
+        model_md: Moondream model instance with .detect() and .query() methods.
+
+    Returns:
+        A tuple of (annotated_image_base64_png, detected_names)
+    """
+    # Ensure RGBA for semi-transparent fills
+    if base_img.mode != "RGBA":
+        base_img = base_img.convert("RGBA")
+    W, H = base_img.width, base_img.height
+
+    # Detect objects
+    try:
+        detect_result = model_md.detect(base_img, "all ui elements")
+        objects = detect_result.get("objects", []) if isinstance(detect_result, dict) else []
+    except Exception:
+        objects = []
+
+    draw = ImageDraw.Draw(base_img)
+    try:
+        font = ImageFont.load_default()
+    except Exception:
+        font = None
+
+    detected_names: List[str] = []
+
+    for i, obj in enumerate(objects):
+        try:
+            # Clamp normalized coords and crop
+            x_min = max(0.0, min(1.0, float(obj.get("x_min", 0.0))))
+            y_min = max(0.0, min(1.0, float(obj.get("y_min", 0.0))))
+            x_max = max(0.0, min(1.0, float(obj.get("x_max", 0.0))))
+            y_max = max(0.0, min(1.0, float(obj.get("y_max", 0.0))))
+            left, top, right, bottom = (
+                int(x_min * W),
+                int(y_min * H),
+                int(x_max * W),
+                int(y_max * H),
+            )
+            left, top = max(0, left), max(0, top)
+            right, bottom = min(W - 1, right), min(H - 1, bottom)
+            crop = base_img.crop((left, top, right, bottom))
+
+            # Prompted short caption
+            try:
+                result = model_md.query(crop, "Caption this UI element in few words.")
+                caption_text = (result or {}).get("answer", "")
+            except Exception:
+                caption_text = ""
+
+            name = (caption_text or "").strip() or f"element_{i+1}"
+            detected_names.append(name)
+
+            # Draw bbox
+            draw.rectangle([left, top, right, bottom], outline=(255, 215, 0, 255), width=2)
+
+            # Label background with padding and rounded corners
+            label = f"{i+1}. {name}"
+            padding = 3
+            if font:
+                text_bbox = draw.textbbox((0, 0), label, font=font)
+            else:
+                text_bbox = draw.textbbox((0, 0), label)
+            text_w = text_bbox[2] - text_bbox[0]
+            text_h = text_bbox[3] - text_bbox[1]
+
+            tx = left + 3
+            ty = top - (text_h + 2 * padding + 4)
+            if ty < 0:
+                ty = top + 3
+
+            bg_left = tx - padding
+            bg_top = ty - padding
+            bg_right = tx + text_w + padding
+            bg_bottom = ty + text_h + padding
+            try:
+                draw.rounded_rectangle(
+                    [bg_left, bg_top, bg_right, bg_bottom],
+                    radius=4,
+                    fill=(0, 0, 0, 160),
+                    outline=(255, 215, 0, 200),
+                    width=1,
+                )
+            except Exception:
+                draw.rectangle(
+                    [bg_left, bg_top, bg_right, bg_bottom],
+                    fill=(0, 0, 0, 160),
+                    outline=(255, 215, 0, 200),
+                    width=1,
+                )
+
+            text_fill = (255, 255, 255, 255)
+            if font:
+                draw.text((tx, ty), label, fill=text_fill, font=font)
+            else:
+                draw.text((tx, ty), label, fill=text_fill)
+        except Exception:
+            continue
+
+    # Encode PNG base64
+    annotated = base_img
+    if annotated.mode not in ("RGBA", "RGB"):
+        annotated = annotated.convert("RGBA")
+    annotated_b64 = _image_to_b64(annotated)
+    return annotated_b64, detected_names
+
+
+GROUNDED_COMPUTER_TOOL_SCHEMA = {
+    "type": "function",
+    "function": {
+        "name": "computer",
+        "description": (
+            "Control a computer by taking screenshots and interacting with UI elements. "
+            "The screenshot action will include a list of detected form UI element names when available. "
+            "Use element descriptions to locate and interact with UI elements on the screen."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "type": "string",
+                    "enum": [
+                        "screenshot",
+                        "click",
+                        "double_click",
+                        "drag",
+                        "type",
+                        "keypress",
+                        "scroll",
+                        "move",
+                        "wait",
+                        "get_current_url",
+                        "get_dimensions",
+                        "get_environment",
+                    ],
+                    "description": "The action to perform (required for all actions)",
+                },
+                "element_description": {
+                    "type": "string",
+                    "description": "Description of the element to interact with (required for click/double_click/move/scroll)",
+                },
+                "start_element_description": {
+                    "type": "string",
+                    "description": "Description of the element to start dragging from (required for drag)",
+                },
+                "end_element_description": {
+                    "type": "string",
+                    "description": "Description of the element to drag to (required for drag)",
+                },
+                "text": {
+                    "type": "string",
+                    "description": "The text to type (required for type)",
+                },
+                "keys": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                    "description": "Key(s) to press (required for keypress)",
+                },
+                "button": {
+                    "type": "string",
+                    "enum": ["left", "right", "wheel", "back", "forward"],
+                    "description": "The mouse button to use for click/double_click",
+                },
+                "scroll_x": {
+                    "type": "integer",
+                    "description": "Horizontal scroll amount (required for scroll)",
+                },
+                "scroll_y": {
+                    "type": "integer",
+                    "description": "Vertical scroll amount (required for scroll)",
+                },
+            },
+            "required": ["action"],
+        },
+    },
+}
+
+
+@register_agent(r"moondream3\+.*", priority=2)
+class Moondream3PlusConfig(AsyncAgentConfig):
+    def __init__(self):
+        self.desc2xy: Dict[str, Tuple[float, float]] = {}
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        # Parse composed model: moondream3+<thinking_model>
+        if "+" not in model:
+            raise ValueError(f"Composed model must be 'moondream3+<thinking_model>', got: {model}")
+        _, thinking_model = model.split("+", 1)
+
+        pre_output_items: List[Dict[str, Any]] = []
+
+        # Acquire last screenshot; if missing, take one
+        last_image_b64: Optional[str] = None
+        for message in reversed(messages):
+            if (
+                isinstance(message, dict)
+                and message.get("type") == "computer_call_output"
+                and isinstance(message.get("output"), dict)
+                and message["output"].get("type") == "input_image"
+            ):
+                image_url = message["output"].get("image_url", "")
+                if image_url.startswith("data:image/png;base64,"):
+                    last_image_b64 = image_url.split(",", 1)[1]
+                    break
+
+        if last_image_b64 is None and computer_handler is not None:
+            # Take a screenshot
+            screenshot_b64 = await computer_handler.screenshot()  # type: ignore
+            if screenshot_b64:
+                call_id = uuid.uuid4().hex
+                pre_output_items += [
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "type": "output_text",
+                                "text": "Taking a screenshot to analyze the current screen.",
+                            }
+                        ],
+                    },
+                    {
+                        "type": "computer_call",
+                        "call_id": call_id,
+                        "status": "completed",
+                        "action": {"type": "screenshot"},
+                    },
+                    {
+                        "type": "computer_call_output",
+                        "call_id": call_id,
+                        "output": {
+                            "type": "input_image",
+                            "image_url": f"data:image/png;base64,{screenshot_b64}",
+                        },
+                    },
+                ]
+                last_image_b64 = screenshot_b64
+                if _on_screenshot:
+                    await _on_screenshot(screenshot_b64)
+
+        # If we have a last screenshot, run Moondream detection and labeling
+        detected_names: List[str] = []
+        if last_image_b64 is not None:
+            base_img = _decode_image_b64(last_image_b64)
+            model_md = get_moondream_model()
+            annotated_b64, detected_names = _annotate_detect_and_label_ui(base_img, model_md)
+            if _on_screenshot:
+                await _on_screenshot(annotated_b64, "annotated_form_ui")
+
+            # Also push a user message listing all detected names
+            if detected_names:
+                names_text = "\n".join(f"- {n}" for n in detected_names)
+                pre_output_items.append(
+                    {
+                        "type": "message",
+                        "role": "user",
+                        "content": [
+                            {"type": "input_text", "text": "Detected form UI elements on screen:"},
+                            {"type": "input_text", "text": names_text},
+                            {
+                                "type": "input_text",
+                                "text": "Please continue with the next action needed to perform your task.",
+                            },
+                        ],
+                    }
+                )
+
+        tool_schemas = []
+        for schema in tools or []:
+            if schema.get("type") == "computer":
+                tool_schemas.append(GROUNDED_COMPUTER_TOOL_SCHEMA)
+            else:
+                tool_schemas.append(schema)
+
+        # Step 1: Convert computer calls from xy to descriptions
+        input_messages = messages + pre_output_items
+        messages_with_descriptions = convert_computer_calls_xy2desc(input_messages, self.desc2xy)
+
+        # Step 2: Convert responses items to completion messages
+        completion_messages = convert_responses_items_to_completion_messages(
+            messages_with_descriptions,
+            allow_images_in_tool_results=False,
+        )
+
+        # Optionally filter images if model lacks vision
+        if not _supports_vision(thinking_model):
+            completion_messages = _filter_images_from_completion_messages(completion_messages)
+
+        # Step 3: Call thinking model with litellm.acompletion
+        api_kwargs = {
+            "model": thinking_model,
+            "messages": completion_messages,
+            "tools": tool_schemas,
+            "max_retries": max_retries,
+            "stream": stream,
+            **kwargs,
+        }
+        if use_prompt_caching:
+            api_kwargs["use_prompt_caching"] = use_prompt_caching
+
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        response = await litellm.acompletion(**api_kwargs)
+
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        usage = {
+            **response.usage.model_dump(),  # type: ignore
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        # Step 4: Convert completion messages back to responses items format
+        response_dict = response.model_dump()  # type: ignore
+        choice_messages = [choice["message"] for choice in response_dict["choices"]]
+        thinking_output_items: List[Dict[str, Any]] = []
+        for choice_message in choice_messages:
+            thinking_output_items.extend(
+                convert_completion_messages_to_responses_items([choice_message])
+            )
+
+        # Step 5: Use Moondream to get coordinates for each description
+        element_descriptions = get_all_element_descriptions(thinking_output_items)
+        if element_descriptions and last_image_b64:
+            for desc in element_descriptions:
+                for _ in range(3):  # try 3 times
+                    coords = await self.predict_click(
+                        model=model,
+                        image_b64=last_image_b64,
+                        instruction=desc,
+                    )
+                    if coords:
+                        self.desc2xy[desc] = coords
+                        break
+
+        # Step 6: Convert computer calls from descriptions back to xy coordinates
+        final_output_items = convert_computer_calls_desc2xy(thinking_output_items, self.desc2xy)
+
+        # Step 7: Return output and usage
+        return {"output": pre_output_items + final_output_items, "usage": usage}
+
+    async def predict_click(
+        self,
+        model: str,
+        image_b64: str,
+        instruction: str,
+        **kwargs,
+    ) -> Optional[Tuple[float, float]]:
+        """Predict click coordinates using Moondream3's point API.
+
+        Returns pixel coordinates (x, y) as floats.
+        """
+        img = _decode_image_b64(image_b64)
+        W, H = img.width, img.height
+        model_md = get_moondream_model()
+        try:
+            result = model_md.point(img, instruction, settings={"max_objects": 1})
+        except Exception:
+            return None
+
+        try:
+            pt = (result or {}).get("points", [])[0]
+            x_norm = float(pt.get("x", 0.0))
+            y_norm = float(pt.get("y", 0.0))
+            x_px = max(0.0, min(float(W - 1), x_norm * W))
+            y_px = max(0.0, min(float(H - 1), y_norm * H))
+            return (x_px, y_px)
+        except Exception:
+            return None
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        return ["click", "step"]

libs/python/agent/cua_agent/loops/omniparser.py

@@ -0,0 +1,533 @@
+"""
+OpenAI computer-use-preview agent loop implementation using liteLLM
+Paper: https://arxiv.org/abs/2408.00203
+Code: https://github.com/microsoft/OmniParser
+"""
+
+import asyncio
+import base64
+import inspect
+import json
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+
+import litellm
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_completion_messages_to_responses_items,
+    convert_responses_items_to_completion_messages,
+)
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+
+SOM_TOOL_SCHEMA = {
+    "type": "function",
+    "function": {
+        "name": "computer",
+        "description": "Control a computer by taking screenshots and interacting with UI elements. This tool shows screenshots with numbered elements overlaid on them. Each UI element has been assigned a unique ID number that you can see in the image. Use the element's ID number to interact with any element instead of pixel coordinates.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "type": "string",
+                    "enum": [
+                        "screenshot",
+                        "click",
+                        "double_click",
+                        "drag",
+                        "type",
+                        "keypress",
+                        "scroll",
+                        "move",
+                        "wait",
+                        "get_current_url",
+                        "get_dimensions",
+                        "get_environment",
+                    ],
+                    "description": "The action to perform",
+                },
+                "element_id": {
+                    "type": "integer",
+                    "description": "The ID of the element to interact with (required for click, double_click, move, scroll actions, and as start/end for drag)",
+                },
+                "start_element_id": {
+                    "type": "integer",
+                    "description": "The ID of the element to start dragging from (required for drag action)",
+                },
+                "end_element_id": {
+                    "type": "integer",
+                    "description": "The ID of the element to drag to (required for drag action)",
+                },
+                "text": {
+                    "type": "string",
+                    "description": "The text to type (required for type action)",
+                },
+                "keys": {
+                    "type": "string",
+                    "description": "Key combination to press (required for keypress action). Single key for individual key press, multiple keys for combinations (e.g., 'ctrl+c')",
+                },
+                "button": {
+                    "type": "string",
+                    "description": "The mouse button to use for click action (left, right, wheel, back, forward) Default: left",
+                },
+                "scroll_x": {
+                    "type": "integer",
+                    "description": "Horizontal scroll amount for scroll action (positive for right, negative for left)",
+                },
+                "scroll_y": {
+                    "type": "integer",
+                    "description": "Vertical scroll amount for scroll action (positive for down, negative for up)",
+                },
+            },
+            "required": ["action", "element_id"],
+        },
+    },
+}
+
+OMNIPARSER_AVAILABLE = False
+try:
+    from som import OmniParser
+
+    OMNIPARSER_AVAILABLE = True
+except ImportError:
+    pass
+OMNIPARSER_SINGLETON = None
+
+
+def get_parser():
+    global OMNIPARSER_SINGLETON
+    if OMNIPARSER_SINGLETON is None:
+        OMNIPARSER_SINGLETON = OmniParser()
+    return OMNIPARSER_SINGLETON
+
+
+def get_last_computer_call_output(messages: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Get the last computer_call_output message from a messages list.
+
+    Args:
+        messages: List of messages to search through
+
+    Returns:
+        The last computer_call_output message dict, or None if not found
+    """
+    for message in reversed(messages):
+        if isinstance(message, dict) and message.get("type") == "computer_call_output":
+            return message
+    return None
+
+
+def _prepare_tools_for_omniparser(tool_schemas: List[Dict[str, Any]]) -> Tuple[Tools, dict]:
+    """Prepare tools for OpenAI API format"""
+    omniparser_tools = []
+    id2xy = dict()
+
+    for schema in tool_schemas:
+        if schema["type"] == "computer":
+            omniparser_tools.append(SOM_TOOL_SCHEMA)
+            if "id2xy" in schema:
+                id2xy = schema["id2xy"]
+            else:
+                schema["id2xy"] = id2xy
+        elif schema["type"] == "function":
+            # Function tools use OpenAI-compatible schema directly (liteLLM expects this format)
+            # Schema should be: {type, name, description, parameters}
+            omniparser_tools.append({"type": "function", **schema["function"]})
+
+    return omniparser_tools, id2xy
+
+
+async def replace_function_with_computer_call(
+    item: Dict[str, Any], id2xy: Dict[int, Tuple[float, float]]
+):
+    item_type = item.get("type")
+
+    def _get_xy(element_id: Optional[int]) -> Union[Tuple[float, float], Tuple[None, None]]:
+        if element_id is None:
+            return (None, None)
+        return id2xy.get(element_id, (None, None))
+
+    if item_type == "function_call":
+        fn_name = item.get("name")
+        fn_args = json.loads(item.get("arguments", "{}"))
+
+        item_id = item.get("id")
+        call_id = item.get("call_id")
+
+        if fn_name == "computer":
+            action = fn_args.get("action")
+            element_id = fn_args.get("element_id")
+            start_element_id = fn_args.get("start_element_id")
+            end_element_id = fn_args.get("end_element_id")
+            text = fn_args.get("text")
+            keys = fn_args.get("keys")
+            button = fn_args.get("button")
+            scroll_x = fn_args.get("scroll_x")
+            scroll_y = fn_args.get("scroll_y")
+
+            x, y = _get_xy(element_id)
+            start_x, start_y = _get_xy(start_element_id)
+            end_x, end_y = _get_xy(end_element_id)
+
+            action_args = {
+                "type": action,
+                "x": x,
+                "y": y,
+                "start_x": start_x,
+                "start_y": start_y,
+                "end_x": end_x,
+                "end_y": end_y,
+                "text": text,
+                "keys": keys,
+                "button": button,
+                "scroll_x": scroll_x,
+                "scroll_y": scroll_y,
+            }
+            # Remove None values to keep the JSON clean
+            action_args = {k: v for k, v in action_args.items() if v is not None}
+
+            return [
+                {
+                    "type": "computer_call",
+                    "action": action_args,
+                    "id": item_id,
+                    "call_id": call_id,
+                    "status": "completed",
+                }
+            ]
+
+    return [item]
+
+
+async def replace_computer_call_with_function(
+    item: Dict[str, Any], xy2id: Dict[Tuple[float, float], int]
+):
+    """
+    Convert computer_call back to function_call format.
+    Also handles computer_call_output -> function_call_output conversion.
+
+    Args:
+        item: The item to convert
+        xy2id: Mapping from (x, y) coordinates to element IDs
+    """
+    item_type = item.get("type")
+
+    def _get_element_id(x: Optional[float], y: Optional[float]) -> Optional[int]:
+        """Get element ID from coordinates, return None if coordinates are None"""
+        if x is None or y is None:
+            return None
+        return xy2id.get((x, y))
+
+    if item_type == "computer_call":
+        action_data = item.get("action", {})
+
+        # Extract coordinates and convert back to element IDs
+        element_id = _get_element_id(action_data.get("x"), action_data.get("y"))
+        start_element_id = _get_element_id(action_data.get("start_x"), action_data.get("start_y"))
+        end_element_id = _get_element_id(action_data.get("end_x"), action_data.get("end_y"))
+
+        # Build function arguments
+        fn_args = {
+            "action": action_data.get("type"),
+            "element_id": element_id,
+            "start_element_id": start_element_id,
+            "end_element_id": end_element_id,
+            "text": action_data.get("text"),
+            "keys": action_data.get("keys"),
+            "button": action_data.get("button"),
+            "scroll_x": action_data.get("scroll_x"),
+            "scroll_y": action_data.get("scroll_y"),
+        }
+
+        # Remove None values to keep the JSON clean
+        fn_args = {k: v for k, v in fn_args.items() if v is not None}
+
+        return [
+            {
+                "type": "function_call",
+                "name": "computer",
+                "arguments": json.dumps(fn_args),
+                "id": item.get("id"),
+                "call_id": item.get("call_id"),
+                "status": "completed",
+            }
+        ]
+
+    elif item_type == "computer_call_output":
+        output = item.get("output")
+
+        if isinstance(output, dict):
+            output = [output]
+
+        return [
+            {
+                "type": "function_call_output",
+                "call_id": item.get("call_id"),
+                "output": item.get("output"),
+                "id": item.get("id"),
+                "status": "completed",
+            }
+        ]
+
+    return [item]
+
+
+@register_agent(models=r"omniparser\+.*|omni\+.*", priority=2)
+class OmniparserConfig(AsyncAgentConfig):
+    """Omniparser agent configuration implementing AsyncAgentConfig protocol."""
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        OpenAI computer-use-preview agent loop using liteLLM responses.
+
+        Supports OpenAI's computer use preview models.
+        """
+        if not OMNIPARSER_AVAILABLE:
+            raise ValueError(
+                "omniparser loop requires som to be installed. Install it with `pip install cua-som`."
+            )
+
+        tools = tools or []
+
+        llm_model = model.split("+")[-1]
+
+        # Get screen dimensions from computer handler
+        try:
+            width, height = await computer_handler.get_dimensions()
+        except Exception:
+            # Fallback to default dimensions if method fails
+            width, height = 1024, 768
+
+        # Prepare tools for OpenAI API
+        openai_tools, id2xy = _prepare_tools_for_omniparser(tools)
+
+        # Build per-screenshot element mappings for historical consistency
+        screenshot_mappings = []  # (message_index, xy2id)
+
+        parser = get_parser()
+
+        for idx, message in enumerate(messages):
+            if not isinstance(message, dict):
+                message = message.__dict__
+
+            if message.get("type") == "computer_call_output":
+                image_url = message.get("output", {}).get("image_url", "")
+                if not image_url:
+                    continue
+
+                image_data = image_url.split(",")[-1]
+                if not image_data:
+                    continue
+
+                result = parser.parse(image_data)
+
+                if _on_screenshot:
+                    await _on_screenshot(result.annotated_image_base64, "annotated_image")
+
+                local_id2xy = {}
+
+                for element in result.elements:
+                    norm_x = (element.bbox.x1 + element.bbox.x2) / 2
+                    norm_y = (element.bbox.y1 + element.bbox.y2) / 2
+                    pixel_x = int(norm_x * width)
+                    pixel_y = int(norm_y * height)
+                    local_id2xy[element.id] = (pixel_x, pixel_y)
+
+                xy2id = {v: k for k, v in local_id2xy.items()}
+                screenshot_mappings.append((idx, xy2id))
+
+                # Replace screenshot with annotated image
+                message["output"][
+                    "image_url"
+                ] = f"data:image/png;base64,{result.annotated_image_base64}"
+
+        def get_mapping_for_index(index):
+            applicable = [m for i, m in screenshot_mappings if i <= index]
+            return applicable[-1] if applicable else {}
+
+        messages_with_element_ids = []
+
+        for i, message in enumerate(messages):
+            if not isinstance(message, dict):
+                message = message.__dict__
+
+            xy2id = get_mapping_for_index(i)
+            converted = await replace_computer_call_with_function(message, xy2id)
+            messages_with_element_ids.extend(converted)
+
+        completion_messages = convert_responses_items_to_completion_messages(
+            messages_with_element_ids, allow_images_in_tool_results=False
+        )
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": llm_model,
+            "messages": completion_messages,
+            "tools": openai_tools if openai_tools else None,
+            "stream": stream,
+            "num_retries": max_retries,
+            **kwargs,
+        }
+
+        # Add Vertex AI specific parameters if using vertex_ai models
+        if llm_model.startswith("vertex_ai/"):
+            import os
+
+            # Pass vertex_project and vertex_location to liteLLM
+            if "vertex_project" not in api_kwargs:
+                api_kwargs["vertex_project"] = os.getenv("GOOGLE_CLOUD_PROJECT")
+            if "vertex_location" not in api_kwargs:
+                api_kwargs["vertex_location"] = "global"
+
+            # Pass through Gemini 3-specific parameters if provided
+            if "thinking_level" in kwargs:
+                api_kwargs["thinking_level"] = kwargs["thinking_level"]
+            if "media_resolution" in kwargs:
+                api_kwargs["media_resolution"] = kwargs["media_resolution"]
+
+        # Call API start hook
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        print(str(api_kwargs)[:1000])
+
+        # Use liteLLM completion
+        response = await litellm.acompletion(**api_kwargs)
+
+        # Call API end hook
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        # Extract usage information
+        usage = {
+            **response.usage.model_dump(),  # type: ignore
+            "response_cost": response._hidden_params.get("response_cost", 0.0),  # type: ignore
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        response_dict = response.model_dump()  # type: ignore
+        choice_messages = [choice["message"] for choice in response_dict["choices"]]
+        responses_items = []
+        for choice_message in choice_messages:
+            responses_items.extend(convert_completion_messages_to_responses_items([choice_message]))
+
+        # Convert element_id → x,y (similar to moondream's convert_computer_calls_desc2xy)
+        final_output = []
+        for item in responses_items:
+            if item.get("type") == "computer_call" and "action" in item:
+                action = item["action"].copy()
+
+                # Handle single element_id
+                if "element_id" in action:
+                    element_id = action["element_id"]
+                    if element_id in id2xy:
+                        x, y = id2xy[element_id]
+                        action["x"] = x
+                        action["y"] = y
+                        del action["element_id"]
+
+                # Handle start_element_id and end_element_id for drag operations
+                elif "start_element_id" in action and "end_element_id" in action:
+                    start_id = action["start_element_id"]
+                    end_id = action["end_element_id"]
+                    if start_id in id2xy and end_id in id2xy:
+                        start_x, start_y = id2xy[start_id]
+                        end_x, end_y = id2xy[end_id]
+                        action["path"] = [{"x": start_x, "y": start_y}, {"x": end_x, "y": end_y}]
+                        del action["start_element_id"]
+                        del action["end_element_id"]
+
+                converted_item = item.copy()
+                converted_item["action"] = action
+                final_output.append(converted_item)
+            else:
+                final_output.append(item)
+
+        return {"output": final_output, "usage": usage}
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[float, float]]:
+        """
+        Predict click coordinates using OmniParser and LLM.
+
+        Uses OmniParser to annotate the image with element IDs, then uses LLM
+        to identify the correct element ID based on the instruction.
+        """
+        if not OMNIPARSER_AVAILABLE:
+            return None
+
+        # Parse the image with OmniParser to get annotated image and elements
+        parser = get_parser()
+        result = parser.parse(image_b64)
+
+        # Extract the LLM model from composed model string
+        llm_model = model.split("+")[-1]
+
+        # Create system prompt for element ID prediction
+        SYSTEM_PROMPT = """
+You are an expert UI element locator. Given a GUI image annotated with numerical IDs over each interactable element, along with a user's element description, provide the ID of the specified element.
+
+The image shows UI elements with numbered overlays. Each number corresponds to a clickable/interactable element.
+
+Output only the element ID as a single integer.
+""".strip()
+
+        # Prepare messages for LLM
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{result.annotated_image_base64}"
+                        },
+                    },
+                    {"type": "text", "text": f"Find the element: {instruction}"},
+                ],
+            },
+        ]
+
+        # Call LLM to predict element ID
+        response = await litellm.acompletion(
+            model=llm_model, messages=messages, max_tokens=10, temperature=0.1
+        )
+
+        # Extract element ID from response
+        response_text = response.choices[0].message.content.strip()  # type: ignore
+
+        # Try to parse the element ID
+        try:
+            element_id = int(response_text)
+
+            # Find the element with this ID and return its center coordinates
+            for element in result.elements:
+                if element.id == element_id:
+                    center_x = (element.bbox.x1 + element.bbox.x2) / 2
+                    center_y = (element.bbox.y1 + element.bbox.y2) / 2
+                    return (center_x, center_y)
+        except ValueError:
+            # If we can't parse the ID, return None
+            pass
+
+        return None
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """Return the capabilities supported by this agent."""
+        return ["step"]

libs/python/agent/cua_agent/loops/openai.py

@@ -0,0 +1,426 @@
+"""
+OpenAI computer-use-preview agent loop implementation using liteLLM
+"""
+
+import asyncio
+import base64
+import json
+from io import BytesIO
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+
+import litellm
+from PIL import Image
+
+from ..decorators import register_agent
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+
+
+async def _map_computer_tool_to_openai(
+    computer_handler: Any, use_native_tool: bool = True
+) -> Dict[str, Any]:
+    """Map a computer tool to OpenAI's tool schema.
+
+    Args:
+        computer_handler: The computer handler instance
+        use_native_tool: If True, use native computer_use_preview format (for computer-use-preview model).
+                        If False, use standard function calling format (for GPT-5.4 etc).
+    """
+    # Get dimensions from the computer handler
+    try:
+        width, height = await computer_handler.get_dimensions()
+    except Exception:
+        # Fallback to default dimensions if method fails
+        width, height = 1024, 768
+
+    # Get environment from the computer handler
+    try:
+        environment = await computer_handler.get_environment()
+    except Exception:
+        # Fallback to default environment if method fails
+        environment = "linux"
+
+    if use_native_tool:
+        # Native computer_use_preview format (for computer-use-preview model)
+        return {
+            "type": "computer_use_preview",
+            "display_width": width,
+            "display_height": height,
+            "environment": environment,  # mac, windows, linux, browser
+        }
+    else:
+        # Standard function calling format (for GPT-5.4 etc)
+        # Responses API requires: {type, name, description, parameters} at root level
+        return {
+            "type": "function",
+            "name": "computer",
+            "description": (
+                f"Use a mouse and keyboard to interact with a computer, and take screenshots.\n"
+                f"Screen resolution: {width}x{height} pixels.\n"
+                f"Environment: {environment}."
+            ),
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "action": {
+                        "description": "The action to perform.",
+                        "type": "string",
+                        "enum": [
+                            "click",
+                            "double_click",
+                            "right_click",
+                            "type",
+                            "keypress",
+                            "scroll",
+                            "move",
+                            "drag",
+                            "screenshot",
+                            "wait",
+                            "terminate",
+                        ],
+                    },
+                    "x": {
+                        "description": "X coordinate for click/move/scroll actions.",
+                        "type": "integer",
+                    },
+                    "y": {
+                        "description": "Y coordinate for click/move/scroll actions.",
+                        "type": "integer",
+                    },
+                    "text": {
+                        "description": "Text to type (for action=type).",
+                        "type": "string",
+                    },
+                    "keys": {
+                        "description": "Keys to press (for action=keypress). Example: ['ctrl', 'c']",
+                        "type": "array",
+                        "items": {"type": "string"},
+                    },
+                    "scroll_x": {
+                        "description": "Horizontal scroll amount. Positive=right, negative=left.",
+                        "type": "integer",
+                    },
+                    "scroll_y": {
+                        "description": "Vertical scroll amount. Positive=down, negative=up.",
+                        "type": "integer",
+                    },
+                    "button": {
+                        "description": "Mouse button for click action.",
+                        "type": "string",
+                        "enum": ["left", "right", "middle"],
+                    },
+                    "start_x": {
+                        "description": "Starting X coordinate for drag action.",
+                        "type": "integer",
+                    },
+                    "start_y": {
+                        "description": "Starting Y coordinate for drag action.",
+                        "type": "integer",
+                    },
+                    "end_x": {
+                        "description": "Ending X coordinate for drag action.",
+                        "type": "integer",
+                    },
+                    "end_y": {
+                        "description": "Ending Y coordinate for drag action.",
+                        "type": "integer",
+                    },
+                    "status": {
+                        "description": "Status for terminate action.",
+                        "type": "string",
+                        "enum": ["success", "failure"],
+                    },
+                },
+                "required": ["action"],
+            },
+        }
+
+
+def _is_native_computer_use_model(model: str) -> bool:
+    """Check if the model supports native computer_use_preview tool format."""
+    import re
+
+    # Only computer-use-preview models support native computer_use_preview tool
+    # GPT 5.4 does NOT support computer_use_preview - it uses function calling
+    return bool(re.search(r"computer-use-preview", model, re.IGNORECASE))
+
+
+async def _prepare_tools_for_openai(tool_schemas: List[Dict[str, Any]], model: str = "") -> Tools:
+    """Prepare tools for OpenAI API format.
+
+    Args:
+        tool_schemas: List of tool schemas to prepare
+        model: Model name to determine tool format
+    """
+    openai_tools = []
+    use_native = _is_native_computer_use_model(model)
+
+    for schema in tool_schemas:
+        if schema["type"] == "computer":
+            # Map computer tool to OpenAI format (native or function based on model)
+            computer_tool = await _map_computer_tool_to_openai(
+                schema["computer"], use_native_tool=use_native
+            )
+            openai_tools.append(computer_tool)
+        elif schema["type"] == "function":
+            # Function tools for Responses API need: {type, name, description, parameters}
+            # Note: parameters are at the root level, NOT nested under 'function'
+            func = schema["function"]
+            openai_tools.append(
+                {
+                    "type": "function",
+                    "name": func["name"],
+                    "description": func.get("description", ""),
+                    "parameters": func.get("parameters", {}),
+                }
+            )
+
+    return openai_tools
+
+
+@register_agent(models=r".*(computer-use-preview|gpt-?5\.?4)")
+class OpenAIComputerUseConfig:
+    """
+    OpenAI computer-use-preview agent configuration using liteLLM responses.
+
+    Supports OpenAI's computer use preview models.
+    """
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Predict the next step based on input items.
+
+        Args:
+            messages: Input items following Responses format
+            model: Model name to use
+            tools: Optional list of tool schemas
+            max_retries: Maximum number of retries
+            stream: Whether to stream responses
+            computer_handler: Computer handler instance
+            _on_api_start: Callback for API start
+            _on_api_end: Callback for API end
+            _on_usage: Callback for usage tracking
+            _on_screenshot: Callback for screenshot events
+            **kwargs: Additional arguments
+
+        Returns:
+            Dictionary with "output" (output items) and "usage" array
+        """
+        tools = tools or []
+
+        # Prepare tools for OpenAI API
+        openai_tools = await _prepare_tools_for_openai(tools, model=model)
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "input": messages,
+            "tools": openai_tools if openai_tools else None,
+            "stream": stream,
+            "reasoning": {"summary": "concise"},
+            "truncation": "auto",
+            "num_retries": max_retries,
+            "request_timeout": kwargs.pop("request_timeout", 120),
+            **kwargs,
+        }
+
+        # Call API start hook
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        # Use liteLLM responses
+        response = await litellm.aresponses(**api_kwargs)
+
+        # Call API end hook
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        # Extract usage information - handle both dict and Pydantic model responses
+        if isinstance(response, dict):
+            response_usage = response.get("usage", {})
+            usage = response_usage if isinstance(response_usage, dict) else {}
+            output_dict = response
+        else:
+            # Response is a Pydantic model - but usage might be dict or model
+            response_usage = response.usage
+            if hasattr(response_usage, "model_dump"):
+                usage = response_usage.model_dump()
+            elif isinstance(response_usage, dict):
+                usage = response_usage
+            else:
+                usage = {}
+            output_dict = response.model_dump()
+
+        # Add response cost if available
+        if hasattr(response, "_hidden_params"):
+            usage["response_cost"] = response._hidden_params.get("response_cost", 0.0)
+        elif isinstance(response, dict):
+            usage["response_cost"] = response.get("_hidden_params", {}).get("response_cost", 0.0)
+
+        if _on_usage:
+            await _on_usage(usage)
+
+        # Return in the expected format
+        output_dict["usage"] = usage
+        return output_dict
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates based on image and instruction.
+
+        Uses OpenAI computer-use-preview with manually constructed input items
+        and a prompt that instructs the agent to only output clicks.
+
+        Args:
+            model: Model name to use
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        # TODO: use computer tool to get dimensions + environment
+        # Manually construct input items with image and click instruction
+        input_items = [
+            {
+                "role": "user",
+                "content": f"""You are a UI grounding expert. 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.
+Task: Click {instruction}. Output ONLY a click action on the target element.""",
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "input_image", "image_url": f"data:image/png;base64,{image_b64}"}
+                ],
+            },
+        ]
+
+        # Get image dimensions from base64 data
+        try:
+            image_data = base64.b64decode(image_b64)
+            image = Image.open(BytesIO(image_data))
+            display_width, display_height = image.size
+        except Exception:
+            # Fallback to default dimensions if image parsing fails
+            display_width, display_height = 1024, 768
+
+        # Prepare computer tool for click actions - use native format only for models that support it
+        use_native = _is_native_computer_use_model(model)
+        if use_native:
+            # Native computer_use_preview format (for computer-use-preview model)
+            computer_tool = {
+                "type": "computer_use_preview",
+                "display_width": display_width,
+                "display_height": display_height,
+                "environment": "windows",
+            }
+        else:
+            # Standard function calling format (for GPT-5.4 etc)
+            computer_tool = {
+                "type": "function",
+                "name": "computer",
+                "description": (
+                    f"Use a mouse and keyboard to interact with a computer, and take screenshots.\n"
+                    f"Screen resolution: {display_width}x{display_height} pixels.\n"
+                    f"Environment: windows."
+                ),
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "action": {
+                            "description": "The action to perform.",
+                            "type": "string",
+                            "enum": ["click"],
+                        },
+                        "x": {
+                            "description": "X coordinate for click action.",
+                            "type": "integer",
+                        },
+                        "y": {
+                            "description": "Y coordinate for click action.",
+                            "type": "integer",
+                        },
+                    },
+                    "required": ["action", "x", "y"],
+                },
+            }
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "input": input_items,
+            "tools": [computer_tool],
+            "stream": False,
+            "reasoning": {"summary": "concise"},
+            "truncation": "auto",
+            "max_tokens": 200,  # Keep response short for click prediction
+            "request_timeout": kwargs.pop("request_timeout", 120),
+            **kwargs,
+        }
+
+        # Use liteLLM responses
+        response = await litellm.aresponses(**api_kwargs)
+
+        # Extract click coordinates from response output - handle both dict and Pydantic model
+        output_dict = response if isinstance(response, dict) else response.model_dump()
+        output_items = output_dict.get("output", [])
+
+        # Look for click coordinates in the response
+        for item in output_items:
+            if not isinstance(item, dict):
+                continue
+
+            # Native format: computer_call with action dict
+            if item.get("type") == "computer_call" and isinstance(item.get("action"), dict):
+                action = item["action"]
+                if action.get("x") is not None and action.get("y") is not None:
+                    return (int(action.get("x")), int(action.get("y")))
+
+            # Function calling format: function_call with arguments
+            if item.get("type") == "function_call" and item.get("name") == "computer":
+                try:
+                    arguments = item.get("arguments", "{}")
+                    if isinstance(arguments, str):
+                        args = json.loads(arguments)
+                    else:
+                        args = arguments
+                    if args.get("x") is not None and args.get("y") is not None:
+                        return (int(args.get("x")), int(args.get("y")))
+                except (json.JSONDecodeError, TypeError):
+                    continue
+
+        return None
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """
+        Get list of capabilities supported by this agent config.
+
+        Returns:
+            List of capability strings
+        """
+        return ["click", "step"]

libs/python/agent/cua_agent/loops/opencua.py

@@ -0,0 +1,435 @@
+"""
+OpenCUA agent loop implementation for click prediction and step execution using litellm.acompletion.
+Based on OpenCUA model for GUI grounding tasks.
+"""
+
+import base64
+import io
+import json
+import re
+import uuid
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from litellm.responses.litellm_completion_transformation.transformation import (
+    LiteLLMCompletionResponsesConfig,
+)
+from PIL import Image
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_completion_messages_to_responses_items,
+    convert_responses_items_to_completion_messages,
+    make_reasoning_item,
+)
+from ..types import AgentCapability
+from .composed_grounded import ComposedGroundedConfig
+from .generic_vlm import (
+    QWEN3_COMPUTER_TOOL,
+    _build_nous_system,
+    _parse_tool_call_from_text,
+    convert_qwen_tool_args_to_computer_action,
+)
+
+
+def extract_coordinates_from_click(text: str) -> Optional[Tuple[int, int]]:
+    """Extract coordinates from click(x=..., y=...) or pyautogui.click(x=..., y=...) format.
+
+    This function supports parsing both generic click() and legacy pyautogui.click() formats
+    for backwards compatibility with models that may still output pyautogui format.
+    """
+    try:
+        # Look for click(x=1443, y=343) or pyautogui.click(x=1443, y=343) pattern
+        pattern = r"(?:pyautogui\.)?click\(x=(\d+),\s*y=(\d+)\)"
+        match = re.search(pattern, text)
+        if match:
+            x, y = int(match.group(1)), int(match.group(2))
+            return (x, y)
+        return None
+    except Exception:
+        return None
+
+
+def _rescale_coordinate(
+    x: int,
+    y: int,
+    orig_w: int,
+    orig_h: int,
+    resized_w: int,
+    resized_h: int,
+) -> Tuple[int, int]:
+    """Rescale coordinates from resized image space back to original image space."""
+    if resized_w == 0 or resized_h == 0:
+        return (x, y)
+    return (round(x * orig_w / resized_w), round(y * orig_h / resized_h))
+
+
+@register_agent(models=r"(?i).*OpenCUA.*")
+class OpenCUAConfig(ComposedGroundedConfig):
+    """OpenCUA agent configuration implementing AsyncAgentConfig protocol for click prediction and step execution."""
+
+    def __init__(self):
+        super().__init__()
+        self.current_model = None
+        self.last_screenshot_b64 = None
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Predict the next step using the OpenCUA model with smart resize (factor=28)."""
+
+        # Convert responses items to completion messages
+        converted_msgs = convert_responses_items_to_completion_messages(
+            messages,
+            allow_images_in_tool_results=False,
+        )
+
+        # Build function schemas from tools array
+        function_schemas: List[Dict[str, Any]] = []
+        if tools:
+            from ..computers import is_agent_computer
+
+            for tool in tools:
+                tool_type = tool.get("type")
+                if tool_type == "computer":
+                    computer = tool.get("computer")
+                    if computer and is_agent_computer(computer):
+                        function_schemas.append(QWEN3_COMPUTER_TOOL["function"])
+                elif tool_type == "function":
+                    function_schema = tool.get("function")
+                    if function_schema:
+                        function_schemas.append(function_schema)
+
+        if not function_schemas:
+            function_schemas = [QWEN3_COMPUTER_TOOL["function"]]
+
+        # Prepend Nous-generated system prompt with tool schema
+        nous_system = _build_nous_system(function_schemas)
+        completion_messages = ([nous_system] if nous_system else []) + converted_msgs
+
+        # ------------------------------------------------------------------
+        # If there are no screenshots in the conversation, take one now
+        # ------------------------------------------------------------------
+        def _has_any_image(msgs: List[Dict[str, Any]]) -> bool:
+            for m in msgs:
+                content = m.get("content")
+                if isinstance(content, list):
+                    for p in content:
+                        if isinstance(p, dict) and p.get("type") == "image_url":
+                            return True
+            return False
+
+        def _has_screenshot_message(msgs: List[Dict[str, Any]]) -> bool:
+            screenshot_text = "Taking a screenshot to see the current computer screen."
+            for m in msgs:
+                content = m.get("content")
+                if isinstance(content, str) and screenshot_text in content:
+                    return True
+                if isinstance(content, list):
+                    for p in content:
+                        if isinstance(p, dict) and p.get("type") == "text":
+                            if screenshot_text in (p.get("text") or ""):
+                                return True
+            return False
+
+        pre_output_items: List[Dict[str, Any]] = []
+        if not _has_any_image(completion_messages):
+            if computer_handler is None or not hasattr(computer_handler, "screenshot"):
+                raise RuntimeError(
+                    "No screenshots present and computer_handler.screenshot is not available."
+                )
+            screenshot_b64 = await computer_handler.screenshot()
+            if not screenshot_b64:
+                raise RuntimeError("Failed to capture screenshot from computer_handler.")
+            completion_messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/png;base64,{screenshot_b64}"},
+                        },
+                        {"type": "text", "text": "Current screen"},
+                    ],
+                }
+            )
+            if not _has_screenshot_message(messages):
+                pre_output_items.append(
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": "Taking a screenshot to see the current computer screen.",
+                            }
+                        ],
+                    }
+                )
+
+        # ------------------------------------------------------------------
+        # Smart-resize all screenshots with factor=28
+        # Unlike generic_vlm (which sets min/max pixel hints for the provider),
+        # OpenCUA uses an OpenAI-compatible endpoint that does not honour those
+        # hints, so we actually resize the image before sending.
+        # ------------------------------------------------------------------
+        MIN_PIXELS = 3136
+        MAX_PIXELS = 12845056
+        FACTOR = 28
+
+        try:
+            from qwen_vl_utils import smart_resize  # type: ignore
+        except ImportError:
+            raise ImportError(
+                "qwen-vl-utils not installed. Please install it with: pip install qwen-vl-utils"
+            )
+
+        last_orig_w: Optional[int] = None
+        last_orig_h: Optional[int] = None
+        last_rw: Optional[int] = None
+        last_rh: Optional[int] = None
+
+        for msg in completion_messages:
+            content = msg.get("content")
+            if not isinstance(content, list):
+                continue
+            for part in content:
+                if isinstance(part, dict) and part.get("type") == "image_url":
+                    url = ((part.get("image_url") or {}).get("url")) or ""
+                    if url.startswith("data:") and "," in url:
+                        b64 = url.split(",", 1)[1]
+                        img_bytes = base64.b64decode(b64)
+                        im = Image.open(io.BytesIO(img_bytes))
+                        orig_h, orig_w = im.height, im.width
+                        rh, rw = smart_resize(
+                            orig_h,
+                            orig_w,
+                            factor=FACTOR,
+                            min_pixels=MIN_PIXELS,
+                            max_pixels=MAX_PIXELS,
+                        )
+
+                        # Actually resize the image
+                        resized_im = im.resize((rw, rh))
+                        buf = io.BytesIO()
+                        resized_im.save(buf, format="PNG")
+                        new_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+                        part["image_url"]["url"] = f"data:image/png;base64,{new_b64}"
+
+                        last_orig_w, last_orig_h = orig_w, orig_h
+                        last_rw, last_rh = rw, rh
+
+        # ------------------------------------------------------------------
+        # Call litellm
+        # ------------------------------------------------------------------
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": completion_messages,
+            "max_retries": max_retries,
+            "stream": stream,
+            **{k: v for k, v in kwargs.items()},
+        }
+        if use_prompt_caching:
+            api_kwargs["use_prompt_caching"] = use_prompt_caching
+
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        response = await litellm.acompletion(**api_kwargs)
+
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        usage = {
+            **LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(  # type: ignore
+                response.usage
+            ).model_dump(),
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        # ------------------------------------------------------------------
+        # Parse response
+        # ------------------------------------------------------------------
+        resp_dict = response.model_dump()  # type: ignore
+        choice = (resp_dict.get("choices") or [{}])[0]
+        message = choice.get("message") or {}
+        content_text = message.get("content") or ""
+        tool_calls_array = message.get("tool_calls") or []
+        reasoning_text = message.get("reasoning") or ""
+
+        output_items: List[Dict[str, Any]] = []
+
+        if reasoning_text:
+            output_items.append(make_reasoning_item(reasoning_text))
+
+        # Helper: rescale coordinates from resized space to original space
+        def _rescale(x: int, y: int) -> Tuple[int, int]:
+            if last_orig_w and last_orig_h and last_rw and last_rh:
+                return _rescale_coordinate(x, y, last_orig_w, last_orig_h, last_rw, last_rh)
+            return (x, y)
+
+        # Priority 1: OpenCUA native click(x=..., y=...) format
+        coords = extract_coordinates_from_click(content_text)
+        if coords:
+            x, y = _rescale(coords[0], coords[1])
+            fake_cm: Dict[str, Any] = {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "id": "call_0",
+                        "function": {
+                            "name": "computer",
+                            "arguments": json.dumps({"action": "left_click", "x": x, "y": y}),
+                        },
+                    }
+                ],
+            }
+            output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+
+        # Priority 2: <tool_call>...</tool_call> XML format
+        elif not tool_calls_array:
+            tool_call = _parse_tool_call_from_text(content_text)
+            if tool_call and isinstance(tool_call, dict):
+                fn_name = tool_call.get("name") or "computer"
+                raw_args = tool_call.get("arguments") or {}
+
+                # Rescale any coordinate field
+                coord = raw_args.get("coordinate")
+                if coord and isinstance(coord, (list, tuple)) and len(coord) >= 2:
+                    rx, ry = _rescale(int(round(float(coord[0]))), int(round(float(coord[1]))))
+                    raw_args = {**raw_args, "coordinate": [rx, ry]}
+
+                fake_cm = {
+                    "role": "assistant",
+                    "tool_calls": [
+                        {
+                            "type": "function",
+                            "id": "call_0",
+                            "function": {
+                                "name": fn_name,
+                                "arguments": json.dumps(raw_args),
+                            },
+                        }
+                    ],
+                }
+                output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+            else:
+                # Plain text response
+                fake_cm = {"role": "assistant", "content": content_text}
+                output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+
+        # Priority 3: tool_calls array from response
+        else:
+            processed_tool_calls = []
+            for tc in tool_calls_array:
+                function = tc.get("function", {})
+                fn_name = function.get("name", "computer")
+                args_str = function.get("arguments", "{}")
+
+                try:
+                    args = json.loads(args_str)
+
+                    # Rescale coordinates if present
+                    coord = args.get("coordinate")
+                    if coord and isinstance(coord, (list, tuple)) and len(coord) >= 2:
+                        rx, ry = _rescale(int(round(float(coord[0]))), int(round(float(coord[1]))))
+                        args = {**args, "coordinate": [rx, ry]}
+
+                    # Convert Qwen format to Computer Calls format
+                    if fn_name == "computer":
+                        converted_action = convert_qwen_tool_args_to_computer_action(args)
+                        if converted_action:
+                            args = converted_action
+
+                    processed_tool_calls.append(
+                        {
+                            "type": tc.get("type", "function"),
+                            "id": tc.get("id", "call_0"),
+                            "function": {
+                                "name": fn_name,
+                                "arguments": json.dumps(args),
+                            },
+                        }
+                    )
+                except json.JSONDecodeError:
+                    processed_tool_calls.append(tc)
+
+            fake_cm = {
+                "role": "assistant",
+                "content": content_text if content_text else "",
+                "tool_calls": processed_tool_calls,
+            }
+            output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+
+        return {"output": (pre_output_items + output_items), "usage": usage}
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using OpenCUA model via litellm.acompletion.
+
+        Args:
+            model: The OpenCUA model name
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        # Prepare system message
+        system_prompt = (
+            "You are a GUI agent. You are given a task and a screenshot of the screen. "
+            "You need to perform a series of click actions to complete the task."
+        )
+
+        system_message = {"role": "system", "content": system_prompt}
+
+        # Prepare user message with image and instruction
+        user_message = {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}},
+                {"type": "text", "text": f"Click on {instruction}"},
+            ],
+        }
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "messages": [system_message, user_message],
+            "max_new_tokens": 2056,
+            "temperature": 0,
+            **kwargs,
+        }
+
+        # Use liteLLM acompletion
+        response = await litellm.acompletion(**api_kwargs)
+
+        # Extract response text
+        output_text = response.choices[0].message.content
+
+        # Extract coordinates from click format
+        coordinates = extract_coordinates_from_click(output_text)
+
+        return coordinates
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """Return the capabilities supported by this agent."""
+        return ["click", "step"]

libs/python/agent/cua_agent/loops/qwen35.py

@@ -0,0 +1,688 @@
+"""
+Qwen3-VL agent loop implementation using litellm with function/tool calling.
+- Passes a ComputerUse tool schema to acompletion
+- Converts between Responses items and completion messages using helpers
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from litellm.responses.litellm_completion_transformation.transformation import (
+    LiteLLMCompletionResponsesConfig,
+)
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_completion_messages_to_responses_items,
+    convert_responses_items_to_completion_messages,
+    make_reasoning_item,
+)
+from ..types import AgentCapability
+
+# ComputerUse tool schema (OpenAI function tool format)
+QWEN3_5_COMPUTER_TOOL: Dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "computer",
+        "description": (
+            "* `key`: Performs key down presses on the arguments passed in order, then performs key releases in reverse order.\n"
+            "* `type`: Type a string of text on the keyboard.\n"
+            "* `mouse_move`: Move the cursor to a specified (x, y) pixel coordinate on the screen.\n"
+            '* `left_click`: Click the left mouse button at a specified (x, y) pixel coordinate on the screen. Optional `text` parameter can specify modifier keys (e.g., "ctrl", "shift", "ctrl+shift") that will be held during the click.\n'
+            "* `left_click_drag`: Click and drag the cursor to a specified (x, y) pixel coordinate on the screen.\n"
+            "* `right_click`: Click the right mouse button at a specified (x, y) pixel coordinate on the screen. Optional `text` parameter can specify modifier keys that will be held during the click.\n"
+            "* `middle_click`: Click the middle mouse button at a specified (x, y) pixel coordinate on the screen. Optional `text` parameter can specify modifier keys that will be held during the click.\n"
+            "* `double_click`: Double-click the left mouse button at a specified (x, y) pixel coordinate on the screen. Optional `text` parameter can specify modifier keys that will be held during the click.\n"
+            "* `triple_click`: Triple-click the left mouse button at a specified (x, y) pixel coordinate on the screen (simulated as double-click since it's the closest action). Optional `text` parameter can specify modifier keys that will be held during the click.\n"
+            '* `scroll`: Performs a scroll of the mouse scroll wheel. Optional `text` parameter can specify a modifier key (e.g., "shift", "ctrl") that will be held during scrolling.\n'
+            "* `hscroll`: Performs a horizontal scroll (mapped to regular scroll). Optional `text` parameter can specify a modifier key that will be held during scrolling.\n"
+            "* `wait`: Wait specified seconds for the change to happen.\n"
+            # "* `terminate`: Terminate the current task and report its completion status.\n"
+            # "* `answer`: Answer a question.\n"
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "description": "The action to perform.",
+                    "enum": [
+                        "key",
+                        "type",
+                        "mouse_move",
+                        "left_click",
+                        "left_click_drag",
+                        "right_click",
+                        "middle_click",
+                        "double_click",
+                        "triple_click",
+                        "scroll",
+                        "hscroll",
+                        # "screenshot",
+                        "wait",
+                        # "terminate",
+                        # "answer",
+                    ],
+                    "type": "string",
+                },
+                "keys": {
+                    "description": "Required only by action=key.",
+                    "type": "array",
+                    "items": {"type": "string"},
+                },
+                "text": {
+                    "description": "Required only by action=type and action=answer.",
+                    "type": "string",
+                },
+                "coordinate": {
+                    "description": "(x, y): Pixel coordinates from top-left.",
+                    "type": "array",
+                    "items": {"type": ["number", "integer"]},
+                    "minItems": 2,
+                    "maxItems": 2,
+                },
+                "pixels": {
+                    "description": "Scroll amount. Positive=up, negative=down. For scroll/hscroll.",
+                    "type": "number",
+                },
+                "time": {
+                    "description": "Seconds to wait (action=wait).",
+                    "type": "number",
+                },
+                # "status": {
+                #     "description": "Task status (action=terminate).",
+                #     "type": "string",
+                #     "enum": ["success", "failure"],
+                # },
+            },
+            "required": ["action"],
+        },
+    },
+}
+
+
+def _build_nous_system(functions: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Use qwen-agent NousFnCallPrompt to generate a system message embedding tool schema."""
+    try:
+        from qwen_agent.llm.fncall_prompts.nous_fncall_prompt import (
+            ContentItem as NousContentItem,
+        )
+        from qwen_agent.llm.fncall_prompts.nous_fncall_prompt import (
+            Message as NousMessage,
+        )
+        from qwen_agent.llm.fncall_prompts.nous_fncall_prompt import (
+            NousFnCallPrompt,
+        )
+    except ImportError:
+        raise ImportError(
+            "qwen-agent not installed. Please install it with `pip install cua-agent[qwen]`."
+        )
+    msgs = NousFnCallPrompt().preprocess_fncall_messages(
+        messages=[
+            NousMessage(
+                role="system", content=[NousContentItem(text="You are a helpful assistant.")]
+            )
+        ],
+        functions=functions,
+        lang="en",
+    )
+    sys = msgs[0].model_dump()
+    # Convert qwen-agent structured content to OpenAI-style content list
+    content = [{"type": "text", "text": c["text"]} for c in sys.get("content", [])]
+    return {"role": "system", "content": content}
+
+
+def _parse_tool_call_from_text(text: str) -> Optional[Dict[str, Any]]:
+    """Extract a tool call from <tool_call>...</tool_call> in model text.
+
+    Handles two formats:
+    1. JSON: ``<tool_call>{"name": "computer", "arguments": {...}}</tool_call>``
+    2. XML-style (qwen35-4b): ``<tool_call><function=computer><parameter=action>left_click</parameter>...</tool_call>``
+    """
+    # --- Format 1: JSON ---
+    m = re.search(r"<tool_call>\s*(\{[\s\S]*?\})\s*</tool_call>", text)
+    if m:
+        try:
+            return json.loads(m.group(1))
+        except Exception:
+            pass
+
+    # --- Format 2: XML-style <function=name><parameter=key>value</parameter> ---
+    fn_match = re.search(
+        r"<tool_call>\s*<function=(\w+)>([\s\S]*?)</function>\s*</tool_call>", text
+    )
+    if fn_match:
+        fn_name = fn_match.group(1)
+        params_block = fn_match.group(2)
+        # Extract all <parameter=key>value</parameter> pairs
+        params: Dict[str, Any] = {}
+        for pm in re.finditer(r"<parameter=(\w+)>\s*([\s\S]*?)\s*</parameter>", params_block):
+            key = pm.group(1)
+            val = pm.group(2).strip()
+            # Try to parse as JSON (for arrays/numbers), fall back to string
+            try:
+                params[key] = json.loads(val)
+            except (json.JSONDecodeError, ValueError):
+                params[key] = val
+        # The XML format uses <parameter=type> for the action field name,
+        # but the Qwen tool schema calls it "action".  Remap if we got
+        # "type" that looks like an action name rather than a literal type.
+        if "type" in params and "action" not in params:
+            params["action"] = params.pop("type")
+        return {"name": fn_name, "arguments": params}
+
+    return None
+
+
+async def _unnormalize_coordinate(args: Dict[str, Any], dims: Tuple[int, int]) -> Dict[str, Any]:
+    """Coordinates appear in 0..1000 space, scale to actual screen size using dims if provided."""
+    coord = args.get("coordinate")
+    if not coord or not isinstance(coord, (list, tuple)) or len(coord) < 2:
+        return args
+    x, y = float(coord[0]), float(coord[1])
+    width, height = float(dims[0]), float(dims[1])
+    x_abs = max(0.0, min(width, (x / 1000.0) * width))
+    y_abs = max(0.0, min(height, (y / 1000.0) * height))
+    args = {**args, "coordinate": [round(x_abs), round(y_abs)]}
+    return args
+
+
+def convert_qwen_tool_args_to_computer_action(args: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """
+    Convert Qwen computer tool arguments to the Computer Calls action schema.
+
+    Qwen (example):
+        {"action": "left_click", "coordinate": [114, 68]}
+
+    Target (example):
+        {"action": "left_click", "x": 114, "y": 68}
+
+    Other mappings:
+    - right_click, middle_click, double_click (triple_click -> double_click)
+    - mouse_move -> { action: "move", x, y }
+    - key -> { action: "keypress", keys: [...] }
+    - type -> { action: "type", text }
+    - scroll/hscroll -> { action: "scroll", scroll_x, scroll_y, x, y }
+    - wait -> { action: "wait" }
+    - terminate/answer are not direct UI actions; return None for now
+    """
+    if not isinstance(args, dict):
+        return None
+
+    action = args.get("action")
+    if not isinstance(action, str):
+        return None
+
+    # Coordinates helper
+    coord = args.get("coordinate")
+    x = y = None
+    if isinstance(coord, (list, tuple)) and len(coord) >= 2:
+        try:
+            x = int(round(float(coord[0])))
+            y = int(round(float(coord[1])))
+        except Exception:
+            x = y = None
+
+    # Map actions
+    a = action.lower()
+    if a in {"left_click", "right_click", "middle_click", "double_click"}:
+        if x is None or y is None:
+            return None
+        return {"action": a, "x": x, "y": y}
+    if a == "triple_click":
+        # Approximate as double_click
+        if x is None or y is None:
+            return None
+        return {"action": "double_click", "x": x, "y": y}
+    if a == "mouse_move":
+        if x is None or y is None:
+            return None
+        return {"action": "move", "x": x, "y": y}
+    if a == "key":
+        keys = args.get("keys")
+        if isinstance(keys, list) and all(isinstance(k, str) for k in keys):
+            return {"action": "keypress", "keys": keys}
+        return None
+    if a == "type":
+        text = args.get("text")
+        if isinstance(text, str):
+            return {"action": "type", "text": text}
+        return None
+    if a in {"scroll", "hscroll"}:
+        pixels = args.get("pixels") or 0
+        try:
+            pixels_val = int(round(float(pixels)))
+        except Exception:
+            pixels_val = 0
+        scroll_x = pixels_val if a == "hscroll" else 0
+        scroll_y = pixels_val if a == "scroll" else 0
+        # Include cursor position if available (optional)
+        out: Dict[str, Any] = {"action": "scroll", "scroll_x": scroll_x, "scroll_y": scroll_y}
+        if x is not None and y is not None:
+            out.update({"x": x, "y": y})
+        return out
+    if a == "wait":
+        return {"action": "wait"}
+
+    # Non-UI or terminal actions: terminate/answer -> not mapped here
+    return None
+
+
+@register_agent(models=r"(?i).*qwen35.*", priority=1)
+class Qwen35Config(AsyncAgentConfig):
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        # Build messages using NousFnCallPrompt system with tool schema in text
+        # Start with converted conversation (images/text preserved)
+        converted_msgs = convert_responses_items_to_completion_messages(
+            messages,
+            allow_images_in_tool_results=False,
+        )
+
+        # print(f"The number of items in the converted_msgs: {len(converted_msgs)}")
+
+        # Build function schemas from tools array
+        function_schemas = []
+        if tools:
+            from ..computers import is_agent_computer
+
+            for tool in tools:
+                tool_type = tool.get("type")
+
+                if tool_type == "computer":
+                    # For computer tools, use QWEN3_COMPUTER_TOOL schema
+                    computer = tool.get("computer")
+                    if computer and is_agent_computer(computer):
+                        function_schemas.append(QWEN3_5_COMPUTER_TOOL["function"])
+                elif tool_type == "function":
+                    # For function tools, use the provided function schema
+                    function_schema = tool.get("function")
+                    if function_schema:
+                        function_schemas.append(function_schema)
+
+        # If no tools provided or no computer tool found, use default QWEN3_COMPUTER_TOOL
+        if not function_schemas:
+            function_schemas = [QWEN3_5_COMPUTER_TOOL["function"]]
+
+        # print(f"[qwen35] function_schemas: {function_schemas}")
+
+        # Prepend Nous-generated system if available
+        nous_system = _build_nous_system(function_schemas)
+        completion_messages = ([nous_system] if nous_system else []) + converted_msgs
+
+        # If there is no screenshot in the conversation, take one now and inject it.
+        # Also record a pre_output_items assistant message to reflect action.
+        def _has_any_image(msgs: List[Dict[str, Any]]) -> bool:
+            for m in msgs:
+                content = m.get("content")
+                if isinstance(content, list):
+                    for p in content:
+                        if isinstance(p, dict) and p.get("type") == "image_url":
+                            return True
+            return False
+
+        def _has_screenshot_message(msgs: List[Dict[str, Any]]) -> bool:
+            """Check if messages already contain the 'Taking a screenshot' text."""
+            screenshot_text = "Taking a screenshot to see the current computer screen."
+            for m in msgs:
+                content = m.get("content")
+                if isinstance(content, str) and screenshot_text in content:
+                    return True
+                if isinstance(content, list):
+                    for p in content:
+                        if isinstance(p, dict) and p.get("type") == "text":
+                            if screenshot_text in (p.get("text") or ""):
+                                return True
+            return False
+
+        pre_output_items: List[Dict[str, Any]] = []
+        if not _has_any_image(completion_messages):
+            if computer_handler is None or not hasattr(computer_handler, "screenshot"):
+                raise RuntimeError(
+                    "No screenshots present and computer_handler.screenshot is not available."
+                )
+            screenshot_b64 = await computer_handler.screenshot()
+            if not screenshot_b64:
+                raise RuntimeError("Failed to capture screenshot from computer_handler.")
+            # Inject a user message with the screenshot so the model can see current context
+            completion_messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/png;base64,{screenshot_b64}"},
+                        },
+                        {"type": "text", "text": "Current screen"},
+                    ],
+                }
+            )
+            # Add assistant message to outputs to reflect the action, only if not already present
+            if not _has_screenshot_message(messages):
+                pre_output_items.append(
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": "Taking a screenshot to see the current computer screen.",
+                            }
+                        ],
+                    }
+                )
+
+        # Smart-resize all screenshots and attach min/max pixel hints. Fail fast if deps missing.
+        # Also record the last resized width/height to unnormalize coordinates later.
+        last_rw: Optional[int] = None
+        last_rh: Optional[int] = None
+        MIN_PIXELS = 3136
+        MAX_PIXELS = 12845056
+        try:
+            import base64
+            import io
+
+            from PIL import Image  # type: ignore
+            from qwen_vl_utils import smart_resize  # type: ignore
+        except Exception:
+            raise ImportError(
+                "qwen-vl-utils not installed. Please install it with `pip install cua-agent[qwen]`."
+            )
+
+        for msg in completion_messages:
+            content = msg.get("content")
+            if not isinstance(content, list):
+                continue
+            for part in content:
+                if isinstance(part, dict) and part.get("type") == "image_url":
+                    url = ((part.get("image_url") or {}).get("url")) or ""
+                    # Expect data URL like data:image/png;base64,<b64>
+                    if url.startswith("data:") and "," in url:
+                        b64 = url.split(",", 1)[1]
+                        img_bytes = base64.b64decode(b64)
+                        im = Image.open(io.BytesIO(img_bytes))
+                        h, w = im.height, im.width
+                        rh, rw = smart_resize(
+                            h, w, factor=32, min_pixels=MIN_PIXELS, max_pixels=MAX_PIXELS
+                        )
+                        # Attach hints on this image block
+                        part["min_pixels"] = MIN_PIXELS
+                        part["max_pixels"] = MAX_PIXELS
+                        last_rw, last_rh = rw, rh
+
+        for i, msg in enumerate(completion_messages):
+            role = msg.get("role")
+            content = msg.get("content")
+            if isinstance(content, list):
+                step_content = []
+                for item in content:
+                    item_type = item.get("type")
+                    if item_type == "text":
+                        step_content.append(item.get("text"))
+                    elif item_type == "image_url":
+                        step_content.append("Image URL: " + item.get("image_url").get("url")[:100])
+            else:
+                item = content
+                step_content = ""
+                if isinstance(item, dict) and item.get("type") == "image_url":
+                    step_content = "Image URL: " + item.get("image_url").get("url")[:100]
+                else:
+                    step_content = content
+
+            print(f"Step {i}: Role: {role}, Content: {step_content}")
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": completion_messages,
+            "max_retries": max_retries,
+            "stream": stream,
+            **{k: v for k, v in kwargs.items()},
+        }
+        if use_prompt_caching:
+            api_kwargs["use_prompt_caching"] = use_prompt_caching
+
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        response = await litellm.acompletion(**api_kwargs)
+
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        usage = {
+            **LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(  # type: ignore
+                response.usage
+            ).model_dump(),
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        # Extract response data
+        resp_dict = response.model_dump()  # type: ignore
+        choice = (resp_dict.get("choices") or [{}])[0]
+        message = choice.get("message") or {}
+        content_text = message.get("content") or ""
+        tool_calls_array = message.get("tool_calls") or []
+        reasoning_text = message.get("reasoning") or ""
+
+        output_items: List[Dict[str, Any]] = []
+
+        # Add reasoning if present (Ollama Cloud format)
+        if reasoning_text:
+            output_items.append(make_reasoning_item(reasoning_text))
+
+        # Priority 1: Try to parse tool call from content text (OpenRouter format)
+        tool_call = _parse_tool_call_from_text(content_text)
+
+        if tool_call and isinstance(tool_call, dict):
+            fn_name = tool_call.get("name") or "computer"
+            raw_args = tool_call.get("arguments") or {}
+
+            output_items.append(
+                {
+                    "type": "message",
+                    "role": "assistant",
+                    "content": [{"type": "output_text", "text": content_text}],
+                }
+            )
+
+            # Unnormalize coordinates to actual screen size using last resized dims
+            if last_rw is None or last_rh is None:
+                raise RuntimeError(
+                    "No screenshots found to derive dimensions for coordinate unnormalization."
+                )
+            args = await _unnormalize_coordinate(raw_args, (last_rw, last_rh))
+
+            # Convert Qwen format to Computer Calls format if this is a computer tool
+            if fn_name == "computer":
+                converted_action = convert_qwen_tool_args_to_computer_action(args)
+                if converted_action:
+                    args = converted_action
+
+            # Build an OpenAI-style tool call so we can reuse the converter
+            fake_cm = {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "type": "function",
+                        "id": "call_0",
+                        "function": {
+                            "name": fn_name,
+                            "arguments": json.dumps(args),
+                        },
+                    }
+                ],
+            }
+            output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+
+        elif tool_calls_array:
+
+            output_items.append(
+                {
+                    "type": "message",
+                    "role": "assistant",
+                    "content": [{"type": "output_text", "text": content_text}],
+                }
+            )
+
+            processed_tool_calls = []
+            for tc in tool_calls_array:
+                function = tc.get("function", {})
+                fn_name = function.get("name", "computer")
+                args_str = function.get("arguments", "{}")
+
+                try:
+                    args = json.loads(args_str)
+
+                    # Unnormalize coordinates if present
+                    if "coordinate" in args and last_rw is not None and last_rh is not None:
+                        args = await _unnormalize_coordinate(args, (last_rw, last_rh))
+
+                    # Convert Qwen format to Computer Calls format if this is a computer tool
+                    if fn_name == "computer":
+                        converted_action = convert_qwen_tool_args_to_computer_action(args)
+                        if converted_action:
+                            args = converted_action
+
+                    processed_tool_calls.append(
+                        {
+                            "type": tc.get("type", "function"),
+                            "id": tc.get("id", "call_0"),
+                            "function": {
+                                "name": fn_name,
+                                "arguments": json.dumps(args),
+                            },
+                        }
+                    )
+                except json.JSONDecodeError:
+                    processed_tool_calls.append(tc)
+
+            fake_cm = {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": processed_tool_calls,
+            }
+            output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+
+        else:
+            # No tool calls found in either format, return text response
+            fake_cm = {"role": "assistant", "content": content_text}
+            output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+
+        # Prepend any pre_output_items (e.g., simulated screenshot-taking message)
+        return {"output": (pre_output_items + output_items), "usage": usage}
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        return ["click", "step"]
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates using Qwen3-VL via litellm.acompletion.
+
+        Only exposes a reduced tool schema with left_click to bias model to output a single click.
+        Returns (x, y) absolute pixels when screen dimensions can be obtained; otherwise normalized 0..1000 integers.
+        """
+        # Reduced tool
+        reduced_tool = {
+            "type": "function",
+            "function": {
+                **QWEN3_5_COMPUTER_TOOL["function"],
+                "parameters": {
+                    **QWEN3_5_COMPUTER_TOOL["function"]["parameters"],
+                    "type": "object",
+                    "properties": {
+                        "action": {"type": "string", "enum": ["left_click"]},
+                        "coordinate": {
+                            "description": "(x, y) in 0..1000 reference space",
+                            "type": "array",
+                            "items": {"type": ["number", "integer"]},
+                            "minItems": 2,
+                            "maxItems": 2,
+                        },
+                    },
+                    "required": ["action", "coordinate"],
+                },
+            },
+        }
+
+        # Build Nous system (lazy import inside helper already raises clear guidance if missing)
+        nous_system = _build_nous_system([reduced_tool["function"]])
+
+        # Pre-process using smart_resize
+        min_pixels = 3136
+        max_pixels = 12845056
+        try:
+            # Lazy import to avoid hard dependency
+            import base64
+            import io
+
+            # If PIL is available, estimate size from image to derive smart bounds
+            from PIL import Image
+            from qwen_vl_utils import smart_resize  # type: ignore
+
+            img_bytes = base64.b64decode(image_b64)
+            im = Image.open(io.BytesIO(img_bytes))
+            h, w = im.height, im.width
+            # Qwen notebook suggests factor=32 and a wide min/max range
+            rh, rw = smart_resize(h, w, factor=32, min_pixels=min_pixels, max_pixels=max_pixels)
+        except Exception:
+            raise ImportError(
+                "qwen-vl-utils not installed. Please install it with `pip install cua-agent[qwen]`."
+            )
+
+        messages = []
+        if nous_system:
+            messages.append(nous_system)
+        image_block: Dict[str, Any] = {
+            "type": "image_url",
+            "image_url": {"url": f"data:image/png;base64,{image_b64}"},
+            "min_pixels": min_pixels,
+            "max_pixels": max_pixels,
+        }
+        # Single user message with image and instruction, matching OpenAI-style content blocks
+        messages.append(
+            {
+                "role": "user",
+                "content": [
+                    image_block,
+                    {"type": "text", "text": instruction},
+                ],
+            }
+        )
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            **{k: v for k, v in kwargs.items()},
+        }
+        response = await litellm.acompletion(**api_kwargs)
+        resp = response.model_dump()  # type: ignore
+        choice = (resp.get("choices") or [{}])[0]
+        content_text = ((choice.get("message") or {}).get("content")) or ""
+        tool_call = _parse_tool_call_from_text(content_text) or {}
+        args = tool_call.get("arguments") or {}
+        args = await _unnormalize_coordinate(args, (rh, rw))
+        coord = args.get("coordinate")
+        if isinstance(coord, (list, tuple)) and len(coord) >= 2:
+            return int(coord[0]), int(coord[1])
+        return None

libs/python/agent/cua_agent/loops/qwen3vl.py

@@ -0,0 +1,18 @@
+"""
+Qwen3-VL dedicated agent loop configuration.
+Re-exports GenericVlmConfig under a Qwen-specific model pattern so that
+Qwen model strings are matched at normal priority instead of the generic
+catch-all (priority -100).
+"""
+
+from __future__ import annotations
+
+from ..decorators import register_agent
+from .generic_vlm import GenericVlmConfig
+
+
+@register_agent(models=r"(?i).*qwen.*")
+class Qwen3VlConfig(GenericVlmConfig):
+    """Qwen3-VL agent loop using litellm with function/tool calling."""
+
+    pass

libs/python/agent/cua_agent/loops/uiins.py

@@ -0,0 +1,175 @@
+"""
+UI-Ins agent loop implementation for click prediction using litellm.acompletion
+Paper: https://arxiv.org/pdf/2510.202861
+Code: https://github.com/alibaba/UI-Ins
+"""
+
+import asyncio
+import base64
+import json
+import math
+import re
+import uuid
+from io import BytesIO
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+
+import litellm
+from PIL import Image
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+
+SYSTEM_PROMPT = """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task.\n\n## Output Format\nReturn a json object with a reasoning process in  tags, a function name and arguments within  XML tags:\n```\n\n...\n\n\n{"name": "grounding", "arguments": }\n\n```\n represents the following item of the action space:\n## Action Space{"action": "click", "coordinate": [x, y]}\nYour task is to accurately locate a UI element based on the instruction. You should first analyze instruction in  tags and finally output the function in  tags.\n"""
+
+
+def parse_coordinates(raw_string: str) -> tuple[int, int]:
+    matches = re.findall(r"\[(\d+),\s*(\d+)\]", raw_string)
+    if matches:
+        return tuple(map(int, matches[0]))
+    return -1, -1
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = 28,
+    min_pixels: int = 3136,
+    max_pixels: int = 8847360,
+) -> Tuple[int, int]:
+    """Smart resize function similar to qwen_vl_utils."""
+    # Calculate the total pixels
+    total_pixels = height * width
+
+    # If already within bounds, return original dimensions
+    if min_pixels <= total_pixels <= max_pixels:
+        # Round to nearest factor
+        new_height = (height // factor) * factor
+        new_width = (width // factor) * factor
+        return new_height, new_width
+
+    # Calculate scaling factor
+    if total_pixels > max_pixels:
+        scale = (max_pixels / total_pixels) ** 0.5
+    else:
+        scale = (min_pixels / total_pixels) ** 0.5
+
+    # Apply scaling
+    new_height = int(height * scale)
+    new_width = int(width * scale)
+
+    # Round to nearest factor
+    new_height = (new_height // factor) * factor
+    new_width = (new_width // factor) * factor
+
+    # Ensure minimum size
+    new_height = max(new_height, factor)
+    new_width = max(new_width, factor)
+
+    return new_height, new_width
+
+
+@register_agent(models=r".*UI-Ins.*")
+class UIInsConfig(AsyncAgentConfig):
+    """UI-Ins agent configuration implementing AsyncAgentConfig protocol for click prediction."""
+
+    def __init__(self):
+        self.current_model = None
+        self.last_screenshot_b64 = None
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        raise NotImplementedError()
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[float, float]]:
+        """
+        Predict click coordinates using UI-Ins model via litellm.acompletion.
+
+        Args:
+            model: The UI-Ins model name
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+
+        Returns:
+            Tuple of (x, y) coordinates or None if prediction fails
+        """
+        # Decode base64 image
+        image_data = base64.b64decode(image_b64)
+        image = Image.open(BytesIO(image_data))
+        width, height = image.width, image.height
+
+        # Smart resize the image (similar to qwen_vl_utils)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=28,  # Default factor for Qwen models
+            min_pixels=3136,
+            max_pixels=4096 * 2160,
+        )
+        resized_image = image.resize((resized_width, resized_height))
+        scale_x, scale_y = width / resized_width, height / resized_height
+
+        # Convert resized image back to base64
+        buffered = BytesIO()
+        resized_image.save(buffered, format="PNG")
+        resized_image_b64 = base64.b64encode(buffered.getvalue()).decode()
+
+        # Prepare system and user messages
+        system_message = {
+            "role": "system",
+            "content": [
+                {"type": "text", "text": "You are a helpful assistant."},
+                {"type": "text", "text": SYSTEM_PROMPT},
+            ],
+        }
+
+        user_message = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/png;base64,{resized_image_b64}"},
+                },
+                {"type": "text", "text": instruction},
+            ],
+        }
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "messages": [system_message, user_message],
+            "max_tokens": 2056,
+            "temperature": 0.0,
+            **kwargs,
+        }
+
+        # Use liteLLM acompletion
+        response = await litellm.acompletion(**api_kwargs)
+
+        # Extract response text
+        output_text = response.choices[0].message.content  # type: ignore
+
+        # Extract and rescale coordinates
+        pred_x, pred_y = parse_coordinates(output_text)  # type: ignore
+        pred_x *= scale_x
+        pred_y *= scale_y
+
+        return (math.floor(pred_x), math.floor(pred_y))
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """Return the capabilities supported by this agent."""
+        return ["click"]

libs/python/agent/cua_agent/loops/uitars.py

@@ -0,0 +1,873 @@
+"""
+UITARS agent loop implementation using liteLLM for ByteDance-Seed/UI-TARS-1.5-7B
+Paper: https://arxiv.org/abs/2501.12326
+Code: https://github.com/bytedance/UI-TARS
+"""
+
+import ast
+import asyncio
+import base64
+import json
+import math
+import re
+from ctypes import cast
+from io import BytesIO
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+
+import litellm
+from litellm.responses.litellm_completion_transformation.transformation import (
+    LiteLLMCompletionResponsesConfig,
+)
+from litellm.responses.utils import Usage
+from litellm.types.utils import ModelResponse
+from openai.types.responses.response_computer_tool_call_param import (
+    ActionType,
+    ResponseComputerToolCallParam,
+)
+from openai.types.responses.response_input_param import ComputerCallOutput
+from openai.types.responses.response_output_message_param import (
+    ResponseOutputMessageParam,
+)
+from openai.types.responses.response_reasoning_item_param import (
+    ResponseReasoningItemParam,
+    Summary,
+)
+from PIL import Image
+
+from ..decorators import register_agent
+from ..responses import (
+    make_click_item,
+    make_double_click_item,
+    make_drag_item,
+    make_input_image_item,
+    make_keypress_item,
+    make_output_text_item,
+    make_reasoning_item,
+    make_scroll_item,
+    make_type_item,
+    make_wait_item,
+)
+from ..types import AgentCapability, AgentResponse, Messages, Tools
+
+# Constants from reference code
+IMAGE_FACTOR = 28
+MIN_PIXELS = 100 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+FINISH_WORD = "finished"
+WAIT_WORD = "wait"
+ENV_FAIL_WORD = "error_env"
+CALL_USER = "call_user"
+
+# Action space prompt for UITARS
+UITARS_ACTION_SPACE = """
+click(start_box='<|box_start|>(x1,y1)<|box_end|>')
+left_double(start_box='<|box_start|>(x1,y1)<|box_end|>')
+right_single(start_box='<|box_start|>(x1,y1)<|box_end|>')
+drag(start_box='<|box_start|>(x1,y1)<|box_end|>', end_box='<|box_start|>(x3,y3)<|box_end|>')
+hotkey(key='')
+type(content='') #If you want to submit your input, use "\\n" at the end of `content`.
+scroll(start_box='<|box_start|>(x1,y1)<|box_end|>', direction='down or up or right or left')
+wait() #Sleep for 5s and take a screenshot to check for any changes.
+finished(content='xxx') # Use escape characters \\', \\", and \\n in content part to ensure we can parse the content in normal python string format.
+"""
+
+UITARS_PROMPT_TEMPLATE = """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task. 
+
+## Output Format
+```
+Thought: ...
+Action: ...
+```
+
+## Action Space
+{action_space}
+
+## Note
+- Use {language} in `Thought` part.
+- Write a small plan and finally summarize your next action (with its target element) in one sentence in `Thought` part.
+
+## User Instruction
+{instruction}
+"""
+
+GROUNDING_UITARS_PROMPT_TEMPLATE = """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task. 
+
+## Output Format
+
+Action: ...
+
+
+## Action Space
+click(point='<|box_start|>(x1,y1)<|box_end|>')
+
+## User Instruction
+{instruction}"""
+
+
+def round_by_factor(number: float, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: float, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: float, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+    1. Both dimensions (height and width) are divisible by 'factor'.
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def escape_single_quotes(text):
+    """Escape single quotes in text for safe string formatting."""
+    pattern = r"(?<!\\)'"
+    return re.sub(pattern, r"\\'", text)
+
+
+def parse_action(action_str):
+    """Parse action string into structured format."""
+    try:
+        node = ast.parse(action_str, mode="eval")
+        if not isinstance(node, ast.Expression):
+            raise ValueError("Not an expression")
+
+        call = node.body
+        if not isinstance(call, ast.Call):
+            raise ValueError("Not a function call")
+
+        # Get function name
+        if isinstance(call.func, ast.Name):
+            func_name = call.func.id
+        elif isinstance(call.func, ast.Attribute):
+            func_name = call.func.attr
+        else:
+            func_name = None
+
+        # Get keyword arguments
+        kwargs = {}
+        for kw in call.keywords:
+            key = kw.arg
+            if isinstance(kw.value, ast.Constant):
+                value = kw.value.value
+            elif isinstance(kw.value, ast.Str):  # Compatibility with older Python
+                value = kw.value.s
+            else:
+                value = None
+            kwargs[key] = value
+
+        return {"function": func_name, "args": kwargs}
+
+    except Exception as e:
+        print(f"Failed to parse action '{action_str}': {e}")
+        return None
+
+
+def parse_uitars_response(text: str, image_width: int, image_height: int) -> List[Dict[str, Any]]:
+    """Parse UITARS model response into structured actions."""
+    text = text.strip()
+
+    # Extract thought
+    thought = None
+    if text.startswith("Thought:"):
+        thought_match = re.search(r"Thought: (.+?)(?=\s*Action:|$)", text, re.DOTALL)
+        if thought_match:
+            thought = thought_match.group(1).strip()
+
+    # Extract action
+    if "Action:" not in text:
+        raise ValueError("No Action found in response")
+
+    action_str = text.split("Action:")[-1].strip()
+
+    # Handle special case for type actions
+    if "type(content" in action_str:
+
+        def escape_quotes(match):
+            return match.group(1)
+
+        pattern = r"type\(content='(.*?)'\)"
+        content = re.sub(pattern, escape_quotes, action_str)
+        action_str = escape_single_quotes(content)
+        action_str = "type(content='" + action_str + "')"
+
+    # Parse the action
+    parsed_action = parse_action(action_str.replace("\n", "\\n").lstrip())
+    if parsed_action is None:
+        raise ValueError(f"Action can't parse: {action_str}")
+
+    action_type = parsed_action["function"]
+    params = parsed_action["args"]
+
+    # Process parameters
+    action_inputs = {}
+    for param_name, param in params.items():
+        if param == "":
+            continue
+        param = str(param).lstrip()
+        action_inputs[param_name.strip()] = param
+
+        # Handle coordinate parameters
+        if "start_box" in param_name or "end_box" in param_name:
+            # Parse coordinates like '<|box_start|>(x,y)<|box_end|>' or '(x,y)'
+            # First, remove special tokens
+            clean_param = param.replace("<|box_start|>", "").replace("<|box_end|>", "")
+            # Then remove parentheses and split
+            numbers = clean_param.replace("(", "").replace(")", "").split(",")
+
+            try:
+                float_numbers = [
+                    float(num.strip()) / 1000 for num in numbers
+                ]  # Normalize to 0-1 range
+
+                if len(float_numbers) == 2:
+                    # Single point, duplicate for box format
+                    float_numbers = [
+                        float_numbers[0],
+                        float_numbers[1],
+                        float_numbers[0],
+                        float_numbers[1],
+                    ]
+
+                action_inputs[param_name.strip()] = str(float_numbers)
+            except ValueError as e:
+                # If parsing fails, keep the original parameter value
+                print(f"Warning: Could not parse coordinates '{param}': {e}")
+                action_inputs[param_name.strip()] = param
+
+    return [
+        {
+            "thought": thought,
+            "action_type": action_type,
+            "action_inputs": action_inputs,
+            "text": text,
+        }
+    ]
+
+
+def convert_to_computer_actions(
+    parsed_responses: List[Dict[str, Any]], image_width: int, image_height: int
+) -> List[ResponseComputerToolCallParam | ResponseOutputMessageParam]:
+    """Convert parsed UITARS responses to computer actions."""
+    computer_actions = []
+
+    for response in parsed_responses:
+        action_type = response.get("action_type")
+        action_inputs = response.get("action_inputs", {})
+
+        if action_type == "finished":
+            finished_text = action_inputs.get("content", "Task completed successfully.")
+            computer_actions.append(make_output_text_item(finished_text))
+            break
+
+        elif action_type == "wait":
+            computer_actions.append(make_wait_item())
+
+        elif action_type == "call_user":
+            computer_actions.append(
+                make_output_text_item("I need assistance from the user to proceed with this task.")
+            )
+
+        elif action_type in ["click", "left_single"]:
+            start_box = action_inputs.get("start_box")
+            if start_box:
+                coords = eval(start_box)
+                x = int((coords[0] + coords[2]) / 2 * image_width)
+                y = int((coords[1] + coords[3]) / 2 * image_height)
+
+                computer_actions.append(make_click_item(x, y, "left"))
+
+        elif action_type in ["double_click", "left_double"]:
+            start_box = action_inputs.get("start_box")
+            if start_box:
+                coords = eval(start_box)
+                x = int((coords[0] + coords[2]) / 2 * image_width)
+                y = int((coords[1] + coords[3]) / 2 * image_height)
+
+                computer_actions.append(make_double_click_item(x, y))
+
+        elif action_type in ["right_click", "right_single"]:
+            start_box = action_inputs.get("start_box")
+            if start_box:
+                coords = eval(start_box)
+                x = int((coords[0] + coords[2]) / 2 * image_width)
+                y = int((coords[1] + coords[3]) / 2 * image_height)
+
+                computer_actions.append(make_click_item(x, y, "right"))
+
+        elif action_type == "type":
+            content = action_inputs.get("content", "")
+            computer_actions.append(make_type_item(content))
+
+        elif action_type == "hotkey":
+            key = action_inputs.get("key", "")
+            keys = key.split()
+            computer_actions.append(make_keypress_item(keys))
+
+        elif action_type == "press":
+            key = action_inputs.get("key", "")
+            computer_actions.append(make_keypress_item([key]))
+
+        elif action_type == "scroll":
+            start_box = action_inputs.get("start_box")
+            direction = action_inputs.get("direction", "down")
+
+            if start_box:
+                coords = eval(start_box)
+                x = int((coords[0] + coords[2]) / 2 * image_width)
+                y = int((coords[1] + coords[3]) / 2 * image_height)
+            else:
+                x, y = image_width // 2, image_height // 2
+
+            scroll_y = 5 if "up" in direction.lower() else -5
+            computer_actions.append(make_scroll_item(x, y, 0, scroll_y))
+
+        elif action_type == "drag":
+            start_box = action_inputs.get("start_box")
+            end_box = action_inputs.get("end_box")
+
+            if start_box and end_box:
+                start_coords = eval(start_box)
+                end_coords = eval(end_box)
+
+                start_x = int((start_coords[0] + start_coords[2]) / 2 * image_width)
+                start_y = int((start_coords[1] + start_coords[3]) / 2 * image_height)
+                end_x = int((end_coords[0] + end_coords[2]) / 2 * image_width)
+                end_y = int((end_coords[1] + end_coords[3]) / 2 * image_height)
+
+                path = [{"x": start_x, "y": start_y}, {"x": end_x, "y": end_y}]
+                computer_actions.append(make_drag_item(path))
+
+    return computer_actions
+
+
+def pil_to_base64(image: Image.Image) -> str:
+    """Convert PIL image to base64 string."""
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def process_image_for_uitars(
+    image_data: str, max_pixels: int = MAX_PIXELS, min_pixels: int = MIN_PIXELS
+) -> tuple[Image.Image, int, int]:
+    """Process image for UITARS model input."""
+    # Decode base64 image
+    if image_data.startswith("data:image"):
+        image_data = image_data.split(",")[1]
+
+    image_bytes = base64.b64decode(image_data)
+    image = Image.open(BytesIO(image_bytes))
+
+    original_width, original_height = image.size
+
+    # Resize image according to UITARS requirements
+    if image.width * image.height > max_pixels:
+        resize_factor = math.sqrt(max_pixels / (image.width * image.height))
+        width = int(image.width * resize_factor)
+        height = int(image.height * resize_factor)
+        image = image.resize((width, height))
+
+    if image.width * image.height < min_pixels:
+        resize_factor = math.sqrt(min_pixels / (image.width * image.height))
+        width = math.ceil(image.width * resize_factor)
+        height = math.ceil(image.height * resize_factor)
+        image = image.resize((width, height))
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    return image, original_width, original_height
+
+
+def sanitize_message(msg: Any) -> Any:
+    """Return a copy of the message with image_url ommited within content parts"""
+    if isinstance(msg, dict):
+        result = {}
+        for key, value in msg.items():
+            if key == "content" and isinstance(value, list):
+                result[key] = [
+                    (
+                        {k: v for k, v in item.items() if k != "image_url"}
+                        if isinstance(item, dict)
+                        else item
+                    )
+                    for item in value
+                ]
+            else:
+                result[key] = value
+        return result
+    elif isinstance(msg, list):
+        return [sanitize_message(item) for item in msg]
+    else:
+        return msg
+
+
+def convert_uitars_messages_to_litellm(messages: Messages) -> List[Dict[str, Any]]:
+    """
+    Convert UITARS internal message format back to LiteLLM format.
+
+    This function processes reasoning, computer_call, and computer_call_output messages
+    and converts them to the appropriate LiteLLM assistant message format.
+
+    Args:
+        messages: List of UITARS internal messages
+
+    Returns:
+        List of LiteLLM formatted messages
+    """
+    litellm_messages = []
+    current_assistant_content = []
+
+    for message in messages:
+        if isinstance(message, dict):
+            message_type = message.get("type")
+
+            if message_type == "reasoning":
+                # Extract reasoning text from summary
+                summary = message.get("summary", [])
+                if summary and isinstance(summary, list):
+                    for summary_item in summary:
+                        if (
+                            isinstance(summary_item, dict)
+                            and summary_item.get("type") == "summary_text"
+                        ):
+                            reasoning_text = summary_item.get("text", "")
+                            if reasoning_text:
+                                current_assistant_content.append(f"Thought: {reasoning_text}")
+
+            elif message_type == "computer_call":
+                # Convert computer action to UITARS action format
+                action = message.get("action", {})
+                action_type = action.get("type")
+
+                if action_type == "click":
+                    x, y = action.get("x", 0), action.get("y", 0)
+                    button = action.get("button", "left")
+                    if button == "left":
+                        action_text = f"Action: click(start_box='({x},{y})')"
+                    elif button == "right":
+                        action_text = f"Action: right_single(start_box='({x},{y})')"
+                    else:
+                        action_text = f"Action: click(start_box='({x},{y})')"
+
+                elif action_type == "double_click":
+                    x, y = action.get("x", 0), action.get("y", 0)
+                    action_text = f"Action: left_double(start_box='({x},{y})')"
+
+                elif action_type == "drag":
+                    start_x, start_y = action.get("start_x", 0), action.get("start_y", 0)
+                    end_x, end_y = action.get("end_x", 0), action.get("end_y", 0)
+                    action_text = f"Action: drag(start_box='({start_x},{start_y})', end_box='({end_x},{end_y})')"
+
+                elif action_type == "key":
+                    key = action.get("key", "")
+                    action_text = f"Action: hotkey(key='{key}')"
+
+                elif action_type == "type":
+                    text = action.get("text", "")
+                    # Escape single quotes in the text
+                    escaped_text = escape_single_quotes(text)
+                    action_text = f"Action: type(content='{escaped_text}')"
+
+                elif action_type == "scroll":
+                    x, y = action.get("x", 0), action.get("y", 0)
+                    direction = action.get("direction", "down")
+                    action_text = f"Action: scroll(start_box='({x},{y})', direction='{direction}')"
+
+                elif action_type == "wait":
+                    action_text = "Action: wait()"
+
+                else:
+                    # Fallback for unknown action types
+                    action_text = f"Action: {action_type}({action})"
+
+                current_assistant_content.append(action_text)
+
+                # When we hit a computer_call_output, finalize the current assistant message
+                if current_assistant_content:
+                    litellm_messages.append(
+                        {
+                            "role": "assistant",
+                            "content": [
+                                {"type": "text", "text": "\n".join(current_assistant_content)}
+                            ],
+                        }
+                    )
+                    current_assistant_content = []
+
+            elif message_type == "computer_call_output":
+                # Add screenshot from computer call output
+                output = message.get("output", {})
+                if isinstance(output, dict) and output.get("type") == "input_image":
+                    image_url = output.get("image_url", "")
+                    if image_url:
+                        litellm_messages.append(
+                            {
+                                "role": "user",
+                                "content": [{"type": "image_url", "image_url": {"url": image_url}}],
+                            }
+                        )
+
+            elif message.get("role") == "user":
+                # # Handle user messages
+                # content = message.get("content", "")
+                # if isinstance(content, str):
+                #     litellm_messages.append({
+                #         "role": "user",
+                #         "content": content
+                #     })
+                # elif isinstance(content, list):
+                #     litellm_messages.append({
+                #         "role": "user",
+                #         "content": content
+                #     })
+                pass
+
+    # Add any remaining assistant content
+    if current_assistant_content:
+        litellm_messages.append(
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "\n".join(current_assistant_content)}],
+            }
+        )
+
+    return litellm_messages
+
+
+@register_agent(models=r"(?i).*ui-?tars.*", priority=-1)
+class UITARSConfig:
+    """
+    UITARS agent configuration using liteLLM for ByteDance-Seed/UI-TARS-1.5-7B model.
+
+    Supports UITARS vision-language models for computer control.
+    """
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Predict the next step based on input messages.
+
+        Args:
+            messages: Input messages following Responses format
+            model: Model name to use
+            tools: Optional list of tool schemas
+            max_retries: Maximum number of retries
+            stream: Whether to stream responses
+            computer_handler: Computer handler instance
+            _on_api_start: Callback for API start
+            _on_api_end: Callback for API end
+            _on_usage: Callback for usage tracking
+            _on_screenshot: Callback for screenshot events
+            **kwargs: Additional arguments
+
+        Returns:
+            Dictionary with "output" (output items) and "usage" array
+        """
+        tools = tools or []
+
+        # Create response items
+        response_items = []
+
+        # Find computer tool for screen dimensions
+        computer_tool = None
+        for tool_schema in tools:
+            if tool_schema["type"] == "computer":
+                computer_tool = tool_schema["computer"]
+                break
+
+        # Get screen dimensions
+        screen_width, screen_height = 1024, 768
+        if computer_tool:
+            try:
+                screen_width, screen_height = await computer_tool.get_dimensions()
+            except:
+                pass
+
+        # Process messages to extract instruction and image
+        instruction = ""
+        image_data = None
+
+        # Convert messages to list if string
+        if isinstance(messages, str):
+            messages = [{"role": "user", "content": messages}]
+
+        # Extract instruction and latest screenshot
+        for message in reversed(messages):
+            if isinstance(message, dict):
+                content = message.get("content", "")
+
+                # Handle different content formats
+                if isinstance(content, str):
+                    if not instruction and message.get("role") == "user":
+                        instruction = content
+                elif isinstance(content, list):
+                    for item in content:
+                        if isinstance(item, dict):
+                            if item.get("type") == "text" and not instruction:
+                                instruction = item.get("text", "")
+                            elif item.get("type") == "image_url" and not image_data:
+                                image_url = item.get("image_url", {})
+                                if isinstance(image_url, dict):
+                                    image_data = image_url.get("url", "")
+                                else:
+                                    image_data = image_url
+
+            # Also check for computer_call_output with screenshots
+            if message.get("type") == "computer_call_output" and not image_data:
+                output = message.get("output", {})
+                if isinstance(output, dict) and output.get("type") == "input_image":
+                    image_data = output.get("image_url", "")
+
+            if instruction and image_data:
+                break
+
+        if not instruction:
+            instruction = (
+                "Help me complete this task by analyzing the screen and taking appropriate actions."
+            )
+
+        # Create prompt
+        user_prompt = UITARS_PROMPT_TEMPLATE.format(
+            instruction=instruction, action_space=UITARS_ACTION_SPACE, language="English"
+        )
+
+        # Convert conversation history to LiteLLM format
+        history_messages = convert_uitars_messages_to_litellm(messages)
+
+        # Prepare messages for liteLLM
+        litellm_messages = [{"role": "system", "content": "You are a helpful assistant."}]
+
+        # Add current user instruction with screenshot
+        current_user_message = {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": user_prompt},
+            ],
+        }
+        litellm_messages.append(current_user_message)
+
+        # Process image for UITARS
+        if not image_data:
+            # Take screenshot if none found in messages
+            if computer_handler:
+                image_data = await computer_handler.screenshot()
+                await _on_screenshot(image_data, "screenshot_before")
+
+                # Add screenshot to output items so it can be retained in history
+                response_items.append(make_input_image_item(image_data))
+            else:
+                raise ValueError("No screenshot found in messages and no computer_handler provided")
+        processed_image, original_width, original_height = process_image_for_uitars(image_data)
+        encoded_image = pil_to_base64(processed_image)
+
+        # Add conversation history
+        if history_messages:
+            litellm_messages.extend(history_messages)
+        else:
+            litellm_messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/png;base64,{encoded_image}"},
+                        }
+                    ],
+                }
+            )
+
+        # Prepare API call kwargs
+        api_kwargs = {
+            "model": model,
+            "messages": litellm_messages,
+            "max_tokens": kwargs.get("max_tokens", 500),
+            "temperature": kwargs.get("temperature", 0.0),
+            "do_sample": kwargs.get("temperature", 0.0) > 0.0,
+            "num_retries": max_retries,
+            **{k: v for k, v in kwargs.items() if k not in ["max_tokens", "temperature"]},
+        }
+
+        # Call API start hook
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        # Call liteLLM with UITARS model
+        response = await litellm.acompletion(**api_kwargs)
+
+        # Call API end hook
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        # Extract response content
+        response_content = response.choices[0].message.content.strip()  # type: ignore
+
+        # Parse UITARS response
+        parsed_responses = parse_uitars_response(response_content, original_width, original_height)
+
+        # Convert to computer actions
+        computer_actions = convert_to_computer_actions(
+            parsed_responses, original_width, original_height
+        )
+
+        # Add computer actions to response items
+        thought = parsed_responses[0].get("thought", "")
+        if thought:
+            response_items.append(make_reasoning_item(thought))
+        response_items.extend(computer_actions)
+
+        # Extract usage information
+        response_usage = {
+            **LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(
+                response.usage
+            ).model_dump(),
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(response_usage)
+
+        # Create agent response
+        agent_response = {"output": response_items, "usage": response_usage}
+
+        return agent_response
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Predict click coordinates based on image and instruction.
+
+        UITARS supports click prediction through its action parsing.
+
+        Args:
+            model: Model name to use
+            image_b64: Base64 encoded image
+            instruction: Instruction for where to click
+
+        Returns:
+            Tuple with (x, y) coordinates or None
+        """
+        try:
+            # Create prompt using grounding template
+            user_prompt = GROUNDING_UITARS_PROMPT_TEMPLATE.format(instruction=instruction)
+
+            # Process image for UITARS
+            processed_image, original_width, original_height = process_image_for_uitars(image_b64)
+            encoded_image = pil_to_base64(processed_image)
+
+            # Prepare messages for liteLLM
+            litellm_messages = [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": user_prompt},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/png;base64,{encoded_image}"},
+                        },
+                    ],
+                },
+            ]
+
+            # Prepare API call kwargs
+            api_kwargs = {
+                "model": model,
+                "messages": litellm_messages,
+                "max_tokens": 2056,
+                "temperature": 0.0,
+                "do_sample": False,
+            }
+            api_kwargs.update({k: v for k, v in (kwargs or {}).items()})
+
+            # Call liteLLM with UITARS model
+            response = await litellm.acompletion(**api_kwargs)
+
+            # Extract response content
+            response_content = response.choices[0].message.content.strip()  # type: ignore
+
+            print(response_content)
+
+            # Parse the response to extract click coordinates
+            # Look for click action with coordinates (with special tokens)
+            click_pattern = r"click\(point='<\|box_start\|>\((\d+),(\d+)\)<\|box_end\|>'\)"
+            match = re.search(click_pattern, response_content)
+
+            # Fallback: Look for simpler format without special tokens
+            if not match:
+                # Pattern for: click(start_box='(x,y)') or click(point='(x,y)')
+                fallback_pattern = r"click\((?:start_box|point)='\((\d+),(\d+)\)'\)"
+                match = re.search(fallback_pattern, response_content)
+
+            if match:
+                x, y = int(match.group(1)), int(match.group(2))
+                # Scale coordinates back to original image dimensions
+                scale_x = original_width / processed_image.width
+                scale_y = original_height / processed_image.height
+
+                scaled_x = int(x * scale_x)
+                scaled_y = int(y * scale_y)
+
+                return (scaled_x, scaled_y)
+
+            return None
+
+        except Exception as e:
+            # Log error and return None
+            print(f"Error in predict_click: {e}")
+            return None
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        """
+        Get list of capabilities supported by this agent config.
+
+        Returns:
+            List of capability strings
+        """
+        return ["step", "click"]

libs/python/agent/cua_agent/loops/uitars2.py

@@ -0,0 +1,951 @@
+"""
+UITARS-2 agent loop implementation using LiteLLM.
+- Prepends a system prompt modeled after the UI-TARS training prompts
+- Converts Responses items -> completion messages
+- Calls litellm.acompletion
+- Parses <seed:tool_call> ... </seed:tool_call> outputs back into Responses items (computer actions)
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import json
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from litellm.responses.litellm_completion_transformation.transformation import (
+    LiteLLMCompletionResponsesConfig,
+)
+
+from ..decorators import register_agent
+from .omniparser import get_last_computer_call_output  # type: ignore
+
+try:
+    from PIL import Image  # type: ignore
+except Exception:  # pragma: no cover
+    Image = None  # type: ignore
+from ..responses import (
+    convert_responses_items_to_completion_messages,
+    make_click_item,
+    make_double_click_item,
+    make_drag_item,
+    make_function_call_item,
+    make_keypress_item,
+    make_move_item,
+    make_output_text_item,
+    make_reasoning_item,
+    make_screenshot_item,
+    make_scroll_item,
+    make_type_item,
+    make_wait_item,
+)
+from ..types import AgentCapability
+
+TOOL_SCHEMAS: List[Dict[str, Any]] = [
+    {
+        "type": "function",
+        "name": "open_computer",
+        "parameters": {},
+        "description": "Open computer.",
+    },
+    {
+        "type": "function",
+        "name": "click",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "point": {
+                    "type": "string",
+                    "description": "Click coordinates. The format is: <point>x y</point>",
+                }
+            },
+            "required": ["point"],
+        },
+        "description": "Mouse left single click action.",
+    },
+    {
+        "type": "function",
+        "name": "left_double",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "point": {
+                    "type": "string",
+                    "description": "Click coordinates. The format is: <point>x y</point>",
+                }
+            },
+            "required": ["point"],
+        },
+        "description": "Mouse left double click action.",
+    },
+    {
+        "type": "function",
+        "name": "right_single",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "point": {
+                    "type": "string",
+                    "description": "Click coordinates. The format is: <point>x y</point>",
+                }
+            },
+            "required": ["point"],
+        },
+        "description": "Mouse right single click action.",
+    },
+    {
+        "type": "function",
+        "name": "scroll",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "point": {
+                    "type": "string",
+                    "description": "Scroll start position. If not specified, default to execute on the current mouse position. The format is: <point>x y</point>",
+                },
+                "direction": {
+                    "type": "string",
+                    "description": "Scroll direction.",
+                    "enum": ["up", "down", "left", "right"],
+                },
+            },
+            "required": ["direction"],
+        },
+        "description": "Scroll action.",
+    },
+    {
+        "type": "function",
+        "name": "move_to",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "point": {
+                    "type": "string",
+                    "description": "Target coordinates. The format is: <point>x y</point>",
+                }
+            },
+            "required": ["point"],
+        },
+        "description": "Mouse move action.",
+    },
+    {
+        "type": "function",
+        "name": "hotkey",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "key": {
+                    "type": "string",
+                    "description": "Hotkeys you want to press. Split keys with a space and use lowercase.",
+                }
+            },
+            "required": ["key"],
+        },
+        "description": "Press hotkey.",
+    },
+    {
+        "type": "function",
+        "name": "finished",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "content": {
+                    "type": "string",
+                    "description": "Provide the final answer or response to complete the task.",
+                }
+            },
+            "required": [],
+        },
+        "description": "This function is used to indicate the completion of a task by providing the final answer or response.",
+    },
+    {
+        "type": "function",
+        "name": "press",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "key": {
+                    "type": "string",
+                    "description": "Key you want to press. Only one key can be pressed at one time.",
+                }
+            },
+            "required": ["key"],
+        },
+        "description": "Press key.",
+    },
+    {
+        "type": "function",
+        "name": "release",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "key": {
+                    "type": "string",
+                    "description": "Key you want to release. Only one key can be released at one time.",
+                }
+            },
+            "required": ["key"],
+        },
+        "description": "Release key.",
+    },
+    {
+        "type": "function",
+        "name": "mouse_down",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "point": {
+                    "type": "string",
+                    "description": "Mouse down position. If not specified, default to execute on the current mouse position. The format is: <point>x y</point>",
+                },
+                "button": {
+                    "type": "string",
+                    "description": "Down button. Default to left.",
+                    "enum": ["left", "right"],
+                },
+            },
+            "required": [],
+        },
+        "description": "Mouse down action.",
+    },
+    {
+        "type": "function",
+        "name": "mouse_up",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "point": {
+                    "type": "string",
+                    "description": "Mouse up position. If not specified, default to execute on the current mouse position. The format is: <point>x y</point>",
+                },
+                "button": {
+                    "type": "string",
+                    "description": "Up button. Default to left.",
+                    "enum": ["left", "right"],
+                },
+            },
+            "required": [],
+        },
+        "description": "Mouse up action.",
+    },
+    {
+        "type": "function",
+        "name": "call_user",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "content": {
+                    "type": "string",
+                    "description": "Message or information displayed to the user to request their input, feedback, or guidance.",
+                }
+            },
+            "required": [],
+        },
+        "description": "This function is used to interact with the user by displaying a message and requesting their input, feedback, or guidance.",
+    },
+    {
+        "type": "function",
+        "name": "wait",
+        "parameters": {
+            "type": "object",
+            "properties": {"time": {"type": "integer", "description": "Wait time in seconds."}},
+            "required": [],
+        },
+        "description": "Wait for a while.",
+    },
+    {
+        "type": "function",
+        "name": "drag",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "start_point": {
+                    "type": "string",
+                    "description": "Drag start point. The format is: <point>x y</point>",
+                },
+                "end_point": {
+                    "type": "string",
+                    "description": "Drag end point. The format is: <point>x y</point>",
+                },
+            },
+            "required": ["start_point", "end_point"],
+        },
+        "description": "Mouse left button drag action.",
+    },
+    {
+        "type": "function",
+        "name": "type",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "content": {
+                    "type": "string",
+                    "description": "Type content. If you want to submit your input, use \\n at the end of content.",
+                }
+            },
+            "required": ["content"],
+        },
+        "description": "Type content.",
+    },
+    {
+        "type": "function",
+        "name": "take_screenshot",
+        "parameters": {},
+        "description": "Take screenshot.",
+    },
+]
+
+
+def _format_tool_schemas_json_lines(schemas: List[Dict[str, Any]]) -> str:
+    # Nicely formatted: pretty JSON with indentation, separated by blank lines
+    return "\n\n".join(json.dumps(s, ensure_ascii=False, indent=2) for s in schemas) + "\n\n"
+
+
+_PROMPT_PREFIX = (
+    "You should begin by detailing the internal reasoning process, and then present the answer to the user. "
+    "The reasoning process should be enclosed within <think_never_used_51bce0c785ca2f68081bfa7d91973934> "
+    "</think_never_used_51bce0c785ca2f68081bfa7d91973934> tags, as follows:\n"
+    "<think_never_used_51bce0c785ca2f68081bfa7d91973934> reasoning process here "
+    "</think_never_used_51bce0c785ca2f68081bfa7d91973934> answer here.\n\n"
+    "You have different modes of thinking:\n"
+    "Unrestricted think mode: Engage in an internal thinking process with thorough reasoning and reflections. "
+    "You have an unlimited budget for thinking tokens and can continue thinking until you fully solve the problem.\n"
+    "Efficient think mode: Provide a concise internal thinking process with efficient reasoning and reflections. "
+    "You don't have a strict token budget but be less verbose and more direct in your thinking.\n"
+    "No think mode: Respond directly to the question without any internal reasoning process or extra thinking tokens. "
+    "Still follow the template with the minimum required thinking tokens to justify the answer.\n"
+    "Budgeted think mode: Limit your internal reasoning and reflections to stay within the specified token budget\n\n"
+    "Based on the complexity of the problem, select the appropriate mode for reasoning among the provided options listed below.\n\n"
+    "Provided Mode(s):\nEfficient think.\n\n"
+    "You are provided with a task description, a history of previous actions, and corresponding screenshots. "
+    "Your goal is to perform the next action to complete the task. "
+    "If performing the same action multiple times results in a static screen with no changes, attempt a modified or alternative action.\n\n"
+    "## Function Definition\n\n"
+    "- You have access to the following functions:\n\n"
+)
+
+_PROMPT_SUFFIX = (
+    "- To call a function, use the following structure without any suffix:\n\n"
+    "<gui_think> reasoning process </gui_think>\n"
+    "<seed:tool_call><function=example_function_name><parameter=example_parameter_1>value_1</parameter>"
+    "<parameter=example_parameter_2>multiline...\n</parameter></function></seed:tool_call>\n\n"
+    "## Important Notes\n"
+    "- Function calls must begin with <function= and end with </function>.\n"
+    "- All required parameters must be explicitly provided.\n"
+    "\n## Additional Notes\n"
+    "- You can execute multiple actions within a single tool call. For example:\n"
+    "<seed:tool_call><function=example_function_1><parameter=example_parameter_1>value_1</parameter><parameter=example_parameter_2>\n"
+    "This is the value for the second parameter\nthat can span\nmultiple lines\n"
+    "</parameter></function><function=example_function_2><parameter=example_parameter_3>value_4</parameter></function></seed:tool_call>"
+)
+
+
+SYSTEM_PROMPT = _PROMPT_PREFIX + _format_tool_schemas_json_lines(TOOL_SCHEMAS) + _PROMPT_SUFFIX
+
+
+def _extract_function_schemas_from_tools(
+    tools: Optional[List[Dict[str, Any]]],
+) -> List[Dict[str, Any]]:
+    schemas: List[Dict[str, Any]] = []
+    if not tools:
+        return schemas
+    for t in tools:
+        if t.get("type") == "function":
+            fn = t.get("function", {})
+            name = fn.get("name")
+            params = fn.get("parameters", {})
+            desc = fn.get("description", "")
+            if name:
+                schemas.append(
+                    {
+                        "type": "function",
+                        "name": name,
+                        "parameters": params if isinstance(params, dict) else {},
+                        "description": desc,
+                    }
+                )
+    return schemas
+
+
+def _parse_seed_tool_calls(text: str) -> List[Dict[str, Any]]:
+    """Parse <seed:tool_call> blocks into a list of {function, parameters} dicts.
+    Also captures optional <gui_think>...</gui_think> as reasoning.
+    """
+    actions: List[Dict[str, Any]] = []
+    if not text:
+        return actions
+
+    # Extract reasoning if present
+    reasoning_text = None
+    think_match = re.search(r"<gui_think>([\s\S]*?)</gui_think>", text)
+    if think_match:
+        reasoning_text = think_match.group(1).strip()
+
+    # Iterate each seed tool_call block
+    for block in re.finditer(r"<seed:tool_call>([\s\S]*?)</seed:tool_call>", text):
+        content = block.group(1)
+        # One or multiple <function=...>...</function> inside
+        for fmatch in re.finditer(r"<function=([\w_]+)>([\s\S]*?)</function>", content):
+            fname = fmatch.group(1)
+            inner = fmatch.group(2)
+            params: Dict[str, str] = {}
+            for pmatch in re.finditer(r"<parameter=([\w_]+)>([\s\S]*?)</parameter>", inner):
+                pname = pmatch.group(1)
+                pval = pmatch.group(2).strip()
+                params[pname] = pval
+            actions.append({"function": fname, "parameters": params})
+
+    # If we have a global reasoning and at least one action, attach it to first
+    if reasoning_text and actions:
+        actions[0]["reasoning"] = reasoning_text
+    elif reasoning_text:
+        actions.append({"function": "reasoning", "parameters": {"content": reasoning_text}})
+
+    return actions
+
+
+def _normalize_xy_to_uitars(x: int, y: int, width: int, height: int) -> Tuple[int, int]:
+    width = max(1, int(width))
+    height = max(1, int(height))
+    nx = max(0, min(1000, int(round((x / width) * 1000))))
+    ny = max(0, min(1000, int(round((y / height) * 1000))))
+    return nx, ny
+
+
+def _denormalize_xy_from_uitars(nx: float, ny: float, width: int, height: int) -> Tuple[int, int]:
+    width = max(1, int(width))
+    height = max(1, int(height))
+    x = int(round((nx / 1000.0) * width))
+    y = int(round((ny / 1000.0) * height))
+    return x, y
+
+
+def _map_computer_action_to_function(
+    action: Dict[str, Any], width: int, height: int
+) -> Optional[Dict[str, Any]]:
+    """Map a computer action item to a UITARS function + parameters dict of strings.
+    Returns dict like {"function": name, "parameters": {..}} or None if unknown.
+    """
+    atype = action.get("type") or action.get("action")
+    if atype == "click":
+        x, y = action.get("x"), action.get("y")
+        btn = action.get("button", "left")
+        if x is None or y is None:
+            return None
+        nx, ny = _normalize_xy_to_uitars(int(x), int(y), width, height)
+        if btn == "right":
+            return {
+                "function": "right_single",
+                "parameters": {"point": f"<point>{nx} {ny}</point>"},
+            }
+        return {"function": "click", "parameters": {"point": f"<point>{nx} {ny}</point>"}}
+    if atype == "double_click":
+        x, y = action.get("x"), action.get("y")
+        if x is None or y is None:
+            return None
+        nx, ny = _normalize_xy_to_uitars(int(x), int(y), width, height)
+        return {"function": "left_double", "parameters": {"point": f"<point>{nx} {ny}</point>"}}
+    if atype == "move":
+        x, y = action.get("x"), action.get("y")
+        if x is None or y is None:
+            return None
+        nx, ny = _normalize_xy_to_uitars(int(x), int(y), width, height)
+        return {"function": "move_to", "parameters": {"point": f"<point>{nx} {ny}</point>"}}
+    if atype == "keypress":
+        keys = action.get("keys", [])
+        if isinstance(keys, list) and keys:
+            if len(keys) == 1:
+                return {"function": "press", "parameters": {"key": keys[0]}}
+            else:
+                return {"function": "hotkey", "parameters": {"key": " ".join(keys)}}
+        return None
+    if atype == "type":
+        text = action.get("text", "")
+        return {"function": "type", "parameters": {"content": text}}
+    if atype == "scroll":
+        x, y = action.get("x", 512), action.get("y", 512)
+        nx, ny = _normalize_xy_to_uitars(int(x), int(y), width, height)
+        sx, sy = action.get("scroll_x", 0), action.get("scroll_y", 0)
+        # Our parser used positive sy for up
+        direction = (
+            "up"
+            if sy and sy > 0
+            else (
+                "down"
+                if sy and sy < 0
+                else ("right" if sx and sx > 0 else ("left" if sx and sx < 0 else "down"))
+            )
+        )
+        return {
+            "function": "scroll",
+            "parameters": {"direction": direction, "point": f"<point>{nx} {ny}</point>"},
+        }
+    if atype == "drag":
+        path = action.get("path", [])
+        if isinstance(path, list) and len(path) >= 2:
+            sx, sy = path[0].get("x"), path[0].get("y")
+            ex, ey = path[-1].get("x"), path[-1].get("y")
+            if sx is None or sy is None or ex is None or ey is None:
+                return None
+            nsx, nsy = _normalize_xy_to_uitars(int(sx), int(sy), width, height)
+            nex, ney = _normalize_xy_to_uitars(int(ex), int(ey), width, height)
+            return {
+                "function": "drag",
+                "parameters": {
+                    "start_point": f"<point>{nsx} {nsy}</point>",
+                    "end_point": f"<point>{nex} {ney}</point>",
+                },
+            }
+        return None
+    if atype == "wait":
+        return {"function": "wait", "parameters": {}}
+    if atype == "screenshot":
+        return {"function": "take_screenshot", "parameters": {}}
+    # Fallback unknown
+    return None
+
+
+def _to_uitars_messages(
+    messages: List[Dict[str, Any]], width: int, height: int
+) -> List[Dict[str, Any]]:
+    """Convert responses items into completion messages tailored for UI-TARS.
+
+    - User content is passed through similar to convert_responses_items_to_completion_messages
+    - Assistant/tool history is rendered as text with <gui_think> and <seed:tool_call> blocks
+    """
+    uitars_messages: List[Dict[str, Any]] = []
+
+    def flush_seed_block(pending_think: Optional[str], pending_functions: List[Dict[str, Any]]):
+        if not pending_think and not pending_functions:
+            return
+        parts: List[str] = []
+        if pending_think:
+            parts.append(f"<gui_think> {pending_think} </gui_think>")
+        if pending_functions:
+            inner = []
+            for f in pending_functions:
+                fname = f["function"]
+                params = f.get("parameters", {})
+                param_blocks = []
+                for k, v in params.items():
+                    param_blocks.append(f"<parameter={k}>{v}</parameter>")
+                inner.append(f"<function={fname}>{''.join(param_blocks)}</function>")
+            parts.append(f"<seed:tool_call>{''.join(inner)}</seed:tool_call>")
+        uitars_messages.append({"role": "assistant", "content": "".join(parts)})
+
+    # Accumulators for a single assistant seed block
+    pending_think: Optional[str] = None
+    pending_functions: List[Dict[str, Any]] = []
+
+    for msg in messages:
+        mtype = msg.get("type")
+        role = msg.get("role")
+
+        # On any user message, flush current assistant block
+        if role == "user" or mtype == "user":
+            flush_seed_block(pending_think, pending_functions)
+            pending_think, pending_functions = None, []
+
+            content = msg.get("content", "")
+            if isinstance(content, list):
+                completion_content = []
+                for item in content:
+                    if item.get("type") == "input_image":
+                        completion_content.append(
+                            {"type": "image_url", "image_url": {"url": item.get("image_url")}}
+                        )
+                    elif item.get("type") in ("input_text", "text"):
+                        completion_content.append({"type": "text", "text": item.get("text")})
+                uitars_messages.append({"role": "user", "content": completion_content})
+            elif isinstance(content, str):
+                uitars_messages.append({"role": "user", "content": content})
+            continue
+
+        # Reasoning item
+        if mtype == "reasoning":
+            # Responses reasoning stores summary list
+            summary = msg.get("summary", [])
+            texts = [
+                s.get("text", "")
+                for s in summary
+                if isinstance(s, dict) and s.get("type") == "summary_text"
+            ]
+            if texts:
+                pending_think = "\n".join([t for t in texts if t])
+            continue
+
+        # Computer/tool calls -> map to functions
+        if mtype == "computer_call":
+            f = _map_computer_action_to_function(msg.get("action", {}), width, height)
+            if f:
+                pending_functions.append(f)
+            continue
+        if mtype == "function_call":
+            # Include custom tools as-is
+            name = msg.get("name")
+            try:
+                args_obj = json.loads(msg.get("arguments", "{}"))
+            except json.JSONDecodeError:
+                args_obj = {}
+            # Ensure string values
+            params = {k: (str(v) if not isinstance(v, str) else v) for k, v in args_obj.items()}
+            pending_functions.append({"function": name, "parameters": params})
+            continue
+
+        # If assistant message text is given, flush current block and add as plain assistant text
+        if role == "assistant" or mtype == "message":
+            flush_seed_block(pending_think, pending_functions)
+            pending_think, pending_functions = None, []
+            content = msg.get("content", [])
+            if isinstance(content, list):
+                texts = [
+                    c.get("text", "")
+                    for c in content
+                    if isinstance(c, dict) and c.get("type") in ("output_text", "text")
+                ]
+                if texts:
+                    uitars_messages.append(
+                        {"role": "assistant", "content": "\n".join([t for t in texts if t])}
+                    )
+            elif isinstance(content, str) and content:
+                uitars_messages.append({"role": "assistant", "content": content})
+            continue
+
+        # On outputs, flush pending assistant block and send outputs as user messages
+        if mtype in ("function_call_output", "computer_call_output"):
+            flush_seed_block(pending_think, pending_functions)
+            pending_think, pending_functions = None, []
+            output = msg.get("output")
+            if isinstance(output, dict) and output.get("type") == "input_image":
+                img_url = output.get("image_url")
+                if img_url:
+                    uitars_messages.append(
+                        {
+                            "role": "user",
+                            "content": [
+                                {"type": "image_url", "image_url": {"url": img_url}},
+                            ],
+                        }
+                    )
+            elif isinstance(output, str):
+                uitars_messages.append({"role": "user", "content": output})
+            else:
+                # Fallback stringify
+                uitars_messages.append({"role": "user", "content": json.dumps(output)})
+            continue
+
+    # Flush any remaining pending seed block
+    flush_seed_block(pending_think, pending_functions)
+
+    return uitars_messages
+
+
+def _to_response_items(
+    actions: List[Dict[str, Any]],
+    tool_names: Optional[set[str]] = None,
+    width: Optional[int] = None,
+    height: Optional[int] = None,
+) -> List[Any]:
+    """Map parsed actions into Responses items (computer actions + optional reasoning)."""
+    items: List[Any] = []
+    tool_names = tool_names or set()
+
+    # Optional top-level reasoning attached to first
+    if actions and actions[0].get("reasoning"):
+        items.append(make_reasoning_item(actions[0]["reasoning"]))
+
+    # Dimensions default
+    w = int(width) if width else 1024
+    h = int(height) if height else 768
+
+    for a in actions:
+        fn = a.get("function")
+        params = a.get("parameters", {})
+        if fn == "reasoning":
+            items.append(make_reasoning_item(params.get("content", "")))
+        elif fn in ("click", "left_double", "right_single"):
+            # params.point is like: <point>x y</point> or plain "x y"
+            point = params.get("point", "").strip()
+            m = re.search(r"([\-\d\.]+)\s+([\-\d\.]+)", point)
+            if not m:
+                continue
+            nx = float(m.group(1))
+            ny = float(m.group(2))
+            x, y = _denormalize_xy_from_uitars(nx, ny, w, h)
+            if fn == "left_double":
+                items.append(make_double_click_item(x, y))
+            elif fn == "right_single":
+                items.append(make_click_item(x, y, "right"))
+            else:
+                items.append(make_click_item(x, y, "left"))
+        elif fn == "move_to":
+            point = params.get("point", "").strip()
+            m = re.search(r"([\-\d\.]+)\s+([\-\d\.]+)", point)
+            if not m:
+                continue
+            nx = float(m.group(1))
+            ny = float(m.group(2))
+            x, y = _denormalize_xy_from_uitars(nx, ny, w, h)
+            items.append(make_move_item(x, y))
+        elif fn == "drag":
+            sp = params.get("start_point", "").strip()
+            ep = params.get("end_point", "").strip()
+            ms = re.search(r"([\-\d\.]+)\s+([\-\d\.]+)", sp)
+            me = re.search(r"([\-\d\.]+)\s+([\-\d\.]+)", ep)
+            if not (ms and me):
+                continue
+            nsx, nsy = float(ms.group(1)), float(ms.group(2))
+            nex, ney = float(me.group(1)), float(me.group(2))
+            sx, sy = _denormalize_xy_from_uitars(nsx, nsy, w, h)
+            ex, ey = _denormalize_xy_from_uitars(nex, ney, w, h)
+            items.append(make_drag_item([{"x": sx, "y": sy}, {"x": ex, "y": ey}]))
+        elif fn == "hotkey":
+            key = params.get("key", "")
+            keys = key.split()
+            if keys:
+                items.append(make_keypress_item(keys))
+        elif fn == "press":
+            key = params.get("key", "")
+            if key:
+                items.append(make_keypress_item([key]))
+        elif fn == "type":
+            content = params.get("content", "")
+            items.append(make_type_item(content))
+        elif fn == "scroll":
+            # direction: up/down/left/right. Point optional
+            direction = params.get("direction", "down").lower()
+            point = params.get("point", "")
+            m = re.search(r"([\-\d\.]+)\s+([\-\d\.]+)", point)
+            if m:
+                nx = float(m.group(1))
+                ny = float(m.group(2))
+                x, y = _denormalize_xy_from_uitars(nx, ny, w, h)
+            else:
+                x, y = _denormalize_xy_from_uitars(500.0, 500.0, w, h)
+            dy = 5 if direction == "up" else -5
+            dx = 5 if direction == "right" else (-5 if direction == "left" else 0)
+            items.append(make_scroll_item(x, y, dx, dy))
+        elif fn == "wait":
+            items.append(make_wait_item())
+        elif fn == "finished":
+            content = params.get("content", "")
+            items.append(make_output_text_item(content or "Task completed."))
+            break
+        elif fn == "take_screenshot":
+            items.append(make_screenshot_item())
+        elif fn == "open_computer":
+            items.append(make_screenshot_item())
+        else:
+            # If this function name is present in provided tool schemas, emit function_call
+            if fn in tool_names:
+                # Convert simple string params into an arguments object
+                # Parameters are strings; pass through as-is
+                items.append(make_function_call_item(fn, params))
+            else:
+                # Unknown function -> surface as assistant text
+                items.append(make_output_text_item(f"Unknown action: {fn} {params}"))
+
+    return items
+
+
+@register_agent(models=r"(?i).*ui-?tars-?2.*")
+class UITARS2Config:
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        # Determine screen dimensions (prefer computer_handler, fallback to last screenshot)
+        width: Optional[int] = None
+        height: Optional[int] = None
+        if computer_handler is not None and hasattr(computer_handler, "get_dimensions"):
+            try:
+                dims = await computer_handler.get_dimensions()  # type: ignore
+                if isinstance(dims, (list, tuple)) and len(dims) == 2:
+                    width, height = int(dims[0]), int(dims[1])
+            except Exception:
+                pass
+
+        if width is None or height is None:
+            try:
+                last_out = get_last_computer_call_output(messages)  # type: ignore
+                if last_out:
+                    image_url = last_out.get("output", {}).get("image_url", "")
+                    if image_url:
+                        b64 = image_url.split(",")[-1]
+                        img_bytes = base64.b64decode(b64)
+                        if Image is not None:
+                            img = Image.open(io.BytesIO(img_bytes))
+                            width, height = img.size
+            except Exception:
+                pass
+
+        if width is None or height is None:
+            width, height = 1024, 768
+
+        # Convert Responses items to UI-TARS style messages with <seed:tool_call> history
+        completion_messages = _to_uitars_messages(messages, width, height)
+
+        # Build dynamic system prompt by concatenating built-in schemas and provided function tools
+        provided_fn_schemas = _extract_function_schemas_from_tools(tools)
+        combined_schemas = (
+            TOOL_SCHEMAS + provided_fn_schemas if provided_fn_schemas else TOOL_SCHEMAS
+        )
+        dynamic_system_prompt = (
+            _PROMPT_PREFIX + _format_tool_schemas_json_lines(combined_schemas) + _PROMPT_SUFFIX
+        )
+
+        # Prepend system prompt (based on training prompts + provided tools)
+        litellm_messages: List[Dict[str, Any]] = [
+            {"role": "system", "content": dynamic_system_prompt},
+        ]
+        litellm_messages.extend(completion_messages)
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": litellm_messages,
+            "max_retries": max_retries,
+            "stream": stream,
+            **{k: v for k, v in kwargs.items()},
+        }
+        if use_prompt_caching:
+            api_kwargs["use_prompt_caching"] = use_prompt_caching
+
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        response = await litellm.acompletion(**api_kwargs)
+
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        usage = {
+            **LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(  # type: ignore
+                response.usage
+            ).model_dump(),
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        # Extract text content (first choice)
+        response_dict = response.model_dump()  # type: ignore
+        content_text = ""
+        choices = response_dict.get("choices", [])
+        if choices:
+            msg = choices[0].get("message", {})
+            # message.content may be string or array; gather text pieces
+            mc = msg.get("content")
+            if isinstance(mc, str):
+                content_text = mc
+            elif isinstance(mc, list):
+                parts = []
+                for part in mc:
+                    if isinstance(part, dict) and part.get("type") == "text":
+                        parts.append(part.get("text", ""))
+                content_text = "\n".join([p for p in parts if p])
+
+        # Parse the seed tool calls and map to response items
+        actions = _parse_seed_tool_calls(content_text)
+        # Build set of tool names from provided tools to emit function_call items
+        tool_names: set[str] = set()
+        for s in provided_fn_schemas:
+            name = s.get("name")
+            if isinstance(name, str):
+                tool_names.add(name)
+        output_items = _to_response_items(actions, tool_names, width, height)
+
+        return {"output": output_items, "usage": usage}
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        return ["step"]
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        """Predict a single click coordinate using a minimal prompt with a click tool.
+
+        This sends the current screenshot and instruction, asking the model to
+        output a click action in the form:
+            Action: click(point='(x,y)')
+        """
+        # Minimal grounding-style prompt
+        system_text = (
+            "You are a GUI agent. Given the instruction, return a single action on the current screen.\n\n"
+            "## Output Format\n\n"
+            "Action: click(point='(x,y)')\n\n"
+            "## User Instruction\n"
+            f"{instruction}"
+        )
+
+        # Build messages with image
+        litellm_messages: List[Dict[str, Any]] = [
+            {"role": "system", "content": system_text},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Please return a single click action."},
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{image_b64}"},
+                    },
+                ],
+            },
+        ]
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": litellm_messages,
+            "max_tokens": kwargs.get("max_tokens", 512),
+            "temperature": kwargs.get("temperature", 0.0),
+            "do_sample": kwargs.get("temperature", 0.0) > 0.0,
+        }
+        api_kwargs.update(
+            {k: v for k, v in (kwargs or {}).items() if k not in ["max_tokens", "temperature"]}
+        )
+
+        response = await litellm.acompletion(**api_kwargs)
+        # Extract response content
+        response_dict = response.model_dump()  # type: ignore
+        choices = response_dict.get("choices", [])
+        if not choices:
+            return None
+        msg = choices[0].get("message", {})
+        content_text = msg.get("content", "")
+        if isinstance(content_text, list):
+            text_parts = [
+                p.get("text", "")
+                for p in content_text
+                if isinstance(p, dict) and p.get("type") == "text"
+            ]
+            content_text = "\n".join([t for t in text_parts if t])
+        if not isinstance(content_text, str):
+            return None
+
+        # Parse coordinates
+        # Pattern for click(point='(x,y)') or click(start_box='(x,y)')
+        patterns = [
+            r"click\(point='\((\d+),(\d+)\)'\)",
+            r"click\((?:start_box|point)='\((\d+),(\d+)\)'\)",
+        ]
+        for pat in patterns:
+            m = re.search(pat, content_text)
+            if m:
+                try:
+                    x, y = int(m.group(1)), int(m.group(2))
+                    return (x, y)
+                except Exception:
+                    pass
+        return None

libs/python/agent/cua_agent/loops/yutori.py

@@ -0,0 +1,397 @@
+"""
+Yutori n1 agent loop implementation using litellm.
+
+n1 is a browser-use model that outputs actions via tool_calls in OpenAI chat
+completions format. Coordinates are in a 1000x1000 normalized space.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import json
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from litellm.responses.litellm_completion_transformation.transformation import (
+    LiteLLMCompletionResponsesConfig,
+)
+from PIL import Image
+
+from ..decorators import register_agent
+from ..loops.base import AsyncAgentConfig
+from ..responses import (
+    convert_completion_messages_to_responses_items,
+    convert_responses_items_to_completion_messages,
+    make_function_call_item,
+    make_output_text_item,
+    make_reasoning_item,
+)
+from ..types import AgentCapability
+
+# Target resolution for n1 (docs recommend 1280x800 WebP)
+N1_TARGET_WIDTH = 1280
+N1_TARGET_HEIGHT = 800
+N1_COORD_SPACE = 1000
+
+
+def _prepare_image_for_n1(image_b64: str) -> str:
+    """Convert a base64 PNG screenshot to WebP at 1280x800 for optimal n1 performance."""
+    try:
+        img_bytes = base64.b64decode(image_b64)
+        img = Image.open(io.BytesIO(img_bytes))
+
+        # Resize to n1's recommended resolution
+        if img.size != (N1_TARGET_WIDTH, N1_TARGET_HEIGHT):
+            img = img.resize((N1_TARGET_WIDTH, N1_TARGET_HEIGHT), Image.LANCZOS)
+
+        # Convert to WebP
+        buf = io.BytesIO()
+        img.save(buf, format="WEBP", quality=85)
+        return base64.b64encode(buf.getvalue()).decode("utf-8")
+    except Exception:
+        # Fallback: return original image if conversion fails
+        return image_b64
+
+
+def _unnormalize_coordinates(
+    coords: List[int], screen_width: int, screen_height: int
+) -> Tuple[int, int]:
+    """Scale coordinates from n1's 1000x1000 space to actual screen pixels."""
+    x = max(0, min(screen_width, round((coords[0] / N1_COORD_SPACE) * screen_width)))
+    y = max(0, min(screen_height, round((coords[1] / N1_COORD_SPACE) * screen_height)))
+    return x, y
+
+
+def _convert_n1_action_to_computer_action(
+    fn_name: str, args: Dict[str, Any], screen_width: int, screen_height: int
+) -> Optional[Dict[str, Any]]:
+    """
+    Convert an n1 tool call to the internal computer_call action schema.
+
+    Returns None for actions that should be emitted as function_calls instead
+    (goto_url, go_back, refresh).
+    """
+    # Actions with coordinates
+    coords = args.get("coordinates")
+    x, y = None, None
+    if isinstance(coords, (list, tuple)) and len(coords) >= 2:
+        x, y = _unnormalize_coordinates(coords, screen_width, screen_height)
+
+    if fn_name == "left_click":
+        if x is None or y is None:
+            return None
+        return {"action": "left_click", "x": x, "y": y}
+
+    if fn_name == "double_click":
+        if x is None or y is None:
+            return None
+        return {"action": "double_click", "x": x, "y": y}
+
+    if fn_name == "triple_click":
+        # Approximate as double_click
+        if x is None or y is None:
+            return None
+        return {"action": "double_click", "x": x, "y": y}
+
+    if fn_name == "right_click":
+        if x is None or y is None:
+            return None
+        return {"action": "right_click", "x": x, "y": y}
+
+    if fn_name == "hover":
+        if x is None or y is None:
+            return None
+        return {"action": "move", "x": x, "y": y}
+
+    if fn_name == "drag":
+        start_coords = args.get("start_coordinates")
+        if (
+            not isinstance(start_coords, (list, tuple))
+            or len(start_coords) < 2
+            or x is None
+            or y is None
+        ):
+            return None
+        sx, sy = _unnormalize_coordinates(start_coords, screen_width, screen_height)
+        return {
+            "action": "drag",
+            "start_x": sx,
+            "start_y": sy,
+            "end_x": x,
+            "end_y": y,
+        }
+
+    if fn_name == "scroll":
+        direction = args.get("direction", "down")
+        amount = int(args.get("amount", 3))
+        # Convert direction + amount to scroll_x/scroll_y pixels
+        # Use ~100 pixels per scroll unit as a reasonable default
+        pixels_per_unit = 100
+        scroll_x, scroll_y = 0, 0
+        if direction == "down":
+            scroll_y = amount * pixels_per_unit
+        elif direction == "up":
+            scroll_y = -(amount * pixels_per_unit)
+        elif direction == "right":
+            scroll_x = amount * pixels_per_unit
+        elif direction == "left":
+            scroll_x = -(amount * pixels_per_unit)
+        out: Dict[str, Any] = {"action": "scroll", "scroll_x": scroll_x, "scroll_y": scroll_y}
+        if x is not None and y is not None:
+            out["x"] = x
+            out["y"] = y
+        return out
+
+    if fn_name == "type":
+        text = args.get("text", "")
+        if args.get("press_enter_after"):
+            text = text + "\n"
+        # Note: clear_before_typing is not supported by the framework's type action.
+        # n1 rarely emits this flag; when it does, the field may already be empty.
+        return {"action": "type", "text": text}
+
+    if fn_name == "key_press":
+        key_comb = args.get("key_comb", "")
+        # n1 uses Playwright-compatible key combos like "Control+a", "Escape"
+        keys = [k.strip() for k in key_comb.split("+")]
+        return {"action": "keypress", "keys": keys}
+
+    if fn_name == "wait":
+        return {"action": "wait"}
+
+    if fn_name == "go_back":
+        return {"action": "history_back"}
+
+    if fn_name == "refresh":
+        return {"action": "keypress", "keys": ["F5"]}
+
+    if fn_name == "goto_url":
+        return {"action": "visit_url", "url": args.get("url", "")}
+
+    return None
+
+
+def _convert_images_to_n1_format(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Convert all images in messages to WebP format optimized for n1."""
+    for msg in messages:
+        content = msg.get("content")
+        if not isinstance(content, list):
+            continue
+        for part in content:
+            if isinstance(part, dict) and part.get("type") == "image_url":
+                url = ((part.get("image_url") or {}).get("url")) or ""
+                if url.startswith("data:") and "," in url:
+                    b64 = url.split(",", 1)[1]
+                    converted = _prepare_image_for_n1(b64)
+                    part["image_url"]["url"] = f"data:image/webp;base64,{converted}"
+    return messages
+
+
+@register_agent(models=r"(yutori/)?n1(-.*)?$", tool_type="browser")
+class YutoriN1Config(AsyncAgentConfig):
+    """
+    Yutori n1 browser-use agent loop.
+
+    n1 is a browser-only model that outputs actions as tool_calls.
+    Coordinates use a 1000x1000 normalized space.
+    """
+
+    async def predict_step(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        max_retries: Optional[int] = None,
+        stream: bool = False,
+        computer_handler=None,
+        use_prompt_caching: Optional[bool] = False,
+        _on_api_start=None,
+        _on_api_end=None,
+        _on_usage=None,
+        _on_screenshot=None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Predict the next browser action using Yutori n1."""
+        tools = tools or []
+
+        # Get screen dimensions for coordinate denormalization
+        screen_width, screen_height = N1_TARGET_WIDTH, N1_TARGET_HEIGHT
+        if computer_handler:
+            try:
+                screen_width, screen_height = await computer_handler.get_dimensions()
+            except Exception:
+                # BrowserTool doesn't have get_dimensions() but has viewport attrs
+                vw = getattr(computer_handler, "viewport_width", None)
+                vh = getattr(computer_handler, "viewport_height", None)
+                if vw and vh:
+                    screen_width, screen_height = vw, vh
+
+        # Convert messages from Responses API format to chat completions format
+        completion_messages = convert_responses_items_to_completion_messages(
+            messages,
+            allow_images_in_tool_results=True,
+        )
+
+        # Convert images to WebP at 1280x800
+        completion_messages = _convert_images_to_n1_format(completion_messages)
+
+        # If there's no screenshot, take one and inject it
+        def _has_any_image(msgs: List[Dict[str, Any]]) -> bool:
+            for m in msgs:
+                content = m.get("content")
+                if isinstance(content, list):
+                    for p in content:
+                        if isinstance(p, dict) and p.get("type") == "image_url":
+                            return True
+            return False
+
+        pre_output_items: List[Dict[str, Any]] = []
+        if not _has_any_image(completion_messages):
+            if computer_handler is None or not hasattr(computer_handler, "screenshot"):
+                raise RuntimeError(
+                    "No screenshots present and computer_handler.screenshot is not available."
+                )
+            screenshot_b64 = await computer_handler.screenshot()
+            if not screenshot_b64:
+                raise RuntimeError("Failed to capture screenshot from computer_handler.")
+
+            converted = _prepare_image_for_n1(screenshot_b64)
+            completion_messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/webp;base64,{converted}"},
+                        },
+                        {"type": "text", "text": "Current browser screen"},
+                    ],
+                }
+            )
+            pre_output_items.append(
+                {
+                    "type": "message",
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "Taking a screenshot to see the current browser screen.",
+                        }
+                    ],
+                }
+            )
+
+        # Build tool list: pass through any custom function tools
+        n1_tools = []
+        for tool in tools:
+            if tool.get("type") == "function":
+                func = tool.get("function")
+                if func:
+                    n1_tools.append({"type": "function", "function": func})
+            # Skip computer tools — n1 has built-in browser actions
+
+        api_kwargs: Dict[str, Any] = {
+            "model": model,
+            "messages": completion_messages,
+            "max_retries": max_retries,
+            "stream": False,  # n1 does not support streaming
+            "temperature": kwargs.pop("temperature", 0.3),
+        }
+
+        if n1_tools:
+            api_kwargs["tools"] = n1_tools
+
+        # Pass through remaining kwargs (api_key, api_base, etc.)
+        api_kwargs.update({k: v for k, v in kwargs.items()})
+
+        if _on_api_start:
+            await _on_api_start(api_kwargs)
+
+        response = await litellm.acompletion(**api_kwargs)
+
+        if _on_api_end:
+            await _on_api_end(api_kwargs, response)
+
+        # Extract usage
+        usage = {
+            **LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(  # type: ignore
+                response.usage
+            ).model_dump(),
+            "response_cost": response._hidden_params.get("response_cost", 0.0),
+        }
+        if _on_usage:
+            await _on_usage(usage)
+
+        # Parse response
+        resp_dict = response.model_dump()  # type: ignore
+        choice = (resp_dict.get("choices") or [{}])[0]
+        message = choice.get("message") or {}
+        content_text = message.get("content") or ""
+        tool_calls_array = message.get("tool_calls") or []
+        reasoning_text = message.get("reasoning") or ""
+
+        output_items: List[Dict[str, Any]] = []
+
+        # Add reasoning if present
+        if reasoning_text:
+            output_items.append(make_reasoning_item(reasoning_text))
+
+        if tool_calls_array:
+            for tc in tool_calls_array:
+                function = tc.get("function", {})
+                fn_name = function.get("name", "")
+                args_str = function.get("arguments", "{}")
+                tc_id = tc.get("id", "call_0")
+
+                try:
+                    args = json.loads(args_str) if isinstance(args_str, str) else args_str
+                except json.JSONDecodeError:
+                    args = {}
+
+                # Try converting to a computer action
+                computer_action = _convert_n1_action_to_computer_action(
+                    fn_name, args, screen_width, screen_height
+                )
+
+                if computer_action is not None:
+                    # Build a fake completion message for the converter
+                    fake_cm = {
+                        "role": "assistant",
+                        "content": content_text or "",
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "id": tc_id,
+                                "function": {
+                                    "name": "computer",
+                                    "arguments": json.dumps(computer_action),
+                                },
+                            }
+                        ],
+                    }
+                    output_items.extend(convert_completion_messages_to_responses_items([fake_cm]))
+                    # Only use content_text once
+                    content_text = ""
+                else:
+                    # Custom tool — emit as function_call
+                    output_items.append(make_function_call_item(fn_name, args, call_id=tc_id))
+        else:
+            # No tool calls — task is complete
+            if content_text:
+                output_items.append(make_output_text_item(content_text))
+            else:
+                output_items.append(make_output_text_item("Task completed."))
+
+        return {"output": (pre_output_items + output_items), "usage": usage}
+
+    async def predict_click(
+        self, model: str, image_b64: str, instruction: str, **kwargs
+    ) -> Optional[Tuple[int, int]]:
+        raise NotImplementedError(
+            "Yutori n1 does not support standalone click prediction. "
+            "Use predict_step for full browser automation."
+        )
+
+    def get_capabilities(self) -> List[AgentCapability]:
+        return ["step"]

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

@@ -0,0 +1,5 @@
+"""Playground server for Cua agents."""
+
+from .server import PlaygroundServer
+
+__all__ = ["PlaygroundServer"]

libs/python/agent/cua_agent/playground/server.py

@@ -0,0 +1,303 @@
+"""Playground server implementation for Cua agents."""
+
+import asyncio
+import logging
+import os
+import platform
+import socket
+import traceback
+import webbrowser
+from typing import Any, Dict, List, Optional, Union
+from urllib.parse import quote
+
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+logger = logging.getLogger(__name__)
+
+
+class PlaygroundServer:
+    """Playground server for running Cua agents via HTTP API."""
+
+    def __init__(self, agent_instance=None):
+        """
+        Initialize the playground server.
+
+        Args:
+            agent_instance: Optional pre-configured agent instance to use
+        """
+        self.agent_instance = agent_instance
+        self.app = FastAPI(
+            title="Cua Playground Server",
+            description="Playground server for Cua agents",
+            version="0.1.0",
+        )
+        self._setup_middleware()
+        self._setup_routes()
+        self.server = None
+        self.port = None
+
+    def _setup_middleware(self):
+        """Setup CORS middleware."""
+        self.app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+    def _setup_routes(self):
+        """Setup API routes."""
+
+        @self.app.get("/status")
+        async def status():
+            """Health check endpoint."""
+            sys = platform.system().lower()
+            if "darwin" in sys or sys in ("macos", "mac"):
+                os_type = "macos"
+            elif "windows" in sys:
+                os_type = "windows"
+            else:
+                os_type = "linux"
+
+            return {
+                "status": "ok",
+                "os_type": os_type,
+                "features": ["agent", "playground"],
+            }
+
+        @self.app.post("/responses")
+        async def responses_endpoint(request: Request):
+            """
+            Run ComputerAgent for up to 2 turns.
+
+            Body JSON:
+            {
+              "model": "...",                 # required
+              "input": "... or messages[]",   # required
+              "agent_kwargs": { ... },         # optional, passed directly to ComputerAgent
+              "env": { ... }                   # optional env overrides for agent
+            }
+            """
+            # Import here to avoid circular imports
+            try:
+                from cua_agent import ComputerAgent
+            except ImportError:
+                raise HTTPException(status_code=501, detail="ComputerAgent not available")
+
+            # Parse request body
+            try:
+                body = await request.json()
+            except Exception as e:
+                raise HTTPException(status_code=400, detail=f"Invalid JSON body: {str(e)}")
+
+            model = body.get("model")
+            input_data = body.get("input")
+            if not model or input_data is None:
+                raise HTTPException(status_code=400, detail="'model' and 'input' are required")
+
+            agent_kwargs: Dict[str, Any] = body.get("agent_kwargs") or {}
+            env_overrides: Dict[str, str] = body.get("env") or {}
+
+            # Simple env override context
+            class _EnvOverride:
+                def __init__(self, overrides: Dict[str, str]):
+                    self.overrides = overrides
+                    self._original: Dict[str, Optional[str]] = {}
+
+                def __enter__(self):
+                    for k, v in (self.overrides or {}).items():
+                        self._original[k] = os.environ.get(k)
+                        os.environ[k] = str(v)
+
+                def __exit__(self, exc_type, exc, tb):
+                    for k, old in self._original.items():
+                        if old is None:
+                            os.environ.pop(k, None)
+                        else:
+                            os.environ[k] = old
+
+            # Convert input to messages
+            def _to_messages(data: Union[str, List[Dict[str, Any]]]) -> List[Dict[str, Any]]:
+                if isinstance(data, str):
+                    return [{"role": "user", "content": data}]
+                if isinstance(data, list):
+                    return data
+                return []
+
+            messages = _to_messages(input_data)
+
+            error = None
+
+            with _EnvOverride(env_overrides):
+                # Use pre-configured agent only if model matches, otherwise create new one
+                # This ensures the model picker in the UI is respected
+                if self.agent_instance and self.agent_instance.model == model:
+                    agent = self.agent_instance
+                else:
+                    # Model changed or no pre-configured agent - create new agent with requested model
+                    agent = ComputerAgent(model=model, **agent_kwargs)  # type: ignore[arg-type]
+
+                total_output: List[Any] = []
+                total_usage: Dict[str, Any] = {}
+
+                pending_computer_call_ids = set()
+                try:
+                    async for result in agent.run(messages):
+                        total_output += result["output"]
+                        # Try to collect usage if present
+                        if (
+                            isinstance(result, dict)
+                            and "usage" in result
+                            and isinstance(result["usage"], dict)
+                        ):
+                            # Merge usage counters
+                            for k, v in result["usage"].items():
+                                if isinstance(v, (int, float)):
+                                    total_usage[k] = total_usage.get(k, 0) + v
+                                else:
+                                    total_usage[k] = v
+                        for msg in result.get("output", []):
+                            if msg.get("type") == "computer_call":
+                                pending_computer_call_ids.add(msg["call_id"])
+                            elif msg.get("type") == "computer_call_output":
+                                pending_computer_call_ids.discard(msg["call_id"])
+                            elif msg.get("type") == "function_call":
+                                pending_computer_call_ids.add(msg["call_id"])
+                            elif msg.get("type") == "function_call_output":
+                                pending_computer_call_ids.discard(msg["call_id"])
+                        # exit if no pending computer calls
+                        if not pending_computer_call_ids:
+                            break
+                except Exception as e:
+                    logger.error(f"Error running agent: {str(e)}")
+                    logger.error(traceback.format_exc())
+                    error = str(e)
+
+            # Build response payload
+            payload = {
+                "model": model,
+                "error": error,
+                "output": total_output,
+                "usage": total_usage,
+                "status": "completed" if not error else "failed",
+            }
+
+            # CORS: allow any origin
+            headers = {
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+            }
+
+            return JSONResponse(content=payload, headers=headers)
+
+    def _find_available_port(self, start_port: int = 8000, max_attempts: int = 100) -> int:
+        """Find an available port starting from start_port."""
+        for port in range(start_port, start_port + max_attempts):
+            try:
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                    s.bind(("127.0.0.1", port))
+                    return port
+            except OSError:
+                continue
+        raise RuntimeError(
+            f"Could not find an available port in range {start_port}-{start_port + max_attempts}"
+        )
+
+    async def start_async(self, port: Optional[int] = None, open_browser: bool = False):
+        """
+        Start the playground server asynchronously.
+
+        Args:
+            port: Port to run the server on. If None, finds an available port.
+            open_browser: Whether to open the browser automatically.
+        """
+        if port is None:
+            port = self._find_available_port()
+
+        self.port = port
+        host = f"http://localhost:{port}"
+
+        logger.info(f"Starting playground server on {host}")
+
+        if open_browser:
+            # Construct the playground URL
+            encoded_host = quote(host, safe="")
+            encoded_model = quote(self.agent_instance.model, safe="")
+            encoded_vnc_url = quote("http://localhost:8006/?autoconnect=true", safe="")
+
+            # Build URL with custom_model if agent instance is configured
+            playground_url = (
+                # f"http://cua.ai/dashboard/playground"
+                f"http://localhost:3000/dashboard/playground"
+                f"?host={encoded_host}"
+                f"&port={port}"
+                f"&id=localhost"
+                f"&name=localhost"
+                f"&custom_model={encoded_model}"
+                f"&custom_vnc_url={encoded_vnc_url}"
+                f"&vnc_password=null"
+                f"&resize=scale"
+                f"&fullscreen=true"
+            )
+
+            logger.info(f"Opening browser at: {playground_url}")
+            webbrowser.open(playground_url)
+
+        config = uvicorn.Config(
+            self.app,
+            host="0.0.0.0",
+            port=port,
+            log_level="info",
+        )
+        self.server = uvicorn.Server(config)
+        await self.server.serve()
+
+    def start(self, port: Optional[int] = None, open_browser: bool = False):
+        """
+        Start the playground server (blocking).
+
+        Args:
+            port: Port to run the server on. If None, finds an available port.
+            open_browser: Whether to open the browser automatically.
+        """
+        # Check if there's already a running event loop
+        try:
+            loop = asyncio.get_running_loop()
+            # If we're in an async context, schedule as a task
+            import threading
+
+            # Run the server in a separate thread to avoid blocking
+            server_thread = threading.Thread(
+                target=self._run_in_new_loop,
+                args=(port, open_browser),
+                daemon=True,
+            )
+            server_thread.start()
+
+            # Give the server a moment to start and open browser
+            import time
+
+            time.sleep(1)
+
+        except RuntimeError:
+            # No running loop, can use asyncio.run() safely
+            asyncio.run(self.start_async(port=port, open_browser=open_browser))
+
+    def _run_in_new_loop(self, port: Optional[int] = None, open_browser: bool = False):
+        """Helper to run server in a new event loop (for threading)."""
+        new_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(new_loop)
+        try:
+            new_loop.run_until_complete(self.start_async(port=port, open_browser=open_browser))
+        finally:
+            new_loop.close()
+
+    async def stop(self):
+        """Stop the playground server."""
+        if self.server:
+            logger.info("Stopping playground server")
+            await self.server.shutdown()

libs/python/agent/cua_agent/proxy/examples.py

@@ -0,0 +1,190 @@
+"""
+Example usage of the proxy server and client requests.
+"""
+
+import dotenv
+
+dotenv.load_dotenv()
+
+import asyncio
+import json
+import os
+from typing import Any, Dict
+
+import aiohttp
+
+
+async def test_http_endpoint():
+    """Test the HTTP /responses endpoint."""
+
+    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
+    assert isinstance(anthropic_api_key, str), "ANTHROPIC_API_KEY environment variable must be set"
+
+    # Example 1: Simple text request
+    simple_request = {
+        "model": "anthropic/claude-sonnet-4-5-20250929",
+        "input": "Tell me a three sentence bedtime story about a unicorn.",
+        "env": {"ANTHROPIC_API_KEY": anthropic_api_key},
+    }
+
+    # Example 2: Multi-modal request with image
+    multimodal_request = {
+        "model": "anthropic/claude-sonnet-4-5-20250929",
+        "input": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "input_text", "text": "what is in this image?"},
+                    {
+                        "type": "input_image",
+                        "image_url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg",
+                    },
+                ],
+            }
+        ],
+        "env": {"ANTHROPIC_API_KEY": anthropic_api_key},
+    }
+
+    # Example 3: Request with custom agent and computer kwargs
+    custom_request = {
+        "model": "anthropic/claude-sonnet-4-5-20250929",
+        "input": "Take a screenshot and tell me what you see",
+        "env": {"ANTHROPIC_API_KEY": anthropic_api_key},
+    }
+
+    # Test requests
+    base_url = "https://m-linux-96lcxd2c2k.containers.cloud.trycua.com:8443"
+    # base_url = "http://localhost:8000"
+    api_key = os.getenv("CUA_API_KEY")
+    assert isinstance(api_key, str), "CUA_API_KEY environment variable must be set"
+
+    async with aiohttp.ClientSession() as session:
+        for i, request_data in enumerate(
+            [
+                simple_request,
+                # multimodal_request,
+                custom_request,
+            ],
+            1,
+        ):
+            print(f"\n--- Test {i} ---")
+            print(f"Request: {json.dumps(request_data, indent=2)}")
+
+            try:
+                print(f"Sending request to {base_url}/responses")
+                async with session.post(
+                    f"{base_url}/responses",
+                    json=request_data,
+                    headers={"Content-Type": "application/json", "X-API-Key": api_key},
+                ) as response:
+                    result = await response.json()
+                    print(f"Status: {response.status}")
+                    print(f"Response: {json.dumps(result, indent=2)}")
+
+            except Exception as e:
+                print(f"Error: {e}")
+
+
+def curl_examples():
+    """Print curl command examples."""
+
+    print("=== CURL Examples ===\n")
+
+    print("1. Simple text request:")
+    print("""curl http://localhost:8000/responses \\
+  -H "Content-Type: application/json" \\
+  -d '{
+    "model": "anthropic/claude-sonnet-4-5-20250929",
+    "input": "Tell me a three sentence bedtime story about a unicorn."
+  }'""")
+
+    print("\n2. Multi-modal request with image:")
+    print("""curl http://localhost:8000/responses \\
+  -H "Content-Type: application/json" \\
+  -d '{
+    "model": "anthropic/claude-sonnet-4-5-20250929",
+    "input": [
+      {
+        "role": "user",
+        "content": [
+          {"type": "input_text", "text": "what is in this image?"},
+          {
+            "type": "input_image",
+            "image_url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+          }
+        ]
+      }
+    ]
+  }'""")
+
+    print("\n3. Request with custom configuration:")
+    print("""curl http://localhost:8000/responses \\
+  -H "Content-Type: application/json" \\
+  -d '{
+    "model": "anthropic/claude-sonnet-4-5-20250929",
+    "input": "Take a screenshot and tell me what you see",
+    "agent_kwargs": {
+      "save_trajectory": true,
+      "verbosity": 20
+    },
+    "computer_kwargs": {
+      "os_type": "linux",
+      "provider_type": "cloud"
+    }
+  }'""")
+
+
+async def test_p2p_client():
+    """Example P2P client using peerjs-python."""
+    try:
+        from aiortc import RTCConfiguration, RTCIceServer
+        from peerjs import ConnectionEventType, Peer, PeerOptions
+
+        # Set up client peer
+        options = PeerOptions(
+            host="0.peerjs.com",
+            port=443,
+            secure=True,
+            config=RTCConfiguration(iceServers=[RTCIceServer(urls="stun:stun.l.google.com:19302")]),
+        )
+
+        client_peer = Peer(id="test-client", peer_options=options)
+        await client_peer.start()
+
+        # Connect to proxy server
+        connection = client_peer.connect("computer-agent-proxy")
+
+        @connection.on(ConnectionEventType.Open)
+        async def connection_open():
+            print("Connected to proxy server")
+
+            # Send a test request
+            request = {
+                "model": "anthropic/claude-sonnet-4-5-20250929",
+                "input": "Hello from P2P client!",
+            }
+            await connection.send(json.dumps(request))
+
+        @connection.on(ConnectionEventType.Data)
+        async def connection_data(data):
+            print(f"Received response: {data}")
+            await client_peer.destroy()
+
+        # Wait for connection
+        await asyncio.sleep(10)
+
+    except ImportError:
+        print("P2P dependencies not available. Install peerjs-python for P2P testing.")
+    except Exception as e:
+        print(f"P2P test error: {e}")
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) > 1 and sys.argv[1] == "curl":
+        curl_examples()
+    elif len(sys.argv) > 1 and sys.argv[1] == "p2p":
+        asyncio.run(test_p2p_client())
+    else:
+        asyncio.run(test_http_endpoint())

libs/python/agent/cua_agent/proxy/handlers.py

@@ -0,0 +1,247 @@
+"""
+Request handlers for the proxy endpoints.
+"""
+
+import json
+import logging
+import os
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional, Union
+
+try:
+    from computer import Computer
+except ImportError:
+    Computer = None  # type: ignore[assignment,misc]
+
+from ..agent import ComputerAgent
+
+logger = logging.getLogger(__name__)
+
+
+class ResponsesHandler:
+    """Handler for /responses endpoint that processes agent requests."""
+
+    def __init__(self):
+        self.computer = None
+        self.agent = None
+        # Simple in-memory caches
+        self._computer_cache: Dict[str, Any] = {}
+        self._agent_cache: Dict[str, Any] = {}
+
+    async def setup_computer_agent(
+        self,
+        model: str,
+        agent_kwargs: Optional[Dict[str, Any]] = None,
+        computer_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """Set up (and cache) computer and agent instances.
+
+        Caching keys:
+        - Computer cache key: computer_kwargs
+        - Agent cache key: {"model": model, **agent_kwargs}
+        """
+        agent_kwargs = agent_kwargs or {}
+        computer_kwargs = computer_kwargs or {}
+
+        def _stable_key(obj: Dict[str, Any]) -> str:
+            try:
+                return json.dumps(obj, sort_keys=True, separators=(",", ":"))
+            except Exception:
+                # Fallback: stringify non-serializable values
+                safe_obj = {}
+                for k, v in obj.items():
+                    try:
+                        json.dumps(v)
+                        safe_obj[k] = v
+                    except Exception:
+                        safe_obj[k] = str(v)
+                return json.dumps(safe_obj, sort_keys=True, separators=(",", ":"))
+
+        # Determine if custom tools are supplied; if so, skip computer setup entirely
+        has_custom_tools = bool(agent_kwargs.get("tools"))
+
+        computer = None
+        if not has_custom_tools:
+            # ---------- Computer setup (with cache) ----------
+            comp_key = _stable_key(computer_kwargs)
+
+            computer = self._computer_cache.get(comp_key)
+            if computer is None:
+                # Default computer configuration
+                default_c_config = {
+                    "os_type": "linux",
+                    "provider_type": "cloud",
+                    "name": os.getenv("CUA_CONTAINER_NAME"),
+                    "api_key": os.getenv("CUA_API_KEY"),
+                }
+                default_c_config.update(computer_kwargs)
+                computer = Computer(**default_c_config)
+                await computer.__aenter__()
+                self._computer_cache[comp_key] = computer
+                logger.info(
+                    f"Computer created and cached with key={comp_key} config={default_c_config}"
+                )
+            else:
+                logger.info(f"Reusing cached computer for key={comp_key}")
+
+        # Bind current computer reference (None if custom tools supplied)
+        self.computer = computer
+
+        # ---------- Agent setup (with cache) ----------
+        # Build agent cache key from {model} + agent_kwargs (excluding tools unless explicitly passed)
+        agent_kwargs_for_key = dict(agent_kwargs)
+        agent_key_payload = {"model": model, **agent_kwargs_for_key}
+        agent_key = _stable_key(agent_key_payload)
+
+        # Pass Computer as the tool - ComputerAgent handles wrapping internally
+        # based on the model's required tool_type (e.g., FARA auto-wraps to BrowserTool)
+        tool = computer
+
+        agent = self._agent_cache.get(agent_key)
+        if agent is None:
+            # Default agent configuration
+            default_a_config: Dict[str, Any] = {"model": model}
+            if not has_custom_tools:
+                default_a_config["tools"] = [tool]
+            # Apply user overrides, but keep tools unless user explicitly sets
+            if agent_kwargs:
+                if not has_custom_tools:
+                    agent_kwargs.setdefault("tools", [tool])
+                default_a_config.update(agent_kwargs)
+            # JSON-derived kwargs may have loose types; ignore static arg typing here
+            agent = ComputerAgent(**default_a_config)  # type: ignore[arg-type]
+            self._agent_cache[agent_key] = agent
+            logger.info(f"Agent created and cached with key={agent_key} model={model}")
+        else:
+            # Ensure cached agent uses the current tool (in case object differs)
+            # Only update if tools not explicitly provided in agent_kwargs
+            if not has_custom_tools:
+                try:
+                    agent.tools = [tool]
+                except Exception:
+                    pass
+            logger.info(f"Reusing cached agent for key={agent_key}")
+
+        # Bind current agent reference
+        self.agent = agent
+
+    async def process_request(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Process a /responses request and return the result.
+
+        Args:
+            request_data: Dictionary containing model, input, and optional kwargs
+
+        Returns:
+            Dictionary with the agent's response
+        """
+        try:
+            # Extract request parameters
+            model = request_data.get("model")
+            input_data = request_data.get("input")
+            agent_kwargs = request_data.get("agent_kwargs", {})
+            computer_kwargs = request_data.get("computer_kwargs", {})
+            env_overrides = request_data.get("env", {}) or {}
+
+            if not model:
+                raise ValueError("Model is required")
+            if not input_data:
+                raise ValueError("Input is required")
+
+            # Apply env overrides for the duration of this request
+            with self._env_overrides(env_overrides):
+                # Set up (and possibly reuse) computer and agent via caches
+                await self.setup_computer_agent(model, agent_kwargs, computer_kwargs)
+
+                # Defensive: ensure agent is initialized for type checkers
+                agent = self.agent
+                if agent is None:
+                    raise RuntimeError("Agent failed to initialize")
+
+                # Convert input to messages format
+                messages = self._convert_input_to_messages(input_data)
+
+                # Run agent and get first result
+                async for result in agent.run(messages):
+                    # Return the first result and break
+                    return {"success": True, "result": result, "model": model}
+
+            # If no results were yielded
+            return {"success": False, "error": "No results from agent", "model": model}
+
+        except Exception as e:
+            logger.error(f"Error processing request: {e}")
+            return {
+                "success": False,
+                "error": str(e),
+                "model": request_data.get("model", "unknown"),
+            }
+
+    def _convert_input_to_messages(
+        self, input_data: Union[str, List[Dict[str, Any]]]
+    ) -> List[Dict[str, Any]]:
+        """Convert input data to messages format."""
+        if isinstance(input_data, str):
+            # Simple string input
+            return [{"role": "user", "content": input_data}]
+        elif isinstance(input_data, list):
+            # Already in messages format
+            messages = []
+            for msg in input_data:
+                # Convert content array format if needed
+                if isinstance(msg.get("content"), list):
+                    content_parts = []
+                    for part in msg["content"]:
+                        if part.get("type") == "input_text":
+                            content_parts.append({"type": "text", "text": part["text"]})
+                        elif part.get("type") == "input_image":
+                            content_parts.append(
+                                {"type": "image_url", "image_url": {"url": part["image_url"]}}
+                            )
+                        else:
+                            content_parts.append(part)
+                    messages.append({"role": msg["role"], "content": content_parts})
+                else:
+                    messages.append(msg)
+            return messages
+        else:
+            raise ValueError("Input must be string or list of messages")
+
+    async def cleanup(self):
+        """Clean up resources."""
+        if self.computer:
+            try:
+                await self.computer.__aexit__(None, None, None)
+            except Exception as e:
+                logger.error(f"Error cleaning up computer: {e}")
+            finally:
+                self.computer = None
+        self.agent = None
+
+    @staticmethod
+    @contextmanager
+    def _env_overrides(env: Dict[str, str]):
+        """Temporarily apply environment variable overrides for the current process.
+        Restores previous values after the context exits.
+
+        Args:
+            env: Mapping of env var names to override for this request.
+        """
+        if not env:
+            # No-op context
+            yield
+            return
+
+        original: Dict[str, Optional[str]] = {}
+        try:
+            for k, v in env.items():
+                original[k] = os.environ.get(k)
+                os.environ[k] = str(v)
+            yield
+        finally:
+            for k, old in original.items():
+                if old is None:
+                    # Was not set before
+                    os.environ.pop(k, None)
+                else:
+                    os.environ[k] = old

libs/python/agent/cua_agent/responses.py

@@ -0,0 +1,875 @@
+"""
+Functions for making various Responses API items from different types of responses.
+Based on the OpenAI spec for Responses API items.
+"""
+
+import base64
+import json
+import uuid
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from openai.types.responses.easy_input_message_param import EasyInputMessageParam
+from openai.types.responses.response_computer_tool_call_param import (
+    ActionClick,
+    ActionDoubleClick,
+    ActionDrag,
+    ActionDragPath,
+    ActionKeypress,
+    ActionMove,
+    ActionScreenshot,
+    ActionScroll,
+)
+from openai.types.responses.response_computer_tool_call_param import (
+    ActionType as ActionTypeAction,
+)
+from openai.types.responses.response_computer_tool_call_param import (
+    ActionWait,
+    PendingSafetyCheck,
+    ResponseComputerToolCallParam,
+)
+from openai.types.responses.response_function_tool_call_param import (
+    ResponseFunctionToolCallParam,
+)
+from openai.types.responses.response_input_image_param import ResponseInputImageParam
+from openai.types.responses.response_output_message_param import (
+    ResponseOutputMessageParam,
+)
+from openai.types.responses.response_output_text_param import ResponseOutputTextParam
+from openai.types.responses.response_reasoning_item_param import (
+    ResponseReasoningItemParam,
+    Summary,
+)
+
+
+def random_id():
+    return str(uuid.uuid4())
+
+
+# User message items
+def make_input_image_item(image_data: Union[str, bytes]) -> EasyInputMessageParam:
+    return EasyInputMessageParam(
+        content=[
+            ResponseInputImageParam(
+                type="input_image",
+                image_url=f"data:image/png;base64,{base64.b64encode(image_data).decode('utf-8') if isinstance(image_data, bytes) else image_data}",
+            )  # type: ignore
+        ],
+        role="user",
+        type="message",
+    )
+
+
+# Text items
+def make_reasoning_item(reasoning: str) -> ResponseReasoningItemParam:
+    return ResponseReasoningItemParam(
+        id=random_id(), summary=[Summary(text=reasoning, type="summary_text")], type="reasoning"
+    )
+
+
+def make_output_text_item(content: str) -> ResponseOutputMessageParam:
+    return ResponseOutputMessageParam(
+        id=random_id(),
+        content=[ResponseOutputTextParam(text=content, type="output_text", annotations=[])],
+        role="assistant",
+        status="completed",
+        type="message",
+    )
+
+
+# Function call items
+def make_function_call_item(
+    function_name: str, arguments: Dict[str, Any], call_id: Optional[str] = None
+) -> ResponseFunctionToolCallParam:
+    return ResponseFunctionToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        name=function_name,
+        arguments=json.dumps(arguments),
+        status="completed",
+        type="function_call",
+    )
+
+
+# Computer tool call items
+def make_click_item(
+    x: int,
+    y: int,
+    button: Literal["left", "right", "wheel", "back", "forward"] = "left",
+    call_id: Optional[str] = None,
+) -> ResponseComputerToolCallParam:
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionClick(button=button, type="click", x=x, y=y),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+def make_double_click_item(
+    x: int, y: int, call_id: Optional[str] = None
+) -> ResponseComputerToolCallParam:
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionDoubleClick(type="double_click", x=x, y=y),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+def make_drag_item(
+    path: List[Dict[str, int]], call_id: Optional[str] = None
+) -> ResponseComputerToolCallParam:
+    drag_path = [ActionDragPath(x=point["x"], y=point["y"]) for point in path]
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionDrag(path=drag_path, type="drag"),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+def make_keypress_item(
+    keys: List[str], call_id: Optional[str] = None
+) -> ResponseComputerToolCallParam:
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionKeypress(keys=keys, type="keypress"),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+def make_move_item(x: int, y: int, call_id: Optional[str] = None) -> ResponseComputerToolCallParam:
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionMove(type="move", x=x, y=y),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+def make_screenshot_item(call_id: Optional[str] = None) -> ResponseComputerToolCallParam:
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionScreenshot(type="screenshot"),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+def make_scroll_item(
+    x: int, y: int, scroll_x: int, scroll_y: int, call_id: Optional[str] = None
+) -> ResponseComputerToolCallParam:
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionScroll(scroll_x=scroll_x, scroll_y=scroll_y, type="scroll", x=x, y=y),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+def make_type_item(text: str, call_id: Optional[str] = None) -> ResponseComputerToolCallParam:
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionTypeAction(text=text, type="type"),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+def make_wait_item(call_id: Optional[str] = None) -> ResponseComputerToolCallParam:
+    return ResponseComputerToolCallParam(
+        id=random_id(),
+        call_id=call_id if call_id else random_id(),
+        action=ActionWait(type="wait"),
+        pending_safety_checks=[],
+        status="completed",
+        type="computer_call",
+    )
+
+
+# Extra anthropic computer calls
+def make_left_mouse_down_item(
+    x: Optional[int] = None, y: Optional[int] = None, call_id: Optional[str] = None
+) -> Dict[str, Any]:
+    return {
+        "id": random_id(),
+        "call_id": call_id if call_id else random_id(),
+        "action": {"type": "left_mouse_down", "x": x, "y": y},
+        "pending_safety_checks": [],
+        "status": "completed",
+        "type": "computer_call",
+    }
+
+
+def make_left_mouse_up_item(
+    x: Optional[int] = None, y: Optional[int] = None, call_id: Optional[str] = None
+) -> Dict[str, Any]:
+    return {
+        "id": random_id(),
+        "call_id": call_id if call_id else random_id(),
+        "action": {"type": "left_mouse_up", "x": x, "y": y},
+        "pending_safety_checks": [],
+        "status": "completed",
+        "type": "computer_call",
+    }
+
+
+def make_failed_tool_call_items(
+    tool_name: str, tool_kwargs: Dict[str, Any], error_message: str, call_id: Optional[str] = None
+) -> List[Dict[str, Any]]:
+    call_id = call_id if call_id else random_id()
+    return [
+        {
+            "type": "function_call",
+            "id": random_id(),
+            "call_id": call_id,
+            "name": tool_name,
+            "arguments": json.dumps(tool_kwargs),
+        },
+        {
+            "type": "function_call_output",
+            "call_id": call_id,
+            "output": json.dumps({"error": error_message}),
+        },
+    ]
+
+
+def make_tool_error_item(error_message: str, call_id: Optional[str] = None) -> Dict[str, Any]:
+    call_id = call_id if call_id else random_id()
+    return {
+        "type": "function_call_output",
+        "call_id": call_id,
+        "output": json.dumps({"error": error_message}),
+    }
+
+
+def replace_failed_computer_calls_with_function_calls(
+    messages: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    """
+    Replace computer_call items with function_call items if they share a call_id with a function_call_output.
+    This indicates the computer call failed and should be treated as a function call instead.
+    We do this because the computer_call_output items do not support text output.
+
+    Args:
+        messages: List of message items to process
+    """
+    messages = messages.copy()
+
+    # Find all call_ids that have function_call_output items
+    failed_call_ids = set()
+    for msg in messages:
+        if msg.get("type") == "function_call_output":
+            call_id = msg.get("call_id")
+            if call_id:
+                failed_call_ids.add(call_id)
+
+    # Replace computer_call items that have matching call_ids
+    for i, msg in enumerate(messages):
+        if msg.get("type") == "computer_call" and msg.get("call_id") in failed_call_ids:
+
+            # Extract action from computer_call
+            action = msg.get("action", {})
+            call_id = msg.get("call_id")
+
+            # Create function_call replacement
+            messages[i] = {
+                "type": "function_call",
+                "id": msg.get("id", random_id()),
+                "call_id": call_id,
+                "name": "computer",
+                "arguments": json.dumps(action),
+            }
+
+    return messages
+
+
+# Conversion functions between element descriptions and coordinates
+def convert_computer_calls_desc2xy(
+    responses_items: List[Dict[str, Any]], desc2xy: Dict[str, tuple]
+) -> List[Dict[str, Any]]:
+    """
+    Convert computer calls from element descriptions to x,y coordinates.
+
+    Args:
+        responses_items: List of response items containing computer calls with element_description
+        desc2xy: Dictionary mapping element descriptions to (x, y) coordinate tuples
+
+    Returns:
+        List of response items with element_description replaced by x,y coordinates
+    """
+    converted_items = []
+
+    for item in responses_items:
+        if item.get("type") == "computer_call" and "action" in item:
+            action = item["action"].copy()
+
+            # Handle single element_description
+            if "element_description" in action:
+                desc = action["element_description"]
+                if desc in desc2xy:
+                    x, y = desc2xy[desc]
+                    action["x"] = x
+                    action["y"] = y
+                    del action["element_description"]
+
+            # Handle start_element_description and end_element_description for drag operations
+            elif "start_element_description" in action and "end_element_description" in action:
+                start_desc = action["start_element_description"]
+                end_desc = action["end_element_description"]
+
+                if start_desc in desc2xy and end_desc in desc2xy:
+                    start_x, start_y = desc2xy[start_desc]
+                    end_x, end_y = desc2xy[end_desc]
+                    action["path"] = [{"x": start_x, "y": start_y}, {"x": end_x, "y": end_y}]
+                    del action["start_element_description"]
+                    del action["end_element_description"]
+
+            converted_item = item.copy()
+            converted_item["action"] = action
+            converted_items.append(converted_item)
+        else:
+            converted_items.append(item)
+
+    return converted_items
+
+
+def convert_computer_calls_xy2desc(
+    responses_items: List[Dict[str, Any]], desc2xy: Dict[str, tuple]
+) -> List[Dict[str, Any]]:
+    """
+    Convert computer calls from x,y coordinates to element descriptions.
+
+    Args:
+        responses_items: List of response items containing computer calls with x,y coordinates
+        desc2xy: Dictionary mapping element descriptions to (x, y) coordinate tuples
+
+    Returns:
+        List of response items with x,y coordinates replaced by element_description
+    """
+    # Create reverse mapping from coordinates to descriptions
+    xy2desc = {coords: desc for desc, coords in desc2xy.items()}
+
+    converted_items = []
+
+    for item in responses_items:
+        if item.get("type") == "computer_call" and "action" in item:
+            action = item["action"].copy()
+
+            # Handle single x,y coordinates
+            if "x" in action and "y" in action:
+                coords = (action["x"], action["y"])
+                if coords in xy2desc:
+                    action["element_description"] = xy2desc[coords]
+                    del action["x"]
+                    del action["y"]
+
+            # Handle path for drag operations
+            elif "path" in action and isinstance(action["path"], list) and len(action["path"]) == 2:
+                start_point = action["path"][0]
+                end_point = action["path"][1]
+
+                if (
+                    "x" in start_point
+                    and "y" in start_point
+                    and "x" in end_point
+                    and "y" in end_point
+                ):
+
+                    start_coords = (start_point["x"], start_point["y"])
+                    end_coords = (end_point["x"], end_point["y"])
+
+                    if start_coords in xy2desc and end_coords in xy2desc:
+                        action["start_element_description"] = xy2desc[start_coords]
+                        action["end_element_description"] = xy2desc[end_coords]
+                        del action["path"]
+
+            converted_item = item.copy()
+            converted_item["action"] = action
+            converted_items.append(converted_item)
+        else:
+            converted_items.append(item)
+
+    return converted_items
+
+
+def get_all_element_descriptions(responses_items: List[Dict[str, Any]]) -> List[str]:
+    """
+    Extract all element descriptions from computer calls in responses items.
+
+    Args:
+        responses_items: List of response items containing computer calls
+
+    Returns:
+        List of unique element descriptions found in computer calls
+    """
+    descriptions = set()
+
+    for item in responses_items:
+        if item.get("type") == "computer_call" and "action" in item:
+            action = item["action"]
+
+            # Handle single element_description
+            if "element_description" in action:
+                descriptions.add(action["element_description"])
+
+            # Handle start_element_description and end_element_description for drag operations
+            if "start_element_description" in action:
+                descriptions.add(action["start_element_description"])
+
+            if "end_element_description" in action:
+                descriptions.add(action["end_element_description"])
+
+    return list(descriptions)
+
+
+# Conversion functions between responses_items and completion messages formats
+def convert_responses_items_to_completion_messages(
+    messages: List[Dict[str, Any]],
+    allow_images_in_tool_results: bool = True,
+    send_multiple_user_images_per_parallel_tool_results: bool = False,
+    use_xml_tools: bool = False,
+) -> List[Dict[str, Any]]:
+    """Convert responses_items message format to liteLLM completion format.
+
+    Args:
+        messages: List of responses_items format messages
+        allow_images_in_tool_results: If True, include images in tool role messages.
+                                    If False, send tool message + separate user message with image.
+        send_multiple_user_images_per_parallel_tool_results: If True, send multiple user images in parallel tool results.
+        use_xml_tools: If True, use XML-style <tool_call> tags instead of tool_calls array.
+                      Also sends tool results as user messages instead of tool role.
+    """
+    # Assert that allow_images_in_tool_results is False when use_xml_tools is True
+    if use_xml_tools:
+        assert (
+            not allow_images_in_tool_results
+        ), "allow_images_in_tool_results must be False when use_xml_tools is True"
+    completion_messages = []
+
+    for i, message in enumerate(messages):
+        msg_type = message.get("type")
+        role = message.get("role")
+
+        # Handle user messages (both with and without explicit type)
+        if role == "user" or msg_type == "user":
+            content = message.get("content", "")
+            if isinstance(content, list):
+                # Handle list content (images, text blocks)
+                completion_content = []
+                for item in content:
+                    if item.get("type") == "input_image":
+                        completion_content.append(
+                            {"type": "image_url", "image_url": {"url": item.get("image_url")}}
+                        )
+                    elif item.get("type") == "input_text":
+                        completion_content.append({"type": "text", "text": item.get("text")})
+                    elif item.get("type") == "text":
+                        completion_content.append({"type": "text", "text": item.get("text")})
+
+                completion_messages.append({"role": "user", "content": completion_content})
+            elif isinstance(content, str):
+                # Handle string content
+                completion_messages.append({"role": "user", "content": content})
+
+        # Handle assistant messages
+        elif role == "assistant" or msg_type == "message":
+            content = message.get("content", [])
+            if isinstance(content, list):
+                text_parts = []
+                for item in content:
+                    if item.get("type") == "output_text":
+                        text_parts.append(item.get("text", ""))
+                    elif item.get("type") == "text":
+                        text_parts.append(item.get("text", ""))
+
+                if text_parts:
+                    completion_messages.append(
+                        {"role": "assistant", "content": "\n".join(text_parts)}
+                    )
+
+        # Handle reasoning items (convert to assistant message)
+        elif msg_type == "reasoning":
+            summary = message.get("summary", [])
+            text_parts = []
+            for item in summary:
+                if item.get("type") == "summary_text":
+                    text_parts.append(item.get("text", ""))
+
+            if text_parts:
+                completion_messages.append({"role": "assistant", "content": "\n".join(text_parts)})
+
+        # Handle function calls
+        elif msg_type == "function_call":
+            if use_xml_tools:
+                # Use XML format instead of tool_calls array
+                if not completion_messages or completion_messages[-1]["role"] != "assistant":
+                    completion_messages.append({"role": "assistant", "content": ""})
+
+                # Ensure arguments is a JSON string (not a dict)
+                arguments = message.get("arguments")
+                if isinstance(arguments, dict):
+                    arguments = json.dumps(arguments)
+
+                # Format as XML tool call
+                tool_call_xml = f'<tool_call>{{"name": "{message.get("name")}", "arguments": {arguments}}}</tool_call>'
+                if completion_messages[-1]["content"]:
+                    completion_messages[-1]["content"] += "\n" + tool_call_xml
+                else:
+                    completion_messages[-1]["content"] = tool_call_xml
+            else:
+                # Add tool call to last assistant message or create new one
+                if not completion_messages or completion_messages[-1]["role"] != "assistant":
+                    completion_messages.append(
+                        {"role": "assistant", "content": "", "tool_calls": []}
+                    )
+
+                if "tool_calls" not in completion_messages[-1]:
+                    completion_messages[-1]["tool_calls"] = []
+
+                # Ensure arguments is a JSON string (not a dict)
+                arguments = message.get("arguments")
+                if isinstance(arguments, dict):
+                    arguments = json.dumps(arguments)
+
+                completion_messages[-1]["tool_calls"].append(
+                    {
+                        "id": message.get("call_id"),
+                        "type": "function",
+                        "function": {
+                            "name": message.get("name"),
+                            "arguments": arguments,
+                        },
+                    }
+                )
+
+        # Handle computer calls
+        elif msg_type == "computer_call":
+            if use_xml_tools:
+                # Use XML format instead of tool_calls array
+                if not completion_messages or completion_messages[-1]["role"] != "assistant":
+                    completion_messages.append({"role": "assistant", "content": ""})
+
+                action = message.get("action", {})
+                # Format as XML tool call
+                tool_call_xml = f'<tool_call>{{"name": "computer", "arguments": {json.dumps(action)}}}</tool_call>'
+                if completion_messages[-1]["content"]:
+                    completion_messages[-1]["content"] += "\n" + tool_call_xml
+                else:
+                    completion_messages[-1]["content"] = tool_call_xml
+            else:
+                # Add tool call to last assistant message or create new one
+                if not completion_messages or completion_messages[-1]["role"] != "assistant":
+                    completion_messages.append(
+                        {"role": "assistant", "content": "", "tool_calls": []}
+                    )
+
+                if "tool_calls" not in completion_messages[-1]:
+                    completion_messages[-1]["tool_calls"] = []
+
+                action = message.get("action", {})
+                completion_messages[-1]["tool_calls"].append(
+                    {
+                        "id": message.get("call_id"),
+                        "type": "function",
+                        "function": {"name": "computer", "arguments": json.dumps(action)},
+                    }
+                )
+
+        # Handle function/computer call outputs
+        elif msg_type in ["function_call_output", "computer_call_output"]:
+            output = message.get("output")
+            call_id = message.get("call_id")
+
+            if use_xml_tools:
+                # When using XML tools, send all results as user messages
+                if isinstance(output, dict) and output.get("type") == "input_image":
+                    # Send image as user message
+                    completion_messages.append(
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "image_url",
+                                    "image_url": {"url": output.get("image_url")},
+                                }
+                            ],
+                        }
+                    )
+                else:
+                    # Send text result as user message
+                    completion_messages.append(
+                        {
+                            "role": "user",
+                            "content": str(output),
+                        }
+                    )
+            else:
+                # Standard tool message handling
+                if isinstance(output, dict) and output.get("type") == "input_image":
+                    if allow_images_in_tool_results:
+                        # Handle image output as tool response (may not work with all APIs)
+                        completion_messages.append(
+                            {
+                                "role": "tool",
+                                "tool_call_id": call_id,
+                                "content": [
+                                    {
+                                        "type": "image_url",
+                                        "image_url": {"url": output.get("image_url")},
+                                    }
+                                ],
+                            }
+                        )
+                    else:
+                        # Determine if the next message is also a tool call output
+                        next_type = None
+                        if i + 1 < len(messages):
+                            next_msg = messages[i + 1]
+                            next_type = next_msg.get("type")
+                        is_next_message_image_result = next_type in [
+                            "computer_call_output",
+                        ]
+                        # Send tool message + separate user message with image (OpenAI compatible)
+                        completion_messages += (
+                            [
+                                {
+                                    "role": "tool",
+                                    "tool_call_id": call_id,
+                                    "content": "[Execution completed. See screenshot below]",
+                                },
+                                {
+                                    "role": "user",
+                                    "content": [
+                                        {
+                                            "type": "image_url",
+                                            "image_url": {"url": output.get("image_url")},
+                                        }
+                                    ],
+                                },
+                            ]
+                            if send_multiple_user_images_per_parallel_tool_results
+                            or (not is_next_message_image_result)
+                            else [
+                                {
+                                    "role": "tool",
+                                    "tool_call_id": call_id,
+                                    "content": "[Execution completed. See screenshot below]",
+                                },
+                            ]
+                        )
+                else:
+                    # Handle text output as tool response
+                    completion_messages.append(
+                        {"role": "tool", "tool_call_id": call_id, "content": str(output)}
+                    )
+
+    return completion_messages
+
+
+def convert_completion_messages_to_responses_items(
+    completion_messages: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    """Convert completion messages format to responses_items message format."""
+    responses_items = []
+    skip_next = False
+
+    for i, message in enumerate(completion_messages):
+        if skip_next:
+            skip_next = False
+            continue
+
+        role = message.get("role")
+        content = message.get("content")
+        tool_calls = message.get("tool_calls", [])
+
+        # Handle assistant messages with text content
+        if role == "assistant" and content and isinstance(content, str):
+            responses_items.append(
+                {
+                    "type": "message",
+                    "role": "assistant",
+                    "content": [{"type": "output_text", "text": content}],
+                }
+            )
+
+        # Handle tool calls
+        if tool_calls:
+            for tool_call in tool_calls:
+                if tool_call.get("type") == "function":
+                    function = tool_call.get("function", {})
+                    function_name = function.get("name")
+
+                    if function_name in ("computer", "computer_use"):
+                        # Parse computer action
+                        try:
+                            action = json.loads(function.get("arguments", "{}"))
+                            # Change key from "action" -> "type"
+                            if action.get("action"):
+                                action["type"] = action["action"]
+                                del action["action"]
+                            responses_items.append(
+                                {
+                                    "type": "computer_call",
+                                    "call_id": tool_call.get("id"),
+                                    "action": action,
+                                    "status": "completed",
+                                }
+                            )
+                        except json.JSONDecodeError:
+                            # Fallback to function call format
+                            responses_items.append(
+                                {
+                                    "type": "function_call",
+                                    "call_id": tool_call.get("id"),
+                                    "name": function_name,
+                                    "arguments": function.get("arguments", "{}"),
+                                    "status": "completed",
+                                }
+                            )
+                    else:
+                        # Regular function call
+                        responses_items.append(
+                            {
+                                "type": "function_call",
+                                "call_id": tool_call.get("id"),
+                                "name": function_name,
+                                "arguments": function.get("arguments", "{}"),
+                                "status": "completed",
+                            }
+                        )
+
+        # Handle tool messages (function/computer call outputs)
+        elif role == "tool" and content:
+            tool_call_id = message.get("tool_call_id")
+            if isinstance(content, str):
+                # Check if this is the "[Execution completed. See screenshot below]" pattern
+                if content == "[Execution completed. See screenshot below]":
+                    # Look ahead for the next user message with image
+                    next_idx = i + 1
+                    if (
+                        next_idx < len(completion_messages)
+                        and completion_messages[next_idx].get("role") == "user"
+                        and isinstance(completion_messages[next_idx].get("content"), list)
+                    ):
+                        # Found the pattern - extract image from next message
+                        next_content = completion_messages[next_idx]["content"]
+                        for item in next_content:
+                            if item.get("type") == "image_url":
+                                responses_items.append(
+                                    {
+                                        "type": "computer_call_output",
+                                        "call_id": tool_call_id,
+                                        "output": {
+                                            "type": "input_image",
+                                            "image_url": item.get("image_url", {}).get("url"),
+                                        },
+                                    }
+                                )
+                                # Skip the next user message since we processed it
+                                skip_next = True
+                                break
+                    else:
+                        # No matching user message, treat as regular text
+                        responses_items.append(
+                            {
+                                "type": "computer_call_output",
+                                "call_id": tool_call_id,
+                                "output": content,
+                            }
+                        )
+                else:
+                    # Determine if this is a computer call or function call output
+                    try:
+                        # Try to parse as structured output
+                        parsed_content = json.loads(content)
+                        if parsed_content.get("type") == "input_image":
+                            responses_items.append(
+                                {
+                                    "type": "computer_call_output",
+                                    "call_id": tool_call_id,
+                                    "output": parsed_content,
+                                }
+                            )
+                        else:
+                            responses_items.append(
+                                {
+                                    "type": "computer_call_output",
+                                    "call_id": tool_call_id,
+                                    "output": content,
+                                }
+                            )
+                    except json.JSONDecodeError:
+                        # Plain text output - could be function or computer call
+                        responses_items.append(
+                            {
+                                "type": "function_call_output",
+                                "call_id": tool_call_id,
+                                "output": content,
+                            }
+                        )
+            elif isinstance(content, list):
+                # Handle structured content (e.g., images)
+                for item in content:
+                    if item.get("type") == "image_url":
+                        responses_items.append(
+                            {
+                                "type": "computer_call_output",
+                                "call_id": tool_call_id,
+                                "output": {
+                                    "type": "input_image",
+                                    "image_url": item.get("image_url", {}).get("url"),
+                                },
+                            }
+                        )
+                    elif item.get("type") == "text":
+                        responses_items.append(
+                            {
+                                "type": "function_call_output",
+                                "call_id": tool_call_id,
+                                "output": item.get("text"),
+                            }
+                        )
+
+        # Handle actual user messages
+        elif role == "user" and content:
+            if isinstance(content, list):
+                # Handle structured user content (e.g., text + images)
+                user_content = []
+                for item in content:
+                    if item.get("type") == "image_url":
+                        user_content.append(
+                            {
+                                "type": "input_image",
+                                "image_url": item.get("image_url", {}).get("url"),
+                            }
+                        )
+                    elif item.get("type") == "text":
+                        user_content.append({"type": "input_text", "text": item.get("text")})
+
+                if user_content:
+                    responses_items.append(
+                        {"role": "user", "type": "message", "content": user_content}
+                    )
+            elif isinstance(content, str):
+                # Handle simple text user message
+                responses_items.append({"role": "user", "content": content})
+
+    return responses_items

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

@@ -0,0 +1,24 @@
+"""
+Agent tools module.
+Provides base classes and registered tools for agent interactions.
+"""
+
+from .base import (
+    TOOL_REGISTRY,
+    BaseComputerTool,
+    BaseTool,
+    get_registered_tools,
+    get_tool,
+    register_tool,
+)
+from .browser_tool import BrowserTool
+
+__all__ = [
+    "BaseTool",
+    "BaseComputerTool",
+    "register_tool",
+    "get_registered_tools",
+    "get_tool",
+    "TOOL_REGISTRY",
+    "BrowserTool",
+]

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

@@ -0,0 +1,253 @@
+"""
+Base tool interface and registration system for agent tools.
+Provides a protocol for implementing tools that can be registered and discovered.
+"""
+
+import json
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Union
+
+# Global tool registry
+TOOL_REGISTRY: Dict[str, type] = {}
+
+
+def register_tool(name: str, allow_overwrite: bool = False):
+    """
+    Decorator to register a tool class with a given name.
+
+    Args:
+        name: The name to register the tool under
+        allow_overwrite: Whether to allow overwriting an existing tool with the same name
+
+    Returns:
+        Decorator function that registers the class
+
+    Example:
+        @register_tool("my_tool")
+        class MyTool(BaseTool):
+            ...
+    """
+
+    def decorator(cls):
+        if name in TOOL_REGISTRY:
+            if allow_overwrite:
+                print(f"Warning: Tool `{name}` already exists! Overwriting with class {cls}.")
+            else:
+                raise ValueError(
+                    f"Tool `{name}` already exists! Please ensure that the tool name is unique."
+                )
+        if hasattr(cls, "name") and cls.name and (cls.name != name):
+            raise ValueError(
+                f'{cls.__name__}.name="{cls.name}" conflicts with @register_tool(name="{name}").'
+            )
+        cls.name = name
+        TOOL_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def is_tool_schema(obj: dict) -> bool:
+    """
+    Check if obj is a valid JSON schema describing a tool compatible with OpenAI's tool calling.
+
+    Example valid schema:
+    {
+      "name": "get_current_weather",
+      "description": "Get the current weather in a given location",
+      "parameters": {
+        "type": "object",
+        "properties": {
+          "location": {
+            "type": "string",
+            "description": "The city and state, e.g. San Francisco, CA"
+          },
+          "unit": {
+            "type": "string",
+            "enum": ["celsius", "fahrenheit"]
+          }
+        },
+        "required": ["location"]
+      }
+    }
+    """
+    try:
+        # Basic structure validation
+        assert set(obj.keys()) == {"name", "description", "parameters"}
+        assert isinstance(obj["name"], str)
+        assert obj["name"].strip()
+        assert isinstance(obj["description"], str)
+        assert isinstance(obj["parameters"], dict)
+
+        # Parameters structure validation
+        assert "type" in obj["parameters"]
+        assert obj["parameters"]["type"] == "object"
+        assert "properties" in obj["parameters"]
+        assert isinstance(obj["parameters"]["properties"], dict)
+
+        if "required" in obj["parameters"]:
+            assert isinstance(obj["parameters"]["required"], list)
+            assert set(obj["parameters"]["required"]).issubset(
+                set(obj["parameters"]["properties"].keys())
+            )
+
+        return True
+    except (AssertionError, KeyError, TypeError):
+        return False
+
+
+class BaseTool(ABC):
+    """
+    Base class for all agent tools.
+
+    Tools must implement:
+    - name: str - The tool name (set by @register_tool decorator)
+    - description: property that returns str - Tool description
+    - parameters: property that returns dict - JSON schema for tool parameters
+    - call: method - Execute the tool with given parameters
+    """
+
+    name: str = ""
+
+    def __init__(self, cfg: Optional[dict] = None):
+        """
+        Initialize the tool.
+
+        Args:
+            cfg: Optional configuration dictionary
+        """
+        self.cfg = cfg or {}
+
+        if not self.name:
+            raise ValueError(
+                f"You must set {self.__class__.__name__}.name, either by "
+                f"@register_tool(name=...) or explicitly setting "
+                f"{self.__class__.__name__}.name"
+            )
+
+        # Validate schema if parameters is a dict
+        if isinstance(self.parameters, dict):
+            if not is_tool_schema(
+                {"name": self.name, "description": self.description, "parameters": self.parameters}
+            ):
+                raise ValueError(
+                    "The parameters, when provided as a dict, must conform to a "
+                    "valid openai-compatible JSON schema."
+                )
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Return the tool description."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def parameters(self) -> dict:
+        """Return the JSON schema for tool parameters."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def call(self, params: Union[str, dict], **kwargs) -> Union[str, list, dict]:
+        """
+        Execute the tool with the given parameters.
+
+        Args:
+            params: The parameters for the tool call (JSON string or dict)
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            The result of the tool execution
+        """
+        raise NotImplementedError
+
+    def _verify_json_format_args(self, params: Union[str, dict]) -> dict:
+        """
+        Verify and parse the parameters as JSON.
+
+        Args:
+            params: Parameters as string or dict
+
+        Returns:
+            Parsed parameters as dict
+
+        Raises:
+            ValueError: If parameters are not valid JSON or don't match schema
+        """
+        if isinstance(params, str):
+            try:
+                params_json: dict = json.loads(params)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Parameters must be formatted as valid JSON: {e}")
+        else:
+            params_json: dict = params
+
+        # Validate against schema if using dict parameters
+        if isinstance(self.parameters, dict):
+            try:
+                # Basic validation of required fields
+                if "required" in self.parameters:
+                    for field in self.parameters["required"]:
+                        if field not in params_json:
+                            raise ValueError(f'Required parameter "{field}" is missing')
+            except (KeyError, TypeError) as e:
+                raise ValueError(f"Invalid parameters: {e}")
+
+        return params_json
+
+    @property
+    def function(self) -> dict:
+        """
+        Return the function information for this tool.
+
+        Returns:
+            Dict with tool metadata
+        """
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": self.parameters,
+        }
+
+
+def get_registered_tools() -> Dict[str, type]:
+    """
+    Get all registered tools.
+
+    Returns:
+        Dictionary mapping tool names to tool classes
+    """
+    return TOOL_REGISTRY.copy()
+
+
+def get_tool(name: str) -> Optional[type]:
+    """
+    Get a registered tool by name.
+
+    Args:
+        name: The tool name
+
+    Returns:
+        The tool class, or None if not found
+    """
+    return TOOL_REGISTRY.get(name)
+
+
+class BaseComputerTool(BaseTool):
+    """
+    Base class for computer tools that can provide screenshots.
+
+    Computer tools must implement:
+    - All BaseTool requirements (name, description, parameters, call)
+    - screenshot() method that returns screenshot as base64 string
+    """
+
+    @abstractmethod
+    async def screenshot(self) -> str:
+        """
+        Take a screenshot of the computer/browser.
+
+        Returns:
+            Screenshot image data as base64-encoded string
+        """
+        raise NotImplementedError

libs/python/agent/cua_agent/tools/browser_tool.py

@@ -0,0 +1,624 @@
+"""
+Browser Tool for agent interactions.
+Allows agents to control a browser programmatically via Playwright.
+Implements the computer_use action interface for comprehensive browser control.
+"""
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Optional, Union
+
+from .base import BaseComputerTool, register_tool
+
+if TYPE_CHECKING:
+    from computer.interface import GenericComputerInterface
+
+logger = logging.getLogger(__name__)
+
+
+@register_tool("computer_use")
+class BrowserTool(BaseComputerTool):
+    """
+    Browser tool that uses the computer SDK's interface to control a browser.
+    Implements a comprehensive computer_use action interface for browser control.
+    """
+
+    def __init__(self, interface: "GenericComputerInterface", cfg: Optional[dict] = None):
+        """
+        Initialize the BrowserTool.
+
+        Args:
+            interface: A GenericComputerInterface instance that provides playwright_exec
+            cfg: Optional configuration dictionary
+        """
+        self.interface = interface
+        self._facts = []  # Store memorized facts
+        self._automation = None  # Cached automation interface
+
+        # Get initial screenshot to determine dimensions
+        self.viewport_width = None
+        self.viewport_height = None
+        self.resized_width = None
+        self.resized_height = None
+
+        # Try to initialize dimensions synchronously
+        try:
+            import asyncio
+
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # If we're in an async context, dimensions will be lazy-loaded
+                pass
+            else:
+                loop.run_until_complete(self._initialize_dimensions())
+        except Exception:
+            # Dimensions will be lazy-loaded on first use
+            pass
+
+        super().__init__(cfg)
+
+    @property
+    def automation(self):
+        """
+        Get the automation interface for keyboard/mouse actions.
+
+        Handles both interface structures:
+        - Nested: interface.interface (wrapper with .interface property)
+        - Direct: interface itself IS the automation handler
+        """
+        if self._automation is not None:
+            return self._automation
+
+        # Try nested structure first (interface.interface)
+        if hasattr(self.interface, "interface") and self.interface.interface is not None:
+            self._automation = self.interface.interface
+        else:
+            # Direct structure - interface IS the automation handler
+            self._automation = self.interface
+
+        return self._automation
+
+    async def _initialize_dimensions(self):
+        """Initialize viewport and resized dimensions from screenshot."""
+        try:
+            import base64
+            import io
+
+            from PIL import Image
+            from qwen_vl_utils import smart_resize
+
+            # Take a screenshot to get actual dimensions
+            screenshot_b64 = await self.screenshot()
+            img_bytes = base64.b64decode(screenshot_b64)
+            im = Image.open(io.BytesIO(img_bytes))
+
+            # Store actual viewport size
+            self.viewport_width = im.width
+            self.viewport_height = im.height
+
+            # Calculate resized dimensions using smart_resize with factor=28
+            MIN_PIXELS = 3136
+            MAX_PIXELS = 12845056
+            rh, rw = smart_resize(
+                im.height, im.width, factor=28, min_pixels=MIN_PIXELS, max_pixels=MAX_PIXELS
+            )
+            self.resized_width = rw
+            self.resized_height = rh
+
+        except Exception as e:
+            # Fall back to defaults if initialization fails
+            logger.warning(f"Failed to initialize dimensions: {e}")
+            self.viewport_width = 1024
+            self.viewport_height = 768
+            self.resized_width = 1024
+            self.resized_height = 768
+
+    async def _proc_coords(self, x: float, y: float) -> tuple:
+        """
+        Process coordinates by converting from resized space to viewport space.
+
+        Args:
+            x: X coordinate in resized space (0 to resized_width)
+            y: Y coordinate in resized space (0 to resized_height)
+
+        Returns:
+            Tuple of (viewport_x, viewport_y) in actual viewport pixels
+        """
+        # Ensure dimensions are initialized
+        if self.resized_width is None or self.resized_height is None:
+            await self._initialize_dimensions()
+
+        # Convert from resized space to viewport space
+        # Normalize by resized dimensions, then scale to viewport dimensions
+        viewport_x = (x / self.resized_width) * self.viewport_width
+        viewport_y = (y / self.resized_height) * self.viewport_height
+
+        return int(round(viewport_x)), int(round(viewport_y))
+
+    @property
+    def description(self) -> str:
+        # Use resized dimensions if available, otherwise use defaults
+        width = self.resized_width if self.resized_width is not None else 1024
+        height = self.resized_height if self.resized_height is not None else 768
+
+        return f"Use a mouse and keyboard to interact with a computer, and take screenshots.\
+* This is an interface to a desktop GUI. You do not have access to a terminal or applications menu. You must click on desktop icons to start applications.\
+* Some applications may take time to start or process actions, so you may need to wait and take successive screenshots to see the results of your actions. E.g. if you click on Firefox and a window doesn't open, try wait and taking another screenshot.\
+* The screen's resolution is {width}x{height}.\
+* Whenever you intend to move the cursor to click on an element like an icon, you should consult a screenshot to determine the coordinates of the element before moving the cursor.\
+* If you tried clicking on a program or link but it failed to load, even after waiting, try adjusting your cursor position so that the tip of the cursor visually falls on the element that you want to click.\
+* Make sure to click any buttons, links, icons, etc with the cursor tip in the center of the element. Don't click boxes on their edges unless asked.\
+* When a separate scrollable container prominently overlays the webpage, if you want to scroll within it, you typically need to mouse_move() over it first and then scroll().\
+* If a popup window appears that you want to close, if left_click() on the 'X' or close button doesn't work, try key(keys=['Escape']) to close it.\
+* On some search bars, when you type(), you may need to press_enter=False and instead separately call left_click() on the search button to submit the search query. This is especially true of search bars that have auto-suggest popups for e.g. locations\
+* For calendar widgets, you usually need to left_click() on arrows to move between months and left_click() on dates to select them; type() is not typically used to input dates there.".strip()
+
+    @property
+    def parameters(self) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "description": """The action to perform. The available actions are:
+* key: Performs key down presses on the arguments passed in order, then performs key releases in reverse order. Includes 'Enter', 'Alt', 'Shift', 'Tab', 'Control', 'Backspace', 'Delete', 'Escape', 'ArrowUp', 'ArrowDown', 'ArrowLeft', 'ArrowRight', 'PageDown', 'PageUp', 'Shift', etc.
+* type: Type a string of text on the keyboard.
+* mouse_move: Move the cursor to a specified (x, y) pixel coordinate on the screen.
+* left_click: Click the left mouse button.
+* scroll: Performs a scroll of the mouse scroll wheel.
+* visit_url: Visit a specified URL.
+* web_search: Perform a web search with a specified query.
+* history_back: Go back to the previous page in the browser history.
+* pause_and_memorize_fact: Pause and memorize a fact for future reference.
+* wait: Wait specified seconds for the change to happen.
+* terminate: Terminate the current task and report its completion status.
+* screenshot: Take a screenshot of the current screen.""",
+                    "enum": [
+                        "key",
+                        "type",
+                        "mouse_move",
+                        "left_click",
+                        "scroll",
+                        "visit_url",
+                        "web_search",
+                        "history_back",
+                        "pause_and_memorize_fact",
+                        "wait",
+                        "terminate",
+                        "screenshot",
+                    ],
+                    "type": "string",
+                },
+                "keys": {
+                    "description": "Required only by action=key.",
+                    "type": "array",
+                    "items": {"type": "string"},
+                },
+                "text": {"description": "Required only by action=type.", "type": "string"},
+                "coordinate": {
+                    "description": "(x, y) coordinates for mouse actions. Required only by action=left_click, action=mouse_move, and action=type.",
+                    "type": "array",
+                    "items": {"type": "number"},
+                },
+                "pixels": {
+                    "description": "Amount of scrolling. Positive = up, Negative = down. Required only by action=scroll.",
+                    "type": "number",
+                },
+                "url": {
+                    "description": "The URL to visit. Required only by action=visit_url.",
+                    "type": "string",
+                },
+                "query": {
+                    "description": "The query to search for. Required only by action=web_search.",
+                    "type": "string",
+                },
+                "fact": {
+                    "description": "The fact to remember for the future. Required only by action=pause_and_memorize_fact.",
+                    "type": "string",
+                },
+                "time": {
+                    "description": "Seconds to wait. Required only by action=wait.",
+                    "type": "number",
+                },
+                "status": {
+                    "description": "Status of the task. Required only by action=terminate.",
+                    "type": "string",
+                    "enum": ["success", "failure"],
+                },
+            },
+            "required": ["action"],
+        }
+
+    def call(self, params: Union[str, dict], **kwargs) -> Union[str, dict]:
+        """
+        Execute a browser action.
+
+        Args:
+            params: Action parameters (JSON string or dict)
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Result of the action execution
+        """
+        # Verify and parse parameters
+        params_dict = self._verify_json_format_args(params)
+        action = params_dict.get("action")
+
+        if not action:
+            return {"success": False, "error": "action parameter is required"}
+
+        # Execute action synchronously by running async method in event loop
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # If we're already in an async context, we can't use run_until_complete
+                # Create a task and wait for it
+                import concurrent.futures
+
+                with concurrent.futures.ThreadPoolExecutor() as executor:
+                    future = executor.submit(asyncio.run, self._execute_action(action, params_dict))
+                    result = future.result()
+            else:
+                result = loop.run_until_complete(self._execute_action(action, params_dict))
+            return result
+        except Exception as e:
+            logger.error(f"Error executing action {action}: {e}")
+            return {"success": False, "error": str(e)}
+
+    async def _execute_action(self, action: str, params: dict) -> dict:
+        """Execute the specific action asynchronously."""
+        try:
+            if action == "key":
+                return await self._action_key(params)
+            elif action == "type":
+                return await self._action_type(params)
+            elif action == "mouse_move":
+                return await self._action_mouse_move(params)
+            elif action == "left_click":
+                return await self._action_left_click(params)
+            elif action == "scroll":
+                return await self._action_scroll(params)
+            elif action == "visit_url":
+                return await self._action_visit_url(params)
+            elif action == "web_search":
+                return await self._action_web_search(params)
+            elif action == "history_back":
+                return await self._action_history_back(params)
+            elif action == "pause_and_memorize_fact":
+                return await self._action_pause_and_memorize_fact(params)
+            elif action == "wait":
+                return await self._action_wait(params)
+            elif action == "terminate":
+                return await self._action_terminate(params)
+            elif action == "screenshot":
+                return await self._action_screenshot(params)
+            else:
+                return {"success": False, "error": f"Unknown action: {action}"}
+        except Exception as e:
+            logger.error(f"Error in action {action}: {e}")
+            return {"success": False, "error": str(e)}
+
+    async def _action_key(self, params: dict) -> dict:
+        """Press keys in sequence."""
+        keys = params.get("keys", [])
+        if not keys:
+            return {"success": False, "error": "keys parameter is required"}
+
+        # Convert keys to proper format and press via hotkey
+        try:
+            await self.automation.hotkey(*keys)
+            return {"success": True, "message": f"Pressed keys: {keys}"}
+        except Exception as e:
+            return {"success": False, "error": str(e)}
+
+    async def _action_type(self, params: dict) -> dict:
+        """Type text."""
+        text = params.get("text")
+        if not text:
+            return {"success": False, "error": "text parameter is required"}
+
+        # If coordinate is provided, click there first
+        coordinate = params.get("coordinate")
+        if coordinate and len(coordinate) == 2:
+            await self.interface.playwright_exec("click", {"x": coordinate[0], "y": coordinate[1]})
+
+        result = await self.interface.playwright_exec("type", {"text": text})
+        return result
+
+    async def _action_mouse_move(self, params: dict) -> dict:
+        """Move mouse to coordinates."""
+        coordinate = params.get("coordinate")
+        if not coordinate or len(coordinate) != 2:
+            return {"success": False, "error": "coordinate parameter [x, y] is required"}
+
+        await self.automation.move_cursor(coordinate[0], coordinate[1])
+        return {"success": True, "message": f"Moved cursor to {coordinate}"}
+
+    async def _action_left_click(self, params: dict) -> dict:
+        """Click at coordinates."""
+        coordinate = params.get("coordinate")
+        if not coordinate or len(coordinate) != 2:
+            return {"success": False, "error": "coordinate parameter [x, y] is required"}
+
+        result = await self.interface.playwright_exec(
+            "click", {"x": coordinate[0], "y": coordinate[1]}
+        )
+        return result
+
+    async def _action_scroll(self, params: dict) -> dict:
+        """Scroll the page."""
+        pixels = params.get("pixels")
+        # Handle None explicitly - default to 0 means "no scroll requested"
+        if pixels is None or pixels == 0:
+            return {"success": False, "error": "pixels parameter is required"}
+
+        # Positive = up (negative delta_y), Negative = down (positive delta_y)
+        result = await self.interface.playwright_exec("scroll", {"delta_x": 0, "delta_y": -pixels})
+        return result
+
+    async def _action_visit_url(self, params: dict) -> dict:
+        """Visit a URL."""
+        url = params.get("url")
+        if not url:
+            return {"success": False, "error": "url parameter is required"}
+
+        result = await self.interface.playwright_exec("visit_url", {"url": url})
+        return result
+
+    async def _action_web_search(self, params: dict) -> dict:
+        """Perform web search."""
+        query = params.get("query")
+        if not query:
+            return {"success": False, "error": "query parameter is required"}
+
+        result = await self.interface.playwright_exec("web_search", {"query": query})
+        return result
+
+    async def _action_history_back(self, params: dict) -> dict:
+        """Go back in browser history."""
+        # Press Alt+Left arrow key combination
+        try:
+            await self.automation.hotkey("Alt", "ArrowLeft")
+            return {"success": True, "message": "Navigated back in history"}
+        except Exception as e:
+            return {"success": False, "error": str(e)}
+
+    async def _action_pause_and_memorize_fact(self, params: dict) -> dict:
+        """Memorize a fact."""
+        fact = params.get("fact")
+        if not fact:
+            return {"success": False, "error": "fact parameter is required"}
+
+        self._facts.append(fact)
+        return {
+            "success": True,
+            "message": f"Memorized fact: {fact}",
+            "total_facts": len(self._facts),
+        }
+
+    async def _action_wait(self, params: dict) -> dict:
+        """Wait for specified seconds."""
+        time = params.get("time")
+        # Handle None or missing time - default to 3 seconds (matches FARA behavior)
+        if time is None:
+            time = 3
+        if time <= 0:
+            return {"success": False, "error": "time parameter must be positive"}
+
+        await asyncio.sleep(time)
+        return {"success": True, "message": f"Waited {time} seconds"}
+
+    async def _action_terminate(self, params: dict) -> dict:
+        """Terminate and report status."""
+        status = params.get("status")
+        # Handle None or missing status - default to "success"
+        if status is None:
+            status = "success"
+        message = f"Task terminated with status: {status}"
+
+        if self._facts:
+            message += f"\nMemorized facts: {self._facts}"
+
+        return {"success": True, "status": status, "message": message, "terminated": True}
+
+    async def _action_screenshot(self, params: dict) -> dict:
+        """Take a screenshot of the current screen."""
+        try:
+            screenshot_b64 = await self.screenshot()
+            return {"success": True, "screenshot": screenshot_b64}
+        except Exception as e:
+            return {"success": False, "error": str(e)}
+
+    # Legacy methods for backward compatibility
+    async def visit_url(self, url: str) -> dict:
+        """Navigate to a URL."""
+        return await self._action_visit_url({"url": url})
+
+    async def click(self, x: int = None, y: int = None, button: str = "left", **kwargs) -> dict:
+        """Click at coordinates. Supports both positional (x, y) and kwargs (button, x, y).
+
+        This is compatible with the normalized format from OperatorNormalizerCallback
+        which transforms actions like {"type": "left_click", "coordinate": [x, y]}
+        into {"type": "click", "button": "left", "x": x, "y": y}.
+        """
+        if x is None or y is None:
+            return {"success": False, "error": "x and y coordinates are required"}
+        if button == "right":
+            return await self.interface.playwright_exec(
+                "click", {"x": x, "y": y, "button": "right"}
+            )
+        elif button == "middle" or button == "wheel":
+            return await self.interface.playwright_exec(
+                "click", {"x": x, "y": y, "button": "middle"}
+            )
+        else:
+            # Default to left click
+            return await self._action_left_click({"coordinate": [x, y]})
+
+    async def type(self, text: str) -> dict:
+        """Type text into the focused element."""
+        return await self._action_type({"text": text})
+
+    async def scroll(
+        self,
+        delta_x: int = None,
+        delta_y: int = None,
+        scroll_x: int = None,
+        scroll_y: int = None,
+        x: int = None,
+        y: int = None,
+        pixels: int = None,
+        coordinate=None,
+        **kwargs,
+    ) -> dict:
+        """Scroll the page. Supports multiple formats:
+        - Legacy: scroll(delta_x, delta_y)
+        - Normalized: scroll(scroll_x=0, scroll_y=100, x=500, y=300)
+        - FARA: scroll(pixels=100, coordinate=[500, 300])
+        """
+        # Determine scroll amounts from various input formats
+        dx = scroll_x or delta_x or 0
+        dy = scroll_y or delta_y or (-(pixels or 0))  # pixels: positive=up, negative=down
+
+        result = await self.interface.playwright_exec("scroll", {"delta_x": dx, "delta_y": dy})
+        return result
+
+    async def web_search(self, query: str) -> dict:
+        """Navigate to a Google search for the query."""
+        return await self._action_web_search({"query": query})
+
+    async def screenshot(self) -> str:
+        """Take a screenshot of the current browser page."""
+        result = await self.interface.playwright_exec("screenshot", {})
+        if result.get("success") and result.get("screenshot"):
+            screenshot_b64 = result["screenshot"]
+            return screenshot_b64
+        else:
+            error = result.get("error", "Unknown error")
+            raise RuntimeError(f"Failed to take screenshot: {error}")
+
+    async def get_current_url(self) -> str:
+        """Get the current URL of the browser page."""
+        result = await self.interface.playwright_exec("get_current_url", {})
+        if result.get("success") and result.get("url"):
+            return result["url"]
+        else:
+            error = result.get("error", "Unknown error")
+            raise RuntimeError(f"Failed to get current URL: {error}")
+
+    # FARA-compatible action methods
+    # These methods accept parameters in the format that FARA model outputs
+    # and agent.py passes via **action_args
+
+    async def left_click(self, coordinate=None, x: int = None, y: int = None, **kwargs) -> dict:
+        """Left click at coordinates. Supports coordinate array or x/y kwargs."""
+        # Accept either coordinate array or x/y kwargs
+        if coordinate and len(coordinate) >= 2:
+            x, y = coordinate[0], coordinate[1]
+        if x is None or y is None:
+            return {"success": False, "error": "coordinate parameter [x, y] or x/y kwargs required"}
+        return await self._action_left_click({"coordinate": [x, y]})
+
+    async def right_click(self, coordinate=None, x: int = None, y: int = None, **kwargs) -> dict:
+        """Right click at coordinates. Supports coordinate array or x/y kwargs."""
+        # Accept either coordinate array or x/y kwargs
+        if coordinate and len(coordinate) >= 2:
+            x, y = coordinate[0], coordinate[1]
+        if x is None or y is None:
+            return {"success": False, "error": "coordinate parameter [x, y] or x/y kwargs required"}
+        result = await self.interface.playwright_exec("click", {"x": x, "y": y, "button": "right"})
+        return result
+
+    async def middle_click(self, coordinate=None, x: int = None, y: int = None, **kwargs) -> dict:
+        """Middle click at coordinates. Supports coordinate array or x/y kwargs."""
+        # Accept either coordinate array or x/y kwargs
+        if coordinate and len(coordinate) >= 2:
+            x, y = coordinate[0], coordinate[1]
+        if x is None or y is None:
+            return {"success": False, "error": "coordinate parameter [x, y] or x/y kwargs required"}
+        result = await self.interface.playwright_exec("click", {"x": x, "y": y, "button": "middle"})
+        return result
+
+    async def double_click(self, coordinate=None, x: int = None, y: int = None, **kwargs) -> dict:
+        """Double click at coordinates. Supports coordinate array or x/y kwargs."""
+        # Accept either coordinate array or x/y kwargs
+        if coordinate and len(coordinate) >= 2:
+            x, y = coordinate[0], coordinate[1]
+        if x is None or y is None:
+            return {"success": False, "error": "coordinate parameter [x, y] or x/y kwargs required"}
+        result = await self.interface.playwright_exec("dblclick", {"x": x, "y": y})
+        return result
+
+    async def triple_click(
+        self, coordinate=None, x: int = None, y: int = None, button: str = None, **kwargs
+    ) -> dict:
+        """Triple click at coordinates. Supports coordinate array or x/y kwargs."""
+        # Accept either coordinate array or x/y kwargs
+        if coordinate and len(coordinate) >= 2:
+            x, y = coordinate[0], coordinate[1]
+        if x is None or y is None:
+            return {"success": False, "error": "coordinate parameter [x, y] or x/y kwargs required"}
+        # Triple click is approximated as double click
+        return await self.double_click(x=x, y=y)
+
+    async def mouse_move(self, coordinate=None, x: int = None, y: int = None, **kwargs) -> dict:
+        """Move mouse to coordinates. Supports coordinate array or x/y kwargs."""
+        # Accept either coordinate array or x/y kwargs
+        if coordinate and len(coordinate) >= 2:
+            x, y = coordinate[0], coordinate[1]
+        if x is None or y is None:
+            return {"success": False, "error": "coordinate parameter [x, y] or x/y kwargs required"}
+        return await self._action_mouse_move({"coordinate": [x, y]})
+
+    async def move(self, x: int = None, y: int = None, **kwargs) -> dict:
+        """Move mouse to coordinates. Alias for mouse_move with x/y kwargs."""
+        return await self.mouse_move(x=x, y=y)
+
+    async def left_click_drag(
+        self, coordinate=None, start_coordinate=None, end_coordinate=None, **kwargs
+    ) -> dict:
+        """Drag from start to end coordinates. FARA-compatible."""
+        if start_coordinate and end_coordinate:
+            # Use start/end coordinates if provided
+            await self.automation.move_cursor(start_coordinate[0], start_coordinate[1])
+            await self.automation.mouse_down(start_coordinate[0], start_coordinate[1])
+            await self.automation.move_cursor(end_coordinate[0], end_coordinate[1])
+            await self.automation.mouse_up(end_coordinate[0], end_coordinate[1])
+            return {
+                "success": True,
+                "message": f"Dragged from {start_coordinate} to {end_coordinate}",
+            }
+        elif coordinate:
+            # Just move to coordinate
+            await self.automation.move_cursor(coordinate[0], coordinate[1])
+            return {"success": True, "message": f"Moved to {coordinate}"}
+        return {
+            "success": False,
+            "error": "start_coordinate and end_coordinate or coordinate required",
+        }
+
+    async def key(self, keys=None, **kwargs) -> dict:
+        """Press keys. FARA-compatible."""
+        return await self._action_key({"keys": keys})
+
+    async def keypress(self, keys=None, **kwargs) -> dict:
+        """Press keys. Alias for key() - used by OperatorNormalizerCallback."""
+        return await self._action_key({"keys": keys})
+
+    async def hscroll(self, pixels=None, coordinate=None, **kwargs) -> dict:
+        """Horizontal scroll. FARA-compatible."""
+        if pixels is None:
+            return {"success": False, "error": "pixels parameter is required"}
+        result = await self.interface.playwright_exec("scroll", {"delta_x": pixels, "delta_y": 0})
+        return result
+
+    async def wait(self, time=None, **kwargs) -> dict:
+        """Wait for specified seconds. FARA-compatible."""
+        return await self._action_wait({"time": time})
+
+    async def history_back(self, **kwargs) -> dict:
+        """Go back in browser history. FARA-compatible."""
+        return await self._action_history_back({})
+
+    async def terminate(self, status=None, **kwargs) -> dict:
+        """Terminate and report status. FARA-compatible."""
+        return await self._action_terminate({"status": status or "success"})

libs/python/agent/cua_agent/types.py

@@ -0,0 +1,45 @@
+"""
+Type definitions for agent
+"""
+
+import re
+from collections.abc import Iterable
+from typing import Any, Callable, Dict, List, Literal, Optional, Protocol
+
+from litellm import ResponseInputParam, ResponsesAPIResponse, ToolParam
+from pydantic import BaseModel
+
+# Agent input types
+Messages = str | ResponseInputParam | List[Dict[str, Any]]
+Tools = Optional[Iterable[ToolParam]]
+
+# Agent output types
+AgentResponse = ResponsesAPIResponse
+AgentCapability = Literal["step", "click"]
+
+
+# Exception types
+class ToolError(RuntimeError):
+    """Base exception for tool-related errors"""
+
+    pass
+
+
+class IllegalArgumentError(ToolError):
+    """Exception raised when function arguments are invalid"""
+
+    pass
+
+
+# Agent config registration
+class AgentConfigInfo(BaseModel):
+    """Information about a registered agent config"""
+
+    agent_class: type
+    models_regex: str
+    priority: int = 0
+    tool_type: Optional[str] = None  # "browser" | "mobile" | None (flexible)
+
+    def matches_model(self, model: str) -> bool:
+        """Check if this agent config matches the given model"""
+        return bool(re.match(self.models_regex, model))

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

@@ -0,0 +1,7 @@
+"""
+UI components for agent
+"""
+
+from .gradio import create_gradio_ui, launch_ui
+
+__all__ = ["launch_ui", "create_gradio_ui"]

libs/python/agent/cua_agent/ui/__main__.py

@@ -0,0 +1,4 @@
+from .gradio import launch_ui
+
+if __name__ == "__main__":
+    launch_ui()

libs/python/agent/cua_agent/ui/gradio/__init__.py

@@ -0,0 +1,8 @@
+"""
+Gradio UI for agent
+"""
+
+from .app import launch_ui
+from .ui_components import create_gradio_ui
+
+__all__ = ["launch_ui", "create_gradio_ui"]

libs/python/agent/cua_agent/ui/gradio/app.py

@@ -0,0 +1,262 @@
+"""
+Advanced Gradio UI for Computer-Use Agent (cua-agent)
+
+This is a Gradio interface for the Computer-Use Agent v0.4.x (cua-agent)
+with an advanced UI for model selection and configuration.
+
+Supported Agent Models:
+- OpenAI: openai/computer-use-preview
+- Anthropic: anthropic/claude-sonnet-4-5-20250929, anthropic/claude-3-7-sonnet-20250219
+- UI-TARS: huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B
+- Omniparser: omniparser+anthropic/claude-sonnet-4-5-20250929, omniparser+ollama_chat/gemma3
+
+Requirements:
+    - Mac with Apple Silicon (M1/M2/M3/M4), Linux, or Windows
+    - macOS 14 (Sonoma) or newer / Ubuntu 20.04+
+    - Python 3.11+
+    - Lume CLI installed (https://github.com/trycua/cua)
+    - OpenAI or Anthropic API key
+"""
+
+import asyncio
+import json
+import logging
+import os
+import platform
+from pathlib import Path
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union, cast
+
+import gradio as gr
+
+# Import from agent package
+from cua_agent import ComputerAgent
+from cua_agent.types import AgentResponse, Messages
+
+try:
+    from computer import Computer
+except ImportError:
+    Computer = None  # type: ignore[assignment,misc]
+from gradio.components.chatbot import MetadataDict
+
+# Global variables
+global_agent = None
+global_computer = None
+SETTINGS_FILE = Path(".gradio_settings.json")
+
+logging.basicConfig(level=logging.INFO)
+
+import dotenv
+
+if dotenv.load_dotenv():
+    print(f"DEBUG - Loaded environment variables from {dotenv.find_dotenv()}")
+else:
+    print("DEBUG - No .env file found")
+
+
+# --- Settings Load/Save Functions ---
+def load_settings() -> Dict[str, Any]:
+    """Loads settings from the JSON file."""
+    if SETTINGS_FILE.exists():
+        try:
+            with open(SETTINGS_FILE, "r") as f:
+                settings = json.load(f)
+                if isinstance(settings, dict):
+                    print(f"DEBUG - Loaded settings from {SETTINGS_FILE}")
+                    return settings
+        except (json.JSONDecodeError, IOError) as e:
+            print(f"Warning: Could not load settings from {SETTINGS_FILE}: {e}")
+    return {}
+
+
+def save_settings(settings: Dict[str, Any]):
+    """Saves settings to the JSON file."""
+    settings.pop("provider_api_key", None)
+    try:
+        with open(SETTINGS_FILE, "w") as f:
+            json.dump(settings, f, indent=4)
+        print(f"DEBUG - Saved settings to {SETTINGS_FILE}")
+    except IOError as e:
+        print(f"Warning: Could not save settings to {SETTINGS_FILE}: {e}")
+
+
+# # Custom Screenshot Handler for Gradio chat
+# class GradioChatScreenshotHandler:
+#     """Custom handler that adds screenshots to the Gradio chatbot."""
+
+#     def __init__(self, chatbot_history: List[gr.ChatMessage]):
+#         self.chatbot_history = chatbot_history
+#         print("GradioChatScreenshotHandler initialized")
+
+#     async def on_screenshot(self, screenshot_base64: str, action_type: str = "") -> None:
+#         """Add screenshot to chatbot when a screenshot is taken."""
+#         image_markdown = f"![Screenshot after {action_type}](data:image/png;base64,{screenshot_base64})"
+
+#         if self.chatbot_history is not None:
+#             self.chatbot_history.append(
+#                 gr.ChatMessage(
+#                     role="assistant",
+#                     content=image_markdown,
+#                     metadata={"title": f"🖥️ Screenshot - {action_type}", "status": "done"},
+#                 )
+#             )
+
+
+# Detect platform capabilities
+is_mac = platform.system().lower() == "darwin"
+is_lume_available = is_mac
+
+print("is_mac: ", is_mac)
+print("Lume available: ", is_lume_available)
+
+# Map model names to agent model strings
+MODEL_MAPPINGS = {
+    "openai": {
+        "default": "openai/computer-use-preview",
+        "OpenAI: Computer-Use Preview": "openai/computer-use-preview",
+    },
+    "anthropic": {
+        "default": "anthropic/claude-3-7-sonnet-20250219",
+        "Anthropic: Claude 4 Opus (20250514)": "anthropic/claude-opus-4-20250514",
+        "Anthropic: Claude 4 Sonnet (20250514)": "anthropic/claude-sonnet-4-20250514",
+        "Anthropic: Claude 3.7 Sonnet (20250219)": "anthropic/claude-3-7-sonnet-20250219",
+    },
+    "omni": {
+        "default": "omniparser+openai/gpt-4o",
+        "OMNI: OpenAI GPT-4o": "omniparser+openai/gpt-4o",
+        "OMNI: OpenAI GPT-4o mini": "omniparser+openai/gpt-4o-mini",
+        "OMNI: Claude 3.7 Sonnet (20250219)": "omniparser+anthropic/claude-3-7-sonnet-20250219",
+    },
+    "uitars": {
+        "default": "huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B" if is_mac else "ui-tars",
+        "huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B": "huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B",
+    },
+}
+
+
+def get_model_string(model_name: str, loop_provider: str) -> str:
+    """Determine the agent model string based on the input."""
+    if model_name == "Custom model (OpenAI compatible API)":
+        return "custom_oaicompat"
+    elif model_name == "Custom model (ollama)":
+        return "custom_ollama"
+    elif loop_provider == "OMNI-OLLAMA" or model_name.startswith("OMNI: Ollama "):
+        if model_name.startswith("OMNI: Ollama "):
+            ollama_model = model_name.split("OMNI: Ollama ", 1)[1]
+            return f"omniparser+ollama_chat/{ollama_model}"
+        return "omniparser+ollama_chat/llama3"
+
+    # Map based on loop provider
+    mapping = MODEL_MAPPINGS.get(loop_provider.lower(), MODEL_MAPPINGS["openai"])
+    return mapping.get(model_name, mapping["default"])
+
+
+def get_ollama_models() -> List[str]:
+    """Get available models from Ollama if installed."""
+    try:
+        import subprocess
+
+        result = subprocess.run(["ollama", "list"], capture_output=True, text=True)
+        if result.returncode == 0:
+            lines = result.stdout.strip().split("\n")
+            if len(lines) < 2:
+                return []
+            models = []
+            for line in lines[1:]:
+                parts = line.split()
+                if parts:
+                    model_name = parts[0]
+                    models.append(f"OMNI: Ollama {model_name}")
+            return models
+        return []
+    except Exception as e:
+        logging.error(f"Error getting Ollama models: {e}")
+        return []
+
+
+def create_computer_instance(
+    verbosity: int = logging.INFO,
+    os_type: str = "macos",
+    provider_type: str = "lume",
+    name: Optional[str] = None,
+    api_key: Optional[str] = None,
+) -> Computer:
+    """Create or get the global Computer instance."""
+    global global_computer
+    if global_computer is None:
+        if provider_type == "localhost":
+            global_computer = Computer(
+                verbosity=verbosity, os_type=os_type, use_host_computer_server=True
+            )
+        else:
+            global_computer = Computer(
+                verbosity=verbosity,
+                os_type=os_type,
+                provider_type=provider_type,
+                name=name if name else "",
+                api_key=api_key,
+            )
+    return global_computer
+
+
+def create_agent(
+    model_string: str,
+    save_trajectory: bool = True,
+    only_n_most_recent_images: int = 3,
+    verbosity: int = logging.INFO,
+    custom_model_name: Optional[str] = None,
+    computer_os: str = "macos",
+    computer_provider: str = "lume",
+    computer_name: Optional[str] = None,
+    computer_api_key: Optional[str] = None,
+    max_trajectory_budget: Optional[float] = None,
+) -> ComputerAgent:
+    """Create or update the global agent with the specified parameters."""
+    global global_agent
+
+    # Create the computer
+    computer = create_computer_instance(
+        verbosity=verbosity,
+        os_type=computer_os,
+        provider_type=computer_provider,
+        name=computer_name,
+        api_key=computer_api_key,
+    )
+
+    # Handle custom models
+    if model_string == "custom_oaicompat" and custom_model_name:
+        model_string = custom_model_name
+    elif model_string == "custom_ollama" and custom_model_name:
+        model_string = f"omniparser+ollama_chat/{custom_model_name}"
+
+    # Create agent kwargs
+    agent_kwargs = {
+        "model": model_string,
+        "tools": [computer],
+        "only_n_most_recent_images": only_n_most_recent_images,
+        "verbosity": verbosity,
+    }
+
+    if save_trajectory:
+        agent_kwargs["trajectory_dir"] = "trajectories"
+
+    if max_trajectory_budget:
+        agent_kwargs["max_trajectory_budget"] = {
+            "max_budget": max_trajectory_budget,
+            "raise_error": True,
+        }
+
+    global_agent = ComputerAgent(**agent_kwargs)
+    return global_agent
+
+
+def launch_ui():
+    """Standalone function to launch the Gradio app."""
+    from cua_agent.ui.gradio.ui_components import create_gradio_ui
+
+    print("Starting Gradio app for Cua Agent...")
+    demo = create_gradio_ui()
+    demo.launch(share=False, inbrowser=True)
+
+
+if __name__ == "__main__":
+    launch_ui()

libs/python/agent/cua_agent/ui/gradio/ui_components.py

@@ -0,0 +1,897 @@
+"""
+UI Components for the Gradio interface
+"""
+
+import asyncio
+import json
+import logging
+import os
+import platform
+from pathlib import Path
+from typing import Any, Dict, List, Optional, cast
+
+import gradio as gr
+from gradio.components.chatbot import MetadataDict
+
+from .app import (
+    create_agent,
+    get_model_string,
+    get_ollama_models,
+    global_agent,
+    global_computer,
+    load_settings,
+    save_settings,
+)
+
+# Global messages array to maintain conversation history
+global_messages = []
+
+
+def create_gradio_ui() -> gr.Blocks:
+    """Create a Gradio UI for the Computer-Use Agent."""
+
+    # Load settings
+    saved_settings = load_settings()
+
+    # Check for API keys
+    openai_api_key = os.environ.get("OPENAI_API_KEY", "")
+    anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY", "")
+    cua_api_key = os.environ.get("CUA_API_KEY", "")
+
+    # Model choices
+    openai_models = ["OpenAI: Computer-Use Preview"]
+    anthropic_models = [
+        "Anthropic: Claude 4 Opus (20250514)",
+        "Anthropic: Claude 4 Sonnet (20250514)",
+        "Anthropic: Claude 3.7 Sonnet (20250219)",
+    ]
+    omni_models = [
+        "OMNI: OpenAI GPT-4o",
+        "OMNI: OpenAI GPT-4o mini",
+        "OMNI: Claude 3.7 Sonnet (20250219)",
+    ]
+
+    # Check if API keys are available
+    has_openai_key = bool(openai_api_key)
+    has_anthropic_key = bool(anthropic_api_key)
+    has_cua_key = bool(cua_api_key)
+
+    # Get Ollama models for OMNI
+    ollama_models = get_ollama_models()
+    if ollama_models:
+        omni_models += ollama_models
+
+    # Detect platform
+    is_mac = platform.system().lower() == "darwin"
+
+    # Format model choices
+    provider_to_models = {
+        "OPENAI": openai_models,
+        "ANTHROPIC": anthropic_models,
+        "OMNI": omni_models + ["Custom model (OpenAI compatible API)", "Custom model (ollama)"],
+        "UITARS": (
+            [
+                "huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B",
+            ]
+            if is_mac
+            else []
+        )
+        + ["Custom model (OpenAI compatible API)"],
+    }
+
+    # Apply saved settings
+    initial_loop = saved_settings.get("agent_loop", "OMNI")
+    available_models_for_loop = provider_to_models.get(initial_loop, [])
+    saved_model_choice = saved_settings.get("model_choice")
+    if saved_model_choice and saved_model_choice in available_models_for_loop:
+        initial_model = saved_model_choice
+    else:
+        if initial_loop == "OPENAI":
+            initial_model = openai_models[0] if openai_models else "No models available"
+        elif initial_loop == "ANTHROPIC":
+            initial_model = anthropic_models[0] if anthropic_models else "No models available"
+        else:  # OMNI
+            initial_model = (
+                omni_models[0] if omni_models else "Custom model (OpenAI compatible API)"
+            )
+
+    initial_custom_model = saved_settings.get("custom_model", "Qwen2.5-VL-7B-Instruct")
+    initial_provider_base_url = saved_settings.get("provider_base_url", "http://localhost:1234/v1")
+    initial_save_trajectory = saved_settings.get("save_trajectory", True)
+    initial_recent_images = saved_settings.get("recent_images", 3)
+
+    # Example prompts
+    example_messages = [
+        "Create a Python virtual environment, install pandas and matplotlib, then plot stock data",
+        "Open a PDF in Preview, add annotations, and save it as a compressed version",
+        "Open Safari, search for 'macOS automation tools', and save the first three results as bookmarks",
+        "Configure SSH keys and set up a connection to a remote server",
+    ]
+
+    def generate_python_code(
+        agent_loop_choice,
+        model_name,
+        tasks,
+        recent_images=3,
+        save_trajectory=True,
+        computer_os="linux",
+        computer_provider="cloud",
+        container_name="",
+        cua_cloud_api_key="",
+        max_budget=None,
+    ):
+        """Generate Python code for the current configuration and tasks."""
+        tasks_str = ""
+        for task in tasks:
+            if task and task.strip():
+                tasks_str += f'            "{task}",\n'
+
+        model_string = get_model_string(model_name, agent_loop_choice)
+
+        computer_args = []
+        if computer_os != "macos":
+            computer_args.append(f'os_type="{computer_os}"')
+        if computer_provider != "lume":
+            computer_args.append(f'provider_type="{computer_provider}"')
+        if container_name:
+            computer_args.append(f'name="{container_name}"')
+        if cua_cloud_api_key:
+            computer_args.append(f'api_key="{cua_cloud_api_key}"')
+
+        computer_args_str = ", ".join(computer_args)
+        if computer_args_str:
+            computer_args_str = f"({computer_args_str})"
+        else:
+            computer_args_str = "()"
+
+        code = f"""import asyncio
+try:
+    from computer import Computer
+except ImportError:
+    Computer = None  # type: ignore[assignment,misc]
+from cua_agent import ComputerAgent
+
+async def main():
+    async with Computer{computer_args_str} as computer:
+        agent = ComputerAgent(
+            model="{model_string}",
+            tools=[computer],
+            only_n_most_recent_images={recent_images},"""
+
+        if save_trajectory:
+            code += """
+            trajectory_dir="trajectories","""
+
+        if max_budget:
+            code += f"""
+            max_trajectory_budget={{"max_budget": {max_budget}, "raise_error": True}},"""
+
+        code += """
+        )
+        """
+
+        if tasks_str:
+            code += f"""
+        # Prompts for the computer-use agent
+        tasks = [
+{tasks_str.rstrip()}
+        ]
+
+        for task in tasks:
+            print(f"Executing task: {{task}}")
+            messages = [{{"role": "user", "content": task}}]
+            async for result in agent.run(messages):
+                for item in result["output"]:
+                    if item["type"] == "message":
+                        print(item["content"][0]["text"])"""
+        else:
+            code += """
+        # Execute a single task
+        task = "Search for information about Cua on GitHub"
+        print(f"Executing task: {task}")
+        messages = [{"role": "user", "content": task}]
+        async for result in agent.run(messages):
+            for item in result["output"]:
+                if item["type"] == "message":
+                    print(item["content"][0]["text"])"""
+
+        code += """
+
+if __name__ == "__main__":
+    asyncio.run(main())"""
+
+        return code
+
+    # Create the Gradio interface
+    with gr.Blocks(title="Computer-Use Agent") as demo:
+        with gr.Row():
+            # Left column for settings
+            with gr.Column(scale=1):
+                # Logo
+                gr.HTML("""
+                    <div style="display: flex; justify-content: center; margin-bottom: 0.5em">
+                        <img alt="Cua Logo" style="width: 80px;"
+                             src="https://github.com/trycua/cua/blob/main/img/logo_white.png?raw=true" />
+                    </div>
+                    """)
+
+                # Python code accordion
+                with gr.Accordion("Python Code", open=False):
+                    code_display = gr.Code(
+                        language="python",
+                        value=generate_python_code(initial_loop, "gpt-4o", []),
+                        interactive=False,
+                    )
+
+                with gr.Accordion("Computer Configuration", open=True):
+                    is_windows = platform.system().lower() == "windows"
+                    is_mac = platform.system().lower() == "darwin"
+
+                    providers = ["cloud", "localhost", "docker"]
+                    if is_mac:
+                        providers += ["lume"]
+                    if is_windows:
+                        providers += ["winsandbox"]
+
+                    # Remove unavailable options
+                    # MacOS is unavailable if Lume is not available
+                    # Windows is unavailable if Winsandbox is not available
+                    # Linux is always available
+                    # This should be removed once we support macOS and Windows on the cloud provider
+                    computer_choices = ["macos", "linux", "windows"]
+                    if not is_mac or "lume" not in providers:
+                        computer_choices.remove("macos")
+                    if not is_windows or "winsandbox" not in providers:
+                        computer_choices.remove("windows")
+
+                    computer_os = gr.Radio(
+                        choices=computer_choices,
+                        label="Operating System",
+                        value=computer_choices[0],
+                        info="Select the operating system for the computer",
+                    )
+
+                    computer_provider = gr.Radio(
+                        choices=providers,
+                        label="Provider",
+                        value="lume" if is_mac else "cloud",
+                        info="Select the computer provider",
+                    )
+
+                    container_name = gr.Textbox(
+                        label="Container Name",
+                        placeholder="Enter container name (optional)",
+                        value=os.environ.get("CUA_CONTAINER_NAME", ""),
+                        info="Optional name for the container",
+                    )
+
+                    cua_cloud_api_key = gr.Textbox(
+                        label="Cua Cloud API Key",
+                        placeholder="Enter your Cua Cloud API key",
+                        value=os.environ.get("CUA_API_KEY", ""),
+                        type="password",
+                        info="Required for cloud provider",
+                        visible=(not has_cua_key),
+                    )
+
+                with gr.Accordion("Agent Configuration", open=True):
+                    agent_loop = gr.Dropdown(
+                        choices=["OPENAI", "ANTHROPIC", "OMNI", "UITARS"],
+                        label="Agent Loop",
+                        value=initial_loop,
+                        info="Select the agent loop provider",
+                    )
+
+                    # Model selection dropdowns
+                    with gr.Group() as model_selection_group:
+                        openai_model_choice = gr.Dropdown(
+                            choices=openai_models,
+                            label="OpenAI Model",
+                            value=openai_models[0] if openai_models else "No models available",
+                            info="Select OpenAI model",
+                            interactive=True,
+                            visible=(initial_loop == "OPENAI"),
+                        )
+
+                        anthropic_model_choice = gr.Dropdown(
+                            choices=anthropic_models,
+                            label="Anthropic Model",
+                            value=(
+                                anthropic_models[0] if anthropic_models else "No models available"
+                            ),
+                            info="Select Anthropic model",
+                            interactive=True,
+                            visible=(initial_loop == "ANTHROPIC"),
+                        )
+
+                        omni_model_choice = gr.Dropdown(
+                            choices=omni_models
+                            + ["Custom model (OpenAI compatible API)", "Custom model (ollama)"],
+                            label="OMNI Model",
+                            value=(
+                                omni_models[0]
+                                if omni_models
+                                else "Custom model (OpenAI compatible API)"
+                            ),
+                            info="Select OMNI model or choose a custom model option",
+                            interactive=True,
+                            visible=(initial_loop == "OMNI"),
+                        )
+
+                        uitars_model_choice = gr.Dropdown(
+                            choices=provider_to_models.get("UITARS", ["No models available"]),
+                            label="UITARS Model",
+                            value=(
+                                provider_to_models.get("UITARS", ["No models available"])[0]
+                                if provider_to_models.get("UITARS")
+                                else "No models available"
+                            ),
+                            info="Select UITARS model",
+                            interactive=True,
+                            visible=(initial_loop == "UITARS"),
+                        )
+
+                        model_choice = gr.Textbox(visible=False)
+
+                    # API key inputs
+                    with gr.Group(
+                        visible=not has_openai_key
+                        and (initial_loop == "OPENAI" or initial_loop == "OMNI")
+                    ) as openai_key_group:
+                        openai_api_key_input = gr.Textbox(
+                            label="OpenAI API Key",
+                            placeholder="Enter your OpenAI API key",
+                            value=os.environ.get("OPENAI_API_KEY", ""),
+                            interactive=True,
+                            type="password",
+                            info="Required for OpenAI models",
+                        )
+
+                    with gr.Group(
+                        visible=not has_anthropic_key
+                        and (initial_loop == "ANTHROPIC" or initial_loop == "OMNI")
+                    ) as anthropic_key_group:
+                        anthropic_api_key_input = gr.Textbox(
+                            label="Anthropic API Key",
+                            placeholder="Enter your Anthropic API key",
+                            value=os.environ.get("ANTHROPIC_API_KEY", ""),
+                            interactive=True,
+                            type="password",
+                            info="Required for Anthropic models",
+                        )
+
+                    # API key handlers
+                    def set_openai_api_key(key):
+                        if key and key.strip():
+                            os.environ["OPENAI_API_KEY"] = key.strip()
+                            print("DEBUG - Set OpenAI API key environment variable")
+                        return key
+
+                    def set_anthropic_api_key(key):
+                        if key and key.strip():
+                            os.environ["ANTHROPIC_API_KEY"] = key.strip()
+                            print("DEBUG - Set Anthropic API key environment variable")
+                        return key
+
+                    openai_api_key_input.change(
+                        fn=set_openai_api_key,
+                        inputs=[openai_api_key_input],
+                        outputs=[openai_api_key_input],
+                        queue=False,
+                    )
+
+                    anthropic_api_key_input.change(
+                        fn=set_anthropic_api_key,
+                        inputs=[anthropic_api_key_input],
+                        outputs=[anthropic_api_key_input],
+                        queue=False,
+                    )
+
+                    # UI update function
+                    def update_ui(
+                        loop=None,
+                        openai_model=None,
+                        anthropic_model=None,
+                        omni_model=None,
+                        uitars_model=None,
+                    ):
+                        loop = loop or agent_loop.value
+
+                        model_value = None
+                        if loop == "OPENAI" and openai_model:
+                            model_value = openai_model
+                        elif loop == "ANTHROPIC" and anthropic_model:
+                            model_value = anthropic_model
+                        elif loop == "OMNI" and omni_model:
+                            model_value = omni_model
+                        elif loop == "UITARS" and uitars_model:
+                            model_value = uitars_model
+
+                        openai_visible = loop == "OPENAI"
+                        anthropic_visible = loop == "ANTHROPIC"
+                        omni_visible = loop == "OMNI"
+                        uitars_visible = loop == "UITARS"
+
+                        show_openai_key = not has_openai_key and (
+                            loop == "OPENAI"
+                            or (
+                                loop == "OMNI"
+                                and model_value
+                                and "OpenAI" in model_value
+                                and "Custom" not in model_value
+                            )
+                        )
+                        show_anthropic_key = not has_anthropic_key and (
+                            loop == "ANTHROPIC"
+                            or (
+                                loop == "OMNI"
+                                and model_value
+                                and "Claude" in model_value
+                                and "Custom" not in model_value
+                            )
+                        )
+
+                        is_custom_openai_api = model_value == "Custom model (OpenAI compatible API)"
+                        is_custom_ollama = model_value == "Custom model (ollama)"
+                        is_any_custom = is_custom_openai_api or is_custom_ollama
+
+                        model_choice_value = model_value if model_value else ""
+
+                        return [
+                            gr.update(visible=openai_visible),
+                            gr.update(visible=anthropic_visible),
+                            gr.update(visible=omni_visible),
+                            gr.update(visible=uitars_visible),
+                            gr.update(visible=show_openai_key),
+                            gr.update(visible=show_anthropic_key),
+                            gr.update(visible=is_any_custom),
+                            gr.update(visible=is_custom_openai_api),
+                            gr.update(visible=is_custom_openai_api),
+                            gr.update(value=model_choice_value),
+                        ]
+
+                    # Custom model inputs
+                    custom_model = gr.Textbox(
+                        label="Custom Model Name",
+                        placeholder="Enter custom model name (e.g., Qwen2.5-VL-7B-Instruct or llama3)",
+                        value=initial_custom_model,
+                        visible=(
+                            initial_model == "Custom model (OpenAI compatible API)"
+                            or initial_model == "Custom model (ollama)"
+                        ),
+                        interactive=True,
+                    )
+
+                    provider_base_url = gr.Textbox(
+                        label="Provider Base URL",
+                        placeholder="Enter provider base URL (e.g., http://localhost:1234/v1)",
+                        value=initial_provider_base_url,
+                        visible=(initial_model == "Custom model (OpenAI compatible API)"),
+                        interactive=True,
+                    )
+
+                    provider_api_key = gr.Textbox(
+                        label="Provider API Key",
+                        placeholder="Enter provider API key (if required)",
+                        value="",
+                        visible=(initial_model == "Custom model (OpenAI compatible API)"),
+                        interactive=True,
+                        type="password",
+                    )
+
+                    # Provider visibility update function
+                    def update_provider_visibility(provider):
+                        """Update visibility of container name and API key based on selected provider."""
+                        is_localhost = provider == "localhost"
+                        return [
+                            gr.update(visible=not is_localhost),  # container_name
+                            gr.update(
+                                visible=not is_localhost and not has_cua_key
+                            ),  # cua_cloud_api_key
+                        ]
+
+                    # Connect provider change event
+                    computer_provider.change(
+                        fn=update_provider_visibility,
+                        inputs=[computer_provider],
+                        outputs=[container_name, cua_cloud_api_key],
+                        queue=False,
+                    )
+
+                    # Connect UI update events
+                    for dropdown in [
+                        agent_loop,
+                        omni_model_choice,
+                        uitars_model_choice,
+                        openai_model_choice,
+                        anthropic_model_choice,
+                    ]:
+                        dropdown.change(
+                            fn=update_ui,
+                            inputs=[
+                                agent_loop,
+                                openai_model_choice,
+                                anthropic_model_choice,
+                                omni_model_choice,
+                                uitars_model_choice,
+                            ],
+                            outputs=[
+                                openai_model_choice,
+                                anthropic_model_choice,
+                                omni_model_choice,
+                                uitars_model_choice,
+                                openai_key_group,
+                                anthropic_key_group,
+                                custom_model,
+                                provider_base_url,
+                                provider_api_key,
+                                model_choice,
+                            ],
+                            queue=False,
+                        )
+
+                    save_trajectory = gr.Checkbox(
+                        label="Save Trajectory",
+                        value=initial_save_trajectory,
+                        info="Save the agent's trajectory for debugging",
+                        interactive=True,
+                    )
+
+                    recent_images = gr.Slider(
+                        label="Recent Images",
+                        minimum=1,
+                        maximum=10,
+                        value=initial_recent_images,
+                        step=1,
+                        info="Number of recent images to keep in context",
+                        interactive=True,
+                    )
+
+                    max_budget = gr.Number(
+                        label="Max Budget ($)",
+                        value=lambda: None,
+                        minimum=-1,
+                        maximum=100.0,
+                        step=0.1,
+                        info="Optional budget limit for trajectory (0 = no limit)",
+                        interactive=True,
+                    )
+
+            # Right column for chat interface
+            with gr.Column(scale=2):
+                gr.Markdown(
+                    "Ask me to perform tasks in a virtual environment.<br>Built with <a href='https://github.com/trycua/cua' target='_blank'>github.com/trycua/cua</a>."
+                )
+
+                chatbot_history = gr.Chatbot()
+                msg = gr.Textbox(placeholder="Ask me to perform tasks in a virtual environment")
+                clear = gr.Button("Clear")
+                cancel_button = gr.Button("Cancel", variant="stop")
+
+                # Add examples
+                example_group = gr.Examples(examples=example_messages, inputs=msg)
+
+                # Chat submission function
+                def chat_submit(message, history):
+                    history.append(gr.ChatMessage(role="user", content=message))
+                    return "", history
+
+                # Cancel function
+                async def cancel_agent_task(history):
+                    global global_agent
+                    if global_agent:
+                        print("DEBUG - Cancelling agent task")
+                        history.append(
+                            gr.ChatMessage(
+                                role="assistant",
+                                content="Task cancelled by user",
+                                metadata={"title": "❌ Cancelled"},
+                            )
+                        )
+                    else:
+                        history.append(
+                            gr.ChatMessage(
+                                role="assistant",
+                                content="No active agent task to cancel",
+                                metadata={"title": "ℹ️ Info"},
+                            )
+                        )
+                    return history
+
+                # Process response function
+                async def process_response(
+                    history,
+                    openai_model_value,
+                    anthropic_model_value,
+                    omni_model_value,
+                    uitars_model_value,
+                    custom_model_value,
+                    agent_loop_choice,
+                    save_traj,
+                    recent_imgs,
+                    custom_url_value=None,
+                    custom_api_key=None,
+                    openai_key_input=None,
+                    anthropic_key_input=None,
+                    computer_os="linux",
+                    computer_provider="cloud",
+                    container_name="",
+                    cua_cloud_api_key="",
+                    max_budget_value=None,
+                ):
+                    if not history:
+                        yield history
+                        return
+
+                    # Get the last user message
+                    last_user_message = history[-1]["content"]
+
+                    # Get the appropriate model value based on the agent loop
+                    if agent_loop_choice == "OPENAI":
+                        model_choice_value = openai_model_value
+                    elif agent_loop_choice == "ANTHROPIC":
+                        model_choice_value = anthropic_model_value
+                    elif agent_loop_choice == "OMNI":
+                        model_choice_value = omni_model_value
+                    elif agent_loop_choice == "UITARS":
+                        model_choice_value = uitars_model_value
+                    else:
+                        model_choice_value = "No models available"
+
+                    # Determine if this is a custom model selection
+                    is_custom_model_selected = model_choice_value in [
+                        "Custom model (OpenAI compatible API)",
+                        "Custom model (ollama)",
+                    ]
+
+                    # Determine the model name string to analyze
+                    if is_custom_model_selected:
+                        model_string_to_analyze = custom_model_value
+                    else:
+                        model_string_to_analyze = model_choice_value
+
+                    try:
+                        # Get the model string
+                        model_string = get_model_string(model_string_to_analyze, agent_loop_choice)
+
+                        # Set API keys if provided
+                        if openai_key_input:
+                            os.environ["OPENAI_API_KEY"] = openai_key_input
+                        if anthropic_key_input:
+                            os.environ["ANTHROPIC_API_KEY"] = anthropic_key_input
+                        if cua_cloud_api_key:
+                            os.environ["CUA_API_KEY"] = cua_cloud_api_key
+
+                        # Save settings
+                        current_settings = {
+                            "agent_loop": agent_loop_choice,
+                            "model_choice": model_choice_value,
+                            "custom_model": custom_model_value,
+                            "provider_base_url": custom_url_value,
+                            "save_trajectory": save_traj,
+                            "recent_images": recent_imgs,
+                            "computer_os": computer_os,
+                            "computer_provider": computer_provider,
+                            "container_name": container_name,
+                        }
+                        save_settings(current_settings)
+
+                        # Create agent
+                        global_agent = create_agent(
+                            model_string=model_string,
+                            save_trajectory=save_traj,
+                            only_n_most_recent_images=recent_imgs,
+                            custom_model_name=(
+                                custom_model_value if is_custom_model_selected else None
+                            ),
+                            computer_os=computer_os,
+                            computer_provider=computer_provider,
+                            computer_name=container_name,
+                            computer_api_key=cua_cloud_api_key,
+                            verbosity=logging.DEBUG,
+                            max_trajectory_budget=(
+                                max_budget_value
+                                if max_budget_value and max_budget_value > 0
+                                else None
+                            ),
+                        )
+
+                        if global_agent is None:
+                            history.append(
+                                gr.ChatMessage(
+                                    role="assistant",
+                                    content="Failed to create agent. Check API keys and configuration.",
+                                )
+                            )
+                            yield history
+                            return
+
+                        # Add user message to global history
+                        global global_messages
+                        global_messages.append({"role": "user", "content": last_user_message})
+
+                        # Stream responses from the agent
+                        async for result in global_agent.run(global_messages):
+                            global_messages += result.get("output", [])
+                            # print(f"DEBUG - Agent response ------- START")
+                            # from pprint import pprint
+                            # pprint(result)
+                            # print(f"DEBUG - Agent response ------- END")
+
+                            # Process the result output
+                            for item in result.get("output", []):
+                                if item.get("type") == "message":
+                                    content = item.get("content", [])
+                                    for content_part in content:
+                                        if content_part.get("text"):
+                                            history.append(
+                                                gr.ChatMessage(
+                                                    role=item.get("role", "assistant"),
+                                                    content=content_part.get("text", ""),
+                                                    metadata=content_part.get("metadata", {}),
+                                                )
+                                            )
+                                elif item.get("type") == "computer_call":
+                                    action = item.get("action", {})
+                                    action_type = action.get("type", "")
+                                    if action_type:
+                                        action_title = f"🛠️ Performing {action_type}"
+                                        if action.get("x") and action.get("y"):
+                                            action_title += f" at ({action['x']}, {action['y']})"
+                                        history.append(
+                                            gr.ChatMessage(
+                                                role="assistant",
+                                                content=f"```json\n{json.dumps(action)}\n```",
+                                                metadata={"title": action_title},
+                                            )
+                                        )
+                                elif item.get("type") == "function_call":
+                                    function_name = item.get("name", "")
+                                    arguments = item.get("arguments", "{}")
+                                    history.append(
+                                        gr.ChatMessage(
+                                            role="assistant",
+                                            content=f"🔧 Calling function: {function_name}\n```json\n{arguments}\n```",
+                                            metadata={"title": f"Function Call: {function_name}"},
+                                        )
+                                    )
+                                elif item.get("type") == "function_call_output":
+                                    output = item.get("output", "")
+                                    history.append(
+                                        gr.ChatMessage(
+                                            role="assistant",
+                                            content=f"📤 Function output:\n```\n{output}\n```",
+                                            metadata={"title": "Function Output"},
+                                        )
+                                    )
+                                elif item.get("type") == "computer_call_output":
+                                    output = item.get("output", {}).get("image_url", "")
+                                    image_markdown = f"![Computer output]({output})"
+                                    history.append(
+                                        gr.ChatMessage(
+                                            role="assistant",
+                                            content=image_markdown,
+                                            metadata={"title": "🖥️ Computer Output"},
+                                        )
+                                    )
+
+                            yield history
+
+                    except Exception as e:
+                        import traceback
+
+                        traceback.print_exc()
+                        history.append(gr.ChatMessage(role="assistant", content=f"Error: {str(e)}"))
+                        yield history
+
+                # Connect the submit button
+                submit_event = msg.submit(
+                    fn=chat_submit,
+                    inputs=[msg, chatbot_history],
+                    outputs=[msg, chatbot_history],
+                    queue=False,
+                ).then(
+                    fn=process_response,
+                    inputs=[
+                        chatbot_history,
+                        openai_model_choice,
+                        anthropic_model_choice,
+                        omni_model_choice,
+                        uitars_model_choice,
+                        custom_model,
+                        agent_loop,
+                        save_trajectory,
+                        recent_images,
+                        provider_base_url,
+                        provider_api_key,
+                        openai_api_key_input,
+                        anthropic_api_key_input,
+                        computer_os,
+                        computer_provider,
+                        container_name,
+                        cua_cloud_api_key,
+                        max_budget,
+                    ],
+                    outputs=[chatbot_history],
+                    queue=True,
+                )
+
+                # Clear button functionality
+                def clear_chat():
+                    global global_messages
+                    global_messages.clear()
+                    return None
+
+                clear.click(clear_chat, None, chatbot_history, queue=False)
+
+                # Connect cancel button
+                cancel_button.click(
+                    cancel_agent_task, [chatbot_history], [chatbot_history], queue=False
+                )
+
+                # Code display update function
+                def update_code_display(
+                    agent_loop,
+                    model_choice_val,
+                    custom_model_val,
+                    chat_history,
+                    recent_images_val,
+                    save_trajectory_val,
+                    computer_os,
+                    computer_provider,
+                    container_name,
+                    cua_cloud_api_key,
+                    max_budget_val,
+                ):
+                    messages = []
+                    if chat_history:
+                        for msg in chat_history:
+                            if isinstance(msg, dict) and msg.get("role") == "user":
+                                messages.append(msg.get("content", ""))
+
+                    return generate_python_code(
+                        agent_loop,
+                        model_choice_val or custom_model_val or "gpt-4o",
+                        messages,
+                        recent_images_val,
+                        save_trajectory_val,
+                        computer_os,
+                        computer_provider,
+                        container_name,
+                        cua_cloud_api_key,
+                        max_budget_val,
+                    )
+
+                # Update code display when configuration changes
+                for component in [
+                    agent_loop,
+                    model_choice,
+                    custom_model,
+                    chatbot_history,
+                    recent_images,
+                    save_trajectory,
+                    computer_os,
+                    computer_provider,
+                    container_name,
+                    cua_cloud_api_key,
+                    max_budget,
+                ]:
+                    component.change(
+                        update_code_display,
+                        inputs=[
+                            agent_loop,
+                            model_choice,
+                            custom_model,
+                            chatbot_history,
+                            recent_images,
+                            save_trajectory,
+                            computer_os,
+                            computer_provider,
+                            container_name,
+                            cua_cloud_api_key,
+                            max_budget,
+                        ],
+                        outputs=[code_display],
+                    )
+
+    return demo

libs/python/agent/example.py

@@ -0,0 +1,153 @@
+"""
+Example usage of the agent library with docstring-based tool definitions.
+"""
+
+import asyncio
+import logging
+
+from computer import Computer
+from computer.helpers import sandboxed
+from cua_agent import ComputerAgent
+
+
+@sandboxed()
+def read_file(location: str) -> str:
+    """Read contents of a file
+
+    Parameters
+    ----------
+    location : str
+        Path to the file to read
+
+    Returns
+    -------
+    str
+        Contents of the file or error message
+    """
+    try:
+        with open(location, "r") as f:
+            return f.read()
+    except Exception as e:
+        return f"Error reading file: {str(e)}"
+
+
+def save_note(content: str, filename: str = "note.txt") -> str:
+    """Save content to a note file
+
+    Parameters
+    ----------
+    content : str
+        Content to save to the file
+    filename : str, optional
+        Name of the file to save to (default is "note.txt")
+
+    Returns
+    -------
+    str
+        Success or error message
+    """
+    try:
+        with open(filename, "w") as f:
+            f.write(content)
+        return f"Saved note to {filename}"
+    except Exception as e:
+        return f"Error saving note: {str(e)}"
+
+
+def calculate(a: int, b: int) -> int:
+    """Calculate the sum of two integers
+
+    Parameters
+    ----------
+    a : int
+        First integer
+    b : int
+        Second integer
+
+    Returns
+    -------
+    int
+        Sum of the two integers
+    """
+    return a + b
+
+
+async def main():
+    """Example usage of ComputerAgent with different models"""
+
+    # Example 1: Using Claude with computer and custom tools
+    print("=== Example 1: Claude with Computer ===")
+
+    import json
+    import os
+
+    import dotenv
+
+    dotenv.load_dotenv()
+
+    assert os.getenv("CUA_CONTAINER_NAME") is not None, "CUA_CONTAINER_NAME is not set"
+    assert os.getenv("CUA_API_KEY") is not None, "CUA_API_KEY is not set"
+
+    async with Computer(
+        os_type="linux",
+        provider_type="cloud",
+        name=os.getenv("CUA_CONTAINER_NAME") or "",
+        api_key=os.getenv("CUA_API_KEY") or "",
+    ) as computer:
+        agent = ComputerAgent(
+            # Supported models:
+            # == OpenAI Cua (computer-use-preview) ==
+            model="openai/computer-use-preview",
+            # == Anthropic Cua (Claude > 3.5) ==
+            # model="anthropic/claude-opus-4-20250514",
+            # model="anthropic/claude-sonnet-4-20250514",
+            # model="anthropic/claude-3-7-sonnet-20250219",
+            # model="anthropic/claude-sonnet-4-5-20250929",
+            # == UI-TARS ==
+            # model="huggingface-local/ByteDance-Seed/UI-TARS-1.5-7B",
+            # TODO: add local mlx provider
+            # model="mlx-community/UI-TARS-1.5-7B-6bit",
+            # model="ollama_chat/0000/ui-tars-1.5-7b",
+            # == Omniparser + Any LLM ==
+            # model="omniparser+..."
+            # model="omniparser+anthropic/claude-opus-4-20250514",
+            tools=[computer],
+            only_n_most_recent_images=3,
+            verbosity=logging.INFO,
+            trajectory_dir="trajectories",
+            use_prompt_caching=True,
+            max_trajectory_budget={
+                "max_budget": 1.0,
+                "raise_error": True,
+                "reset_after_each_run": False,
+            },
+        )
+
+        history = []
+        while True:
+            user_input = input("> ")
+            history.append({"role": "user", "content": user_input})
+
+            # Non-streaming usage
+            async for result in agent.run(history, stream=False):
+                history += result["output"]
+
+                # # Print output
+                # for item in result["output"]:
+                #     if item["type"] == "message":
+                #         print(item["content"][0]["text"])
+                #     elif item["type"] == "computer_call":
+                #         action = item["action"]
+                #         action_type = action["type"]
+                #         action_args = {k: v for k, v in action.items() if k != "type"}
+                #         print(f"{action_type}({action_args})")
+                #     elif item["type"] == "function_call":
+                #         action = item["name"]
+                #         action_args = item["arguments"]
+                #         print(f"{action}({action_args})")
+                #     elif item["type"] == "function_call_output":
+                #         print("===>", item["output"])
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

libs/python/agent/pyproject.toml

@@ -0,0 +1,134 @@
+[build-system]
+requires = ["pdm-backend"]
+build-backend = "pdm.backend"
+
+[project]
+name = "cua-agent"
+version = "0.8.4"
+description = "Cua (Computer Use) Agent for AI-driven computer interaction"
+readme = "README.md"
+authors = [
+    { name = "TryCua", email = "gh@trycua.com" }
+]
+dependencies = [
+    "httpx>=0.27.0",
+    "aiohttp>=3.9.3",
+    "asyncio",
+    "anyio>=4.4.1",
+    "typing-extensions>=4.12.2",
+    "pydantic>=2.6.4",
+    "rich>=13.7.1",
+    "python-dotenv>=1.0.1",
+    "cua-core>=0.3.0,<0.4.0",
+    "certifi>=2024.2.2",
+    "litellm==1.86.2",
+    "Pillow>=10.0.0"
+]
+requires-python = ">=3.11,<3.14"
+
+[project.optional-dependencies]
+computer = [
+    "cua-computer>=0.5.0,<0.6.0",
+]
+openai = []
+anthropic = []
+qwen = [
+    "qwen-vl-utils",
+    "qwen-agent",
+    "Pillow>=10.0.0",
+]
+omni = [
+    "cua-som>=0.1.0,<0.2.0",
+]
+uitars = []
+uitars-mlx = [
+    "mlx-vlm>=0.1.27; sys_platform == 'darwin'"
+]
+uitars-hf = [
+    "accelerate",
+    "torch",
+    "transformers>=4.54.0"
+]
+glm45v-hf = [
+    "accelerate",
+    "torch",
+    "transformers>=5.0.0rc3"
+]
+opencua-hf = [
+    "accelerate",
+    "torch",
+    "transformers>=4.53.0",
+    "tiktoken>=0.11.0",
+    "blobfile>=3.0.0"
+]
+internvl-hf = [
+    "accelerate",
+    "torch",
+    "transformers>=4.55.0",
+    "einops",
+    "timm"
+]
+moondream3 = [
+    "accelerate",
+    "torch",
+    "transformers>=4.55.0"
+]
+ui = [
+    "gradio>=6.0.0",
+    "python-dotenv>=1.0.1",
+]
+cli = [
+    "yaspin>=3.1.0",
+]
+hud = [
+    "hud-python==0.4.52",
+    "fastmcp>=3.2.0,<3.3.0",
+]
+gemini = [
+    "google-genai>=1.41.0",
+]
+# API-only extras for cloud/container deployments — excludes torch/transformers/local models.
+# Supports: openai, anthropic, gemini, uitars (API mode).
+cloud = [
+    # gemini API SDK
+    "google-genai>=1.41.0",
+    # cli spinner
+    "yaspin>=3.1.0",
+    # qwen tool-call formatting (used by generic_vlm.py for API-based Qwen calls)
+    "qwen-agent",
+    "qwen-vl-utils",
+    "Pillow>=10.0.0",
+]
+all = [
+    # uitars requirements
+    "mlx-vlm>=0.1.27; sys_platform == 'darwin'",
+    "accelerate",
+    "torch",
+    "transformers>=4.55.0",
+    # internvl requirements,
+    "einops",
+    "timm",
+    # opencua requirements
+    "tiktoken>=0.11.0",
+    "blobfile>=3.0.0",
+    # ui requirements
+    "gradio>=6.0.0",
+    "python-dotenv>=1.0.1",
+    # cli requirements
+    "yaspin>=3.1.0",
+    # gemini requirements
+    "google-genai>=1.41.0",
+    # qwen requirements
+    "qwen-vl-utils",
+    "qwen-agent",
+    "Pillow>=10.0.0",
+]
+
+[tool.uv]
+constraint-dependencies = ["fastrtc>0.43.0", "mlx-audio>0.2.3"]
+
+[tool.pdm]
+distribution = true
+
+[tool.pdm.build]
+includes = ["cua_agent/"]

libs/python/agent/tests/conftest.py

@@ -0,0 +1,84 @@
+"""Pytest configuration and shared fixtures for agent package tests.
+
+This file contains shared fixtures and configuration for all agent tests.
+Following SRP: This file ONLY handles test setup/teardown.
+"""
+
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+
+import pytest
+
+
+@pytest.fixture
+def mock_litellm():
+    """Mock liteLLM completion calls.
+
+    Use this fixture to avoid making real LLM API calls during tests.
+    Returns a mock that simulates LLM responses.
+    """
+    with patch("litellm.acompletion") as mock_completion:
+
+        async def mock_response(*args, **kwargs):
+            """Simulate a typical LLM response."""
+            return {
+                "id": "chatcmpl-test123",
+                "object": "chat.completion",
+                "created": 1234567890,
+                "model": kwargs.get("model", "anthropic/claude-sonnet-4-5-20250929"),
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": "This is a mocked response for testing.",
+                        },
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 10,
+                    "completion_tokens": 20,
+                    "total_tokens": 30,
+                },
+            }
+
+        mock_completion.side_effect = mock_response
+        yield mock_completion
+
+
+@pytest.fixture
+def mock_computer():
+    """Mock Computer interface for agent tests.
+
+    Use this fixture to test agent logic without requiring a real Computer instance.
+    """
+    computer = AsyncMock()
+    computer.interface = AsyncMock()
+    computer.interface.screenshot = AsyncMock(return_value=b"fake_screenshot_data")
+    computer.interface.left_click = AsyncMock()
+    computer.interface.type = AsyncMock()
+    computer.interface.key = AsyncMock()
+
+    # Mock context manager
+    computer.__aenter__ = AsyncMock(return_value=computer)
+    computer.__aexit__ = AsyncMock()
+
+    return computer
+
+
+@pytest.fixture
+def disable_telemetry(monkeypatch):
+    """Disable telemetry for tests.
+
+    Use this fixture to ensure no telemetry is sent during tests.
+    """
+    monkeypatch.setenv("CUA_TELEMETRY_ENABLED", "false")
+
+
+@pytest.fixture
+def sample_messages():
+    """Provide sample messages for testing.
+
+    Returns a list of messages in the expected format.
+    """
+    return [{"role": "user", "content": "Take a screenshot and tell me what you see"}]

libs/python/agent/tests/test_computer_agent.py

@@ -0,0 +1,139 @@
+"""Unit tests for ComputerAgent class.
+
+This file tests ONLY the ComputerAgent initialization and basic functionality.
+Following SRP: This file tests ONE class (ComputerAgent).
+All external dependencies (liteLLM, Computer) are mocked.
+"""
+
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+
+import pytest
+
+
+class TestComputerAgentInitialization:
+    """Test ComputerAgent initialization (SRP: Only tests initialization)."""
+
+    @patch("cua_agent.agent.litellm")
+    def test_agent_initialization_with_model(self, mock_litellm, disable_telemetry):
+        """Test that agent can be initialized with a model string."""
+        from cua_agent import ComputerAgent
+
+        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929")
+
+        assert agent is not None
+        assert hasattr(agent, "model")
+        assert agent.model == "anthropic/claude-sonnet-4-5-20250929"
+
+    @patch("cua_agent.agent.litellm")
+    def test_agent_initialization_with_tools(self, mock_litellm, disable_telemetry, mock_computer):
+        """Test that agent can be initialized with tools."""
+        from cua_agent import ComputerAgent
+
+        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929", tools=[mock_computer])
+
+        assert agent is not None
+        assert hasattr(agent, "tools")
+
+    @patch("cua_agent.agent.litellm")
+    def test_agent_initialization_with_max_budget(self, mock_litellm, disable_telemetry):
+        """Test that agent can be initialized with max trajectory budget."""
+        from cua_agent import ComputerAgent
+
+        budget = 5.0
+        agent = ComputerAgent(
+            model="anthropic/claude-sonnet-4-5-20250929", max_trajectory_budget=budget
+        )
+
+        assert agent is not None
+
+    @patch("cua_agent.agent.litellm")
+    def test_agent_requires_model(self, mock_litellm, disable_telemetry):
+        """Test that agent requires a model parameter."""
+        from cua_agent import ComputerAgent
+
+        with pytest.raises(TypeError):
+            # Should fail without model parameter - intentionally missing required argument
+            ComputerAgent()  # type: ignore[call-arg]
+
+
+class TestComputerAgentRun:
+    """Test ComputerAgent.run() method (SRP: Only tests run logic)."""
+
+    @pytest.mark.asyncio
+    @patch("cua_agent.agent.litellm")
+    async def test_agent_run_with_messages(self, mock_litellm, disable_telemetry, sample_messages):
+        """Test that agent.run() works with valid messages."""
+        from cua_agent import ComputerAgent
+
+        # Mock liteLLM response
+        mock_response = {
+            "id": "chatcmpl-test",
+            "choices": [
+                {
+                    "message": {"role": "assistant", "content": "Test response"},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
+        }
+
+        mock_litellm.acompletion = AsyncMock(return_value=mock_response)
+
+        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929")
+
+        # Run should return an async generator
+        result_generator = agent.run(sample_messages)
+
+        assert result_generator is not None
+        # Check it's an async generator
+        assert hasattr(result_generator, "__anext__")
+
+    def test_agent_has_run_method(self, disable_telemetry):
+        """Test that agent has run method available."""
+        from cua_agent import ComputerAgent
+
+        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929")
+
+        # Verify run method exists
+        assert hasattr(agent, "run")
+        assert callable(agent.run)
+
+    def test_agent_has_agent_loop(self, disable_telemetry):
+        """Test that agent has agent_loop initialized."""
+        from cua_agent import ComputerAgent
+
+        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929")
+
+        # Verify agent_loop is initialized
+        assert hasattr(agent, "agent_loop")
+        assert agent.agent_loop is not None
+
+
+class TestComputerAgentTypes:
+    """Test AgentResponse and Messages types (SRP: Only tests type definitions)."""
+
+    def test_messages_type_exists(self):
+        """Test that Messages type is exported."""
+        from cua_agent import Messages
+
+        assert Messages is not None
+
+    def test_agent_response_type_exists(self):
+        """Test that AgentResponse type is exported."""
+        from cua_agent import AgentResponse
+
+        assert AgentResponse is not None
+
+
+class TestComputerAgentIntegration:
+    """Test ComputerAgent integration with Computer tool (SRP: Integration within package)."""
+
+    def test_agent_accepts_computer_tool(self, disable_telemetry, mock_computer):
+        """Test that agent can be initialized with Computer tool."""
+        from cua_agent import ComputerAgent
+
+        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929", tools=[mock_computer])
+
+        # Verify agent accepted the tool
+        assert agent is not None
+        assert hasattr(agent, "tools")

libs/python/agent/tests/test_predict_click_zero_coords.py

@@ -0,0 +1,96 @@
+"""Tests for predict_click zero-coordinate fix (issue #1400).
+
+These tests verify the coordinate-extraction logic in
+AnthropicHostedToolsConfig.predict_click directly, without requiring the full
+cua_agent import chain (which needs cua-computer, cua-core, etc.).
+"""
+
+import pytest
+
+
+def _extract_click_coords_old(responses_items):
+    """Buggy implementation: falsy check silently drops x=0 or y=0."""
+    for item in responses_items:
+        if (
+            isinstance(item, dict)
+            and item.get("type") == "computer_call"
+            and isinstance(item.get("action"), dict)
+        ):
+            action = item["action"]
+            if action.get("x") and action.get("y"):  # BUG: 0 is falsy
+                return (int(action.get("x")), int(action.get("y")))
+    return None
+
+
+def _extract_click_coords_fixed(responses_items):
+    """Fixed implementation: explicit None check preserves zero coordinates."""
+    for item in responses_items:
+        if (
+            isinstance(item, dict)
+            and item.get("type") == "computer_call"
+            and isinstance(item.get("action"), dict)
+        ):
+            action = item["action"]
+            if action.get("x") is not None and action.get("y") is not None:
+                return (int(action.get("x")), int(action.get("y")))
+    return None
+
+
+def _make_items(x, y):
+    return [{"type": "computer_call", "action": {"type": "click", "x": x, "y": y}}]
+
+
+# --- Regression tests: these all FAIL with the old code, PASS with the fix ---
+
+
+def test_zero_x_was_broken_before_fix():
+    """Old code returns None for x=0; new code returns (0, y)."""
+    items = _make_items(0, 100)
+    assert _extract_click_coords_old(items) is None, "confirm old bug"
+    assert _extract_click_coords_fixed(items) == (0, 100)
+
+
+def test_zero_y_was_broken_before_fix():
+    """Old code returns None for y=0; new code returns (x, 0)."""
+    items = _make_items(200, 0)
+    assert _extract_click_coords_old(items) is None, "confirm old bug"
+    assert _extract_click_coords_fixed(items) == (200, 0)
+
+
+def test_zero_zero_was_broken_before_fix():
+    """Old code returns None for (0, 0); new code returns (0, 0)."""
+    items = _make_items(0, 0)
+    assert _extract_click_coords_old(items) is None, "confirm old bug"
+    assert _extract_click_coords_fixed(items) == (0, 0)
+
+
+# --- Positive tests: non-zero coordinates work in both old and new code ---
+
+
+def test_nonzero_coordinates_still_work():
+    items = _make_items(512, 384)
+    assert _extract_click_coords_fixed(items) == (512, 384)
+
+
+def test_returns_none_when_no_computer_call():
+    items = [{"type": "text", "text": "no click"}]
+    assert _extract_click_coords_fixed(items) is None
+
+
+# --- Verify the actual fix is present in the source file ---
+
+
+def test_source_uses_is_not_none_check():
+    """Confirm the fix is applied in the real anthropic.py source."""
+    import pathlib
+
+    src = (
+        pathlib.Path(__file__).parent.parent / "cua_agent" / "loops" / "anthropic.py"
+    ).read_text()
+    assert (
+        'action.get("x") is not None and action.get("y") is not None' in src
+    ), "Fix not found in anthropic.py — the 'is not None' check is missing"
+    # Ensure the old buggy pattern is gone
+    assert (
+        'if action.get("x") and action.get("y"):' not in src
+    ), "Old buggy truthiness check still present in anthropic.py"

libs/python/agent/tests/test_telemetry_events.py

@@ -0,0 +1,145 @@
+"""
+Test script to verify telemetry events are emitted correctly.
+"""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+
+class TestAgentTelemetryEvents:
+    """Test telemetry events emitted by ComputerAgent."""
+
+    @patch("cua_agent.agent.record_event")
+    @patch("cua_agent.agent.is_telemetry_enabled", return_value=True)
+    def test_agent_init_event(self, mock_telemetry_enabled, mock_record_event):
+        """Test that agent_init event is emitted with correct args_provided."""
+        from cua_agent.agent import ComputerAgent
+
+        # Create agent with various args
+        agent = ComputerAgent(
+            model="anthropic/claude-sonnet-4-5-20250929",
+            instructions="Test instructions",
+            max_retries=5,  # non-default
+            trajectory_dir="/tmp/test",
+        )
+
+        # Find the agent_init call
+        agent_init_calls = [
+            call for call in mock_record_event.call_args_list if call[0][0] == "agent_init"
+        ]
+
+        assert len(agent_init_calls) == 1, "agent_init should be called once"
+
+        event_name, event_data = agent_init_calls[0][0]
+        assert event_name == "agent_init"
+        assert event_data["model"] == "anthropic/claude-sonnet-4-5-20250929"
+        assert "instructions" in event_data["args_provided"]
+        assert "max_retries" in event_data["args_provided"]
+        assert "trajectory_dir" in event_data["args_provided"]
+
+    @patch("cua_agent.agent.record_event")
+    @patch("cua_agent.agent.is_telemetry_enabled", return_value=True)
+    def test_agent_init_minimal_args(self, mock_telemetry_enabled, mock_record_event):
+        """Test agent_init with minimal args (defaults)."""
+        from cua_agent.agent import ComputerAgent
+
+        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929")
+
+        agent_init_calls = [
+            call for call in mock_record_event.call_args_list if call[0][0] == "agent_init"
+        ]
+
+        assert len(agent_init_calls) == 1
+        event_name, event_data = agent_init_calls[0][0]
+
+        # With defaults, only model-related things should be tracked
+        # instructions, trajectory_dir, etc. should NOT be in args_provided
+        assert "instructions" not in event_data["args_provided"]
+        assert "trajectory_dir" not in event_data["args_provided"]
+        assert "max_retries" not in event_data["args_provided"]  # default is 3
+
+    @patch("cua_agent.agent.record_event")
+    @patch("cua_agent.agent.is_telemetry_enabled", return_value=False)
+    def test_no_events_when_telemetry_disabled(self, mock_telemetry_enabled, mock_record_event):
+        """Test that no events are emitted when telemetry is disabled."""
+        from cua_agent.agent import ComputerAgent
+
+        agent = ComputerAgent(
+            model="anthropic/claude-sonnet-4-5-20250929",
+            telemetry_enabled=False,
+        )
+
+        # No agent_init should be called (telemetry disabled)
+        agent_init_calls = [
+            call for call in mock_record_event.call_args_list if call[0][0] == "agent_init"
+        ]
+
+        assert len(agent_init_calls) == 0
+
+
+class TestActionTelemetryEvents:
+    """Test telemetry events for computer actions."""
+
+    @pytest.mark.asyncio
+    @patch("cua_agent.agent.record_event")
+    @patch("cua_agent.agent.is_telemetry_enabled", return_value=True)
+    async def test_computer_action_executed_event(self, mock_telemetry_enabled, mock_record_event):
+        """Test that computer_action_executed is emitted for computer calls."""
+        from cua_agent.agent import ComputerAgent
+
+        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929")
+        agent.telemetry_enabled = True
+
+        # Mock computer handler
+        mock_computer = MagicMock()
+        mock_computer.click = AsyncMock(return_value=None)
+        mock_computer.screenshot = AsyncMock(return_value="base64screenshot")
+
+        # Create a mock computer_call item
+        item = {
+            "type": "computer_call",
+            "call_id": "test-call-id",
+            "action": {
+                "type": "click",
+                "x": 100,
+                "y": 200,
+            },
+        }
+
+        # Process the item (this would normally happen in the agent loop)
+        # Note: We can't easily test this without running the full agent loop
+        # This is more of an integration test
+
+        # For unit testing, we verify the event structure
+        expected_event = {
+            "action_type": "click",
+        }
+
+        # Verify event structure is correct
+        assert "action_type" in expected_event
+
+
+class TestToolExecutedEvents:
+    """Test telemetry events for tool execution."""
+
+    def test_event_structure(self):
+        """Test that agent_tool_executed event has correct structure."""
+        expected_computer_tool_event = {
+            "tool_type": "computer",
+            "tool_name": "click",
+        }
+
+        expected_function_tool_event = {
+            "tool_type": "function",
+            "tool_name": "my_custom_function",
+        }
+
+        # Verify expected structure
+        assert "tool_type" in expected_computer_tool_event
+        assert "tool_name" in expected_computer_tool_event
+        assert expected_computer_tool_event["tool_type"] in ["computer", "function"]
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

libs/python/agent/tests/test_tool_resolution.py

@@ -0,0 +1,262 @@
+"""
+Tests for centralized tool resolution in ComputerAgent.
+
+Tests that:
+1. FARA (specialized model) auto-wraps Computer to BrowserTool with warning
+2. FARA accepts explicit BrowserTool without warning
+3. Claude (general model) accepts any tool without wrapping
+4. Custom function tools pass through unchanged
+"""
+
+import warnings
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from cua_agent.computers import is_agent_computer
+from cua_agent.tools.browser_tool import BrowserTool
+from cua_agent.types import AgentConfigInfo
+
+
+# Mock agent config class for testing
+class MockAgentConfig:
+    """Mock agent config class for testing"""
+
+    async def predict_step(self, **kwargs):
+        return {"output": [], "usage": {}}
+
+    async def predict_click(self, **kwargs):
+        return None
+
+    def get_capabilities(self):
+        return ["step"]
+
+
+class TestToolResolution:
+    """Tests for ComputerAgent._resolve_tools()"""
+
+    @pytest.fixture
+    def mock_computer(self):
+        """Create a mock Computer object"""
+        computer = Mock()
+        computer.interface = Mock()
+        computer.interface.interface = Mock()  # For hotkey, etc.
+        computer.interface.playwright_exec = Mock()
+        return computer
+
+    @pytest.fixture
+    def mock_browser_tool(self):
+        """Create a mock BrowserTool"""
+        tool = Mock(spec=BrowserTool)
+        return tool
+
+    @pytest.mark.asyncio
+    async def test_fara_auto_wraps_computer_to_browser_tool(self, mock_computer):
+        """FARA auto-wraps Computer to BrowserTool with warning."""
+        from cua_agent.agent import ComputerAgent
+
+        # Patch find_agent_config to return FARA config
+        fara_config = AgentConfigInfo(
+            agent_class=MockAgentConfig,
+            models_regex=r"(?i).*fara.*",
+            priority=0,
+            tool_type="browser",
+        )
+
+        # Patch is_agent_computer to recognize our mock as a computer
+        def mock_is_agent_computer(tool):
+            return tool is mock_computer
+
+        with patch("cua_agent.agent.find_agent_config", return_value=fara_config):
+            with patch("cua_agent.agent.is_agent_computer", side_effect=mock_is_agent_computer):
+                agent = ComputerAgent(model="cua/microsoft/fara-7b", tools=[mock_computer])
+                # Tool resolution happens in _resolve_tools, which is async
+                with pytest.warns(UserWarning, match="Auto-wrapping Computer to BrowserTool"):
+                    resolved = await agent._resolve_tools(agent.tools, "browser")
+
+        # Should have wrapped to BrowserTool
+        assert len(resolved) == 1
+        assert isinstance(resolved[0], BrowserTool)
+
+    def test_fara_accepts_explicit_browser_tool_no_warning(self, mock_browser_tool):
+        """FARA accepts explicit BrowserTool without warning."""
+        from cua_agent.agent import ComputerAgent
+
+        # Patch find_agent_config to return FARA config
+        fara_config = AgentConfigInfo(
+            agent_class=MockAgentConfig,
+            models_regex=r"(?i).*fara.*",
+            priority=0,
+            tool_type="browser",
+        )
+
+        with patch("cua_agent.agent.find_agent_config", return_value=fara_config):
+            # Should not raise any warnings
+            with warnings.catch_warnings():
+                warnings.simplefilter("error")
+                agent = ComputerAgent(model="cua/microsoft/fara-7b", tools=[mock_browser_tool])
+
+        # Should keep the original BrowserTool
+        assert len(agent.tools) == 1
+        assert agent.tools[0] is mock_browser_tool
+
+    def test_claude_accepts_any_tool_no_wrapping(self, mock_computer):
+        """Claude (general model) accepts any tool without wrapping."""
+        from cua_agent.agent import ComputerAgent
+
+        # Patch find_agent_config to return Claude config (no tool_type)
+        claude_config = AgentConfigInfo(
+            agent_class=MockAgentConfig,
+            models_regex=r".*claude.*",
+            priority=0,
+            tool_type=None,  # General model, no tool type requirement
+        )
+
+        with patch("cua_agent.agent.find_agent_config", return_value=claude_config):
+            with warnings.catch_warnings():
+                warnings.simplefilter("error")  # Fail if any warning
+                agent = ComputerAgent(model="claude-sonnet-4", tools=[mock_computer])
+
+        # Should keep original Computer, not wrapped
+        assert len(agent.tools) == 1
+        assert agent.tools[0] is mock_computer
+
+    @pytest.mark.asyncio
+    async def test_custom_tools_pass_through(self, mock_computer):
+        """Custom function tools pass through unchanged."""
+        from cua_agent.agent import ComputerAgent
+
+        def custom_tool():
+            """A custom tool function."""
+            pass
+
+        # Patch find_agent_config to return FARA config
+        fara_config = AgentConfigInfo(
+            agent_class=MockAgentConfig,
+            models_regex=r"(?i).*fara.*",
+            priority=0,
+            tool_type="browser",
+        )
+
+        # Patch is_agent_computer to recognize our mock as a computer
+        def mock_is_agent_computer(tool):
+            return tool is mock_computer
+
+        with patch("cua_agent.agent.find_agent_config", return_value=fara_config):
+            with patch("cua_agent.agent.is_agent_computer", side_effect=mock_is_agent_computer):
+                agent = ComputerAgent(
+                    model="cua/microsoft/fara-7b", tools=[mock_computer, custom_tool]
+                )
+                # Tool resolution happens in _resolve_tools, which is async
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore")
+                    resolved = await agent._resolve_tools(agent.tools, "browser")
+
+        # Should have wrapped computer but kept custom tool unchanged
+        assert len(resolved) == 2
+        assert isinstance(resolved[0], BrowserTool)
+        assert resolved[1] is custom_tool
+
+
+class TestAgentConfigInfo:
+    """Tests for AgentConfigInfo with tool_type field"""
+
+    def test_agent_config_info_with_tool_type(self):
+        """AgentConfigInfo accepts tool_type parameter"""
+        config = AgentConfigInfo(
+            agent_class=MockAgentConfig,
+            models_regex=r".*fara.*",
+            priority=0,
+            tool_type="browser",
+        )
+        assert config.tool_type == "browser"
+
+    def test_agent_config_info_without_tool_type(self):
+        """AgentConfigInfo defaults tool_type to None"""
+        config = AgentConfigInfo(
+            agent_class=MockAgentConfig,
+            models_regex=r".*claude.*",
+            priority=0,
+        )
+        assert config.tool_type is None
+
+    def test_agent_config_info_matches_model(self):
+        """AgentConfigInfo.matches_model works correctly"""
+        config = AgentConfigInfo(
+            agent_class=MockAgentConfig,
+            models_regex=r"(?i).*fara-7b.*",
+            priority=0,
+            tool_type="browser",
+        )
+        assert config.matches_model("cua/microsoft/fara-7b")
+        assert config.matches_model("FARA-7B")
+        assert not config.matches_model("claude-sonnet-4")
+
+
+class TestRegisterAgentDecorator:
+    """Tests for @register_agent decorator with tool_type"""
+
+    def test_register_agent_with_tool_type(self):
+        """@register_agent accepts tool_type parameter"""
+        from cua_agent.decorators import _agent_configs, register_agent
+
+        # Clear registry for test
+        original_configs = _agent_configs.copy()
+        _agent_configs.clear()
+
+        try:
+
+            @register_agent(models=r"test-model.*", tool_type="browser")
+            class TestAgentConfig:
+                async def predict_step(self, **kwargs):
+                    pass
+
+                async def predict_click(self, **kwargs):
+                    pass
+
+                def get_capabilities(self):
+                    return ["step"]
+
+            # Find the registered config
+            from cua_agent.decorators import find_agent_config
+
+            config = find_agent_config("test-model-123")
+            assert config is not None
+            assert config.tool_type == "browser"
+
+        finally:
+            # Restore original registry
+            _agent_configs.clear()
+            _agent_configs.extend(original_configs)
+
+    def test_register_agent_without_tool_type(self):
+        """@register_agent without tool_type defaults to None"""
+        from cua_agent.decorators import _agent_configs, register_agent
+
+        # Clear registry for test
+        original_configs = _agent_configs.copy()
+        _agent_configs.clear()
+
+        try:
+
+            @register_agent(models=r"general-model.*")
+            class GeneralAgentConfig:
+                async def predict_step(self, **kwargs):
+                    pass
+
+                async def predict_click(self, **kwargs):
+                    pass
+
+                def get_capabilities(self):
+                    return ["step"]
+
+            # Find the registered config
+            from cua_agent.decorators import find_agent_config
+
+            config = find_agent_config("general-model-123")
+            assert config is not None
+            assert config.tool_type is None
+
+        finally:
+            # Restore original registry
+            _agent_configs.clear()
+            _agent_configs.extend(original_configs)

libs/python/bench-ui/.bumpversion.cfg

@@ -0,0 +1,10 @@
+[bumpversion]
+current_version = 0.7.0
+commit = True
+tag = True
+tag_name = bench-ui-v{new_version}
+message = chore: bump cua-bench-ui version to {new_version}
+
+[bumpversion:file:pyproject.toml]
+search = version = "{current_version}"
+replace = version = "{new_version}"