sentienceapi 0.92.2__py3-none-any.whl → 0.98.0__py3-none-any.whl

sentience/snapshot.py

@@ -12,6 +12,7 @@ import requests
 
 from .browser import AsyncSentienceBrowser, SentienceBrowser
 from .browser_evaluator import BrowserEvaluator
+from .constants import SENTIENCE_API_URL
 from .models import Snapshot, SnapshotOptions
 from .sentience_methods import SentienceMethod
 
@@ -19,6 +20,206 @@ from .sentience_methods import SentienceMethod
 MAX_PAYLOAD_BYTES = 10 * 1024 * 1024
 
 
+def _is_execution_context_destroyed_error(e: Exception) -> bool:
+    """
+    Playwright can throw while a navigation is in-flight, invalidating the JS execution context.
+
+    Common symptoms:
+    - "Execution context was destroyed, most likely because of a navigation"
+    - "Cannot find context with specified id"
+    """
+    msg = str(e).lower()
+    return (
+        "execution context was destroyed" in msg
+        or "most likely because of a navigation" in msg
+        or "cannot find context with specified id" in msg
+    )
+
+
+async def _page_evaluate_with_nav_retry(
+    page: Any,
+    expression: str,
+    arg: Any = None,
+    *,
+    retries: int = 2,
+    settle_timeout_ms: int = 10000,
+) -> Any:
+    """
+    Evaluate JS with a small retry loop if the page is mid-navigation.
+
+    This prevents flaky crashes when callers snapshot right after triggering a navigation
+    (e.g., pressing Enter on Google).
+    """
+    last_err: Exception | None = None
+    for attempt in range(retries + 1):
+        try:
+            if arg is None:
+                return await page.evaluate(expression)
+            return await page.evaluate(expression, arg)
+        except Exception as e:
+            last_err = e
+            if not _is_execution_context_destroyed_error(e) or attempt >= retries:
+                raise
+            try:
+                await page.wait_for_load_state("domcontentloaded", timeout=settle_timeout_ms)
+            except Exception:
+                pass
+            await asyncio.sleep(0.25)
+    raise last_err if last_err else RuntimeError("Page.evaluate failed")
+
+
+async def _wait_for_function_with_nav_retry(
+    page: Any,
+    expression: str,
+    *,
+    timeout_ms: int,
+    retries: int = 2,
+) -> None:
+    last_err: Exception | None = None
+    for attempt in range(retries + 1):
+        try:
+            await page.wait_for_function(expression, timeout=timeout_ms)
+            return
+        except Exception as e:
+            last_err = e
+            if not _is_execution_context_destroyed_error(e) or attempt >= retries:
+                raise
+            try:
+                await page.wait_for_load_state("domcontentloaded", timeout=timeout_ms)
+            except Exception:
+                pass
+            await asyncio.sleep(0.25)
+    raise last_err if last_err else RuntimeError("wait_for_function failed")
+
+
+def _build_snapshot_payload(
+    raw_result: dict[str, Any],
+    options: SnapshotOptions,
+) -> dict[str, Any]:
+    """
+    Build payload dict for gateway snapshot API.
+
+    Shared helper used by both sync and async snapshot implementations.
+    """
+    diagnostics = raw_result.get("diagnostics") or {}
+    client_metrics = None
+    try:
+        client_metrics = diagnostics.get("metrics")
+    except Exception:
+        client_metrics = None
+
+    return {
+        "raw_elements": raw_result.get("raw_elements", []),
+        "url": raw_result.get("url", ""),
+        "viewport": raw_result.get("viewport"),
+        "goal": options.goal,
+        "options": {
+            "limit": options.limit,
+            "filter": options.filter.model_dump() if options.filter else None,
+        },
+        "client_metrics": client_metrics,
+    }
+
+
+def _validate_payload_size(payload_json: str) -> None:
+    """
+    Validate payload size before sending to gateway.
+
+    Raises ValueError if payload exceeds server limit.
+    """
+    payload_size = len(payload_json.encode("utf-8"))
+    if payload_size > MAX_PAYLOAD_BYTES:
+        raise ValueError(
+            f"Payload size ({payload_size / 1024 / 1024:.2f}MB) exceeds server limit "
+            f"({MAX_PAYLOAD_BYTES / 1024 / 1024:.0f}MB). "
+            f"Try reducing the number of elements on the page or filtering elements."
+        )
+
+
+def _post_snapshot_to_gateway_sync(
+    payload: dict[str, Any],
+    api_key: str,
+    api_url: str = SENTIENCE_API_URL,
+) -> dict[str, Any]:
+    """
+    Post snapshot payload to gateway (synchronous).
+
+    Used by sync snapshot() function.
+    """
+    payload_json = json.dumps(payload)
+    _validate_payload_size(payload_json)
+
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+
+    response = requests.post(
+        f"{api_url}/v1/snapshot",
+        data=payload_json,
+        headers=headers,
+        timeout=30,
+    )
+    response.raise_for_status()
+    return response.json()
+
+
+async def _post_snapshot_to_gateway_async(
+    payload: dict[str, Any],
+    api_key: str,
+    api_url: str = SENTIENCE_API_URL,
+) -> dict[str, Any]:
+    """
+    Post snapshot payload to gateway (asynchronous).
+
+    Used by async backend snapshot() function.
+    """
+    # Lazy import httpx - only needed for async API calls
+    import httpx
+
+    payload_json = json.dumps(payload)
+    _validate_payload_size(payload_json)
+
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        response = await client.post(
+            f"{api_url}/v1/snapshot",
+            content=payload_json,
+            headers=headers,
+        )
+        response.raise_for_status()
+        return response.json()
+
+
+def _merge_api_result_with_local(
+    api_result: dict[str, Any],
+    raw_result: dict[str, Any],
+) -> dict[str, Any]:
+    """
+    Merge API result with local data (screenshot, etc.).
+
+    Shared helper used by both sync and async snapshot implementations.
+    """
+    return {
+        "status": api_result.get("status", "success"),
+        "timestamp": api_result.get("timestamp"),
+        "url": api_result.get("url", raw_result.get("url", "")),
+        "viewport": api_result.get("viewport", raw_result.get("viewport")),
+        "elements": api_result.get("elements", []),
+        "screenshot": raw_result.get("screenshot"),  # Keep local screenshot
+        "screenshot_format": raw_result.get("screenshot_format"),
+        "error": api_result.get("error"),
+        # Phase 2: Runtime stability/debug info
+        "diagnostics": api_result.get("diagnostics", raw_result.get("diagnostics")),
+        # Phase 2: Ordinal support - dominant group key from Gateway
+        "dominant_group_key": api_result.get("dominant_group_key"),
+    }
+
+
 def _save_trace_to_file(raw_elements: list[dict[str, Any]], trace_path: str | None = None) -> None:
     """
     Save raw_elements to a JSON file for benchmarking/training
@@ -72,14 +273,18 @@ def snapshot(
     if options is None:
         options = SnapshotOptions()
 
+    # Resolve API key: options.sentience_api_key takes precedence, then browser.api_key
+    # This allows browser-use users to pass api_key via options without SentienceBrowser
+    effective_api_key = options.sentience_api_key or browser.api_key
+
     # Determine if we should use server-side API
     should_use_api = (
-        options.use_api if options.use_api is not None else (browser.api_key is not None)
+        options.use_api if options.use_api is not None else (effective_api_key is not None)
     )
 
-    if should_use_api and browser.api_key:
+    if should_use_api and effective_api_key:
         # Use server-side API (Pro/Enterprise tier)
-        return _snapshot_via_api(browser, options)
+        return _snapshot_via_api(browser, options, effective_api_key)
     else:
         # Use local extension (Free tier)
         return _snapshot_via_extension(browser, options)
@@ -127,10 +332,15 @@ def _snapshot_via_extension(
     if options.save_trace:
         _save_trace_to_file(result.get("raw_elements", []), options.trace_path)
 
+    # Validate and parse with Pydantic
+    snapshot_obj = Snapshot(**result)
+
     # Show visual overlay if requested
     if options.show_overlay:
-        raw_elements = result.get("raw_elements", [])
-        if raw_elements:
+        # Prefer processed semantic elements for overlay (have bbox/importance/visual_cues).
+        # raw_elements may not match the overlay renderer's expected shape.
+        elements_for_overlay = result.get("elements") or result.get("raw_elements") or []
+        if elements_for_overlay:
             browser.page.evaluate(
                 """
                 (elements) => {
@@ -139,27 +349,46 @@ def _snapshot_via_extension(
                     }
                 }
                 """,
-                raw_elements,
+                elements_for_overlay,
+            )
+
+    # Show grid overlay if requested
+    if options.show_grid:
+        # Get all grids (don't filter by grid_id here - we want to show all but highlight the target)
+        grids = snapshot_obj.get_grid_bounds(grid_id=None)
+        if grids:
+            # Convert GridInfo to dict for JavaScript
+            grid_dicts = [grid.model_dump() for grid in grids]
+            # Pass grid_id as targetGridId to highlight it in red
+            target_grid_id = options.grid_id if options.grid_id is not None else None
+            browser.page.evaluate(
+                """
+                (grids, targetGridId) => {
+                    if (window.sentience && window.sentience.showGrid) {
+                        window.sentience.showGrid(grids, targetGridId);
+                    } else {
+                        console.warn('[SDK] showGrid not available in extension');
+                    }
+                }
+                """,
+                grid_dicts,
+                target_grid_id,
             )
 
-    # Validate and parse with Pydantic
-    snapshot_obj = Snapshot(**result)
     return snapshot_obj
 
 
 def _snapshot_via_api(
     browser: SentienceBrowser,
     options: SnapshotOptions,
+    api_key: str,
 ) -> Snapshot:
     """Take snapshot using server-side API (Pro/Enterprise tier)"""
     if not browser.page:
         raise RuntimeError("Browser not started. Call browser.start() first.")
 
-    if not browser.api_key:
-        raise ValueError("API key required for server-side processing")
-
-    if not browser.api_url:
-        raise ValueError("API URL required for server-side processing")
+    # Use browser.api_url if set, otherwise default
+    api_url = browser.api_url or SENTIENCE_API_URL
 
     # CRITICAL: Wait for extension injection to complete (CSP-resistant architecture)
     # Even for API mode, we need the extension to collect raw data locally
@@ -169,6 +398,14 @@ def _snapshot_via_api(
     raw_options: dict[str, Any] = {}
     if options.screenshot is not False:
         raw_options["screenshot"] = options.screenshot
+    # Important: also pass limit/filter to extension to keep raw_elements payload bounded.
+    # Without this, large pages (e.g. Amazon) can exceed gateway request size limits (HTTP 413).
+    if options.limit != 50:
+        raw_options["limit"] = options.limit
+    if options.filter is not None:
+        raw_options["filter"] = (
+            options.filter.model_dump() if hasattr(options.filter, "model_dump") else options.filter
+        )
 
     raw_result = BrowserEvaluator.invoke(browser.page, SentienceMethod.SNAPSHOT, **raw_options)
 
@@ -179,54 +416,16 @@ def _snapshot_via_api(
     # Step 2: Send to server for smart ranking/filtering
     # Use raw_elements (raw data) instead of elements (processed data)
     # Server validates API key and applies proprietary ranking logic
-    payload = {
-        "raw_elements": raw_result.get("raw_elements", []),  # Raw data needed for server processing
-        "url": raw_result.get("url", ""),
-        "viewport": raw_result.get("viewport"),
-        "goal": options.goal,  # Optional goal/task description
-        "options": {
-            "limit": options.limit,
-            "filter": options.filter.model_dump() if options.filter else None,
-        },
-    }
-
-    # Check payload size before sending (server has 10MB limit)
-    payload_json = json.dumps(payload)
-    payload_size = len(payload_json.encode("utf-8"))
-    if payload_size > MAX_PAYLOAD_BYTES:
-        raise ValueError(
-            f"Payload size ({payload_size / 1024 / 1024:.2f}MB) exceeds server limit "
-            f"({MAX_PAYLOAD_BYTES / 1024 / 1024:.0f}MB). "
-            f"Try reducing the number of elements on the page or filtering elements."
-        )
-
-    headers = {
-        "Authorization": f"Bearer {browser.api_key}",
-        "Content-Type": "application/json",
-    }
+    payload = _build_snapshot_payload(raw_result, options)
 
     try:
-        response = requests.post(
-            f"{browser.api_url}/v1/snapshot",
-            data=payload_json,  # Reuse already-serialized JSON
-            headers=headers,
-            timeout=30,
-        )
-        response.raise_for_status()
-
-        api_result = response.json()
+        api_result = _post_snapshot_to_gateway_sync(payload, api_key, api_url)
 
         # Merge API result with local data (screenshot, etc.)
-        snapshot_data = {
-            "status": api_result.get("status", "success"),
-            "timestamp": api_result.get("timestamp"),
-            "url": api_result.get("url", raw_result.get("url", "")),
-            "viewport": api_result.get("viewport", raw_result.get("viewport")),
-            "elements": api_result.get("elements", []),
-            "screenshot": raw_result.get("screenshot"),  # Keep local screenshot
-            "screenshot_format": raw_result.get("screenshot_format"),
-            "error": api_result.get("error"),
-        }
+        snapshot_data = _merge_api_result_with_local(api_result, raw_result)
+
+        # Create snapshot object
+        snapshot_obj = Snapshot(**snapshot_data)
 
         # Show visual overlay if requested (use API-ranked elements)
         if options.show_overlay:
@@ -243,9 +442,31 @@ def _snapshot_via_api(
                     elements,
                 )
 
-        return Snapshot(**snapshot_data)
+        # Show grid overlay if requested
+        if options.show_grid:
+            # Get all grids (don't filter by grid_id here - we want to show all but highlight the target)
+            grids = snapshot_obj.get_grid_bounds(grid_id=None)
+            if grids:
+                grid_dicts = [grid.model_dump() for grid in grids]
+                # Pass grid_id as targetGridId to highlight it in red
+                target_grid_id = options.grid_id if options.grid_id is not None else None
+                browser.page.evaluate(
+                    """
+                    (grids, targetGridId) => {
+                        if (window.sentience && window.sentience.showGrid) {
+                            window.sentience.showGrid(grids, targetGridId);
+                        } else {
+                            console.warn('[SDK] showGrid not available in extension');
+                        }
+                    }
+                    """,
+                    grid_dicts,
+                    target_grid_id,
+                )
+
+        return snapshot_obj
     except requests.exceptions.RequestException as e:
-        raise RuntimeError(f"API request failed: {e}")
+        raise RuntimeError(f"API request failed: {e}") from e
 
 
 # ========== Async Snapshot Functions ==========
@@ -281,14 +502,18 @@ async def snapshot_async(
     if options is None:
         options = SnapshotOptions()
 
+    # Resolve API key: options.sentience_api_key takes precedence, then browser.api_key
+    # This allows browser-use users to pass api_key via options without SentienceBrowser
+    effective_api_key = options.sentience_api_key or browser.api_key
+
     # Determine if we should use server-side API
     should_use_api = (
-        options.use_api if options.use_api is not None else (browser.api_key is not None)
+        options.use_api if options.use_api is not None else (effective_api_key is not None)
     )
 
-    if should_use_api and browser.api_key:
+    if should_use_api and effective_api_key:
         # Use server-side API (Pro/Enterprise tier)
-        return await _snapshot_via_api_async(browser, options)
+        return await _snapshot_via_api_async(browser, options, effective_api_key)
     else:
         # Use local extension (Free tier)
         return await _snapshot_via_extension_async(browser, options)
@@ -304,18 +529,20 @@ async def _snapshot_via_extension_async(
 
     # Wait for extension injection to complete
     try:
-        await browser.page.wait_for_function(
+        await _wait_for_function_with_nav_retry(
+            browser.page,
             "typeof window.sentience !== 'undefined'",
-            timeout=5000,
+            timeout_ms=5000,
         )
     except Exception as e:
         try:
-            diag = await browser.page.evaluate(
+            diag = await _page_evaluate_with_nav_retry(
+                browser.page,
                 """() => ({
                     sentience_defined: typeof window.sentience !== 'undefined',
                     extension_id: document.documentElement.dataset.sentienceExtensionId || 'not set',
                     url: window.location.href
-                })"""
+                })""",
             )
         except Exception:
             diag = {"error": "Could not gather diagnostics"}
@@ -341,7 +568,8 @@ async def _snapshot_via_extension_async(
         )
 
     # Call extension API
-    result = await browser.page.evaluate(
+    result = await _page_evaluate_with_nav_retry(
+        browser.page,
         """
         (options) => {
             return window.sentience.snapshot(options);
@@ -356,11 +584,26 @@ async def _snapshot_via_extension_async(
     if options.save_trace:
         _save_trace_to_file(result.get("raw_elements", []), options.trace_path)
 
+    # Extract screenshot_format from data URL if not provided by extension
+    if result.get("screenshot") and not result.get("screenshot_format"):
+        screenshot_data_url = result.get("screenshot", "")
+        if screenshot_data_url.startswith("data:image/"):
+            # Extract format from "data:image/jpeg;base64,..." or "data:image/png;base64,..."
+            format_match = screenshot_data_url.split(";")[0].split("/")[-1]
+            if format_match in ["jpeg", "jpg", "png"]:
+                result["screenshot_format"] = "jpeg" if format_match in ["jpeg", "jpg"] else "png"
+
+    # Validate and parse with Pydantic
+    snapshot_obj = Snapshot(**result)
+
     # Show visual overlay if requested
     if options.show_overlay:
-        raw_elements = result.get("raw_elements", [])
-        if raw_elements:
-            await browser.page.evaluate(
+        # Prefer processed semantic elements for overlay (have bbox/importance/visual_cues).
+        # raw_elements may not match the overlay renderer's expected shape.
+        elements_for_overlay = result.get("elements") or result.get("raw_elements") or []
+        if elements_for_overlay:
+            await _page_evaluate_with_nav_retry(
+                browser.page,
                 """
                 (elements) => {
                     if (window.sentience && window.sentience.showOverlay) {
@@ -368,41 +611,53 @@ async def _snapshot_via_extension_async(
                     }
                 }
                 """,
-                raw_elements,
+                elements_for_overlay,
             )
 
-    # Extract screenshot_format from data URL if not provided by extension
-    if result.get("screenshot") and not result.get("screenshot_format"):
-        screenshot_data_url = result.get("screenshot", "")
-        if screenshot_data_url.startswith("data:image/"):
-            # Extract format from "data:image/jpeg;base64,..." or "data:image/png;base64,..."
-            format_match = screenshot_data_url.split(";")[0].split("/")[-1]
-            if format_match in ["jpeg", "jpg", "png"]:
-                result["screenshot_format"] = "jpeg" if format_match in ["jpeg", "jpg"] else "png"
+    # Show grid overlay if requested
+    if options.show_grid:
+        # Get all grids (don't filter by grid_id here - we want to show all but highlight the target)
+        grids = snapshot_obj.get_grid_bounds(grid_id=None)
+        if grids:
+            grid_dicts = [grid.model_dump() for grid in grids]
+            # Pass grid_id as targetGridId to highlight it in red
+            target_grid_id = options.grid_id if options.grid_id is not None else None
+            await _page_evaluate_with_nav_retry(
+                browser.page,
+                """
+                (args) => {
+                    const [grids, targetGridId] = args;
+                    if (window.sentience && window.sentience.showGrid) {
+                        window.sentience.showGrid(grids, targetGridId);
+                    } else {
+                        console.warn('[SDK] showGrid not available in extension');
+                    }
+                }
+                """,
+                [grid_dicts, target_grid_id],
+            )
 
-    # Validate and parse with Pydantic
-    snapshot_obj = Snapshot(**result)
     return snapshot_obj
 
 
 async def _snapshot_via_api_async(
     browser: AsyncSentienceBrowser,
     options: SnapshotOptions,
+    api_key: str,
 ) -> Snapshot:
     """Take snapshot using server-side API (Pro/Enterprise tier) - async"""
     if not browser.page:
         raise RuntimeError("Browser not started. Call await browser.start() first.")
 
-    if not browser.api_key:
-        raise ValueError("API key required for server-side processing")
-
-    if not browser.api_url:
-        raise ValueError("API URL required for server-side processing")
+    # Use browser.api_url if set, otherwise default
+    api_url = browser.api_url or SENTIENCE_API_URL
 
     # Wait for extension injection
     try:
-        await browser.page.wait_for_function(
-            "typeof window.sentience !== 'undefined'", timeout=5000
+        await _wait_for_function_with_nav_retry(
+            browser.page,
+            "typeof window.sentience !== 'undefined'",
+            timeout_ms=5000,
         )
     except Exception as e:
         raise RuntimeError(
@@ -419,8 +674,17 @@ async def _snapshot_via_api_async(
             raw_options["screenshot"] = options.screenshot.model_dump()
         else:
             raw_options["screenshot"] = options.screenshot
+    # Important: also pass limit/filter to extension to keep raw_elements payload bounded.
+    # Without this, large pages (e.g. Amazon) can exceed gateway request size limits (HTTP 413).
+    if options.limit != 50:
+        raw_options["limit"] = options.limit
+    if options.filter is not None:
+        raw_options["filter"] = (
+            options.filter.model_dump() if hasattr(options.filter, "model_dump") else options.filter
+        )
 
-    raw_result = await browser.page.evaluate(
+    raw_result = await _page_evaluate_with_nav_retry(
+        browser.page,
         """
         (options) => {
             return window.sentience.snapshot(options);
@@ -466,7 +730,7 @@ async def _snapshot_via_api_async(
         )
 
     headers = {
-        "Authorization": f"Bearer {browser.api_key}",
+        "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json",
     }
 
@@ -476,7 +740,7 @@ async def _snapshot_via_api_async(
 
         async with httpx.AsyncClient(timeout=30.0) as client:
             response = await client.post(
-                f"{browser.api_url}/v1/snapshot",
+                f"{api_url}/v1/snapshot",
                 content=payload_json,
                 headers=headers,
             )
@@ -502,11 +766,15 @@ async def _snapshot_via_api_async(
             "error": api_result.get("error"),
         }
 
+        # Create snapshot object
+        snapshot_obj = Snapshot(**snapshot_data)
+
         # Show visual overlay if requested
         if options.show_overlay:
             elements = api_result.get("elements", [])
             if elements:
-                await browser.page.evaluate(
+                await _page_evaluate_with_nav_retry(
+                    browser.page,
                     """
                     (elements) => {
                         if (window.sentience && window.sentience.showOverlay) {
@@ -517,7 +785,30 @@ async def _snapshot_via_api_async(
                     elements,
                 )
 
-        return Snapshot(**snapshot_data)
+        # Show grid overlay if requested
+        if options.show_grid:
+            # Get all grids (don't filter by grid_id here - we want to show all but highlight the target)
+            grids = snapshot_obj.get_grid_bounds(grid_id=None)
+            if grids:
+                grid_dicts = [grid.model_dump() for grid in grids]
+                # Pass grid_id as targetGridId to highlight it in red
+                target_grid_id = options.grid_id if options.grid_id is not None else None
+                await _page_evaluate_with_nav_retry(
+                    browser.page,
+                    """
+                    (args) => {
+                        const [grids, targetGridId] = args;
+                        if (window.sentience && window.sentience.showGrid) {
+                            window.sentience.showGrid(grids, targetGridId);
+                        } else {
+                            console.warn('[SDK] showGrid not available in extension');
+                        }
+                    }
+                    """,
+                    [grid_dicts, target_grid_id],
+                )
+
+        return snapshot_obj
     except ImportError:
         # Fallback to requests if httpx not available (shouldn't happen in async context)
         raise RuntimeError(