feat: collect a Router deadline 504 under the same Idempotency-Key by default

README.md

@@ -413,26 +413,31 @@ sends the same `Idempotency-Key`**, and a new call mints a new one — that is
 what lets a server tell a retry apart from a second order, on a surface where
 one call is a billed generation.
 
-That one key is also what decides *which* failures are worth retrying. The API
-contract makes `Idempotency-Key` **single-use, reject-on-duplicate, with no
-response replay**: the first request to present a key is processed, and a later
-one presenting the same key is rejected `422 idempotency_key_reuse` rather than
-re-run. The contract also says when the key is released instead of claimed — a
-request that definitively failed without starting work frees its key, while one
-whose outcome the server could not characterise (a 5xx, an upstream timeout)
-keeps it. Retrying under one key is only safe where that key is still unspent,
-so that is exactly what the default policy retries.
+That one key is also what decides *which* failures are worth retrying. The v2
+jobs contract makes its `Idempotency-Key` **single-use, reject-on-duplicate,
+with no response replay**: the first request to present a key is processed, and
+a later one presenting the same key is rejected `422 idempotency_key_reuse`
+rather than re-run. (That rule governs `submit()`. On the router surface a
+resend of the same key with the same body *collects* the generation the first
+request started instead of dispatching another; a resend with a different body
+is the one that is rejected `422`.) The contract also says when the key is
+released instead of claimed — a request that definitively failed without
+starting work frees its key, while one whose outcome the server could not
+characterise (a 5xx, an upstream timeout) keeps it. Retrying under one key is
+only safe where that key is still spendable, so that is exactly what the default
+policy retries.
 
 The default policy:
 
 | Condition | Behaviour |
 |---|---|
 | Retried | connect-phase transport failures (connection refused, connect timeout, no pooled connection, proxy error) — the request never reached the server, so the key was never claimed |
 | Retried, at the server's pace | a `429` carrying `Retry-After` (queue full, out of credits, a concurrency limit) — a reject that started no work, so the key is released. The delay is the one the server named, not a guess |
-| Not retried | every other 4xx — `400`/`content_policy_violation`, `404`, `409`, `422`, `401`, `402` — because asking again cannot change a deterministic refusal. A `429` with no `Retry-After` is not asking to be asked again either |
-| Not retried by default | anything whose outcome is unknown: a **5xx response** (including the router's `service_unavailable` `503`, which asks a caller to retry with backoff but says nothing about the key), and a client-side timeout where the server may still be generating. The key stays claimed for these, so a same-key retry comes back `422 idempotency_key_reuse` and hides the real error — while a fresh-key retry is the second billed generation the one-key rule exists to prevent |
-| Budget | 60 seconds of **total elapsed time** from the first attempt, not a number of attempts |
-| Backoff | 0.5s doubling to a 15s ceiling, with full jitter (each wait is drawn from `[0, ceiling]`), clamped to whatever is left of the budget |
+| Retried, at the server's pace | the answers that pace a resend of the *same* key for work already running: a `deadline_exceeded` `504` carrying `Retry-After` (Comfy stopped holding the connection at its own bound; the contract says to retry with the same key, which collects that generation rather than dispatching another), and a `generation_in_progress` `409` carrying `Retry-After` (the same key, asked for again before the generation finished). One `run()` rides that loop to the finished result |
+| Not retried | every other 4xx — `400`/`content_policy_violation`, `404`, any `409` that is not the paced `generation_in_progress` one above (`hash_mismatch` carries a `Retry-After` and is still deterministic), `422`, `401`, `402` — because asking again cannot change a deterministic refusal. A `429` with no `Retry-After` is not asking to be asked again either |
+| Not retried by default | anything whose outcome is unknown: any **other 5xx response** — including the router's `service_unavailable` `503` (which asks a caller to retry with backoff but says nothing about the key), a `504` carrying no `Retry-After` (the router sends it only when it holds a generation to collect), and a `504` that is `provider_timeout` rather than `deadline_exceeded` — and a client-side timeout where the server may still be generating. The key stays claimed for these, so a same-key retry comes back `422 idempotency_key_reuse` and hides the real error — while a fresh-key retry is the second billed generation the one-key rule exists to prevent |
+| Budget | 60 seconds of **total elapsed time** from the first attempt, not a number of attempts. The collect loop gets its own, longer budget: 1200 seconds, two server deadline windows, so it can outlast the deadline that started it |
+| Backoff | 0.5s doubling to a 15s ceiling, with full jitter (each wait is drawn from `[0, ceiling]`), clamped to whatever is left of the budget. A `Retry-After` the server named is used as given instead |
 
 The budget bounds when the *last* attempt may **start**; an attempt already
 running is never interrupted by it, so a slow generation is never abandoned
@@ -444,26 +449,46 @@ Tune or disable it per client:
 ```python
 from comfy_sdk import Comfy, NO_RETRY, RetryPolicy
 
-Comfy(retry=NO_RETRY)                          # exactly one attempt, ever
-Comfy(retry=RetryPolicy(max_elapsed=300.0))    # keep trying for five minutes
+Comfy(retry=NO_RETRY)                              # exactly one attempt, ever
+Comfy(retry=RetryPolicy(max_elapsed=300.0))        # fast classes: five minutes
+Comfy(retry=RetryPolicy(collect_max_elapsed=60.0)) # bound the collect loop
+Comfy(retry=RetryPolicy(retry_collectable=False))  # raise the 504/409 instead
 
-client.models.retry                            # the policy in force, read-only
+client.models.retry                                # the policy in force, read-only
 ```
 
-5xx responses and client-side timeouts are the cases left out by default, and
-for the same reason. `run` holds the connection open while the server
-generates, so neither one tells you whether the generation happened — retrying
-either starts a second generation unless the server replays the repeated key
-rather than re-running it, and against a server that *rejects* it instead the
-retry simply cannot succeed. Against a deployment that does replay, opt in:
+There are two budgets because the classes have two shapes. `max_elapsed` governs
+the fast ones — a connect failure, a paced `429` — which resolve in seconds or
+not at all, so an unreachable host gives up in a minute rather than pinning a
+caller (a whole thread, on the sync client) for longer. `collect_max_elapsed`
+governs the collect loop alone, and is twenty minutes because that is the one
+class that has to outlast a *server-side* bound: a `deadline_exceeded` `504`
+arrives at Comfy's own ten-minute deadline, so a single-window budget would
+already be spent when it lands and the collect attempt it exists for would never
+start. Nothing else pays for that room.
+
+A note on what the default trades: `POST /models/run` is in neither vendored
+spec, so a deployment may apply the v2 rule instead and keep the key claimed
+across the `504`. There the collect resend comes back `422
+idempotency_key_reuse` in place of the real `504`. Set
+`retry_collectable=False` on such a deployment.
+
+Other 5xx responses and client-side timeouts are the cases left out by default,
+and for the same reason. `run` holds the connection open while the server
+generates, so neither one tells you whether the generation happened, and no
+contract says the key survives them — retrying either starts a second generation
+unless the server replays the repeated key rather than re-running it, and
+against a server that *rejects* it instead the retry simply cannot succeed.
+Against a deployment that does replay, opt in:
 
 ```python
 Comfy(retry=RetryPolicy(max_elapsed=1200.0, retry_possibly_in_flight=True))
 ```
 
-Raise `max_elapsed` when you do: one full-length client timeout on a run can
-spend the default 60-second budget on its own, leaving no room for the retry
-you just asked for.
+Raise `max_elapsed` when you do: one full-length client timeout on a run spends
+many times the default 60-second budget on its own, leaving no room for the
+retry you just asked for. `collect_max_elapsed` does not help here — that budget
+is the collect class's alone.
 
 `retry` governs `client.models` only. `submit()`/`run()` on the client keep
 their own 429 handling, which follows the server's `Retry-After`.

src/comfy_low/errors.py

@@ -98,25 +98,67 @@ class Forbidden(ApiError):
 }
 
 
+def _clean(value: Any) -> str | None:
+    """``value`` as a non-empty string, or ``None``.
+
+    Anything else — a missing key, a number, Router's ``detail[]`` list form —
+    reads as absent rather than being coerced, so a malformed body degrades to
+    the status-derived default instead of producing a nonsense code.
+    """
+    if not isinstance(value, str):
+        return None
+    stripped = value.strip()
+    return stripped or None
+
+
 def error_from_envelope(
     http_status: int,
     body: dict[str, Any] | None,
     *,
     retry_after: int | None = None,
+    error_type: str | None = None,
 ) -> ApiError:
     """Build the typed exception for an error response.
 
     Falls back to a status-derived code when the body is missing or not a
     well-formed envelope (so a bare ``401`` with no JSON still maps to
     ``Unauthorized``).
+
+    Not every route answers in the envelope shape. ``POST /api/v2/models/run``
+    is fronted by Router, whose error body is ``{detail, error_type}`` and which
+    repeats the same coarse bucket on the ``X-Comfy-Error-Type`` header
+    (``spec/router-openapi.yaml``). Reading only ``error["code"]`` would collapse
+    every one of those to the status-derived default — for a ``504`` that
+    default is the meaningless ``"error"``, and ``comfy_sdk.retry`` keys its
+    default-on collect rule on the bucket, so the rule would be a silent no-op
+    against every real Router ``504``. ``error_type`` (the header, passed by the
+    caller) and the body's own top-level ``error_type`` are read to prevent that.
+
+    They are consulted in one narrow place, though: *after* the envelope's
+    ``code``, which always wins, and *after* :data:`_CODE_BY_STATUS`, which
+    already determines a code for every status where this API has a documented
+    one. So a Router ``429`` still maps to ``queue_full`` and a Router ``401``
+    still to ``unauthorized`` — reordering those would silently retype
+    exceptions that integrators already catch. What is left is exactly the 5xx
+    range, where the status determines nothing and the bucket is the only name
+    the response has.
     """
     err = (body or {}).get("error") if isinstance(body, dict) else None
     code = (err or {}).get("code") if isinstance(err, dict) else None
     message = (err or {}).get("message") if isinstance(err, dict) else None
     details = (err or {}).get("details") if isinstance(err, dict) else None
 
     if code is None:
-        code = _CODE_BY_STATUS.get(http_status, "error")
+        code = _CODE_BY_STATUS.get(http_status)
+    if code is None:
+        code = _clean(error_type) or _clean(
+            (body or {}).get("error_type") if isinstance(body, dict) else None
+        )
+    if code is None:
+        code = "error"
+    if not message:
+        # Router names its human-readable string `detail`, not `error.message`.
+        message = _clean((body or {}).get("detail") if isinstance(body, dict) else None)
     if not message:
         message = f"HTTP {http_status}"
 

src/comfy_low/transport.py

@@ -220,7 +220,16 @@ def parse_or_raise(self, resp: httpx.Response, ok: tuple[int, ...]) -> dict[str,
             body = resp.json()
         except Exception:
             body = None
-        raise error_from_envelope(resp.status_code, body, retry_after=_retry_after(resp))
+        raise error_from_envelope(
+            resp.status_code,
+            body,
+            retry_after=_retry_after(resp),
+            # Router repeats its coarse bucket on this header, and the model-run
+            # route answers in Router's body shape rather than the envelope's.
+            # Passing it here is what keeps the bucket alive across the layer
+            # boundary; see `error_from_envelope`.
+            error_type=resp.headers.get("X-Comfy-Error-Type"),
+        )
 
 
 def parse_expiry(url: str) -> datetime | None:

src/comfy_sdk/models.py

@@ -35,6 +35,7 @@
 import asyncio
 import time
 from collections.abc import Mapping
+from copy import deepcopy
 from typing import Any, cast
 
 import httpx
@@ -134,17 +135,29 @@ def run(
 
         An ``Idempotency-Key`` is sent on every run; a fresh one is minted per
         call unless ``idempotency_key`` is given, so an accidental exact resend
-        is the server's to reject rather than a second charged generation.
+        is the server's to deduplicate rather than a second charged generation.
 
         A failed attempt is retried under the client's policy, backed off with
         jitter and bounded by total elapsed time. **Every attempt of this one
         call reuses the one key**, which is what stops a retry from being
         billed as a second generation; calling ``run`` again is a new call and
-        mints a new key. Because the key is single-use and reject-on-duplicate,
-        only failures that leave it unclaimed are retried by default —
-        connect-phase failures, and a ``429`` that names a ``Retry-After``. A
-        completed 5xx or a client-side timeout leaves the outcome unknown, and
-        with it the key claimed, so those need
+        mints a new key. A key presented twice is *not* re-run: on the router
+        surface a resend of the same key with the same body collects the
+        generation the first request started rather than dispatching another,
+        and a resend with a *different* body is rejected ``422``
+        ``idempotency_key_reuse`` — which is why the body is snapshotted before
+        the first attempt. (The single-use, reject-on-duplicate rule
+        ``spec/openapi.yaml`` states is the **v2 jobs API**'s, and governs
+        ``submit()``, not this route.)
+
+        Retried by default: connect-phase failures, a ``429`` that names a
+        ``Retry-After``, and the answers that name a pace for collecting work
+        already running — a ``deadline_exceeded`` ``504`` and a
+        ``generation_in_progress`` ``409``, each carrying ``Retry-After``. One
+        ``run()`` can therefore ride the collect loop through a server-side
+        deadline to the finished generation. Not retried by default: a completed
+        5xx that named no such pace, and a client-side timeout — those leave the
+        outcome genuinely unknown and need
         ``RetryPolicy(retry_possibly_in_flight=True)``. See
         :mod:`comfy_sdk.retry`, and ``Comfy(retry=NO_RETRY)`` to switch it off.
         """
@@ -156,7 +169,13 @@ def run(
         # mapping inside each attempt would let a mutation between attempts
         # send a different body under the *same* key, which is precisely the
         # same-key-different-body case the contract rejects outright.
-        payload = dict(arguments)
+        #
+        # Deep, not shallow: a shallow copy leaves every nested list and dict
+        # shared with the caller, so mutating `arguments["config"]["steps"]`
+        # during the retry window would still change the body under the one key
+        # and earn the 422 this snapshot exists to prevent. The body is JSON on
+        # the wire, so everything legal in it is deep-copyable.
+        payload = deepcopy(dict(arguments))
         retrier = Retrier(self._retry, now=_now)
         with translating():
             while True:
@@ -192,10 +211,11 @@ async def run(
         """
         low = cast(AsyncComfyLow, self._low)
         key = idempotency_key or new_idempotency_key()
-        # Snapshotted before the first attempt — see :meth:`Models.run`. The
-        # window is wider here: the caller's coroutine can mutate `arguments`
-        # while the retry sleeps.
-        payload = dict(arguments)
+        # Snapshotted deeply before the first attempt — see :meth:`Models.run`.
+        # The window is wider here: the retry sleeps inside the caller's own
+        # event loop, so another task is free to run and mutate `arguments`,
+        # nested values included.
+        payload = deepcopy(dict(arguments))
         retrier = Retrier(self._retry, now=_now)
         with translating():
             while True: