diff --git a/artemis/agents/validator/validator.py b/artemis/agents/validator/validator.py index 382c814..3970339 100644 --- a/artemis/agents/validator/validator.py +++ b/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)) diff --git a/tests/unit/agents/test_validator.py b/tests/unit/agents/test_validator.py index 5f4e8b8..ebf6ef5 100644 --- a/tests/unit/agents/test_validator.py +++ b/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" +