From f543ae16f2b2f78c667dde865d1e01eff5238aff Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 17:58:05 -0700
Subject: [PATCH 01/16] Add poll_async_job helper

---
 setup.py                 |  2 +-
 surge/async_jobs.py      | 62 ++++++++++++++++++++++++++++++++++++++
 tests/test_async_jobs.py | 64 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 127 insertions(+), 1 deletion(-)
 create mode 100644 surge/async_jobs.py
 create mode 100644 tests/test_async_jobs.py

diff --git a/setup.py b/setup.py
index 869afb4..980af90 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-VERSION = "1.5.23"
+VERSION = "1.5.24"
 
 with open("requirements.txt") as f:
     requirements = f.read().splitlines()
diff --git a/surge/async_jobs.py b/surge/async_jobs.py
new file mode 100644
index 0000000..323ce32
--- /dev/null
+++ b/surge/async_jobs.py
@@ -0,0 +1,62 @@
+"""Polling helper for async-job status endpoints."""
+
+import time
+
+DEFAULT_IN_PROGRESS_STATUSES = ("IN_PROGRESS", "CREATING")
+DEFAULT_COMPLETED_STATUSES = ("COMPLETED", "READY")
+DEFAULT_ERROR_STATUSES = ("ERROR", "FAILED")
+
+
+class AsyncJobError(Exception):
+    """Raised when an async job returns a terminal error status."""
+
+    def __init__(self, status):
+        self.status = status
+        error = status.get("error") or status.get("type") or "async job failed"
+        super().__init__(error)
+
+
+class AsyncJobTimeoutError(Exception):
+    """Raised when an async job does not reach a terminal status in time."""
+
+
+def poll_async_job(
+    check_status,
+    *,
+    poll_time,
+    poll_interval,
+    in_progress_statuses=DEFAULT_IN_PROGRESS_STATUSES,
+    completed_statuses=DEFAULT_COMPLETED_STATUSES,
+    error_statuses=DEFAULT_ERROR_STATUSES,
+):
+    """Poll ``check_status()`` until it returns a terminal status.
+
+    Arguments:
+        check_status: zero-arg callable returning a dict with a ``status`` key.
+        poll_time: maximum seconds to wait before raising AsyncJobTimeoutError.
+        poll_interval: seconds to sleep between polls.
+        in_progress_statuses: status strings that mean "keep polling".
+        completed_statuses: status strings that mean "return the status dict".
+        error_statuses: status strings that mean "raise AsyncJobError".
+
+    Returns the terminal status dict.
+    """
+    in_progress = set(in_progress_statuses)
+    completed = set(completed_statuses)
+    error = set(error_statuses)
+
+    deadline = time.monotonic() + poll_time
+    while True:
+        status = check_status()
+        value = status.get("status")
+        if value in completed:
+            return status
+        if value in error:
+            raise AsyncJobError(status)
+        if value not in in_progress:
+            raise AsyncJobError({**status, "error": f"unexpected status {value!r}"})
+        if time.monotonic() >= deadline:
+            raise AsyncJobTimeoutError(
+                f"async job did not complete within {poll_time} seconds"
+            )
+        time.sleep(poll_interval)
diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py
new file mode 100644
index 0000000..b2d22a0
--- /dev/null
+++ b/tests/test_async_jobs.py
@@ -0,0 +1,64 @@
+from unittest.mock import patch
+
+import pytest
+
+from surge.async_jobs import (
+    AsyncJobError,
+    AsyncJobTimeoutError,
+    poll_async_job,
+)
+
+
+def test_returns_on_first_completed_status():
+    statuses = iter([{"status": "COMPLETED", "url": "u"}])
+    result = poll_async_job(lambda: next(statuses), poll_time=10, poll_interval=0)
+    assert result == {"status": "COMPLETED", "url": "u"}
+
+
+def test_polls_through_in_progress_then_returns_completed():
+    statuses = iter(
+        [
+            {"status": "IN_PROGRESS"},
+            {"status": "IN_PROGRESS"},
+            {"status": "COMPLETED", "url": "u"},
+        ]
+    )
+    with patch("time.sleep"):
+        result = poll_async_job(lambda: next(statuses), poll_time=10, poll_interval=0)
+    assert result == {"status": "COMPLETED", "url": "u"}
+
+
+def test_accepts_creating_and_ready_status_vocabulary():
+    statuses = iter(
+        [
+            {"status": "CREATING"},
+            {"status": "READY", "url": "u"},
+        ]
+    )
+    with patch("time.sleep"):
+        result = poll_async_job(lambda: next(statuses), poll_time=10, poll_interval=0)
+    assert result == {"status": "READY", "url": "u"}
+
+
+def test_raises_async_job_error_on_error_status():
+    with pytest.raises(AsyncJobError, match="kaboom"):
+        poll_async_job(
+            lambda: {"status": "ERROR", "error": "kaboom"},
+            poll_time=10,
+            poll_interval=0,
+        )
+
+
+def test_raises_async_job_error_on_unknown_status():
+    with pytest.raises(AsyncJobError, match="weird"):
+        poll_async_job(lambda: {"status": "weird"}, poll_time=10, poll_interval=0)
+
+
+def test_raises_timeout_when_poll_time_elapses():
+    with patch("time.sleep"):
+        with pytest.raises(AsyncJobTimeoutError):
+            poll_async_job(
+                lambda: {"status": "IN_PROGRESS"},
+                poll_time=0,
+                poll_interval=0,
+            )

From a6123b0ccc74b8dac652c3fae75bc672f6c9b63c Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 18:07:27 -0700
Subject: [PATCH 02/16] Check deadline before next poll

---
 surge/async_jobs.py      | 10 ++++++----
 tests/test_async_jobs.py | 27 +++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/surge/async_jobs.py b/surge/async_jobs.py
index 323ce32..d5d55da 100644
--- a/surge/async_jobs.py
+++ b/surge/async_jobs.py
@@ -46,7 +46,13 @@ def poll_async_job(
     error = set(error_statuses)
 
     deadline = time.monotonic() + poll_time
+    first_iteration = True
     while True:
+        if not first_iteration and time.monotonic() >= deadline:
+            raise AsyncJobTimeoutError(
+                f"async job did not complete within {poll_time} seconds"
+            )
+        first_iteration = False
         status = check_status()
         value = status.get("status")
         if value in completed:
@@ -55,8 +61,4 @@ def poll_async_job(
             raise AsyncJobError(status)
         if value not in in_progress:
             raise AsyncJobError({**status, "error": f"unexpected status {value!r}"})
-        if time.monotonic() >= deadline:
-            raise AsyncJobTimeoutError(
-                f"async job did not complete within {poll_time} seconds"
-            )
         time.sleep(poll_interval)
diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py
index b2d22a0..435fa99 100644
--- a/tests/test_async_jobs.py
+++ b/tests/test_async_jobs.py
@@ -62,3 +62,30 @@ def test_raises_timeout_when_poll_time_elapses():
                 poll_time=0,
                 poll_interval=0,
             )
+
+
+def test_raises_timeout_instead_of_returning_completed_after_deadline():
+    """When poll_interval > poll_time the loop sleeps past the deadline.
+    The second status must not be honored — even if it would be COMPLETED."""
+
+    def fake_sleep(_):
+        # Make the deadline elapse during the first sleep without actually
+        # waiting in the test.
+        import time as _time
+
+        _time.monotonic = lambda: 1e9
+
+    statuses = iter(
+        [
+            {"status": "IN_PROGRESS"},
+            {"status": "COMPLETED", "url": "u"},
+        ]
+    )
+
+    with patch("time.sleep", side_effect=fake_sleep):
+        with pytest.raises(AsyncJobTimeoutError):
+            poll_async_job(
+                lambda: next(statuses),
+                poll_time=1,
+                poll_interval=2,
+            )

From 6de3056d2f6ee7b377904ce4c2c745562d5e6cc6 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 18:12:58 -0700
Subject: [PATCH 03/16] Use poll_async_job in Report.save_report

---
 surge/reports.py      | 70 ++++++++++++++++++++---------------
 tests/test_reports.py | 86 +++++++++++++++++++++----------------------
 2 files changed, 82 insertions(+), 74 deletions(-)

diff --git a/surge/reports.py b/surge/reports.py
index 06579d0..62bfc38 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -1,5 +1,4 @@
 import gzip
-from time import sleep, monotonic
 import urllib
 import tempfile
 import shutil
@@ -8,17 +7,17 @@
 import warnings
 
 from surge.api_resource import REPORTS_ENDPOINT, APIResource
+from surge.async_jobs import AsyncJobError, AsyncJobTimeoutError, poll_async_job
 from surge.errors import SurgeRequestError
 
 
 class Report(APIResource):
-
     def __init__(self, **kwargs):
         super().__init__()
         self.__dict__.update(kwargs)
 
     def __str__(self):
-        return f"<surge.Report>"
+        return "<surge.Report>"
 
     def __repr__(self):
         return f"<surge.Report {self.attrs_repr()}>"
@@ -58,34 +57,48 @@ def save_report(
             poll_time (int): Maximum number of seconds to wait for the report to be generated
             poll_interval (int or float): Seconds to wait between status checks
         """
-        response = cls.request(project_id=project_id,
-                               type=type,
-                               api_key=api_key)
-        # Capture the job_id from the initial CREATING response and
-        # reuse it across polls. check_status's IN_PROGRESS response
-        # does not include job_id; only RETRYING does (and means the
-        # server kicked off a new underlying job).
-        job_id = getattr(response, "job_id", None)
-        deadline = monotonic() + poll_time
-        while response.status in ("CREATING", "IN_PROGRESS", "RETRYING"):
-            if monotonic() >= deadline:
+        initial = cls.request(project_id=project_id, type=type, api_key=api_key)
+        if initial.status == "READY":
+            url = initial.url
+        else:
+            # Capture the initial CREATING job_id; check_status's IN_PROGRESS
+            # response does not include one. RETRYING updates job_id mid-poll
+            # (server kicked off a new underlying job).
+            state = {"job_id": getattr(initial, "job_id", None)}
+
+            def _check():
+                r = cls.check_status(project_id, state["job_id"], api_key=api_key)
+                if r.status == "RETRYING":
+                    state["job_id"] = r.job_id
+                return vars(r)
+
+            try:
+                terminal = poll_async_job(
+                    _check,
+                    poll_time=poll_time,
+                    poll_interval=poll_interval,
+                    in_progress_statuses=("IN_PROGRESS", "CREATING", "RETRYING"),
+                )
+            except AsyncJobTimeoutError:
                 raise Exception(
-                    "Report failed to generate within {poll_time} seconds".
-                    format(poll_time=poll_time))
-            sleep(poll_interval)
-            response = cls.check_status(project_id, job_id, api_key=api_key)
-            if response.status == "RETRYING":
-                job_id = response.job_id
-
-        if response.status not in ("READY", "COMPLETED"):
-            raise ValueError("Report failed to generate with status {}".format(
-                response.status))
+                    "Report failed to generate within {poll_time} seconds".format(
+                        poll_time=poll_time
+                    )
+                )
+            except AsyncJobError as e:
+                raise ValueError(
+                    "Report failed to generate with status {}".format(
+                        e.status.get("status")
+                    )
+                )
+            url = terminal["url"]
 
         file_ext = "csv" if "csv" in type else "json"
         default_file_name = "project_{project_id}_results.{file_ext}".format(
-            project_id=project_id, file_ext=file_ext)
+            project_id=project_id, file_ext=file_ext
+        )
         target = filepath or default_file_name
-        with urllib.request.urlopen(response.url) as remote:
+        with urllib.request.urlopen(url) as remote:
             with tempfile.NamedTemporaryFile() as tmp_file:
                 shutil.copyfileobj(remote, tmp_file)
                 tmp_file.flush()
@@ -103,10 +116,7 @@ def save_report(
         return data
 
     @classmethod
-    def download_json(cls,
-                      project_id: str,
-                      poll_time=5 * 60,
-                      api_key: str = None):
+    def download_json(cls, project_id: str, poll_time=5 * 60, api_key: str = None):
         """
         Download and parse the results JSON for a project
 
diff --git a/tests/test_reports.py b/tests/test_reports.py
index 2b0341a..ac651b2 100644
--- a/tests/test_reports.py
+++ b/tests/test_reports.py
@@ -17,8 +17,7 @@ def test_save_report_on_empty_project_raises_an_error():
     with mock.patch.object(Report, "post") as mock_post:
         mock_post.return_value = {"error": "Project has no responses"}
         with pytest.raises(SurgeRequestError):
-            Report.save_report("fake_project_id", "export_csv",
-                               "my_report.csv")
+            Report.save_report("fake_project_id", "export_csv", "my_report.csv")
 
 
 def test_save_report_downloads_when_request_returns_ready():
@@ -26,12 +25,12 @@ def test_save_report_downloads_when_request_returns_ready():
     payload = b'[{"a": 1}]'
     ready = Report(status="READY", url="https://signed.example/report.gz")
     sink = io.BytesIO()
-    with mock.patch.object(Report, "request",
-                           return_value=ready) as mock_request, \
-            mock.patch.object(Report, "check_status") as mock_check, \
-            mock.patch("urllib.request.urlopen") as mock_urlopen:
-        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(
-            _gzipped(payload))
+    with (
+        mock.patch.object(Report, "request", return_value=ready) as mock_request,
+        mock.patch.object(Report, "check_status") as mock_check,
+        mock.patch("urllib.request.urlopen") as mock_urlopen,
+    ):
+        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(_gzipped(payload))
         Report.save_report("proj-123", "export_json", filepath=sink)
     mock_request.assert_called_once()
     mock_check.assert_not_called()
@@ -48,42 +47,43 @@ def test_save_report_polls_check_status_for_returned_job_id():
     payload = b"a,b\n1,2\n"
     creating = Report(status="CREATING", job_id="job-abc")
     in_progress = Report(status="IN_PROGRESS")
-    completed = Report(status="COMPLETED",
-                       url="https://signed.example/report.gz")
+    completed = Report(status="COMPLETED", url="https://signed.example/report.gz")
     sink = io.BytesIO()
-    with mock.patch.object(Report, "request",
-                           return_value=creating) as mock_request, \
-            mock.patch.object(Report, "check_status",
-                              side_effect=[in_progress, completed]) as mock_check, \
-            mock.patch("urllib.request.urlopen") as mock_urlopen, \
-            mock.patch("surge.reports.sleep") as mock_sleep:
-        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(
-            _gzipped(payload))
+    with (
+        mock.patch.object(Report, "request", return_value=creating) as mock_request,
+        mock.patch.object(
+            Report, "check_status", side_effect=[in_progress, completed]
+        ) as mock_check,
+        mock.patch("urllib.request.urlopen") as mock_urlopen,
+        mock.patch("time.sleep") as mock_sleep,
+    ):
+        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(_gzipped(payload))
         Report.save_report("proj-123", "export_csv", filepath=sink)
     mock_request.assert_called_once()
     assert mock_check.call_count == 2
     for call in mock_check.call_args_list:
         assert call.args[:2] == ("proj-123", "job-abc")
-    assert mock_sleep.call_count == 2
+    assert mock_sleep.call_count == 1
     assert sink.getvalue() == payload
 
 
 def test_save_report_switches_job_id_on_retrying():
     """RETRYING responses include a new job_id; subsequent polls use it."""
-    payload = b'[]'
+    payload = b"[]"
     creating = Report(status="CREATING", job_id="job-abc")
     retrying = Report(status="RETRYING", job_id="job-xyz")
     in_progress = Report(status="IN_PROGRESS")
-    completed = Report(status="COMPLETED",
-                       url="https://signed.example/report.gz")
+    completed = Report(status="COMPLETED", url="https://signed.example/report.gz")
     sink = io.BytesIO()
-    with mock.patch.object(Report, "request", return_value=creating), \
-            mock.patch.object(Report, "check_status",
-                              side_effect=[retrying, in_progress, completed]) as mock_check, \
-            mock.patch("urllib.request.urlopen") as mock_urlopen, \
-            mock.patch("surge.reports.sleep"):
-        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(
-            _gzipped(payload))
+    with (
+        mock.patch.object(Report, "request", return_value=creating),
+        mock.patch.object(
+            Report, "check_status", side_effect=[retrying, in_progress, completed]
+        ) as mock_check,
+        mock.patch("urllib.request.urlopen") as mock_urlopen,
+        mock.patch("time.sleep"),
+    ):
+        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(_gzipped(payload))
         Report.save_report("proj-123", "export_json", filepath=sink)
     job_ids = [call.args[1] for call in mock_check.call_args_list]
     assert job_ids == ["job-abc", "job-xyz", "job-xyz"]
@@ -92,25 +92,23 @@ def test_save_report_switches_job_id_on_retrying():
 def test_save_report_raises_on_unexpected_status():
     creating = Report(status="CREATING", job_id="job-abc")
     error = Report(status="ERROR", type="Report generation error")
-    with mock.patch.object(Report, "request", return_value=creating), \
-            mock.patch.object(Report, "check_status", return_value=error), \
-            mock.patch("surge.reports.sleep"):
+    with (
+        mock.patch.object(Report, "request", return_value=creating),
+        mock.patch.object(Report, "check_status", return_value=error),
+        mock.patch("time.sleep"),
+    ):
         with pytest.raises(ValueError, match="ERROR"):
-            Report.save_report("proj-123",
-                               "export_json",
-                               filepath=io.BytesIO())
+            Report.save_report("proj-123", "export_json", filepath=io.BytesIO())
 
 
 def test_save_report_times_out_when_job_never_completes():
     creating = Report(status="CREATING", job_id="job-abc")
     in_progress = Report(status="IN_PROGRESS", job_id="job-abc")
-    with mock.patch.object(Report, "request", return_value=creating), \
-            mock.patch.object(Report, "check_status",
-                              return_value=in_progress), \
-            mock.patch("surge.reports.sleep"), \
-            mock.patch("surge.reports.monotonic",
-                       side_effect=[0.0, 1000.0, 2000.0]):
+    with (
+        mock.patch.object(Report, "request", return_value=creating),
+        mock.patch.object(Report, "check_status", return_value=in_progress),
+        mock.patch("time.sleep"),
+        mock.patch("time.monotonic", side_effect=[0.0, 1000.0, 2000.0]),
+    ):
         with pytest.raises(Exception, match="within 300 seconds"):
-            Report.save_report("proj-123",
-                               "export_json",
-                               filepath=io.BytesIO())
+            Report.save_report("proj-123", "export_json", filepath=io.BytesIO())

From c6766fbda4daabe970f4d3d8d0827ffbadf5f27e Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 18:22:12 -0700
Subject: [PATCH 04/16] Cap sleep to remaining poll_time; yapf

---
 surge/async_jobs.py      | 12 +++--
 surge/reports.py         | 30 +++++++------
 tests/test_async_jobs.py | 96 ++++++++++++++++++++++++++++++----------
 tests/test_reports.py    | 75 ++++++++++++++++++-------------
 4 files changed, 142 insertions(+), 71 deletions(-)

diff --git a/surge/async_jobs.py b/surge/async_jobs.py
index d5d55da..07506a6 100644
--- a/surge/async_jobs.py
+++ b/surge/async_jobs.py
@@ -50,8 +50,7 @@ def poll_async_job(
     while True:
         if not first_iteration and time.monotonic() >= deadline:
             raise AsyncJobTimeoutError(
-                f"async job did not complete within {poll_time} seconds"
-            )
+                f"async job did not complete within {poll_time} seconds")
         first_iteration = False
         status = check_status()
         value = status.get("status")
@@ -60,5 +59,10 @@ def poll_async_job(
         if value in error:
             raise AsyncJobError(status)
         if value not in in_progress:
-            raise AsyncJobError({**status, "error": f"unexpected status {value!r}"})
-        time.sleep(poll_interval)
+            raise AsyncJobError({
+                **status, "error":
+                f"unexpected status {value!r}"
+            })
+        remaining = deadline - time.monotonic()
+        if remaining > 0:
+            time.sleep(min(poll_interval, remaining))
diff --git a/surge/reports.py b/surge/reports.py
index 62bfc38..e30dd2b 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -12,6 +12,7 @@
 
 
 class Report(APIResource):
+
     def __init__(self, **kwargs):
         super().__init__()
         self.__dict__.update(kwargs)
@@ -57,7 +58,9 @@ def save_report(
             poll_time (int): Maximum number of seconds to wait for the report to be generated
             poll_interval (int or float): Seconds to wait between status checks
         """
-        initial = cls.request(project_id=project_id, type=type, api_key=api_key)
+        initial = cls.request(project_id=project_id,
+                              type=type,
+                              api_key=api_key)
         if initial.status == "READY":
             url = initial.url
         else:
@@ -67,7 +70,9 @@ def save_report(
             state = {"job_id": getattr(initial, "job_id", None)}
 
             def _check():
-                r = cls.check_status(project_id, state["job_id"], api_key=api_key)
+                r = cls.check_status(project_id,
+                                     state["job_id"],
+                                     api_key=api_key)
                 if r.status == "RETRYING":
                     state["job_id"] = r.job_id
                 return vars(r)
@@ -77,26 +82,22 @@ def _check():
                     _check,
                     poll_time=poll_time,
                     poll_interval=poll_interval,
-                    in_progress_statuses=("IN_PROGRESS", "CREATING", "RETRYING"),
+                    in_progress_statuses=("IN_PROGRESS", "CREATING",
+                                          "RETRYING"),
                 )
             except AsyncJobTimeoutError:
                 raise Exception(
-                    "Report failed to generate within {poll_time} seconds".format(
-                        poll_time=poll_time
-                    )
-                )
+                    "Report failed to generate within {poll_time} seconds".
+                    format(poll_time=poll_time))
             except AsyncJobError as e:
                 raise ValueError(
                     "Report failed to generate with status {}".format(
-                        e.status.get("status")
-                    )
-                )
+                        e.status.get("status")))
             url = terminal["url"]
 
         file_ext = "csv" if "csv" in type else "json"
         default_file_name = "project_{project_id}_results.{file_ext}".format(
-            project_id=project_id, file_ext=file_ext
-        )
+            project_id=project_id, file_ext=file_ext)
         target = filepath or default_file_name
         with urllib.request.urlopen(url) as remote:
             with tempfile.NamedTemporaryFile() as tmp_file:
@@ -116,7 +117,10 @@ def _check():
         return data
 
     @classmethod
-    def download_json(cls, project_id: str, poll_time=5 * 60, api_key: str = None):
+    def download_json(cls,
+                      project_id: str,
+                      poll_time=5 * 60,
+                      api_key: str = None):
         """
         Download and parse the results JSON for a project
 
diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py
index 435fa99..19fe797 100644
--- a/tests/test_async_jobs.py
+++ b/tests/test_async_jobs.py
@@ -11,39 +11,56 @@
 
 def test_returns_on_first_completed_status():
     statuses = iter([{"status": "COMPLETED", "url": "u"}])
-    result = poll_async_job(lambda: next(statuses), poll_time=10, poll_interval=0)
+    result = poll_async_job(lambda: next(statuses),
+                            poll_time=10,
+                            poll_interval=0)
     assert result == {"status": "COMPLETED", "url": "u"}
 
 
 def test_polls_through_in_progress_then_returns_completed():
-    statuses = iter(
-        [
-            {"status": "IN_PROGRESS"},
-            {"status": "IN_PROGRESS"},
-            {"status": "COMPLETED", "url": "u"},
-        ]
-    )
+    statuses = iter([
+        {
+            "status": "IN_PROGRESS"
+        },
+        {
+            "status": "IN_PROGRESS"
+        },
+        {
+            "status": "COMPLETED",
+            "url": "u"
+        },
+    ])
     with patch("time.sleep"):
-        result = poll_async_job(lambda: next(statuses), poll_time=10, poll_interval=0)
+        result = poll_async_job(lambda: next(statuses),
+                                poll_time=10,
+                                poll_interval=0)
     assert result == {"status": "COMPLETED", "url": "u"}
 
 
 def test_accepts_creating_and_ready_status_vocabulary():
-    statuses = iter(
-        [
-            {"status": "CREATING"},
-            {"status": "READY", "url": "u"},
-        ]
-    )
+    statuses = iter([
+        {
+            "status": "CREATING"
+        },
+        {
+            "status": "READY",
+            "url": "u"
+        },
+    ])
     with patch("time.sleep"):
-        result = poll_async_job(lambda: next(statuses), poll_time=10, poll_interval=0)
+        result = poll_async_job(lambda: next(statuses),
+                                poll_time=10,
+                                poll_interval=0)
     assert result == {"status": "READY", "url": "u"}
 
 
 def test_raises_async_job_error_on_error_status():
     with pytest.raises(AsyncJobError, match="kaboom"):
         poll_async_job(
-            lambda: {"status": "ERROR", "error": "kaboom"},
+            lambda: {
+                "status": "ERROR",
+                "error": "kaboom"
+            },
             poll_time=10,
             poll_interval=0,
         )
@@ -51,7 +68,9 @@ def test_raises_async_job_error_on_error_status():
 
 def test_raises_async_job_error_on_unknown_status():
     with pytest.raises(AsyncJobError, match="weird"):
-        poll_async_job(lambda: {"status": "weird"}, poll_time=10, poll_interval=0)
+        poll_async_job(lambda: {"status": "weird"},
+                       poll_time=10,
+                       poll_interval=0)
 
 
 def test_raises_timeout_when_poll_time_elapses():
@@ -64,6 +83,32 @@ def test_raises_timeout_when_poll_time_elapses():
             )
 
 
+def test_caps_sleep_to_remaining_poll_time():
+    """Sleep should never exceed the remaining budget — even with a
+    poll_interval much larger than poll_time."""
+
+    statuses = iter([
+        {
+            "status": "IN_PROGRESS"
+        },
+        {
+            "status": "COMPLETED",
+            "url": "u"
+        },
+    ])
+
+    with patch("time.sleep") as mock_sleep:
+        poll_async_job(
+            lambda: next(statuses),
+            poll_time=5,
+            poll_interval=60,
+        )
+
+    assert mock_sleep.call_count == 1
+    (slept_seconds, ) = mock_sleep.call_args.args
+    assert slept_seconds <= 5
+
+
 def test_raises_timeout_instead_of_returning_completed_after_deadline():
     """When poll_interval > poll_time the loop sleeps past the deadline.
     The second status must not be honored — even if it would be COMPLETED."""
@@ -75,12 +120,15 @@ def fake_sleep(_):
 
         _time.monotonic = lambda: 1e9
 
-    statuses = iter(
-        [
-            {"status": "IN_PROGRESS"},
-            {"status": "COMPLETED", "url": "u"},
-        ]
-    )
+    statuses = iter([
+        {
+            "status": "IN_PROGRESS"
+        },
+        {
+            "status": "COMPLETED",
+            "url": "u"
+        },
+    ])
 
     with patch("time.sleep", side_effect=fake_sleep):
         with pytest.raises(AsyncJobTimeoutError):
diff --git a/tests/test_reports.py b/tests/test_reports.py
index ac651b2..24aa665 100644
--- a/tests/test_reports.py
+++ b/tests/test_reports.py
@@ -17,7 +17,8 @@ def test_save_report_on_empty_project_raises_an_error():
     with mock.patch.object(Report, "post") as mock_post:
         mock_post.return_value = {"error": "Project has no responses"}
         with pytest.raises(SurgeRequestError):
-            Report.save_report("fake_project_id", "export_csv", "my_report.csv")
+            Report.save_report("fake_project_id", "export_csv",
+                               "my_report.csv")
 
 
 def test_save_report_downloads_when_request_returns_ready():
@@ -26,11 +27,13 @@ def test_save_report_downloads_when_request_returns_ready():
     ready = Report(status="READY", url="https://signed.example/report.gz")
     sink = io.BytesIO()
     with (
-        mock.patch.object(Report, "request", return_value=ready) as mock_request,
-        mock.patch.object(Report, "check_status") as mock_check,
-        mock.patch("urllib.request.urlopen") as mock_urlopen,
+            mock.patch.object(Report, "request", return_value=ready) as
+            mock_request,
+            mock.patch.object(Report, "check_status") as mock_check,
+            mock.patch("urllib.request.urlopen") as mock_urlopen,
     ):
-        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(_gzipped(payload))
+        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(
+            _gzipped(payload))
         Report.save_report("proj-123", "export_json", filepath=sink)
     mock_request.assert_called_once()
     mock_check.assert_not_called()
@@ -47,17 +50,21 @@ def test_save_report_polls_check_status_for_returned_job_id():
     payload = b"a,b\n1,2\n"
     creating = Report(status="CREATING", job_id="job-abc")
     in_progress = Report(status="IN_PROGRESS")
-    completed = Report(status="COMPLETED", url="https://signed.example/report.gz")
+    completed = Report(status="COMPLETED",
+                       url="https://signed.example/report.gz")
     sink = io.BytesIO()
     with (
-        mock.patch.object(Report, "request", return_value=creating) as mock_request,
-        mock.patch.object(
-            Report, "check_status", side_effect=[in_progress, completed]
-        ) as mock_check,
-        mock.patch("urllib.request.urlopen") as mock_urlopen,
-        mock.patch("time.sleep") as mock_sleep,
+            mock.patch.object(Report, "request",
+                              return_value=creating) as mock_request,
+            mock.patch.object(Report,
+                              "check_status",
+                              side_effect=[in_progress, completed]) as
+            mock_check,
+            mock.patch("urllib.request.urlopen") as mock_urlopen,
+            mock.patch("time.sleep") as mock_sleep,
     ):
-        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(_gzipped(payload))
+        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(
+            _gzipped(payload))
         Report.save_report("proj-123", "export_csv", filepath=sink)
     mock_request.assert_called_once()
     assert mock_check.call_count == 2
@@ -73,17 +80,20 @@ def test_save_report_switches_job_id_on_retrying():
     creating = Report(status="CREATING", job_id="job-abc")
     retrying = Report(status="RETRYING", job_id="job-xyz")
     in_progress = Report(status="IN_PROGRESS")
-    completed = Report(status="COMPLETED", url="https://signed.example/report.gz")
+    completed = Report(status="COMPLETED",
+                       url="https://signed.example/report.gz")
     sink = io.BytesIO()
     with (
-        mock.patch.object(Report, "request", return_value=creating),
-        mock.patch.object(
-            Report, "check_status", side_effect=[retrying, in_progress, completed]
-        ) as mock_check,
-        mock.patch("urllib.request.urlopen") as mock_urlopen,
-        mock.patch("time.sleep"),
+            mock.patch.object(Report, "request", return_value=creating),
+            mock.patch.object(Report,
+                              "check_status",
+                              side_effect=[retrying, in_progress, completed])
+            as mock_check,
+            mock.patch("urllib.request.urlopen") as mock_urlopen,
+            mock.patch("time.sleep"),
     ):
-        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(_gzipped(payload))
+        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(
+            _gzipped(payload))
         Report.save_report("proj-123", "export_json", filepath=sink)
     job_ids = [call.args[1] for call in mock_check.call_args_list]
     assert job_ids == ["job-abc", "job-xyz", "job-xyz"]
@@ -93,22 +103,27 @@ def test_save_report_raises_on_unexpected_status():
     creating = Report(status="CREATING", job_id="job-abc")
     error = Report(status="ERROR", type="Report generation error")
     with (
-        mock.patch.object(Report, "request", return_value=creating),
-        mock.patch.object(Report, "check_status", return_value=error),
-        mock.patch("time.sleep"),
+            mock.patch.object(Report, "request", return_value=creating),
+            mock.patch.object(Report, "check_status", return_value=error),
+            mock.patch("time.sleep"),
     ):
         with pytest.raises(ValueError, match="ERROR"):
-            Report.save_report("proj-123", "export_json", filepath=io.BytesIO())
+            Report.save_report("proj-123",
+                               "export_json",
+                               filepath=io.BytesIO())
 
 
 def test_save_report_times_out_when_job_never_completes():
     creating = Report(status="CREATING", job_id="job-abc")
     in_progress = Report(status="IN_PROGRESS", job_id="job-abc")
     with (
-        mock.patch.object(Report, "request", return_value=creating),
-        mock.patch.object(Report, "check_status", return_value=in_progress),
-        mock.patch("time.sleep"),
-        mock.patch("time.monotonic", side_effect=[0.0, 1000.0, 2000.0]),
+            mock.patch.object(Report, "request", return_value=creating),
+            mock.patch.object(Report, "check_status",
+                              return_value=in_progress),
+            mock.patch("time.sleep"),
+            mock.patch("time.monotonic", side_effect=[0.0, 1000.0, 2000.0]),
     ):
         with pytest.raises(Exception, match="within 300 seconds"):
-            Report.save_report("proj-123", "export_json", filepath=io.BytesIO())
+            Report.save_report("proj-123",
+                               "export_json",
+                               filepath=io.BytesIO())

From 9209b86e6aca5ef0e15ab87a7e7267bdf1e2c5b4 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 18:33:44 -0700
Subject: [PATCH 05/16] Add extra_params kwarg to Report.request and
 save_report

---
 surge/reports.py      | 22 +++++++++++++++++++---
 tests/test_reports.py | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/surge/reports.py b/surge/reports.py
index e30dd2b..98ec79c 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -35,6 +35,8 @@ def save_report(
         poll_time=5 * 60,
         api_key: str = None,
         poll_interval: float = 2,
+        *,
+        extra_params: dict = None,
     ):
         """
         Request creation of a report, poll until the report is generated, and save the data to a file all in one call.
@@ -57,10 +59,13 @@ def save_report(
             filepath (string or IO or None): Location to save the results file. If not specified, will save to "project_{project_id}_results.{csv/json}
             poll_time (int): Maximum number of seconds to wait for the report to be generated
             poll_interval (int or float): Seconds to wait between status checks
+            extra_params (dict, optional): Additional params merged into the
+                initial report request body. See ``request``.
         """
         initial = cls.request(project_id=project_id,
                               type=type,
-                              api_key=api_key)
+                              api_key=api_key,
+                              extra_params=extra_params)
         if initial.status == "READY":
             url = initial.url
         else:
@@ -143,7 +148,14 @@ def download_json(cls,
         return json.load(bytesio)
 
     @classmethod
-    def request(cls, project_id: str, type: str, api_key: str = None):
+    def request(
+        cls,
+        project_id: str,
+        type: str,
+        api_key: str = None,
+        *,
+        extra_params: dict = None,
+    ):
         """
         Request creation of a report for the given type. Note that reports are generated
         asychronously so the response may include a `job_id` which needs to be used with
@@ -178,12 +190,16 @@ def request(cls, project_id: str, type: str, api_key: str = None):
         Arguments:
             project_id (str): ID of project.
             report_type (str): report type
+            extra_params (dict, optional): Additional params merged into the
+                report request body. Use for backend-supported params that
+                aren't first-class kwargs on this method; refer to the API
+                reference for valid keys.
 
         Returns:
             status: Report status object which includes report id
         """
         endpoint = f"{REPORTS_ENDPOINT}/{project_id}/report"
-        params = {"report_type": type}
+        params = {"report_type": type, **(extra_params or {})}
         response_json = cls.post(endpoint, params, api_key=api_key)
         if "error" in response_json:
             raise SurgeRequestError(response_json["error"])
diff --git a/tests/test_reports.py b/tests/test_reports.py
index 24aa665..f0ac81d 100644
--- a/tests/test_reports.py
+++ b/tests/test_reports.py
@@ -127,3 +127,39 @@ def test_save_report_times_out_when_job_never_completes():
             Report.save_report("proj-123",
                                "export_json",
                                filepath=io.BytesIO())
+
+
+def test_request_merges_extra_params_into_post_body():
+    with mock.patch.object(Report, "post") as mock_post:
+        mock_post.return_value = {"status": "CREATING", "job_id": "j"}
+        Report.request("proj-1", "export_json", extra_params={"item_id": "i"})
+    endpoint, params = mock_post.call_args.args[:2]
+    assert "report" in endpoint
+    assert params == {"report_type": "export_json", "item_id": "i"}
+
+
+def test_request_without_extra_params_sends_only_report_type():
+    with mock.patch.object(Report, "post") as mock_post:
+        mock_post.return_value = {"status": "CREATING", "job_id": "j"}
+        Report.request("proj-1", "export_json")
+    _, params = mock_post.call_args.args[:2]
+    assert params == {"report_type": "export_json"}
+
+
+def test_save_report_forwards_extra_params_to_request():
+    payload = b"[]"
+    ready = Report(status="READY", url="https://signed.example/report.gz")
+    with (
+            mock.patch.object(Report, "request", return_value=ready) as
+            mock_request,
+            mock.patch("urllib.request.urlopen") as mock_urlopen,
+    ):
+        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(
+            _gzipped(payload))
+        Report.save_report(
+            "proj-1",
+            "export_json",
+            filepath=io.BytesIO(),
+            extra_params={"item_id": "i"},
+        )
+    assert mock_request.call_args.kwargs["extra_params"] == {"item_id": "i"}

From f473c955a6a5039a4d1b4201a478d53305db9013 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 18:38:51 -0700
Subject: [PATCH 06/16] Treat COMPLETED as terminal in initial request response

---
 surge/reports.py      |  2 +-
 tests/test_reports.py | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/surge/reports.py b/surge/reports.py
index 98ec79c..86fa3fd 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -66,7 +66,7 @@ def save_report(
                               type=type,
                               api_key=api_key,
                               extra_params=extra_params)
-        if initial.status == "READY":
+        if initial.status in ("READY", "COMPLETED"):
             url = initial.url
         else:
             # Capture the initial CREATING job_id; check_status's IN_PROGRESS
diff --git a/tests/test_reports.py b/tests/test_reports.py
index f0ac81d..1dd88bb 100644
--- a/tests/test_reports.py
+++ b/tests/test_reports.py
@@ -146,6 +146,26 @@ def test_request_without_extra_params_sends_only_report_type():
     assert params == {"report_type": "export_json"}
 
 
+def test_save_report_treats_completed_initial_response_as_terminal():
+    """If request() returns COMPLETED (not just READY), skip polling."""
+    payload = b"[]"
+    completed = Report(status="COMPLETED",
+                       url="https://signed.example/report.gz")
+    sink = io.BytesIO()
+    with (
+            mock.patch.object(Report, "request", return_value=completed) as
+            mock_request,
+            mock.patch.object(Report, "check_status") as mock_check,
+            mock.patch("urllib.request.urlopen") as mock_urlopen,
+    ):
+        mock_urlopen.return_value.__enter__.return_value = io.BytesIO(
+            _gzipped(payload))
+        Report.save_report("proj-c", "export_json", filepath=sink)
+    mock_request.assert_called_once()
+    mock_check.assert_not_called()
+    assert sink.getvalue() == payload
+
+
 def test_save_report_forwards_extra_params_to_request():
     payload = b"[]"
     ready = Report(status="READY", url="https://signed.example/report.gz")

From 47671a62e47a38de9a5f3b68a8089a1e2efc411b Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 18:44:25 -0700
Subject: [PATCH 07/16] Use mock.patch for time.monotonic to avoid leaking
 state

---
 tests/test_async_jobs.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py
index 19fe797..bcaf6ef 100644
--- a/tests/test_async_jobs.py
+++ b/tests/test_async_jobs.py
@@ -113,13 +113,6 @@ def test_raises_timeout_instead_of_returning_completed_after_deadline():
     """When poll_interval > poll_time the loop sleeps past the deadline.
     The second status must not be honored — even if it would be COMPLETED."""
 
-    def fake_sleep(_):
-        # Make the deadline elapse during the first sleep without actually
-        # waiting in the test.
-        import time as _time
-
-        _time.monotonic = lambda: 1e9
-
     statuses = iter([
         {
             "status": "IN_PROGRESS"
@@ -130,7 +123,10 @@ def fake_sleep(_):
         },
     ])
 
-    with patch("time.sleep", side_effect=fake_sleep):
+    # Simulate the deadline elapsing immediately after the first poll:
+    # call 1 sets the deadline, calls 2+ all report well past it.
+    with patch("time.sleep"), patch("time.monotonic",
+                                    side_effect=[0.0, 1e9, 1e9, 1e9]):
         with pytest.raises(AsyncJobTimeoutError):
             poll_async_job(
                 lambda: next(statuses),

From e82b07989580c7ab25227fa89d1e57dc08c79232 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 18:50:15 -0700
Subject: [PATCH 08/16] Raise immediately on terminal error from initial
 request

---
 surge/reports.py      |  3 +++
 tests/test_reports.py | 12 ++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/surge/reports.py b/surge/reports.py
index 86fa3fd..c8f4534 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -68,6 +68,9 @@ def save_report(
                               extra_params=extra_params)
         if initial.status in ("READY", "COMPLETED"):
             url = initial.url
+        elif initial.status not in ("CREATING", "IN_PROGRESS", "RETRYING"):
+            raise ValueError("Report failed to generate with status {}".format(
+                initial.status))
         else:
             # Capture the initial CREATING job_id; check_status's IN_PROGRESS
             # response does not include one. RETRYING updates job_id mid-poll
diff --git a/tests/test_reports.py b/tests/test_reports.py
index 1dd88bb..e472312 100644
--- a/tests/test_reports.py
+++ b/tests/test_reports.py
@@ -146,6 +146,18 @@ def test_request_without_extra_params_sends_only_report_type():
     assert params == {"report_type": "export_json"}
 
 
+def test_save_report_raises_when_initial_request_returns_error():
+    """If request() returns ERROR directly, raise without polling check_status."""
+    error = Report(status="ERROR", type="Report generation error")
+    with (
+            mock.patch.object(Report, "request", return_value=error),
+            mock.patch.object(Report, "check_status") as mock_check,
+    ):
+        with pytest.raises(ValueError, match="ERROR"):
+            Report.save_report("proj-e", "export_json", filepath=io.BytesIO())
+    mock_check.assert_not_called()
+
+
 def test_save_report_treats_completed_initial_response_as_terminal():
     """If request() returns COMPLETED (not just READY), skip polling."""
     payload = b"[]"

From 42a623357d3a9db808035dbb9fdec3cf240c5291 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 18:57:21 -0700
Subject: [PATCH 09/16] Reject report_type override in extra_params

---
 surge/reports.py      |  7 ++++++-
 tests/test_reports.py | 11 +++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/surge/reports.py b/surge/reports.py
index c8f4534..d2f034b 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -201,8 +201,13 @@ def request(
         Returns:
             status: Report status object which includes report id
         """
+        extra = extra_params or {}
+        if "report_type" in extra:
+            raise ValueError(
+                "`report_type` is set via the `type` argument and cannot "
+                "be passed in `extra_params`")
         endpoint = f"{REPORTS_ENDPOINT}/{project_id}/report"
-        params = {"report_type": type, **(extra_params or {})}
+        params = {"report_type": type, **extra}
         response_json = cls.post(endpoint, params, api_key=api_key)
         if "error" in response_json:
             raise SurgeRequestError(response_json["error"])
diff --git a/tests/test_reports.py b/tests/test_reports.py
index e472312..a42343e 100644
--- a/tests/test_reports.py
+++ b/tests/test_reports.py
@@ -178,6 +178,17 @@ def test_save_report_treats_completed_initial_response_as_terminal():
     assert sink.getvalue() == payload
 
 
+def test_request_rejects_report_type_in_extra_params():
+    """report_type is set via the `type` argument; allowing it in
+    extra_params would silently override the explicit `type`."""
+    with mock.patch.object(Report, "post") as mock_post:
+        with pytest.raises(ValueError, match="report_type"):
+            Report.request("proj-1",
+                           "export_csv",
+                           extra_params={"report_type": "export_json"})
+    mock_post.assert_not_called()
+
+
 def test_save_report_forwards_extra_params_to_request():
     payload = b"[]"
     ready = Report(status="READY", url="https://signed.example/report.gz")

From babe9b1c307aeb616445f20824b0bc859be76250 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Tue, 26 May 2026 23:39:33 -0700
Subject: [PATCH 10/16] Document save_report as the recommended entry point

---
 surge/reports.py | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/surge/reports.py b/surge/reports.py
index d2f034b..28e351e 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -160,11 +160,19 @@ def request(
         extra_params: dict = None,
     ):
         """
-        Request creation of a report for the given type. Note that reports are generated
-        asychronously so the response may include a `job_id` which needs to be used with
-        the `status` method to get the job status. In the event that the report has is
-        already generated and current, the report URL will be returned. Note that the URL
-        is a presigned URL which is active for only a limited duration.
+        Request creation of a report for the given type.
+
+        Most callers want `save_report` instead — it kicks off the
+        request, polls the specific job, and saves the result to a
+        file in a single call. Use this raw method for fire-and-forget
+        batches or custom poll loops.
+
+        Reports are generated asychronously so the response may
+        include a `job_id` which needs to be used with `check_status`
+        to get the job status. In the event that the report has is
+        already generated and current, the report URL will be returned.
+        Note that the URL is a presigned URL which is active for only
+        a limited duration.
 
         Type may be one of these types:
           * `export_json`
@@ -259,7 +267,12 @@ def status(cls, project_id: str, job_id: str, api_key: str = None):
     @classmethod
     def check_status(cls, project_id: str, job_id: str, api_key: str = None):
         """
-        Checks the status of a given report job. The response will be of one of these shapes:
+        Checks the status of a given report job.
+
+        Most callers want `save_report` instead — this method is for
+        custom poll loops over a `job_id` returned by the raw `request`.
+
+        The response will be of one of these shapes:
 
         1) Report is still being generated (HTTP 202):
 

From 76169a8cce8e41574105d38ddca4d9b06b354737 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Thu, 28 May 2026 10:44:25 -0700
Subject: [PATCH 11/16] Fold initial request into poll loop

Drops the separate dispatch on the initial request response by treating
it as the first iteration of poll_async_job. Trades one extra
poll_interval sleep on the not-yet-ready path for ~12 fewer lines and a
single status-handling code path.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 surge/reports.py      | 68 ++++++++++++++++++++-----------------------
 tests/test_reports.py |  2 +-
 2 files changed, 32 insertions(+), 38 deletions(-)

diff --git a/surge/reports.py b/surge/reports.py
index 28e351e..294133a 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -62,46 +62,40 @@ def save_report(
             extra_params (dict, optional): Additional params merged into the
                 initial report request body. See ``request``.
         """
-        initial = cls.request(project_id=project_id,
-                              type=type,
-                              api_key=api_key,
-                              extra_params=extra_params)
-        if initial.status in ("READY", "COMPLETED"):
-            url = initial.url
-        elif initial.status not in ("CREATING", "IN_PROGRESS", "RETRYING"):
-            raise ValueError("Report failed to generate with status {}".format(
-                initial.status))
-        else:
-            # Capture the initial CREATING job_id; check_status's IN_PROGRESS
-            # response does not include one. RETRYING updates job_id mid-poll
-            # (server kicked off a new underlying job).
-            state = {"job_id": getattr(initial, "job_id", None)}
-
-            def _check():
+        # job_id arrives on the first request and on RETRYING swaps;
+        # IN_PROGRESS responses omit it, so only update when present.
+        state = {"job_id": None, "first": True}
+
+        def _check():
+            if state["first"]:
+                state["first"] = False
+                r = cls.request(project_id=project_id,
+                                type=type,
+                                api_key=api_key,
+                                extra_params=extra_params)
+            else:
                 r = cls.check_status(project_id,
                                      state["job_id"],
                                      api_key=api_key)
-                if r.status == "RETRYING":
-                    state["job_id"] = r.job_id
-                return vars(r)
-
-            try:
-                terminal = poll_async_job(
-                    _check,
-                    poll_time=poll_time,
-                    poll_interval=poll_interval,
-                    in_progress_statuses=("IN_PROGRESS", "CREATING",
-                                          "RETRYING"),
-                )
-            except AsyncJobTimeoutError:
-                raise Exception(
-                    "Report failed to generate within {poll_time} seconds".
-                    format(poll_time=poll_time))
-            except AsyncJobError as e:
-                raise ValueError(
-                    "Report failed to generate with status {}".format(
-                        e.status.get("status")))
-            url = terminal["url"]
+            if getattr(r, "job_id", None):
+                state["job_id"] = r.job_id
+            return vars(r)
+
+        try:
+            terminal = poll_async_job(
+                _check,
+                poll_time=poll_time,
+                poll_interval=poll_interval,
+                in_progress_statuses=("IN_PROGRESS", "CREATING", "RETRYING"),
+            )
+        except AsyncJobTimeoutError:
+            raise Exception(
+                "Report failed to generate within {poll_time} seconds".format(
+                    poll_time=poll_time))
+        except AsyncJobError as e:
+            raise ValueError("Report failed to generate with status {}".format(
+                e.status.get("status")))
+        url = terminal["url"]
 
         file_ext = "csv" if "csv" in type else "json"
         default_file_name = "project_{project_id}_results.{file_ext}".format(
diff --git a/tests/test_reports.py b/tests/test_reports.py
index a42343e..46f8fc7 100644
--- a/tests/test_reports.py
+++ b/tests/test_reports.py
@@ -70,7 +70,7 @@ def test_save_report_polls_check_status_for_returned_job_id():
     assert mock_check.call_count == 2
     for call in mock_check.call_args_list:
         assert call.args[:2] == ("proj-123", "job-abc")
-    assert mock_sleep.call_count == 1
+    assert mock_sleep.call_count == 2
     assert sink.getvalue() == payload
 
 

From 2d1eed6e5a8669b475bf6513c7ddf43d9674f9af Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Thu, 28 May 2026 11:13:01 -0700
Subject: [PATCH 12/16] Collapse AsyncJobTimeoutError into AsyncJobError

Timeouts now surface as AsyncJobError with a synthetic
status="TIMEOUT" sentinel, letting callers distinguish them from
server-side failures via e.status without a separate exception class.

Also moves the deadline check after each poll instead of at the top of
the next iteration, so a job that completes during the final sleep is
honored instead of being lost to a timeout.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 surge/async_jobs.py      | 31 ++++++++++++++-----------------
 surge/reports.py         |  9 ++-------
 tests/test_async_jobs.py | 12 +++++-------
 tests/test_reports.py    |  4 ++--
 4 files changed, 23 insertions(+), 33 deletions(-)

diff --git a/surge/async_jobs.py b/surge/async_jobs.py
index 07506a6..8ebf829 100644
--- a/surge/async_jobs.py
+++ b/surge/async_jobs.py
@@ -8,7 +8,12 @@
 
 
 class AsyncJobError(Exception):
-    """Raised when an async job returns a terminal error status."""
+    """Raised when an async job fails or times out.
+
+    The ``status`` dict carries the underlying error; timeouts use the
+    synthetic ``status="TIMEOUT"`` sentinel so callers can distinguish
+    them from server-side failures.
+    """
 
     def __init__(self, status):
         self.status = status
@@ -16,10 +21,6 @@ def __init__(self, status):
         super().__init__(error)
 
 
-class AsyncJobTimeoutError(Exception):
-    """Raised when an async job does not reach a terminal status in time."""
-
-
 def poll_async_job(
     check_status,
     *,
@@ -33,7 +34,8 @@ def poll_async_job(
 
     Arguments:
         check_status: zero-arg callable returning a dict with a ``status`` key.
-        poll_time: maximum seconds to wait before raising AsyncJobTimeoutError.
+        poll_time: maximum seconds to wait before raising AsyncJobError
+            with the synthetic ``status="TIMEOUT"`` sentinel.
         poll_interval: seconds to sleep between polls.
         in_progress_statuses: status strings that mean "keep polling".
         completed_statuses: status strings that mean "return the status dict".
@@ -46,12 +48,7 @@ def poll_async_job(
     error = set(error_statuses)
 
     deadline = time.monotonic() + poll_time
-    first_iteration = True
     while True:
-        if not first_iteration and time.monotonic() >= deadline:
-            raise AsyncJobTimeoutError(
-                f"async job did not complete within {poll_time} seconds")
-        first_iteration = False
         status = check_status()
         value = status.get("status")
         if value in completed:
@@ -59,10 +56,10 @@ def poll_async_job(
         if value in error:
             raise AsyncJobError(status)
         if value not in in_progress:
-            raise AsyncJobError({
-                **status, "error":
-                f"unexpected status {value!r}"
-            })
+            message = f"unexpected status {value!r}"
+            raise AsyncJobError({**status, "error": message})
         remaining = deadline - time.monotonic()
-        if remaining > 0:
-            time.sleep(min(poll_interval, remaining))
+        if remaining <= 0:
+            message = f"async job did not complete within {poll_time} seconds"
+            raise AsyncJobError({"status": "TIMEOUT", "error": message})
+        time.sleep(min(poll_interval, remaining))
diff --git a/surge/reports.py b/surge/reports.py
index 294133a..c832d94 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -7,7 +7,7 @@
 import warnings
 
 from surge.api_resource import REPORTS_ENDPOINT, APIResource
-from surge.async_jobs import AsyncJobError, AsyncJobTimeoutError, poll_async_job
+from surge.async_jobs import AsyncJobError, poll_async_job
 from surge.errors import SurgeRequestError
 
 
@@ -88,13 +88,8 @@ def _check():
                 poll_interval=poll_interval,
                 in_progress_statuses=("IN_PROGRESS", "CREATING", "RETRYING"),
             )
-        except AsyncJobTimeoutError:
-            raise Exception(
-                "Report failed to generate within {poll_time} seconds".format(
-                    poll_time=poll_time))
         except AsyncJobError as e:
-            raise ValueError("Report failed to generate with status {}".format(
-                e.status.get("status")))
+            raise ValueError(f"Report failed to generate: {e}") from e
         url = terminal["url"]
 
         file_ext = "csv" if "csv" in type else "json"
diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py
index bcaf6ef..c42c944 100644
--- a/tests/test_async_jobs.py
+++ b/tests/test_async_jobs.py
@@ -2,11 +2,7 @@
 
 import pytest
 
-from surge.async_jobs import (
-    AsyncJobError,
-    AsyncJobTimeoutError,
-    poll_async_job,
-)
+from surge.async_jobs import AsyncJobError, poll_async_job
 
 
 def test_returns_on_first_completed_status():
@@ -75,12 +71,13 @@ def test_raises_async_job_error_on_unknown_status():
 
 def test_raises_timeout_when_poll_time_elapses():
     with patch("time.sleep"):
-        with pytest.raises(AsyncJobTimeoutError):
+        with pytest.raises(AsyncJobError) as excinfo:
             poll_async_job(
                 lambda: {"status": "IN_PROGRESS"},
                 poll_time=0,
                 poll_interval=0,
             )
+    assert excinfo.value.status.get("status") == "TIMEOUT"
 
 
 def test_caps_sleep_to_remaining_poll_time():
@@ -127,9 +124,10 @@ def test_raises_timeout_instead_of_returning_completed_after_deadline():
     # call 1 sets the deadline, calls 2+ all report well past it.
     with patch("time.sleep"), patch("time.monotonic",
                                     side_effect=[0.0, 1e9, 1e9, 1e9]):
-        with pytest.raises(AsyncJobTimeoutError):
+        with pytest.raises(AsyncJobError) as excinfo:
             poll_async_job(
                 lambda: next(statuses),
                 poll_time=1,
                 poll_interval=2,
             )
+    assert excinfo.value.status.get("status") == "TIMEOUT"
diff --git a/tests/test_reports.py b/tests/test_reports.py
index 46f8fc7..4b407f5 100644
--- a/tests/test_reports.py
+++ b/tests/test_reports.py
@@ -107,7 +107,7 @@ def test_save_report_raises_on_unexpected_status():
             mock.patch.object(Report, "check_status", return_value=error),
             mock.patch("time.sleep"),
     ):
-        with pytest.raises(ValueError, match="ERROR"):
+        with pytest.raises(ValueError, match="Report generation error"):
             Report.save_report("proj-123",
                                "export_json",
                                filepath=io.BytesIO())
@@ -153,7 +153,7 @@ def test_save_report_raises_when_initial_request_returns_error():
             mock.patch.object(Report, "request", return_value=error),
             mock.patch.object(Report, "check_status") as mock_check,
     ):
-        with pytest.raises(ValueError, match="ERROR"):
+        with pytest.raises(ValueError, match="Report generation error"):
             Report.save_report("proj-e", "export_json", filepath=io.BytesIO())
     mock_check.assert_not_called()
 

From dc487df3c60aebe8a3a16371b6c275f13e6eec75 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Thu, 28 May 2026 11:27:21 -0700
Subject: [PATCH 13/16] Add initial_status parameter to poll_async_job

Lets callers hand the helper a kick-off response to consume as iter 1
in place of the first check_status() call. save_report uses it so its
_check closure no longer needs to branch between request() and
check_status() based on iteration count.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 surge/async_jobs.py      | 10 +++++++++-
 surge/reports.py         | 24 +++++++++++-------------
 tests/test_async_jobs.py | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/surge/async_jobs.py b/surge/async_jobs.py
index 8ebf829..1bf236a 100644
--- a/surge/async_jobs.py
+++ b/surge/async_jobs.py
@@ -26,6 +26,7 @@ def poll_async_job(
     *,
     poll_time,
     poll_interval,
+    initial_status=None,
     in_progress_statuses=DEFAULT_IN_PROGRESS_STATUSES,
     completed_statuses=DEFAULT_COMPLETED_STATUSES,
     error_statuses=DEFAULT_ERROR_STATUSES,
@@ -37,6 +38,9 @@ def poll_async_job(
         poll_time: maximum seconds to wait before raising AsyncJobError
             with the synthetic ``status="TIMEOUT"`` sentinel.
         poll_interval: seconds to sleep between polls.
+        initial_status: optional status dict consumed in place of the first
+            ``check_status()`` call — useful when the caller already has the
+            response from a kick-off endpoint.
         in_progress_statuses: status strings that mean "keep polling".
         completed_statuses: status strings that mean "return the status dict".
         error_statuses: status strings that mean "raise AsyncJobError".
@@ -48,8 +52,12 @@ def poll_async_job(
     error = set(error_statuses)
 
     deadline = time.monotonic() + poll_time
+    pending = initial_status
     while True:
-        status = check_status()
+        if pending is not None:
+            status, pending = pending, None
+        else:
+            status = check_status()
         value = status.get("status")
         if value in completed:
             return status
diff --git a/surge/reports.py b/surge/reports.py
index c832d94..fdb7bc4 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -62,21 +62,18 @@ def save_report(
             extra_params (dict, optional): Additional params merged into the
                 initial report request body. See ``request``.
         """
-        # job_id arrives on the first request and on RETRYING swaps;
-        # IN_PROGRESS responses omit it, so only update when present.
-        state = {"job_id": None, "first": True}
+        initial = cls.request(project_id=project_id,
+                              type=type,
+                              api_key=api_key,
+                              extra_params=extra_params)
+        # job_id swaps mid-poll on RETRYING; IN_PROGRESS responses
+        # omit it, so only update when present.
+        state = {"job_id": getattr(initial, "job_id", None)}
 
         def _check():
-            if state["first"]:
-                state["first"] = False
-                r = cls.request(project_id=project_id,
-                                type=type,
-                                api_key=api_key,
-                                extra_params=extra_params)
-            else:
-                r = cls.check_status(project_id,
-                                     state["job_id"],
-                                     api_key=api_key)
+            r = cls.check_status(project_id,
+                                 state["job_id"],
+                                 api_key=api_key)
             if getattr(r, "job_id", None):
                 state["job_id"] = r.job_id
             return vars(r)
@@ -86,6 +83,7 @@ def _check():
                 _check,
                 poll_time=poll_time,
                 poll_interval=poll_interval,
+                initial_status=vars(initial),
                 in_progress_statuses=("IN_PROGRESS", "CREATING", "RETRYING"),
             )
         except AsyncJobError as e:
diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py
index c42c944..cb1da8c 100644
--- a/tests/test_async_jobs.py
+++ b/tests/test_async_jobs.py
@@ -50,6 +50,38 @@ def test_accepts_creating_and_ready_status_vocabulary():
     assert result == {"status": "READY", "url": "u"}
 
 
+def test_initial_status_is_consumed_in_place_of_first_check_call():
+    """When given, initial_status acts as iter 1; check_status is invoked
+    starting from iter 2."""
+    check_calls = []
+
+    def check():
+        check_calls.append(True)
+        return {"status": "COMPLETED", "url": "u"}
+
+    result = poll_async_job(
+        check,
+        poll_time=10,
+        poll_interval=0,
+        initial_status={"status": "READY", "url": "from-initial"},
+    )
+    assert result == {"status": "READY", "url": "from-initial"}
+    assert check_calls == []
+
+
+def test_initial_status_falls_through_to_polling_when_in_progress():
+    """In-progress initial_status hands off to check_status from iter 2."""
+    statuses = iter([{"status": "COMPLETED", "url": "u"}])
+    with patch("time.sleep"):
+        result = poll_async_job(
+            lambda: next(statuses),
+            poll_time=10,
+            poll_interval=0,
+            initial_status={"status": "CREATING"},
+        )
+    assert result == {"status": "COMPLETED", "url": "u"}
+
+
 def test_raises_async_job_error_on_error_status():
     with pytest.raises(AsyncJobError, match="kaboom"):
         poll_async_job(

From 7d214efb934b202d37d40ce006e5dd2b174d8d6c Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Thu, 28 May 2026 11:38:11 -0700
Subject: [PATCH 14/16] yapf

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 surge/reports.py         | 4 +---
 tests/test_async_jobs.py | 5 ++++-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/surge/reports.py b/surge/reports.py
index fdb7bc4..ae3dd68 100644
--- a/surge/reports.py
+++ b/surge/reports.py
@@ -71,9 +71,7 @@ def save_report(
         state = {"job_id": getattr(initial, "job_id", None)}
 
         def _check():
-            r = cls.check_status(project_id,
-                                 state["job_id"],
-                                 api_key=api_key)
+            r = cls.check_status(project_id, state["job_id"], api_key=api_key)
             if getattr(r, "job_id", None):
                 state["job_id"] = r.job_id
             return vars(r)
diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py
index cb1da8c..1434b9b 100644
--- a/tests/test_async_jobs.py
+++ b/tests/test_async_jobs.py
@@ -63,7 +63,10 @@ def check():
         check,
         poll_time=10,
         poll_interval=0,
-        initial_status={"status": "READY", "url": "from-initial"},
+        initial_status={
+            "status": "READY",
+            "url": "from-initial"
+        },
     )
     assert result == {"status": "READY", "url": "from-initial"}
     assert check_calls == []

From be2bae68f7bb50d70e024e3cf050556a3fcb9cac Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Thu, 28 May 2026 14:06:23 -0700
Subject: [PATCH 15/16] Address review nits: frozenset constants, simpler poll
 loop

Make the default status collections frozensets so membership checks
don't rebuild sets on every call, and drop the per-call set() wrapping.

Restructure the poll loop to seed the first status and fetch the next
one at the end of the body, removing the pending/swap dance.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 surge/async_jobs.py | 23 ++++++++---------------
 1 file changed, 8 insertions(+), 15 deletions(-)

diff --git a/surge/async_jobs.py b/surge/async_jobs.py
index 1bf236a..1062e49 100644
--- a/surge/async_jobs.py
+++ b/surge/async_jobs.py
@@ -2,9 +2,9 @@
 
 import time
 
-DEFAULT_IN_PROGRESS_STATUSES = ("IN_PROGRESS", "CREATING")
-DEFAULT_COMPLETED_STATUSES = ("COMPLETED", "READY")
-DEFAULT_ERROR_STATUSES = ("ERROR", "FAILED")
+DEFAULT_IN_PROGRESS_STATUSES = frozenset({"IN_PROGRESS", "CREATING"})
+DEFAULT_COMPLETED_STATUSES = frozenset({"COMPLETED", "READY"})
+DEFAULT_ERROR_STATUSES = frozenset({"ERROR", "FAILED"})
 
 
 class AsyncJobError(Exception):
@@ -47,23 +47,15 @@ def poll_async_job(
 
     Returns the terminal status dict.
     """
-    in_progress = set(in_progress_statuses)
-    completed = set(completed_statuses)
-    error = set(error_statuses)
-
     deadline = time.monotonic() + poll_time
-    pending = initial_status
+    status = initial_status if initial_status is not None else check_status()
     while True:
-        if pending is not None:
-            status, pending = pending, None
-        else:
-            status = check_status()
         value = status.get("status")
-        if value in completed:
+        if value in completed_statuses:
             return status
-        if value in error:
+        if value in error_statuses:
             raise AsyncJobError(status)
-        if value not in in_progress:
+        if value not in in_progress_statuses:
             message = f"unexpected status {value!r}"
             raise AsyncJobError({**status, "error": message})
         remaining = deadline - time.monotonic()
@@ -71,3 +63,4 @@ def poll_async_job(
             message = f"async job did not complete within {poll_time} seconds"
             raise AsyncJobError({"status": "TIMEOUT", "error": message})
         time.sleep(min(poll_interval, remaining))
+        status = check_status()

From a5d1270c7b9ba5a1f3a37c5deaea916d522df3b7 Mon Sep 17 00:00:00 2001
From: c-fang <charliefang@surgehq.ai>
Date: Thu, 28 May 2026 14:20:23 -0700
Subject: [PATCH 16/16] Clarify poll_time as a between-polls waiting budget

Reframes poll_time in the docstring as the budget for time spent waiting
between polls rather than a hard wall-clock ceiling, matching the
implementation that honors a terminal status already returned by a poll.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 surge/async_jobs.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/surge/async_jobs.py b/surge/async_jobs.py
index 1062e49..23dd7ca 100644
--- a/surge/async_jobs.py
+++ b/surge/async_jobs.py
@@ -35,8 +35,9 @@ def poll_async_job(
 
     Arguments:
         check_status: zero-arg callable returning a dict with a ``status`` key.
-        poll_time: maximum seconds to wait before raising AsyncJobError
-            with the synthetic ``status="TIMEOUT"`` sentinel.
+        poll_time: seconds budgeted for waiting between polls; when exhausted,
+            AsyncJobError is raised with the synthetic ``status="TIMEOUT"``
+            sentinel.
         poll_interval: seconds to sleep between polls.
         initial_status: optional status dict consumed in place of the first
             ``check_status()`` call — useful when the caller already has the