async playwright backend

Author: ollmer

src/agentlab/backends/browser/base.py

@@ -80,12 +80,15 @@ def run_js(self, js: str):
     def goto(self, url: str) -> str:
         raise NotImplementedError
 
-    def page_snapshot(self) -> str:
+    def page_html(self) -> str:
         raise NotImplementedError
 
     def page_screenshot(self) -> Image:
         raise NotImplementedError
 
+    def page_axtree(self) -> str:
+        raise NotImplementedError
+
     def step(self, action: ToolCallAction) -> str:
         raise NotImplementedError
 

src/agentlab/backends/browser/env.py

@@ -31,12 +31,13 @@ def reset(self, seed: int):
         if setup_js:
             self.goal = self.backend.run_js(setup_js)
             logger.info(f"Task goal: {self.goal}")
-        page_content = self.backend.page_snapshot()
+        html = self.backend.page_html()
         screenshot = self.backend.page_screenshot()
+        axtree = self.backend.page_axtree()
         obs = {
             "goal_object": [{"type": "text", "text": self.goal}],
-            "pruned_html": "",
-            "axtree_txt": page_content,
+            "pruned_html": html,
+            "axtree_txt": axtree,
             "screenshot": screenshot,
             "last_action_error": "",
             "focused_element_bid": "none",
@@ -89,12 +90,13 @@ def step(self, action: ToolCallAction | str) -> tuple[dict, float, bool, bool, d
 
     def _step(self, action: ToolCallAction) -> dict:
         obs_dict = self.backend.step(action)
-        return {
-            "goal_object": [{"type": "text", "text": self.goal}],
-            **obs_dict,
-            "last_action_error": "",
-            "focused_element_bid": "none",
-        }
+        if "goal_object" not in obs_dict:
+            obs_dict["goal_object"] = [{"type": "text", "text": self.goal}]
+        if "last_action_error" not in obs_dict:
+            obs_dict["last_action_error"] = ""
+        if "focused_element_bid" not in obs_dict:
+            obs_dict["focused_element_bid"] = "none"
+        return obs_dict
 
     def validate_task(self, action: ToolCallAction, observation: dict) -> tuple[float, dict]:
         validate_js = self.task.get_step_validate_js()

src/agentlab/backends/browser/mcp_playwright.py

@@ -32,16 +32,20 @@ def step(self, action: ToolCallAction) -> dict:
         tool_result = "\n".join(
             [c.text for c in contents if c.type == "text" and "# Ran Playwright code" not in c.text]
         )
-        snapshot = self.page_snapshot()
+        html = self.page_html()
         screenshot = self.page_screenshot()
+        axtree = self.page_axtree()
         return {
             "tool_result": tool_result,
-            "pruned_html": "",
-            "axtree_txt": snapshot,
+            "pruned_html": html,
+            "axtree_txt": axtree,
             "screenshot": screenshot,
         }
 
-    def page_snapshot(self) -> str:
+    def page_html(self) -> str:
+       return ""
+
+    def page_axtree(self) -> str:
         contents = self.call_tool("browser_snapshot", {})
         return "\n".join([c.text for c in contents if c.type == "text"])
 

src/agentlab/backends/browser/playwright.py

@@ -1,18 +1,20 @@
+import asyncio
 import logging
 from io import BytesIO
 from typing import Any, Callable
 
 from PIL import Image
-from playwright.sync_api import Page, sync_playwright
+from playwright.async_api import Browser, Page, async_playwright
 
 from agentlab.backends.browser.base import BrowserBackend, ToolCallAction, ToolSpec
 
 logger = logging.getLogger(__name__)
 
 
-class PlaywrightSyncBackend(BrowserBackend):
+class AsyncPlaywright(BrowserBackend):
     _actions: dict[str, Callable]
-    _browser: Any
+    _loop: asyncio.AbstractEventLoop
+    _browser: Browser
     _page: Page
 
     def model_post_init(self, __context: Any):
@@ -26,86 +28,97 @@ def model_post_init(self, __context: Any):
             "browser_mouse_click_xy": self.browser_mouse_click_xy,
         }
 
-    def browser_press_key(self, key: str):
+    def initialize(self):
+        self._loop = asyncio.get_event_loop()
+        self._loop.run_until_complete(self.ainitialize())
+
+    async def ainitialize(self):
+        pw = await async_playwright().start()
+        self._browser = await pw.chromium.launch(headless=True, chromium_sandbox=True)
+        self._page = await self._browser.new_page()
+
+    async def browser_press_key(self, key: str):
         """
         Press a key on the keyboard.
         """
-        self._page.keyboard.press(key)
+        await self._page.keyboard.press(key)
 
-    def browser_type(self, text: str):
+    async def browser_type(self, text: str):
         """
         Type text into the focused element.
         """
-        self._page.type(text)
+        await self._page.type(text)
 
-    def browser_click(self, selector: str):
+    async def browser_click(self, selector: str):
         """
         Click on a selector.
         """
-        self._page.click(selector)
+        await self._page.click(selector)
 
-    def browser_drag(self, from_selector: str, to_selector: str):
+    async def browser_drag(self, from_selector: str, to_selector: str):
         """
         Drag and drop from one selector to another.
         """
         from_elem = self._page.locator(from_selector)
-        from_elem.hover(timeout=500)
-        self._page.mouse.down()
+        await from_elem.hover(timeout=500)
+        await self._page.mouse.down()
 
         to_elem = self._page.locator(to_selector)
-        to_elem.hover(timeout=500)
-        self._page.mouse.up()
+        await to_elem.hover(timeout=500)
+        await self._page.mouse.up()
 
-    def browser_hover(self, selector: str):
+    async def browser_hover(self, selector: str):
         """
         Hover over a given element.
         """
-        self._page.hover(selector)
+        await self._page.hover(selector)
 
-    def browser_select_option(self, selector: str):
+    async def browser_select_option(self, selector: str):
         """
         Select an option from a given element.
         """
-        self._page.select_option(selector)
+        await self._page.select_option(selector)
 
-    def browser_mouse_click_xy(self, x: int, y: int):
+    async def browser_mouse_click_xy(self, x: int, y: int):
         """
         Click at a given x, y coordinate using the mouse.
         """
-        self._page.mouse.click(x, y)
-
-    def initialize(self):
-        self._browser = sync_playwright().start().chromium.launch(headless=True, chromium_sandbox=True)
-        self._page = self._browser.new_page()
+        await self._page.mouse.click(x, y)
 
     def run_js(self, js: str):
-        js_result = self._page.evaluate(js)
+        js_result = self._loop.run_until_complete(self._page.evaluate(js))
         logger.info(f"JS result: {js_result}")
         return js_result
 
     def goto(self, url: str):
-        self._page.goto(url)
+        self._loop.run_until_complete(self._page.goto(url))
 
-    def page_snapshot(self):
-        return self._page.content()
+    def page_html(self):
+        return self._loop.run_until_complete(self._page.content())
 
     def page_screenshot(self):
-        scr_bytes = self._page.screenshot()
+        scr_bytes = self._loop.run_until_complete(self._page.screenshot())
         return Image.open(BytesIO(scr_bytes))
 
+    def page_axtree(self):
+        return ""
+
     def step(self, action: ToolCallAction):
         fn = self._actions[action.function.name]
-        action_result = fn(**action.function.arguments)
-        snapshot = self.page_snapshot()
+        action_result = self._loop.run_until_complete(fn(**action.function.arguments))
+        html = self.page_html()
         screenshot = self.page_screenshot()
+        axtree = self.page_axtree()
         return {
-            "pruned_html": f"{action_result or ''}\n{snapshot}",
-            "axtree_txt": snapshot,
+            "tool_result": action_result,
+            "pruned_html": html,
+            "axtree_txt": axtree,
             "screenshot": screenshot,
         }
+
     def actions(self) -> tuple[ToolSpec]:
         specs = [ToolSpec.from_function(fn) for fn in self._actions.values()]
         return tuple(specs)
 
     def close(self):
-        self._browser.close()
+        self._loop.run_until_complete(self._browser.close())