Refactor browser_use MIME type detection

Replace if-elif chain with dictionary-based lookup for cleaner, more maintainable code.
Extract detection logic into reusable detect_image_mime_type() helper function.
Add comprehensive test coverage for JPEG, PNG, GIF, WebP, and unknown formats.

Co-authored-by: openhands <[email protected]>

Author: openhands-agent

openhands-tools/openhands/tools/browser_use/definition.py

@@ -24,6 +24,29 @@
 # Maximum output size for browser observations
 MAX_BROWSER_OUTPUT_SIZE = 50000
 
+# Mapping of base64 prefixes to MIME types for image detection
+BASE64_IMAGE_PREFIXES = {
+    "/9j/": "image/jpeg",
+    "iVBORw0KGgo": "image/png",
+    "R0lGODlh": "image/gif",
+    "UklGR": "image/webp",
+}
+
+
+def detect_image_mime_type(base64_data: str) -> str:
+    """Detect MIME type from base64-encoded image data.
+
+    Args:
+        base64_data: Base64-encoded image data
+
+    Returns:
+        Detected MIME type, defaults to "image/png" if not detected
+    """
+    for prefix, mime_type in BASE64_IMAGE_PREFIXES.items():
+        if base64_data.startswith(prefix):
+            return mime_type
+    return "image/png"
+
 
 class BrowserObservation(Observation):
     """Base observation for browser operations."""
@@ -48,15 +71,7 @@ def to_llm_content(self) -> Sequence[TextContent | ImageContent]:
             )
 
         if self.screenshot_data:
-            mime_type = "image/png"
-            if self.screenshot_data.startswith("/9j/"):
-                mime_type = "image/jpeg"
-            elif self.screenshot_data.startswith("iVBORw0KGgo"):
-                mime_type = "image/png"
-            elif self.screenshot_data.startswith("R0lGODlh"):
-                mime_type = "image/gif"
-            elif self.screenshot_data.startswith("UklGR"):
-                mime_type = "image/webp"
+            mime_type = detect_image_mime_type(self.screenshot_data)
             # Convert base64 to data URL format for ImageContent
             data_url = f"data:{mime_type};base64,{self.screenshot_data}"
             llm_content.append(ImageContent(image_urls=[data_url]))

tests/tools/browser_use/test_browser_observation.py

@@ -112,3 +112,41 @@ def test_browser_observation_empty_screenshot_handling():
     observation = BrowserObservation.from_text(text="Test", screenshot_data=None)
     agent_obs = observation.to_llm_content
     assert len(agent_obs) == 1  # Only text content, no image
+
+
+def test_browser_observation_mime_type_detection():
+    """Test MIME type detection for different image formats."""
+    test_cases = [
+        (
+            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==",  # noqa: E501
+            "image/png",
+        ),
+        (
+            "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwA/",  # noqa: E501
+            "image/jpeg",
+        ),
+        (
+            "R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7",
+            "image/gif",
+        ),
+        (
+            "UklGRiQAAABXRUJQVlA4IBgAAAAwAQCdASoBAAEAAQAcJaQAA3AA/v3AgAA=",
+            "image/webp",
+        ),
+        (
+            "AAAABBBBCCCC",  # Unknown format
+            "image/png",  # Falls back to PNG
+        ),
+    ]
+
+    for screenshot_data, expected_mime_type in test_cases:
+        observation = BrowserObservation.from_text(
+            text="Test", screenshot_data=screenshot_data
+        )
+        agent_obs = observation.to_llm_content
+
+        assert len(agent_obs) == 2
+        assert isinstance(agent_obs[1], ImageContent)
+        assert (
+            agent_obs[1].image_urls[0].startswith(f"data:{expected_mime_type};base64,")
+        )