fix: enhance screenshot retrieval in PythonController

- Added a static method to validate image responses for PNG and JPEG formats using magic bytes.
- Improved error handling in the get_screenshot method to log invalid payloads and retry attempts.
- Updated the requests call to include a timeout for better reliability.

desktop_env/controllers/python.py

@@ -20,17 +20,41 @@ class PythonController:
         self.retry_times = 3
         self.retry_interval = 5
 
+    @staticmethod
+    def _is_valid_image_response(content_type: str, data: Optional[bytes]) -> bool:
+        """Quick validation for PNG/JPEG payload using magic bytes; Content-Type is advisory.
+        Returns True only when bytes look like a real PNG or JPEG.
+        """
+        if not isinstance(data, (bytes, bytearray)) or not data:
+            return False
+        # PNG magic
+        if len(data) >= 8 and data[:8] == b"\x89PNG\r\n\x1a\n":
+            return True
+        # JPEG magic
+        if len(data) >= 3 and data[:3] == b"\xff\xd8\xff":
+            return True
+        # If server explicitly marks as image, accept as a weak fallback (some environments strip magic)
+        if content_type and ("image/png" in content_type or "image/jpeg" in content_type or "image/jpg" in content_type):
+            return True
+        return False
+
     def get_screenshot(self) -> Optional[bytes]:
         """
         Gets a screenshot from the server. With the cursor. None -> no screenshot or unexpected error.
         """
 
-        for _ in range(self.retry_times):
+        for attempt_idx in range(self.retry_times):
             try:
-                response = requests.get(self.http_server + "/screenshot")
+                response = requests.get(self.http_server + "/screenshot", timeout=10)
                 if response.status_code == 200:
-                    logger.info("Got screenshot successfully")
-                    return response.content
+                    content_type = response.headers.get("Content-Type", "")
+                    content = response.content
+                    if self._is_valid_image_response(content_type, content):
+                        logger.info("Got screenshot successfully")
+                        return content
+                    else:
+                        logger.error("Invalid screenshot payload (attempt %d/%d).", attempt_idx + 1, self.retry_times)
+                        logger.info("Retrying to get screenshot.")
                 else:
                     logger.error("Failed to get screenshot. Status code: %d", response.status_code)
                     logger.info("Retrying to get screenshot.")