diff --git a/CHANGELOG.md b/CHANGELOG.md index b7f7c10..daa163a 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/s1a/browser/decision_model.py b/s1a/browser/decision_model.py index b5d785d..7fd548c 100644 --- a/s1a/browser/decision_model.py +++ b/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) diff --git a/tests/test_browser_screenshot.py b/tests/test_browser_screenshot.py index 4f59e5e..bf4c0a0 100644 --- a/tests/test_browser_screenshot.py +++ b/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))