Merge pull request #8 from jhshi1879/fix/validator-screenshot-error-handling

fix(validator): safely handle non-image/error MCP screenshot results in validator

artemis/agents/validator/validator.py

@@ -405,26 +405,16 @@ class ValidatorNode:
                     try:
                         if session is not None:
                             result = await session.call_tool("take_screenshot", {})
-                            success_img, post_screenshot_b64 = self._parse_mcp_result(result)
+                            success_img, content = self._parse_mcp_result(result)
+                            if success_img and content:
+                                post_screenshot_b64 = content
+                            else:
+                                logger.warning(f"Take screenshot via MCP returned non-success: {content}")
                         elif getattr(self, "_local_controller", None):
                             post_screenshot_b64 = await self._local_controller.take_screenshot()
-                            success_img = True
-                        else:
-                            success_img, post_screenshot_b64 = False, None
-
-                        if success_img and post_screenshot_b64:
-                            if self.ctx.data_engine:
-                                post_image_name = self.ctx.data_engine.get_or_create_image(
-                                    base64.b64decode(post_screenshot_b64)
-                                )
-                                screenshot_path = str(
-                                    self.ctx.data_engine.get_image_path(post_image_name)
-                                )
-                                state.latest_screenshot = screenshot_path
                     except Exception as e:
                         logger.error(f"Failed to capture live screenshot for failure analysis: {e}")
                         post_screenshot_b64 = None
-                        post_image_name = None
                 else:
                     # 2. Local execution attempts
                     max_local_retries = 1 if action_name == "launch_app" else 2
@@ -454,20 +444,16 @@ class ValidatorNode:
                         try:
                             if session is not None:
                                 result = await session.call_tool("take_screenshot", {})
-                                success_img, post_screenshot_b64 = self._parse_mcp_result(result)
+                                success_img, content = self._parse_mcp_result(result)
+                                if success_img and content:
+                                    post_screenshot_b64 = content
+                                else:
+                                    logger.warning(f"Take failure screenshot via MCP returned non-success: {content}")
                             elif getattr(self, "_local_controller", None):
                                 post_screenshot_b64 = await self._local_controller.take_screenshot()
-                                success_img = True
-                            else:
-                                success_img, post_screenshot_b64 = False, None
-
-                            if success_img and post_screenshot_b64:
-                                if self.ctx.data_engine:
-                                    post_image_name = self.ctx.data_engine.get_or_create_image(
-                                        base64.b64decode(post_screenshot_b64)
-                                    )
                         except Exception as e:
                             logger.error(f"Failed to capture failure screenshot: {e}")
+                            post_screenshot_b64 = None
 
                 # Enrich and record execution
                 if len(attempts_log) > 1 or not success:
@@ -475,11 +461,11 @@ class ValidatorNode:
                 execution.append(action_item)
 
                 if post_screenshot_b64:
+                    decoded_bytes = base64.b64decode(post_screenshot_b64)
                     last_screenshot_b64 = post_screenshot_b64
+                    logger.info(f"Successfully decoded post_screenshot_b64 ({len(decoded_bytes)} bytes)")
                     if self.ctx.data_engine:
-                        post_image_name = self.ctx.data_engine.get_or_create_image(
-                            base64.b64decode(post_screenshot_b64)
-                        )
+                        post_image_name = self.ctx.data_engine.get_or_create_image(decoded_bytes)
                         last_screenshot_name = post_image_name
 
                         screenshot_path = str(self.ctx.data_engine.get_image_path(post_image_name))

tests/unit/agents/test_validator.py

@@ -911,3 +911,42 @@ async def test_validator_pre_execution_xml_failure_not_overridden_when_pixel_byp
     mock_analyze.assert_called_once()
     args, kwargs = mock_analyze.call_args
     assert "Pre-execution validation failed" in args[2]
+
+
+@pytest.mark.asyncio
+async def test_validator_failure_screenshot_mcp_error_handled_safely(
+    mock_mcp, mock_context, temp_screenshot
+):
+    """Test that when take_screenshot fails on MCP server and returns an error text,
+    the ValidatorNode does not crash with base64 decoding error.
+    """
+    def mock_call_tool(name, args):
+        if name == "take_screenshot":
+            return Mock(content=[Mock(text="Error: ADB connection lost")])
+        if name == "tap":
+            return Mock(content=[Mock(text="Error: Target not found")])
+        return Mock(content=[Mock(text="Success")])
+
+    mock_mcp.call_tool.side_effect = mock_call_tool
+
+    decisions = json.dumps([{"action": "tap", "coordinates": [105, 205]}])
+    state = DummyState(structured_decisions=decisions, latest_screenshot=temp_screenshot)
+    node = ValidatorNode(mock_context)
+
+    with (
+        patch("artemis.agents.validator.failure_analyzer.FailureAnalyzer.analyze") as mock_analyze,
+        patch("artemis.utils.image_diff.check_ui_change", return_value=False),
+    ):
+        mock_analyze.return_value = {
+            "status": "cannot_fix",
+            "analysis": "Action failed and cannot fix.",
+        }
+
+        result = await node(state)
+
+    assert "last_execution_result" in result
+    report = result["last_execution_result"]
+    assert report["status"] == "failed"
+    assert len(report["execution"]) == 1
+    assert report["execution"][0]["action"] == "tap"
+