handling navigation for new tabs - cua

Author: miguelg719

stagehand/agent/anthropic_cua.py

@@ -86,8 +86,16 @@ def __init__(
                     "required": ["url"],
                 },
             },
+            {
+                "name": "navigate_back",
+                "description": "Navigate back to the previous page",
+                "input_schema": {
+                    "type": "object",
+                    "properties": {},
+                },
+            },
         ]
-        self.max_tokens = kwargs.get("max_tokens", 4096)
+        self.max_tokens = kwargs.get("max_tokens", 1024)
         self.last_tool_use_ids = None
         self.logger.info(
             f"AnthropicCUAClient initialized for model: {model}",
@@ -143,7 +151,8 @@ async def run_task(
                 response = self.anthropic_sdk_client.beta.messages.create(
                     model=self.model,
                     max_tokens=self.max_tokens,
-                    system=self.instructions,  # System prompt
+                    system=self.instructions
+                    + "Remember to call the computer tools, and only goto or navigate_back if you need to. Screenshots, clicks, etc, will be parsed from computer tool calls",  # System prompt
                     messages=current_messages,
                     tools=self.tools,
                     betas=["computer-use-2025-01-24"],
@@ -300,9 +309,6 @@ def _process_provider_response(
         task_completed = not bool(
             agent_action
         )  # Task is complete if no tool_use blocks
-        self.logger.info(
-            f"{agent_action}, {model_message_text}, {task_completed}, {raw_assistant_content_blocks}"
-        )
 
         return (
             agent_action,
@@ -314,14 +320,18 @@ def _process_provider_response(
     def _convert_tool_use_to_agent_action(
         self, tool_name: str, tool_input: dict[str, Any]
     ) -> Optional[AgentAction]:
-        if tool_name != "computer" and tool_name != "goto":
+        if (
+            tool_name != "computer"
+            and tool_name != "goto"
+            and tool_name != "navigate_back"
+        ):
             self.logger.warning(
                 f"Unsupported tool name from Anthropic: {tool_name}",
                 category=StagehandFunctionName.AGENT,
             )
             return None
 
-        if tool_name == "goto":
+        if tool_name == "goto" or tool_name == "navigate_back":
             action_type_str = "function"
         else:
             action_type_str = tool_input.get("action")
@@ -542,6 +552,13 @@ def _convert_tool_use_to_agent_action(
                             category=StagehandFunctionName.AGENT,
                         )
                         return None
+                elif tool_name == "navigate_back":
+                    action_model_payload = AgentActionType(
+                        type="function",
+                        name="navigate_back",
+                        arguments=FunctionArguments(),
+                    )
+                    action_type_str = "function"
             else:
                 self.logger.warning(
                     f"Unsupported action type '{action_type_str}' from Anthropic computer tool.",
@@ -573,7 +590,6 @@ def _format_action_feedback(
         new_screenshot_base64: str,
         current_url: Optional[str],
     ) -> list[dict[str, Any]]:
-        self.logger.info(f"Action result: {action_result}")
         content_for_tool_result: list[dict[str, Any]] = []
         is_error_result = not action_result.get("success", False)
 

stagehand/handlers/cua_handler.py

@@ -53,6 +53,9 @@ async def perform_action(self, action: AgentAction) -> ActionExecutionResult:
             }
 
         try:
+            # Store initial URL to detect navigation
+            initial_url = self.page.url
+
             if action_type == "click":
                 # specific_action_model is already an instance of ClickAction
                 x, y = specific_action_model.x, specific_action_model.y
@@ -66,7 +69,9 @@ async def perform_action(self, action: AgentAction) -> ActionExecutionResult:
                     await self._animate_click(x, y)
                     await asyncio.sleep(0.1)  # Ensure animation is visible
                     await self.page.mouse.click(x, y, button=button)
-                # Consider new tab/page handling logic here if needed
+
+                # Check for page navigation
+                await self.handle_page_navigation("click", initial_url)
                 return {"success": True}
 
             elif action_type == "double_click":
@@ -78,6 +83,9 @@ async def perform_action(self, action: AgentAction) -> ActionExecutionResult:
                 await self._animate_click(x, y)
                 await asyncio.sleep(0.1)
                 await self.page.mouse.dblclick(x, y)
+
+                # Check for page navigation
+                await self.handle_page_navigation("double_click", initial_url)
                 return {"success": True}
 
             elif action_type == "type":
@@ -92,6 +100,9 @@ async def perform_action(self, action: AgentAction) -> ActionExecutionResult:
                 for key_str in specific_action_model.keys:
                     playwright_key = self._convert_key_name(key_str)
                     await self.page.keyboard.press(playwright_key)  # Press each key
+
+                # Check for page navigation - keys like Enter can cause navigation
+                await self.handle_page_navigation("keypress", initial_url)
                 return {"success": True}
 
             elif action_type == "scroll":
@@ -110,6 +121,9 @@ async def perform_action(self, action: AgentAction) -> ActionExecutionResult:
                 if name == "goto" and args.url:
                     await self.page.goto(args.url)
                     return {"success": True}
+                elif name == "navigate_back":
+                    await self.page.go_back()
+                    return {"success": True}
                 # Add other function calls like back, forward, reload if needed, similar to TS version
                 self.logger.warning(
                     f"Unsupported function call: {name}",
@@ -131,6 +145,9 @@ async def perform_action(self, action: AgentAction) -> ActionExecutionResult:
                 else:
                     # Use _convert_key_name for consistency if possible, or press directly
                     await self.page.keyboard.press(self._convert_key_name(text))
+
+                # Check for page navigation - Enter and other keys may navigate
+                await self.handle_page_navigation("key", initial_url)
                 return {"success": True}
 
             elif action_type == "wait":
@@ -248,3 +265,68 @@ def _convert_key_name(self, key: str) -> str:
         # Convert to uppercase for case-insensitive matching then check map,
         # default to original key if not found.
         return key_map.get(key.upper(), key)
+
+    async def _handle_page_navigation(self) -> None:
+        """Handle page navigation actions."""
+        pass
+
+    async def handle_page_navigation(
+        self,
+        action_description: str,
+        initial_url: str,
+        dom_settle_timeout_ms: int = 1000,
+    ) -> None:
+        """Handle possible page navigation after an action."""
+        self.logger.debug(
+            f"{action_description} - checking for page navigation",
+            category=StagehandFunctionName.AGENT,
+        )
+
+        # Check for new tab/window
+        new_opened_tab = None
+        try:
+            async with self.page.context.expect_page(timeout=1500) as new_page_info:
+                # Just checking if a page was opened by the action
+                pass
+            new_opened_tab = await new_page_info.value
+        except Exception:
+            new_opened_tab = None
+
+        # Handle new tab if one was opened
+        if new_opened_tab:
+            self.logger.info(
+                f"New tab detected with URL: {new_opened_tab.url}",
+                category=StagehandFunctionName.AGENT,
+            )
+            new_tab_url = new_opened_tab.url
+            await new_opened_tab.close()
+            await self.page.goto(new_tab_url)
+            await self.page.wait_for_load_state("domcontentloaded")
+
+        # Wait for DOM to settle
+        try:
+            await self.page.wait_for_load_state(
+                "domcontentloaded", timeout=dom_settle_timeout_ms
+            )
+            # Additional optional wait for network idle
+            await self.page.wait_for_load_state(
+                "networkidle", timeout=dom_settle_timeout_ms
+            )
+        except Exception as e:
+            self.logger.warning(
+                f"Wait for DOM settle timed out: {str(e)}",
+                category=StagehandFunctionName.AGENT,
+            )
+
+        # Check if URL changed
+        current_url = self.page.url
+        if current_url != initial_url:
+            self.logger.debug(
+                f"Page navigation detected: {initial_url} -> {current_url}",
+                category=StagehandFunctionName.AGENT,
+            )
+
+        self.logger.debug(
+            "Finished checking for page navigation",
+            category=StagehandFunctionName.AGENT,
+        )