[Fix] Read the browser screenshot from the capture's result object, not by pattern search

CHANGELOG.md

@@ -19,6 +19,9 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); ver
 
 ### Fixed
 
+- Browser front: the screenshot for an image-reading decision model is read from the `png` field of the
+  capture's own result, layer by layer, and must decode to a PNG. Before, a pattern search over the whole run-code
+  output could take PNG-like text from elsewhere in the result, such as the page state, for the picture.
 - Laya browser-state compaction preserves the recovery plan for the next decision.
 - Desktop recovery resets its detection window on observed progress without resetting its cumulative budget.
 - Screenshot-enabled desktop recovery detects pixel changes while ignoring capture IDs and PNG metadata.

s1a/browser/decision_model.py

@@ -102,7 +102,8 @@
     "return JSON.stringify({png: png.toString('base64')}); }"
 )
 SCREENSHOT_ATTEMPTS = 2  # a failed capture is tried once more before the step goes on without a picture
-_SCREENSHOT_RE = re.compile(r'png\\*"\s*:\s*\\*"([A-Za-z0-9+/=]+)')
+PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
+_RUN_CODE_LAYERS = 4  # compact RPC envelope, MCP content list, result field, JSON text (possibly encoded twice)
 
 
 @dataclass(frozen=True)
@@ -707,14 +708,14 @@ async def _screenshot(self) -> tuple[Image | None, int]:
                 raw = await executor(SCREENSHOT_JS)
             except Exception:  # noqa: BLE001 - a failed capture degrades to a text-only observation
                 logger.warning("[BrowserDecisionModel] screenshot failed", exc_info=True)
-            found = _SCREENSHOT_RE.search(json.dumps(raw, default=str)) if raw is not None else None
+            found = _png_from_run_code(raw)
             if found is not None:
                 break
         ms = round((time.perf_counter() - started) * 1000)
         if found is None:
             logger.warning("[BrowserDecisionModel] no screenshot in the run-code result")
             return None, ms
-        return Image(base64.b64decode(found.group(1))), ms
+        return Image(found), ms
 
     async def _raw_probe(self, *, settle_ms: int, quiet_ms: int, after: dict[str, Any] | None) -> dict[str, Any]:
         params = {
@@ -1167,6 +1168,38 @@ def _progressed(snapshot: dict[str, Any], before_key: Any) -> bool:
     return not _probe_failed(snapshot) and snapshot.get("page_key") != before_key
 
 
+def _png_from_run_code(raw: Any) -> bytes | None:
+    """The PNG that ``SCREENSHOT_JS`` returned, read from the ``png`` field of its result object.
+
+    The object arrives as is, or wrapped: in the compact RPC envelope (``payload``), as the text of an MCP content
+    list, under ``result``/``text``/``data``, or as JSON text after a ``### Result`` header, possibly encoded twice.
+    Each layer is opened in turn. ``None`` when no layer holds a ``png`` string, or when that string is not base64 of
+    a PNG, so text elsewhere in the result (the page's own content included) is never taken for the picture."""
+    for _layer in range(_RUN_CODE_LAYERS + 1):
+        if isinstance(raw, dict) and "png" in raw:
+            break
+        if isinstance(raw, dict):
+            content = raw.get("content")
+            if isinstance(content, list):
+                raw = "\n".join(
+                    str(i.get("text") or "") for i in content if isinstance(i, dict) and i.get("type") == "text"
+                )
+            else:
+                raw = next((raw[key] for key in ("payload", "result", "text", "data") if key in raw), None)
+        elif isinstance(raw, str):
+            raw = extract_json_object(raw) or None
+        else:
+            return None
+    png = raw.get("png") if isinstance(raw, dict) else None
+    if not isinstance(png, str):
+        return None
+    try:
+        data = base64.b64decode(png, validate=True)
+    except ValueError:
+        return None
+    return data if data.startswith(PNG_SIGNATURE) else None
+
+
 def _json_field(content: Any, key: str) -> Any:
     try:
         return json.loads(content if isinstance(content, str) else "{}").get(key)

tests/test_browser_screenshot.py

@@ -7,9 +7,9 @@
 import json
 from types import SimpleNamespace
 from typing import Any
-from unittest import IsolatedAsyncioTestCase
+from unittest import IsolatedAsyncioTestCase, TestCase
 
-from s1a.browser.decision_model import SCREENSHOT_JS, BrowserDecisionModel
+from s1a.browser.decision_model import SCREENSHOT_JS, BrowserDecisionModel, _png_from_run_code
 
 PNG = b"\x89PNG\r\n\x1a\n fake"
 ENCODED = base64.b64encode(PNG).decode()
@@ -57,3 +57,44 @@ async def test_one_failed_capture_is_retried_once(self) -> None:
     async def test_a_result_without_a_png_is_no_image(self) -> None:
         image, _ms = await _shot(_runtime({"payload": {"result": "{}"}}))
         self.assertIsNone(image)
+
+
+def _mcp_text(result_json: str, page_text: str = "- generic [ref=e1]") -> dict[str, Any]:
+    """A playwright-mcp run-code result: the returned value as JSON text after ``### Result``, then the page state."""
+    text = f"### Result\n{result_json}\n\n### Ran Playwright code\n```js\n...\n```\n\n### Page state\n{page_text}"
+    return {"content": [{"type": "text", "text": text}]}
+
+
+class TestPngFromRunCode(TestCase):
+    def test_every_wrapping_of_the_result_object_is_opened(self) -> None:
+        returned = json.dumps({"png": ENCODED})
+        shapes = {
+            "object": {"png": ENCODED},
+            "json text": returned,
+            "result field": {"result": returned},
+            "compact rpc": {"__browser_compact_rpc__": True, "payload": {"result": returned}},
+            "mcp text, encoded twice": _mcp_text(json.dumps(returned)),
+            "compact rpc around mcp": {"__browser_compact_rpc__": True, "payload": _mcp_text(returned)},
+        }
+        for name, raw in shapes.items():
+            with self.subTest(name):
+                self.assertEqual(_png_from_run_code(raw), PNG)
+
+    def test_png_like_text_outside_the_result_object_is_not_the_picture(self) -> None:
+        # The capture timed out; the page itself shows text that a search over the whole result would take for it.
+        raw = _mcp_text(json.dumps({"error": "timeout"}), page_text=f'- text: png": "{ENCODED}"')
+        self.assertIsNone(_png_from_run_code(raw))
+
+    def test_a_png_field_that_is_not_base64_of_a_png_is_no_image(self) -> None:
+        for name, png in {
+            "not base64": "not base64!",
+            "base64 of other bytes": base64.b64encode(b"GIF89a").decode(),
+            "not a string": 42,
+        }.items():
+            with self.subTest(name):
+                self.assertIsNone(_png_from_run_code({"png": png}))
+
+    def test_an_unknown_shape_is_no_image(self) -> None:
+        for raw in (None, 7, ["png", ENCODED], {"other": ENCODED}):
+            with self.subTest(raw=raw):
+                self.assertIsNone(_png_from_run_code(raw))