fix(crawl): retry a stop on 429, 502 and 504

src/zenrows/crawl/client.py

@@ -49,11 +49,13 @@
 # 429 `too_many_crawls` is a capacity limit, not a transient failure.
 _CREATE_RETRY_STATUSES = _RETRYABLE_STATUSES - {429}
 
-# 503 `crawl_busy` on stop: nothing was saved and the crawl is still
-# running, so the stop is safe to repeat after `Retry-After`.
-_CRAWL_BUSY = "crawl_busy"
-# The longest a crawl_busy `Retry-After` may make a stop wait, so a large
-# server value cannot hang it.
+# A stop is safe to repeat (an ended crawl answers as it ended), so
+# `stop` retries these statuses itself even though it is a POST without
+# an idempotency key. 503 includes `crawl_busy`: nothing was saved and
+# the crawl is still running.
+_STOP_RETRY_STATUSES = _RETRYABLE_STATUSES
+# The longest a `Retry-After` may make a stop wait, so a large server
+# value cannot hang it.
 _MAX_STOP_RETRY_AFTER = 30.0
 
 _POLL_INTERVAL = 2.0
@@ -127,7 +129,7 @@ def __init__(
         `create` that carries an `idempotency_key` (never on 429 for
         `create`). Retries use jittered exponential backoff and honor
         `Retry-After`; set `retries=0` to disable. It also bounds the
-        retries of a `stop` answered 503 `crawl_busy`."""
+        retries of a `stop` answered 429, 502, 503 or 504."""
         if not api_key:
             raise ValueError("ZenRowsCrawlClient: api_key is required.")
         self._retries = max(0, retries)
@@ -276,9 +278,12 @@ def stop(self, crawl_id: str) -> CrawlStop:
         The answer carries no counts; read the final coverage and
         results with `get` once pages already in flight have finished.
 
-        A 503 `crawl_busy` means the stop was not saved and the crawl is
-        still running; `stop` retries it up to `retries` times, waiting
-        `Retry-After` seconds (at most 30) each time, then raises it.
+        Because repeating a stop is safe, `stop` retries a 429, 502,
+        503 (including `crawl_busy`, which means the stop was not saved
+        and the crawl is still running) or 504 up to `retries` times,
+        waiting `Retry-After` seconds (at most 30; 1 when missing or
+        unusable) each time, then raises the last error. Network errors
+        are not retried.
         """
         attempt = 0
         while True:
@@ -287,7 +292,7 @@ def stop(self, crawl_id: str) -> CrawlStop:
                     self._t.request_json("POST", f"/crawls/{crawl_id}/stop")
                 )
             except CrawlAPIError as err:
-                if err.status_code != 503 or err.code != _CRAWL_BUSY or attempt >= self._retries:
+                if err.status_code not in _STOP_RETRY_STATUSES or attempt >= self._retries:
                     raise
                 time.sleep(_stop_retry_wait(err.retry_after))
                 attempt += 1
@@ -351,7 +356,7 @@ def wait(self, crawl_id: str, *, timeout: float = 600.0) -> Crawl:
 
 
 def _stop_retry_wait(retry_after: float | None) -> float:
-    """Seconds `stop` waits after a 503 crawl_busy: the server's
+    """Seconds `stop` waits after a 429/502/503/504: the server's
     `Retry-After`, 1s when it is missing or unusable, at most 30s."""
     if retry_after is None or not math.isfinite(retry_after) or retry_after < 0:
         return 1.0

tests/test_crawl_client.py

@@ -301,11 +301,50 @@ def test_stop_surfaces_crawl_busy_after_retries(no_sleep):
     assert no_sleep == [1.0, 1.0]
 
 
+@pytest.mark.parametrize(
+    ("status", "code"),
+    [(429, "rate_limited"), (502, "bad_gateway"), (503, "service_unavailable"), (504, None)],
+)
+@respx.mock
+def test_stop_retries_transient_status(no_sleep, status: int, code: str | None):
+    client = ZenRowsCrawlClient(api_key=API_KEY, base_url=BASE_URL, retries=2)
+    first = (
+        problem(status, code, headers={"Retry-After": "2"})
+        if code
+        else Response(status, headers={"Retry-After": "2"})
+    )
+    route = respx.post(f"{BASE_URL}/crawls/c_1/stop").mock(
+        side_effect=[
+            first,
+            Response(200, json={"crawl_id": "c_1", "status": "stopped", "stop_reason": "user"}),
+        ]
+    )
+
+    stopped = client.stop("c_1")
+
+    assert stopped.status is CrawlStatus.STOPPED
+    assert route.call_count == 2
+    assert no_sleep == [2.0]
+
+
+@respx.mock
+def test_stop_surfaces_502_after_retries(no_sleep):
+    client = ZenRowsCrawlClient(api_key=API_KEY, base_url=BASE_URL, retries=3)
+    route = respx.post(f"{BASE_URL}/crawls/c_1/stop").mock(return_value=problem(502, "bad_gateway"))
+
+    with pytest.raises(CrawlAPIError) as exc_info:
+        client.stop("c_1")
+
+    assert exc_info.value.status_code == 502
+    assert route.call_count == 4
+    assert no_sleep == [1.0, 1.0, 1.0]
+
+
 @respx.mock
-def test_stop_does_not_retry_other_503(no_sleep):
+def test_stop_does_not_retry_other_status(no_sleep):
     client = ZenRowsCrawlClient(api_key=API_KEY, base_url=BASE_URL, retries=2)
     route = respx.post(f"{BASE_URL}/crawls/c_1/stop").mock(
-        return_value=problem(503, "service_unavailable", headers={"Retry-After": "1"})
+        return_value=problem(500, "internal_error")
     )
 
     with pytest.raises(CrawlAPIError):
@@ -597,12 +636,15 @@ def test_wait_returns_failed_crawl(client: ZenRowsCrawlClient, no_sleep):
 @pytest.mark.parametrize(
     ("retry_after", "slept"), [("3600", 30.0), ("-1", 1.0), ("nan", 1.0), ("2", 2.0)]
 )
+@pytest.mark.parametrize(("status", "code"), [(503, "crawl_busy"), (429, "rate_limited")])
 @respx.mock
-def test_stop_waits_at_most_30_seconds_for_crawl_busy(no_sleep, retry_after: str, slept: float):
+def test_stop_waits_at_most_30_seconds(
+    no_sleep, status: int, code: str, retry_after: str, slept: float
+):
     client = ZenRowsCrawlClient(api_key=API_KEY, base_url=BASE_URL, retries=1)
     respx.post(f"{BASE_URL}/crawls/c_1/stop").mock(
         side_effect=[
-            problem(503, "crawl_busy", headers={"Retry-After": retry_after}),
+            problem(status, code, headers={"Retry-After": retry_after}),
             Response(200, json={"crawl_id": "c_1", "status": "stopped", "stop_reason": "user"}),
         ]
     )