fix(models): raise the original failure when a same-key retry is refused 422

CHANGELOG.md

@@ -147,6 +147,19 @@ notes for each version.
 
 ### Changed
 
+- **Behaviour change: when a `client.models.run()` retry is refused `422`
+  `idempotency_key_reuse`, the failure that *caused* the retry is what is
+  raised** — the `deadline_exceeded` `504` the default collect loop resent
+  under, or the `500` a `retry_possibly_in_flight=True` policy resent under —
+  with the key refusal chained onto it as `__cause__`. The refusal is an
+  artefact of the retry loop rather than an answer about the request, and it
+  used to be the only error the caller saw, so the real failure was lost. **An
+  `except IdempotencyKeyReuse` around `models.run` no longer catches this
+  case**: catch the failure you actually care about (or `ComfyError`) and
+  inspect `exc.__cause__` to tell a rejected resend apart from a first-attempt
+  refusal. Nothing about *which* failures are retried changed, and no other
+  terminal failure is substituted — a `500` followed by a `404` still raises
+  the `404`.
 - **Breaking (wire): `client.models.run` now posts to Comfy Router.** It sends
   `POST {COMFY_ROUTER_BASE_URL}/v1/models/{provider}/{model}` — the route
   `spec/router-openapi.yaml` declares as `runRouterModel` — with the partner

README.md

@@ -474,7 +474,7 @@ The default policy:
 | 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 |
 | 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 an in-flight-key `concurrency_limit_exceeded` `409` carrying `Retry-After` (another call is already running under the same key; the resend collects it — the same bucket on a `429` is plain throttling and takes the ordinary retry path). 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 in-flight `concurrency_limit_exceeded` one above (`hash_mismatch` carries a `Retry-After` and is still deterministic, and the router's `invalid_input` key cases are answered by a NEW key), `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 |
+| 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 only comes back `422 idempotency_key_reuse` (that refusal no longer hides the real error — see below — but it is still a wasted request) — 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 |
 
@@ -510,8 +510,10 @@ A note on what the default trades: the collect rule is
 `spec/router-openapi.yaml`'s, so it binds Comfy Router — but
 `COMFY_ROUTER_BASE_URL` can name a deployment that applies the v2 rule instead
 and keeps 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.
+`422 idempotency_key_reuse` — one wasted request, but not a lost diagnosis:
+`run` raises the real `504` and chains the refusal onto it as `__cause__`, so
+`except IdempotencyKeyReuse` around `run` does *not* catch it. Set
+`retry_collectable=False` on such a deployment to skip the resend entirely.
 
 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
@@ -631,7 +633,10 @@ asset, job, event, and output helpers translate protocol errors, so catches of
   no replay — so if you pass your own `idempotency_key=` and reuse it, the second
   call raises this. After an ambiguous failure (e.g. a timeout where you don't
   know if the job was created), poll or list your jobs rather than resubmitting
-  with the same key.
+  with the same key. One exception: when `models.run`'s *own* retry is refused
+  for key reuse, it raises the failure that caused the retry and chains this
+  exception onto it as `__cause__`, since the refusal says nothing about why the
+  call failed.
 - `InsufficientCredits` — the account can't afford the job.
 - `QueueFull` — backpressure; carries `.retry_after` seconds. `client.submit`
   retries 429 responses with `Retry-After` for a bounded budget (including

src/comfy_sdk/models.py

@@ -49,10 +49,11 @@
 import httpx
 
 from comfy_low.errors import ApiError
+from comfy_low.errors import IdempotencyKeyReuse as ProtocolIdempotencyKeyReuse
 from comfy_low.transport import MODEL_RUN_TIMEOUT, AsyncComfyLow, ComfyLow
 
 from ._core import new_idempotency_key, validate_idempotency_key
-from .exceptions import translating
+from .exceptions import IdempotencyKeyReuse, _stamp, to_sdk_error, translating
 from .retry import DEFAULT_RETRY, Retrier, RetryPolicy
 from .router_exceptions import RouterError
 
@@ -72,9 +73,55 @@
 #: raising them, with no test failing.
 _CANDIDATE_FAILURES = (ApiError, RouterError, httpx.TransportError)
 
+#: The answer that means the resend was never going to work: the server refused
+#: the *key*, not the request. ``422`` is the v2 jobs rule
+#: (single-use, reject-on-duplicate), which a deployment named by
+#: ``COMFY_ROUTER_BASE_URL`` may apply to this route even though the router
+#: contract does not.
+_KEY_REUSE_STATUS = 422
+_KEY_REUSE_CODE = "idempotency_key_reuse"
+
 _now = time.monotonic
 
 
+def _is_key_reuse(exc: BaseException) -> bool:
+    """Whether ``exc`` is the server refusing a repeated ``Idempotency-Key``.
+
+    Matched on the wire facts — ``422`` plus the ``idempotency_key_reuse``
+    code — rather than on a class alone, because the retry loop runs *inside*
+    ``translating()``: what it catches is the raw
+    :class:`comfy_low.errors.ApiError`, and only the copy that leaves the block
+    is the idiomatic :class:`~comfy_sdk.exceptions.IdempotencyKeyReuse`. Either
+    layer's typed class is accepted outright so the answer does not depend on
+    which one raised.
+    """
+    if isinstance(exc, (IdempotencyKeyReuse, ProtocolIdempotencyKeyReuse)):
+        return True
+    return (
+        getattr(exc, "http_status", None) == _KEY_REUSE_STATUS
+        and getattr(exc, "code", None) == _KEY_REUSE_CODE
+    )
+
+
+def _as_sdk_error(exc: BaseException, idempotency_key: str) -> BaseException:
+    """``exc`` on the surface a caller catches, carrying ``idempotency_key``.
+
+    ``translating()`` converts and stamps the exception it is *handed*, and
+    nothing else — so an ``ApiError`` chained onto that one as ``__cause__``
+    would reach the caller as a raw protocol type this SDK otherwise never
+    shows. Both halves of the substitution below therefore go through here
+    first. Anything already idiomatic (a ``RouterError``, an ``httpx`` failure
+    with no response to translate) is only stamped.
+
+    The translated copy inherits the original's traceback, so the chain still
+    points at the attempt that produced each half rather than at the one line
+    of this module that re-raised them.
+    """
+    if not isinstance(exc, ApiError):
+        return _stamp(exc, idempotency_key)
+    return _stamp(to_sdk_error(exc).with_traceback(exc.__traceback__), idempotency_key)
+
+
 class _ModelsBase:
     """Read-only view of the configuration inherited from the host client."""
 
@@ -193,6 +240,15 @@ def run(
         ``RetryPolicy(retry_possibly_in_flight=True)``. See
         :mod:`comfy_sdk.retry`, and ``Comfy(retry=NO_RETRY)`` to switch it off.
 
+        **When a resend is refused because the key was already claimed
+        (``422`` ``idempotency_key_reuse``), what is raised is the failure that
+        caused the retry** — the ``504``, the ``500`` — with the key refusal
+        chained onto it as ``__cause__``. That refusal is an artefact of this
+        retry loop rather than an answer about the request, so surfacing it in
+        place of the real failure would hide the only error worth diagnosing.
+        Every other terminal failure is raised exactly as the server sent it,
+        first or last.
+
         **Every exception this raises carries the key it sent** on
         ``.idempotency_key`` (and the server's ``.request_id`` when the response
         named one), so a caller who lost the response — a ``deadline_exceeded``
@@ -244,6 +300,11 @@ def run(
         # the wire, so everything legal in it is deep-copyable.
         payload = deepcopy(dict(arguments))
         retrier = Retrier(self._retry, now=_now)
+        # The first retryable failure, kept so a key refusal on a later attempt
+        # cannot bury it. Only the first: every later one is the same call
+        # failing again, and the one the caller has to diagnose is the one that
+        # started the retrying.
+        first: BaseException | None = None
         # The key is stamped onto whatever this raises: it is a local of this
         # frame, so an exception that propagates past it would otherwise take
         # the caller's only route back to an already-billed generation with it.
@@ -252,9 +313,18 @@ def run(
                 try:
                     return low.post_model_run(model, payload, idempotency_key=key, timeout=timeout)
                 except _CANDIDATE_FAILURES as exc:
+                    if first is not None and _is_key_reuse(exc):
+                        # The resend could never have succeeded — the server
+                        # refused the key, not the request — so raising it
+                        # would replace the real failure with an artefact of
+                        # this loop. Chained, not discarded: the 422 stays
+                        # reachable on `__cause__` and in the traceback.
+                        raise _as_sdk_error(first, key) from _as_sdk_error(exc, key)
                     delay = retrier.delay_before_retry(exc)
                     if delay is None:
                         raise
+                    if first is None:
+                        first = exc
                     time.sleep(delay)
 
 
@@ -277,8 +347,9 @@ async def run(
 
         This *is* the async form of ``run``: awaiting it on ``AsyncComfy`` is
         the whole difference from the sync client — including the model-id
-        rule, the retry policy, the one-key-per-call rule, and the
-        ``.idempotency_key`` every exception it raises carries for the replay.
+        rule, the retry policy, the one-key-per-call rule, the failure raised
+        when a resend is refused for key reuse, and the ``.idempotency_key``
+        every exception it raises carries for the replay.
         See :meth:`Models.run`.
         """
         low = cast(AsyncComfyLow, self._low)
@@ -293,6 +364,8 @@ async def run(
         # nested values included.
         payload = deepcopy(dict(arguments))
         retrier = Retrier(self._retry, now=_now)
+        # The first retryable failure — see :meth:`Models.run`.
+        first: BaseException | None = None
         # The key is stamped onto whatever this raises: it is a local of this
         # frame, so an exception that propagates past it would otherwise take
         # the caller's only route back to an already-billed generation with it.
@@ -303,7 +376,13 @@ async def run(
                         model, payload, idempotency_key=key, timeout=timeout
                     )
                 except _CANDIDATE_FAILURES as exc:
+                    if first is not None and _is_key_reuse(exc):
+                        # The rejected resend replaces nothing — see
+                        # :meth:`Models.run`.
+                        raise _as_sdk_error(first, key) from _as_sdk_error(exc, key)
                     delay = retrier.delay_before_retry(exc)
                     if delay is None:
                         raise
+                    if first is None:
+                        first = exc
                     await asyncio.sleep(delay)

src/comfy_sdk/retry.py

@@ -81,11 +81,13 @@
    on a run is the important member: the generation-sized client timeout expired
    with no answer, which is precisely when the server is most likely still
    generating). Nothing on the wire says the key survives this, so a same-key
-   retry may be a ``422`` that replaces the real error — and a *fresh*-key retry
+   retry may be refused ``422`` — and a *fresh*-key retry
    is the second billed generation this module exists to prevent. Not retried by
    default. :attr:`RetryPolicy.retry_possibly_in_flight` opts in, and it is
    correct exactly when a deployment replays a repeated key instead of rejecting
-   it.
+   it. That refusal no longer costs the caller the diagnosis: ``models.run``
+   raises the failure that caused the retry and chains the ``422`` onto it as
+   ``__cause__``, so the real error is what surfaces.
 5. **Everything else** — every other 4xx is the server's considered answer
    about *this* request, and asking again spends money to be refused again.
    Never retried.
@@ -273,8 +275,9 @@ def is_unknown_outcome_status(status: int) -> bool:
     ``invalid_input``. That is why
     this class sits behind :attr:`RetryPolicy.retry_possibly_in_flight` rather
     than being retried by default: unless the deployment replays a claimed key,
-    the same-key retry cannot succeed and *replaces* the genuine 5xx with a
-    confusing key-reuse error.
+    the same-key retry cannot succeed at all. What it costs is one wasted
+    request rather than the diagnosis — ``models.run`` raises the genuine 5xx
+    and chains the key-reuse refusal onto it as ``__cause__``.
 
     A ``502``/``504`` from an intermediary belongs here for the same reason:
     the proxy's response completed, which says nothing about whether the origin
@@ -421,7 +424,9 @@ class RetryPolicy:
     #: caution: a key whose recorded answer cannot be replayed is answered
     #: ``409`` ``invalid_input`` on the router surface (use a NEW key), and the
     #: v2 jobs contract makes ``Idempotency-Key`` single-use outright — so a
-    #: same-key retry here can surface a key refusal that hides the real error.
+    #: same-key retry here buys a key refusal rather than an answer. It no
+    #: longer hides the real error: ``models.run`` raises the failure that
+    #: caused the retry, with the refusal chained on as ``__cause__``.
     #: Turn it on for a deployment that replays a repeated
     #: key instead of rejecting it — and raise ``max_elapsed`` when you do, since
     #: one full-length client timeout on a run spends the whole default budget on