fix(agent): route cua/ prefix Gemini models through correct agent loop - #1117
Conversation
|
@sarinali is attempting to deploy a commit to the Cua Team on Vercel. A member of the Team first needs to authorize it. |
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe changes introduce CUA (Cloud Universal API) routing support across the agent system. A helper function strips CUA prefixes from model names to align routed models with bare configurations. The Gemini implementation adds centralized client creation, system instruction support for Gemini 3 models, enhanced reasoning content extraction, and expanded function declarations with coordinate semantics. Changes
Sequence DiagramsequenceDiagram
participant Caller
participant ConfigFinder as Config Finder<br/>(find_agent_config)
participant CUAHandler as CUA Handler<br/>(_strip_cua_prefix)
participant ClientFactory as Client Factory<br/>(_create_gemini_client)
participant GeminiAPI as Gemini API
Caller->>ConfigFinder: find_agent_config(cua/provider/model)
ConfigFinder->>CUAHandler: _strip_cua_prefix(cua/provider/model)
CUAHandler-->>ConfigFinder: bare_model
ConfigFinder->>ConfigFinder: Try original, then stripped
ConfigFinder-->>Caller: agent_config
Caller->>ClientFactory: _create_gemini_client(cua/provider/model, ...)
ClientFactory->>ClientFactory: Detect CUA prefix
ClientFactory->>ClientFactory: Set CUA auth & routing
ClientFactory-->>Caller: (gemini_client, bare_model)
Caller->>GeminiAPI: GenerateContent(with system_instruction)
GeminiAPI-->>Caller: Response with reasoning
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
libs/python/agent/agent/loops/gemini.py (1)
941-958: Consider extractingclick_atdeclaration to reduce duplication.The
click_atfunction declaration here duplicates the one in_build_custom_function_declarations(lines 358-375). If the description or parameters change, both locations need updating.♻️ Extract shared click_at declaration
+def _get_click_at_declaration(types: Any) -> Any: + """Return the click_at function declaration for Gemini 3 models.""" + return 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.", + parameters={ + "type": "object", + "properties": { + "x": { + "type": "integer", + "description": "X coordinate (0-999 normalized). 0 is the left edge, 999 is the right edge.", + }, + "y": { + "type": "integer", + "description": "Y coordinate (0-999 normalized). 0 is the top edge, 999 is the bottom edge.", + }, + }, + "required": ["x", "y"], + }, + )Then use
_get_click_at_declaration(types)in both_build_custom_function_declarationsandpredict_click.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@libs/python/agent/agent/loops/gemini.py` around lines 941 - 958, Extract the duplicated click_at FunctionDeclaration into a single helper function named _get_click_at_declaration(types) and return the FunctionDeclaration object from it; then replace the inline click_at declarations in _build_custom_function_declarations and predict_click with calls to _get_click_at_declaration(types). Ensure the helper accepts the same `types` argument used where FunctionDeclaration is constructed and keep the original description, parameter schema and required fields unchanged so both callers get the identical declaration.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@libs/python/agent/agent/loops/gemini.py`:
- Around line 941-958: Extract the duplicated click_at FunctionDeclaration into
a single helper function named _get_click_at_declaration(types) and return the
FunctionDeclaration object from it; then replace the inline click_at
declarations in _build_custom_function_declarations and predict_click with calls
to _get_click_at_declaration(types). Ensure the helper accepts the same `types`
argument used where FunctionDeclaration is constructed and keep the original
description, parameter schema and required fields unchanged so both callers get
the identical declaration.
d73a2eb to
189516b
Compare
2 participants
Summary
cua/prefix models:cua/google/gemini-3-flash-previewwas falling through toGenericVlmConfiginstead ofGeminiComputerUseConfigbecausefind_agent_config()didn't strip the routing prefix before regex matching. All Gemini-specific features (function declarations, ComputerUse tool, thinking traces) were silently unused._create_gemini_client()helper detects thecua/prefix and configures the Google GenAI SDK to route through{CUA_BASE_URL}/gemini, matching how other loops (Anthropic, OpenAI) transparently route through litellm's CUA adapter. Single code path for both direct-Google and CUA-routed models.Problem
When using
cua/google/gemini-3-flash-preview(CUA inference API routing):Wrong agent loop:
find_agent_config()matched the rawcua/google/gemini-3-flash-previewstring against registered patterns. The Gemini loop's anchored regex^(gemini-2\.5-computer-use-preview.*|gemini-3-flash-preview.*|...)$didn't match, so the model fell through toGenericVlmConfig's catch-all(?i).*at priority -100.No CUA proxy support in Gemini loop: Even after fixing routing, the Gemini loop called the Google SDK directly (
client.models.generate_content()), unlike Anthropic/OpenAI loops which use litellm (which transparently routescua/models throughCUAAdapter). Socua/prefix models would fail authentication sinceGOOGLE_API_KEYcontained a CUA key, not a real Google key.Changes
libs/python/agent/agent/decorators.py_strip_cua_prefix(model): Stripscua/<provider>/routing prefix → bare model name (e.g.cua/google/gemini-3-flash-preview→gemini-3-flash-preview)find_agent_config(model): For each config in priority order, tries original string first, then stripped. This ensures specialized loops match before the catch-all, without breaking direct model strings.libs/python/agent/agent/loops/gemini.py_create_gemini_client(): Shared helper forpredict_stepandpredict_click. Detectscua/prefix → configures Google GenAI SDK withhttp_options={"base_url": f"{CUA_BASE_URL}/gemini"}and CUA API key. Falls back to standardGOOGLE_API_KEY/ Vertex AI for non-CUA models.click_at,type_text_at,hover_at,scroll_at,drag_and_drop) now include explicit 0-999 edge semantics.types.Part(text="[screenshot]")alongside image parts incomputer_call_outputto fix LiteLLM content-part warnings.make_reasoning_itemimport andgetattr(p, "thought", False)check to capture Gemini's thinking traces in trajectory output.Design Decision
All other agent loops pass the raw model string to litellm, which detects the
cua/prefix and routes throughCUAAdapter. The Gemini loop is unique in using the Google GenAI SDK directly (for nativeComputerUsetool support, function declarations, thinking config, etc.). Rather than rewriting it to use litellm, we configure the Google SDK to point at the CUA inference proxy — same endpoint, same format, single code path for both direct and CUA-routed models.Test plan
cua/google/gemini-3-flash-previewroutes toGeminiComputerUseConfig(notGenericVlmConfig)agent_loop: "GeminiComputerUseConfig"in metadataGOOGLE_API_KEYpath still works (nocua/prefix)🤖 Generated with Claude Code
Summary by CodeRabbit
Release Notes
New Features
Refactor