fix(reports): anchor tiled budget clock at overall screenshot start

Follow-up to #42118, implementing the non-blocking review suggestion
left at merge time: the tiled wait budget is derived from the running
Celery task's own time limit, but elapsed time was measured from a
clock started inside take_tiled_screenshot() -- so navigation
(page.goto, bounded 60s), the headstart sleep, element waits, and
dimension probing all ran before the clock started, effectively
granting the tile loop a fresh full budget on top of task time already
spent. On hard-limit-only tasks that overhang can still crest the task
limit -- the SIGKILL #42118 exists to prevent. The non-tiled readiness
wait already avoids this by threading screenshot_started_at from the
top of get_screenshot (#42427); this applies the same pattern to the
tiled call:

- take_tiled_screenshot() accepts screenshot_started_at (optional,
  defaults to "now" for backward compatibility) and anchors the budget
  clock on it;
- WebDriverPlaywright.get_screenshot() passes its existing
  screenshot_started_at into the tiled call, putting both capture paths
  on one clock.

Tests: pre-capture elapsed time reduces the first tile's capped wait
when the anchor is provided; the local-clock default is unchanged when
it is omitted; the existing tiled call-site assertion now pins the new
argument.

Co-Authored-By: Claude <noreply@anthropic.com>

superset/utils/screenshot_utils.py

@@ -281,6 +281,7 @@ def take_tiled_screenshot(  # noqa: C901
     load_wait: int = 60,
     animation_wait: int = 0,
     log_context: str | None = None,
+    screenshot_started_at: float | None = None,
 ) -> bytes | None:
     """
     Take a tiled screenshot of a large dashboard by scrolling and capturing sections.
@@ -294,6 +295,11 @@ def take_tiled_screenshot(  # noqa: C901
         log_context: Optional identifier (e.g. report execution id, or a
             cache key for thumbnails) appended to log lines so a slow/timed-out
             capture can be traced back to the run that produced it.
+        screenshot_started_at: Optional time.monotonic() timestamp taken at
+            the start of the overall screenshot operation (before browser
+            navigation), so time already spent on goto/headstart/element
+            waits counts against the budget -- the same clock the non-tiled
+            readiness wait uses. Falls back to "now" when not provided.
 
     Returns:
         Combined screenshot bytes or None if failed
@@ -323,7 +329,15 @@ def take_tiled_screenshot(  # noqa: C901
     wait_budget_seconds = resolve_screenshot_task_budget_seconds(log_context)
     if wait_budget_seconds is None:
         wait_budget_seconds = float(TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS)
-    start_time = time.monotonic()
+    # Anchor the budget clock at the start of the overall screenshot
+    # operation when the caller provides it: the budget is derived from the
+    # Celery task's own limit, and navigation/headstart/element waits before
+    # tiling spend that same task time. A locally-reset clock would grant
+    # the tile loop a fresh full budget on top of whatever was already
+    # spent, which on hard-limit-only tasks can still crest the task limit.
+    start_time = (
+        screenshot_started_at if screenshot_started_at is not None else time.monotonic()
+    )
     try:
         # Get the target element
         element = page.locator(f".{element_name}")

superset/utils/webdriver.py

@@ -592,6 +592,7 @@ class WebDriverPlaywright(WebDriverProxy):
                             load_wait=self._screenshot_load_wait,
                             animation_wait=selenium_animation_wait,
                             log_context=log_context,
+                            screenshot_started_at=screenshot_started_at,
                         )
                         if not img:
                             # _get_screenshot() has no wait/readiness logic at

tests/unit_tests/utils/test_screenshot_utils.py

@@ -915,6 +915,64 @@ class TestTileWaitBudget:
         assert first_timeout == 96 * 1000
         assert first_timeout < 200 * 1000
 
+    def test_screenshot_started_at_counts_pre_capture_time_against_budget(
+        self, mock_page, monkeypatch
+    ):
+        """Time spent before tiling (navigation, headstart, element waits)
+        must count against the budget when the caller provides the overall
+        screenshot start time -- the budget is derived from the Celery
+        task's own limit, and that same task time was already being spent
+        before the tile loop started. Same clock the non-tiled readiness
+        wait uses via _wait_for_charts_ready."""
+        monkeypatch.setattr(
+            "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS",  # noqa: E501
+            1000,
+        )
+        clock = self._FakeClock()
+        # The overall screenshot started 900s ago (goto/headstart/element
+        # waits consumed it); only 100s of the 1000s budget remains.
+        clock.now = 900.0
+
+        with patch("superset.utils.screenshot_utils.current_task", None):
+            with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
+                with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
+                    take_tiled_screenshot(
+                        mock_page,
+                        "dashboard",
+                        tile_height=2000,
+                        load_wait=500,
+                        screenshot_started_at=0.0,
+                    )
+
+        first_timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"]
+        assert first_timeout == 100 * 1000
+
+    def test_omitted_screenshot_started_at_anchors_clock_locally(
+        self, mock_page, monkeypatch
+    ):
+        """Without the caller-provided anchor the clock starts at entry
+        (backward-compatible default): the same 900s of prior wall-clock
+        time does not reduce the budget."""
+        monkeypatch.setattr(
+            "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS",  # noqa: E501
+            1000,
+        )
+        clock = self._FakeClock()
+        clock.now = 900.0
+
+        with patch("superset.utils.screenshot_utils.current_task", None):
+            with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
+                with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
+                    take_tiled_screenshot(
+                        mock_page,
+                        "dashboard",
+                        tile_height=2000,
+                        load_wait=500,
+                    )
+
+        first_timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"]
+        assert first_timeout == 500 * 1000
+
     def test_fast_dashboard_matches_default_behavior(self, mock_page):
         """Well under budget, waits are not capped and behavior is unchanged."""
         with patch("superset.utils.screenshot_utils.current_task", None):

tests/unit_tests/utils/webdriver_test.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from unittest.mock import MagicMock, patch, PropertyMock
+from unittest.mock import ANY, MagicMock, patch, PropertyMock
 
 import pytest
 
@@ -1681,6 +1681,7 @@ class TestWebDriverPlaywrightAnimationWaitOrder:
             load_wait=30,
             animation_wait=2,
             log_context=None,
+            screenshot_started_at=ANY,
         )
         # The only wait_for_timeout call should be the 0ms headstart; no global
         # animation wait should be issued (handled per-tile by take_tiled_screenshot)