diff --git a/docs/content/docs/cuabench/guide/fundamentals/agent-traces.mdx b/docs/content/docs/cuabench/guide/fundamentals/agent-traces.mdx index 5582dc828a..70ca0d2924 100644 --- a/docs/content/docs/cuabench/guide/fundamentals/agent-traces.mdx +++ b/docs/content/docs/cuabench/guide/fundamentals/agent-traces.mdx @@ -96,21 +96,17 @@ Opens an interactive viewer showing: - Action details and metadata - Agent reasoning (if available) -### View All Traces in a Run +### View Agent Trajectories from a Run ```bash -# View all task traces in a grid -cb trace grid +# Open agent trajectories in the cua.ai trajectory viewer +cb trace traj # Example -cb trace grid 96d41b51 +cb trace traj 96d41b51 ``` -Shows a grid view of all task traces in the run, useful for: - -- Comparing agent performance across variants -- Quickly identifying failures -- Reviewing oracle solutions +Zips and serves the cua-agent trajectory recordings from the run artifacts folder, then opens [cua.ai/trajectory-viewer](https://cua.ai/trajectory-viewer). For runs with multiple tasks, an index page is shown so you can open each session individually. ## Recording Custom Events diff --git a/docs/content/docs/cuabench/reference/cli-reference.mdx b/docs/content/docs/cuabench/reference/cli-reference.mdx index 42e838cdbd..78e8d9aec5 100644 --- a/docs/content/docs/cuabench/reference/cli-reference.mdx +++ b/docs/content/docs/cuabench/reference/cli-reference.mdx @@ -251,10 +251,10 @@ View and manage trace datasets. cb trace ``` -| Command | Description | -| -------------------- | ------------------------------------------------------------- | -| `cb trace view ` | View a single trace in browser (accepts run_id or session_id) | -| `cb trace grid ` | View all traces in a run as a grid (accepts run_id) | +| Command | Description | +| -------------------- | -------------------------------------------------------------------- | +| `cb trace view ` | View a single trace in browser (accepts run_id or session_id) | +| `cb trace traj ` | View agent trajectories from a run in cua.ai/trajectory-viewer | ## Dataset Commands diff --git a/libs/cua-bench/cua_bench/cli/commands/trace.py b/libs/cua-bench/cua_bench/cli/commands/trace.py index fb1f7a11c0..f526812000 100644 --- a/libs/cua-bench/cua_bench/cli/commands/trace.py +++ b/libs/cua-bench/cua_bench/cli/commands/trace.py @@ -2,7 +2,7 @@ Usage: cb trace view # View a single trace (run_id or session_id) - cb trace grid # View all traces in a run as a grid + cb trace traj # View agent trajectories from a run in cua.ai/trajectory-viewer """ from __future__ import annotations @@ -10,9 +10,12 @@ import base64 import json import os +import shutil +import tempfile import threading import urllib.parse import webbrowser +import zipfile as _zipfile from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any, List, Optional, Tuple @@ -125,9 +128,18 @@ def register_parser(subparsers): help='Run ID or session ID (e.g., "30c12572" or "task-30c12572-click-button-v0")', ) - # cb trace grid - grid_parser = trace_subparsers.add_parser("grid", help="View all traces in a run as a grid") - grid_parser.add_argument("identifier", help='Run ID (e.g., "30c12572")') + # cb trace traj + traj_parser = trace_subparsers.add_parser( + "traj", help="View agent trajectories from a run in cua.ai/trajectory-viewer" + ) + traj_parser.add_argument("identifier", help='Run ID (e.g., "30c12572")') + traj_parser.add_argument( + "--port", + "-p", + type=int, + default=8090, + help="Local file server port (default: 8090)", + ) def execute(args): @@ -136,13 +148,13 @@ def execute(args): if trace_command == "view": return cmd_view(args) - elif trace_command == "grid": - return cmd_grid(args) + elif trace_command == "traj": + return cmd_traj(args) else: print(f"{YELLOW}Usage: cb trace {RESET}") print(f"\n{GREY}Commands:{RESET}") print(" view View a single trace (run_id or session_id)") - print(" grid View all traces in a run as a grid") + print(" traj View agent trajectories in cua.ai/trajectory-viewer") return 1 @@ -337,242 +349,187 @@ def log_message(self, format, *args): # ============================================================================ -# cmd_grid - View multiple traces in a grid +# cmd_traj - View agent trajectories via cua.ai/trajectory-viewer # ============================================================================ -def cmd_grid(args) -> int: - """View all traces in a run as a grid layout.""" - identifier = args.identifier +def _collect_run_trajectories(run_dir: Path) -> List[Tuple[str, Path]]: + """Collect all cua-agent trajectory sessions in a run directory. + + Scans for task_*_agent_logs/trajectories// inside each task dir. + + Returns: + List of (task_name, session_dir) tuples sorted by task name. + """ + sessions: List[Tuple[str, Path]] = [] + if not run_dir.exists(): + return sessions + + for task_dir in sorted(run_dir.iterdir()): + if not task_dir.is_dir(): + continue + for agent_logs_dir in sorted(task_dir.glob("task_*_agent_logs")): + traj_root = agent_logs_dir / "trajectories" + if not traj_root.exists(): + continue + for session_dir in sorted(traj_root.iterdir()): + if session_dir.is_dir(): + sessions.append((task_dir.name, session_dir)) + + return sessions - # Resolve identifier to run directory - run_dir = _resolve_trace_path(identifier) - if run_dir is None: +def cmd_traj(args) -> int: + """View cua-agent trajectories from a run using the cua.ai trajectory viewer.""" + identifier = args.identifier + port = getattr(args, "port", 8090) + + run_path = _resolve_trace_path(identifier) + if run_path is None: print(f"{RED}Run not found: {identifier}{RESET}") print(f"{GREY}Try:{RESET}") print(f" cb run list {GREY}# List all runs{RESET}") return 1 - # If it's a specific trace (session_id), error - grid needs a run_id - if (run_dir / "dataset_info.json").exists(): - print(f"{YELLOW}Grid view requires a run_id, not a session_id.{RESET}") - print(f"{GREY}To view a single trace:{RESET}") - print(f" cb trace view {identifier}") + if (run_path / "dataset_info.json").exists(): + print(f"{YELLOW}traj requires a run_id, not a session_id.{RESET}") + print(f"{GREY}To view a single trace: cb trace view {identifier}{RESET}") return 1 - # Collect all traces in the run - traces = _collect_run_traces(run_dir) - if not traces: - print(f"{YELLOW}No traces found in run: {identifier}{RESET}") + sessions = _collect_run_trajectories(run_path) + if not sessions: + print(f"{YELLOW}No agent trajectories found in run: {identifier}{RESET}") + print(f"{GREY}(Trajectories are saved when tasks run with cua-agent){RESET}") return 1 - cards_html: List[str] = [] - for name, p in traces: - ds = load_from_disk(str(p)) - preview_b64 = "" - for row in ds: - imgs = row.get("data_images") or [] - if not imgs: - continue - img0 = imgs[0] - if hasattr(img0, "save"): - import io - - buf = io.BytesIO() - img0.save(buf, format="PNG") - preview_b64 = base64.b64encode(buf.getvalue()).decode("ascii") - elif isinstance(img0, dict): - data = img0.get("bytes") - if data: - preview_b64 = base64.b64encode(data).decode("ascii") - else: - pth = img0.get("path") - if pth and Path(pth).exists(): - data = Path(pth).read_bytes() - preview_b64 = base64.b64encode(data).decode("ascii") - if preview_b64: - break - row_count = len(ds) - href = "/trace?path=" + urllib.parse.quote(str(p)) - img_html = ( - f'' - if preview_b64 - else '
(no image)
' - ) - cards_html.append( - f""" - -
- {img_html} -
-
{name}
-
{row_count} rows
-
{p}
-
-
- - """ - ) - - index_html = f""" - + tmp_dir = Path(tempfile.mkdtemp(prefix="cb-traj-")) + try: + # Zip each session into the temp dir + zip_entries: List[Tuple[str, str, str]] = [] # (task_name, session_name, zip_name) + for task_name, session_dir in sessions: + zip_name = f"{task_name}__{session_dir.name}.zip" + zip_path = tmp_dir / zip_name + with _zipfile.ZipFile(zip_path, "w", _zipfile.ZIP_DEFLATED) as zf: + for child in sorted(session_dir.rglob("*")): + if child.is_file(): + zf.write(child, child.relative_to(session_dir)) + zip_entries.append((task_name, session_dir.name, zip_name)) + + viewer_base = "https://cua.ai/trajectory-viewer" + zip_names = {e[2] for e in zip_entries} + + # Build index page + items_html = "" + for task_name, session_name, zip_name in zip_entries: + zip_url = f"http://localhost:{port}/{zip_name}" + viewer_url = f"{viewer_base}?zip={urllib.parse.quote(zip_url, safe='')}" + items_html += f""" +
+
+
{task_name}
+
{session_name}
+
+ Open Viewer +
""" + + index_html = f""" - - Trace Grid - Run {identifier} + + Trajectories \u2014 Run {identifier} -

Trace Grid - Run {identifier}

-
- {''.join(cards_html)} +

Trajectories \u2014 Run {identifier}

+
{items_html}
- -""" +""" - def render_single(path: Path) -> bytes: - """Render a single trace for the popup window.""" - ds = load_from_disk(str(path)) - ds = ds.cast_column("data_images", ds.features["data_images"]) - ds.set_format(type=None, columns=None) - rows: List[str] = [] - row_meta: List[dict] = [] - row_data: List[dict] = [] - for i, row in enumerate(ds): - imgs_html: List[str] = [] - imgs = row.get("data_images") or [] - for j, img in enumerate(imgs): - if hasattr(img, "save"): - import io + index_bytes = index_html.encode("utf-8") - buf = io.BytesIO() - img.save(buf, format="PNG") - b64 = base64.b64encode(buf.getvalue()).decode("ascii") - imgs_html.append( - f'' - ) - elif isinstance(img, dict): - data = img.get("bytes") - if data: - b64 = base64.b64encode(data).decode("ascii") - imgs_html.append( - f'' - ) - else: - pth = img.get("path") - if pth and Path(pth).exists(): - data = Path(pth).read_bytes() - b64 = base64.b64encode(data).decode("ascii") - imgs_html.append( - f'' - ) - meta = { - "event_name": row.get("event_name"), - "timestamp": row.get("timestamp"), - "trajectory_id": row.get("trajectory_id"), - } - row_meta.append( - { - "timestamp": meta["timestamp"], - "trajectory_id": meta["trajectory_id"], - } - ) - data_json = row.get("data_json") - if isinstance(data_json, (dict, list)): - parsed = data_json - elif isinstance(data_json, str): - try: - parsed = json.loads(data_json) - except Exception: - parsed = data_json - else: - parsed = data_json - - row_data.append(parsed if isinstance(parsed, dict) else {}) - ( - json.dumps(parsed, ensure_ascii=False, indent=2) - if not isinstance(parsed, str) - else str(parsed) - ) - rows.append( - f""" - - - {json.dumps(meta, ensure_ascii=False, indent=2)} - - - {json.dumps(parsed, ensure_ascii=False, indent=2) if not isinstance(parsed, str) else parsed} - - {'
'.join(imgs_html)} - - """ - ) - template_path = Path(__file__).resolve().parents[2] / "www" / "trace_viewer.html" - html_template = template_path.read_text(encoding="utf-8") - row_data_b64 = base64.b64encode(json.dumps(row_data).encode("utf-8")).decode("ascii") - - html = ( - html_template.replace("__PATH_NAME__", str(path.name)) - .replace("__PATH__", str(path)) - .replace("__ROWS__", "".join(rows)) - .replace("__ROW_META__", json.dumps(row_meta)) - .replace("__ROW_DATA__", f'"{row_data_b64}"') - ) - return html.encode("utf-8") + class _Handler(BaseHTTPRequestHandler): + def _send_cors(self): + self.send_header("Access-Control-Allow-Origin", "*") + self.send_header("Access-Control-Allow-Methods", "GET, OPTIONS") + self.send_header("Access-Control-Allow-Headers", "*") - class _Handler(BaseHTTPRequestHandler): - def do_GET(self): - try: - if self.path.startswith("/trace?"): - q = urllib.parse.parse_qs(urllib.parse.urlsplit(self.path).query) - p = Path(q.get("path", [""])[0]) - if not p.exists(): + def do_OPTIONS(self): + self.send_response(204) + self._send_cors() + self.end_headers() + + def do_GET(self): + try: + path = self.path.split("?")[0].lstrip("/") + if not path: + body = index_bytes + ct = "text/html; charset=utf-8" + elif path in zip_names: + body = (tmp_dir / path).read_bytes() + ct = "application/zip" + else: self.send_response(404) self.end_headers() return - body = render_single(p) - else: - body = index_html.encode("utf-8") - self.send_response(200) - self.send_header("Content-Type", "text/html; charset=utf-8") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - except Exception as e: - import traceback - - traceback.print_exc() - print(f"{RED}Error serving request: {e}{RESET}") - self.send_response(500) - self.end_headers() - - def log_message(self, format, *args): - return - - httpd = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) - port = httpd.server_address[1] - url = f"http://127.0.0.1:{port}/" + self.send_response(200) + self.send_header("Content-Type", ct) + self.send_header("Content-Length", str(len(body))) + self._send_cors() + self.end_headers() + self.wfile.write(body) + except Exception: + self.send_response(500) + self.end_headers() + + def log_message(self, format, *args): + return + + try: + httpd = ThreadingHTTPServer(("127.0.0.1", port), _Handler) + except OSError: + httpd = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + port = httpd.server_address[1] + + t = threading.Thread(target=httpd.serve_forever, daemon=True) + t.start() + + # Single session → open viewer directly; multiple → open index + if len(zip_entries) == 1: + task_name, session_name, zip_name = zip_entries[0] + zip_url = f"http://localhost:{port}/{zip_name}" + open_url = f"{viewer_base}?zip={urllib.parse.quote(zip_url, safe='')}" + else: + open_url = f"http://localhost:{port}/" + + webbrowser.open(open_url) + + n = len(zip_entries) + print(f"{CYAN}Serving {n} trajectory session{'s' if n != 1 else ''}:{RESET}") + for task_name, session_name, zip_name in zip_entries: + zip_url = f"http://localhost:{port}/{zip_name}" + viewer_url = f"{viewer_base}?zip={urllib.parse.quote(zip_url, safe='')}" + print(f" {GREY}{task_name}:{RESET} {viewer_url}") + print(f"\n{GREY}Press Enter to stop...{RESET}") + + try: + input() + except KeyboardInterrupt: + pass + finally: + httpd.shutdown() + httpd.server_close() - t = threading.Thread(target=httpd.serve_forever, daemon=True) - t.start() - webbrowser.open(url) - print(f"{CYAN}Serving traces viewer at:{RESET} {url}\n{GREY}Press Enter to stop...{RESET}") - try: - input() - except KeyboardInterrupt: - pass finally: - httpd.shutdown() - httpd.server_close() + shutil.rmtree(tmp_dir, ignore_errors=True) + return 0 diff --git a/libs/python/agent/agent/loops/gemini.py b/libs/python/agent/agent/loops/gemini.py index 3385b796c8..1189332081 100644 --- a/libs/python/agent/agent/loops/gemini.py +++ b/libs/python/agent/agent/loops/gemini.py @@ -7,6 +7,7 @@ - gemini-2.5-computer-use-preview-10-2025 (uses built-in ComputerUse tool) - gemini-3-flash-preview (and variants) (uses custom function declarations) - gemini-3-pro-preview (and variants) (uses custom function declarations) +- gemini-3.1-pro-preview-computer-use (and variants) (uses custom function declarations) Key features: - Lazy import of google.genai @@ -339,9 +340,9 @@ def _denormalize(v: int, size: int) -> int: return 0 -def _is_gemini_3_model(model: str) -> bool: - """Check if the model is a Gemini 3 model (Flash or Pro Preview).""" - return "gemini-3" in model.lower() or "gemini-2.0" in model.lower() +def _has_builtin_computer_use(model: str) -> bool: + """Check if the model has a built-in ComputerUse tool (e.g. gemini-2.5-computer-use-preview).""" + return "computer-use" in model.lower() def _build_custom_function_declarations(types: Any) -> List[Any]: @@ -683,8 +684,9 @@ def _map_gemini_fc_to_computer_call( # - gemini-2.5-computer-use-preview-* : Uses built-in ComputerUse tool # - gemini-3-flash-preview-* : Uses custom function declarations # - gemini-3-pro-preview-* : Uses custom function declarations +# - gemini-3.1-pro-preview-* : Uses custom function declarations @register_agent( - models=r"^(gemini-2\.5-computer-use-preview.*|gemini-3-flash-preview.*|gemini-3-pro-preview.*)$" + models=r"^(gemini-2\.5-computer-use-preview.*|gemini-3(\.\d+)?-flash-preview.*|gemini-3(\.\d+)?-pro-preview.*)$" ) class GeminiComputerUseConfig(AsyncAgentConfig): async def predict_step( @@ -751,18 +753,15 @@ async def predict_step( contents, (screen_w, screen_h) = _convert_messages_to_gemini_contents(messages, types) # Compose tools config based on model type - # Gemini 2.5 Computer Use Preview uses built-in ComputerUse tool - # Gemini 3 Flash/Pro Preview uses custom function declarations - is_gemini_3 = _is_gemini_3_model(model) + # Models with "computer-use" in the name use built-in ComputerUse tool + # All other models use custom function declarations + has_builtin_cu = _has_builtin_computer_use(model) - if is_gemini_3: - # Use custom function declarations for Gemini 3 models + if not has_builtin_cu: custom_functions = _build_custom_function_declarations(types) - print(f"[DEBUG] Using custom function declarations for Gemini 3 model: {model}") + print(f"[DEBUG] Using custom function declarations for model: {model}") print(f"[DEBUG] Number of custom functions: {len(custom_functions)}") - # Build system instruction with coordinate system and screen resolution context - # to improve coordinate precision for Gemini 3 models system_instruction = ( f"You are controlling a computer with screen resolution {screen_w}x{screen_h} pixels. " "When using coordinate-based functions (click_at, type_text_at, hover_at, scroll_at, drag_and_drop), " @@ -801,7 +800,7 @@ async def predict_step( computer_environment.lower(), types.Environment.ENVIRONMENT_BROWSER ) - print(f"[DEBUG] Using built-in ComputerUse tool for Gemini 2.5 model: {model}") + print(f"[DEBUG] Using built-in ComputerUse tool for model: {model}") print(f"[DEBUG] Environment: {resolved_environment}") print(f"[DEBUG] Excluded functions: {excluded}") @@ -948,10 +947,10 @@ async def predict_click( client, model = _create_gemini_client(model, genai, kwargs) # Build tools config based on model type - is_gemini_3 = _is_gemini_3_model(model) + has_builtin_cu = _has_builtin_computer_use(model) - if is_gemini_3: - # For Gemini 3 models, use only click_at function declaration + if not has_builtin_cu: + # Use only click_at function declaration for models without built-in ComputerUse click_function = types.FunctionDeclaration( name="click_at", description="Click at the specified x,y coordinates on the screen. x and y are normalized 0-999 where 0 is the left/top edge and 999 is the right/bottom edge of the screen. Look carefully at the screenshot to identify the exact position of the target element before clicking.",