feat: wait before screenshots

Author: mmilanta

src/playwright_computer_use/async_api.py

@@ -192,17 +192,25 @@ def to_params(self) -> BetaToolComputerUse20241022Param:
         """Params describing the tool. Used by Claude to understand this is a computer use tool."""
         return {"name": self.name, "type": self.api_type, **self.options}
 
-    def __init__(self, page: Page, use_cursor: bool = True):
+    def __init__(
+        self,
+        page: Page,
+        use_cursor: bool = True,
+        screenshot_wait_until: Literal["load", "domcontentloaded", "networkidle"]
+        | None = None,
+    ):
         """Initializes the PlaywrightComputerTool.
 
         Args:
             page: The Async Playwright page to interact with.
             use_cursor: Whether to display the cursor in the screenshots or not.
+            screenshot_wait_until: Optional, wait until the page is in a specific state before taking a screenshot. Default does not wait
         """
         super().__init__()
         self.page = page
         self.use_cursor = use_cursor
         self.mouse_position: tuple[int, int] = (0, 0)
+        self.screenshot_wait_until = screenshot_wait_until
 
     async def __call__(
         self,
@@ -284,6 +292,9 @@ async def __call__(
 
     async def screenshot(self) -> ToolResult:
         """Take a screenshot of the current screen and return the base64 encoded image."""
+        if self.screenshot_wait_until is not None:
+            await self.page.wait_for_timeout(self.screenshot_wait_until)
+        await self.page.wait_for_load_state()
         screenshot = await self.page.screenshot()
         image = Image.open(io.BytesIO(screenshot))
         img_small = image.resize((self.width, self.height), Image.LANCZOS)