feat(agent): add Gemini 3.1 support and improve model routing

[ddupont808]

Improve support for gemini computer-use and refactor built-in ComputerUse detection to key off "computer-use" in the model name instead of hardcoded version checks. Update cuabench trace CLI and docs to support the trajectory viewer web UI.

Summary by CodeRabbit

Release Notes

• Documentation

  • Updated guide and CLI reference to reflect trace command changes and trajectory viewer integration.

• New Features

  • Integrated trajectory viewer for analyzing and visualizing agent execution paths during runs.
  • Extended Gemini model support for additional model variants with computer use capabilities.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
docs	Ready	Preview, Comment	Mar 4, 2026 3:44pm

[github-actions]

📦 Publishable packages changed

• pypi/agent
• pypi/bench

Add release:<service> labels to auto-release on merge (+ optional bump:minor or bump:major, default is patch).
Or add no-release to skip.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: d59d8efe-6587-457a-b01c-e7e730081568

📥 Commits

Reviewing files that changed from the base of the PR and between 6bf7f42 and adfce50.

📒 Files selected for processing (4)

• docs/content/docs/cuabench/guide/fundamentals/agent-traces.mdx
• docs/content/docs/cuabench/reference/cli-reference.mdx
• libs/cua-bench/cua_bench/cli/commands/trace.py
• libs/python/agent/agent/loops/gemini.py

---

📝 Walkthrough

Walkthrough

The pull request updates the trace command from grid-based viewing to a trajectory-based viewer, implementing HTTP serving and ZIP packaging of trajectory data. Additionally, it updates Gemini model detection logic to identify models with built-in ComputerUse capabilities via a new helper function.

Changes

Cohort / File(s)	Summary
Documentation Updates docs/content/docs/cuabench/guide/fundamentals/agent-traces.mdx, docs/content/docs/cuabench/reference/cli-reference.mdx	Replaces cb trace grid command references with cb trace traj and updates descriptions to reflect trajectory viewer functionality instead of grid view.
Trace Command Implementation libs/cua-bench/cua_bench/cli/commands/trace.py	Replaces grid command with new traj command. Adds trajectory collection logic, HTTP server with CORS support, ZIP packaging of trajectory sessions, and viewer integration that launches cua.ai/trajectory-viewer with generated session URLs. Removes prior grid-based UI in favor of index page listing trajectory links.
Gemini Model Detection libs/python/agent/agent/loops/gemini.py	Introduces _has_builtin_computer_use() helper to detect models with built-in ComputerUse by name pattern. Replaces _is_gemini_3_model() checks throughout. Updates model selector regex to support extended Gemini 3 variants (e.g., gemini-3.1-*) and adjusts tool configuration logic accordingly.

Sequence Diagram

sequenceDiagram
    actor User
    participant CLI
    participant TrajCollector
    participant HTTPServer
    participant Browser
    participant TrajectoryViewer

    User->>CLI: cb trace traj <run_id>
    CLI->>TrajCollector: collect_run_trajectories(run_dir)
    TrajCollector->>TrajCollector: discover task_*/agent_logs/trajectories
    TrajCollector-->>CLI: return (task_name, session_dir) pairs

    CLI->>HTTPServer: create temp ZIPs for each session
    HTTPServer->>HTTPServer: generate per-session ZIP files
    HTTPServer->>HTTPServer: start local HTTP server with CORS

    CLI->>Browser: open cua.ai/trajectory-viewer?zip_url=...
    Browser->>TrajectoryViewer: request trajectory viewer
    TrajectoryViewer->>HTTPServer: fetch trajectory ZIP
    HTTPServer-->>TrajectoryViewer: serve ZIP content
    TrajectoryViewer-->>Browser: render trajectory visualization
    TrajectoryViewer-->>User: display trajectory data

    CLI->>HTTPServer: cleanup temp files after serving

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• fix(agent): route cua/ prefix Gemini models through correct agent loop #1117: Modifies libs/python/agent/agent/loops/gemini.py to refactor Gemini model variant detection and ComputerUse routing logic, directly related to the gemini.py changes in this PR.

Poem

🐰 From grid to traj, the traces now flow,
Zipped trajectories ready to show,
Gemini models detect with flair,
Built-in computer-use, oh how rare!
Serve and view with HTTP's grace,
A trajectory viewer finds its place! 🎯✨

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch ddupont/gemini-support

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[sentry]

Codecov Report

❌ Patch coverage is 12.50000% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
libs/python/agent/agent/loops/gemini.py	12.50%	7 Missing ⚠️

📢 Thoughts on this report? Let us know!