feat: carry the Idempotency-Key and request id on every exception models.run raises

CHANGELOG.md

@@ -14,6 +14,50 @@ notes for each version.
 
 ### Added
 
+- Every exception `client.models.run()` raises **for a failed call** now
+  carries the `Idempotency-Key` it was made under, on `.idempotency_key` — the
+  typed `RouterError` buckets, a `RouterError` whose `error_type` this version
+  does not recognise, any other `ComfyError`, and a transport failure with no
+  response at all (a dropped connection, a read timeout), and a cancelled
+  `await` of `AsyncModels.run` — the `asyncio.wait_for` a caller wraps a
+  ten-minute call in abandons a generation that may already be dispatched and
+  billed, and the cancellation still propagates unchanged. "Failed call" is the
+  boundary, not "every exception": a programming error escaping the call is not
+  a failed request, a key means nothing on it, and it reaches you untouched —
+  as does `KeyboardInterrupt`. `run` mints that key itself unless you pass
+  `idempotency_key=`, and it used to be a local of the call: when the call
+  raised, the key went with it. Since collecting a generation you were already
+  billed for after a lost response means asking again under the *same* key,
+  that made the auto-minted case uncollectable — only callers who chose and
+  stored their own key could recover. The recovery idiom is now
+  `client.models.run(model, arguments, idempotency_key=exc.idempotency_key)`;
+  see the README. Nothing about what is retried, or what goes on the wire,
+  changed.
+- `ComfyError.request_id` — the server's `X-Comfy-Request-Id` for the failed
+  call, when the response carried that header, on every SDK exception rather
+  than only on `RouterError`. It is the id to quote in a support request, and it
+  was previously unreachable once the response object was gone. `None` when the
+  response named none, or when there was no response — including on a transport
+  failure, where the attribute now reads as `None` rather than being absent, so
+  a handler never has to guard the access. The header is bounded and filtered
+  before it is stored (it is server-controlled and the id is meant to be
+  displayed and pasted into support tickets), identically on both error
+  surfaces.
+- `ComfyError.retry_after` — seconds the server asked the caller to wait, from
+  `Retry-After`, now forwarded for every error code rather than only for
+  `queue_full`. The replay documented above tells a caller to ask again "after
+  the `Retry-After` the server named", and a `deadline_exceeded` `504` that
+  carried one had nowhere to surface it, so the caller had nothing to wait on.
+  `None` when the server named no pace.
+
+### Fixed
+
+- A success status whose body will not decode (a proxy interstitial served
+  under a `200`, a response truncated mid-stream) now raises a translated SDK
+  error instead of letting `json.JSONDecodeError` escape from outside the
+  translated surface. On `models.run` that is a generation that ran and was
+  billed with the result lost — precisely the failure the `Idempotency-Key`
+  has to ride out on, and it previously carried no key.
 - Automatic retry for `client.models.run`, on by default, with the
   `Idempotency-Key` sent unconditionally on **every** attempt of one logical
   call — a new call mints a new key. That is what keeps a retry from being

README.md

@@ -493,6 +493,64 @@ 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`.
 
+### Collecting a generation after a lost response
+
+Everything above is about the retry the SDK makes *for* you. When it gives up —
+or when you turned it off — the failure that reaches you may still be a
+generation that ran and was billed: a `504` where the server stopped holding the
+connection at its own deadline, or a connection that dropped while it was
+generating. Recovering that generation means asking again under the **same**
+`Idempotency-Key`, and if `run` minted the key for you then that key is the one
+thing you need and the one thing you never saw.
+
+So every exception `models.run` raises carries it:
+
+```python
+import httpx
+from comfy_sdk import Comfy, ComfyError
+
+with Comfy() as client:
+    try:
+        result = client.models.run("acme/flux/dev", {"prompt": "a cat"})
+    # Both, and the second is not optional: a dropped connection or a read
+    # timeout — one of the two cases this section is about — never reached a
+    # response to translate, so it arrives as the `httpx` error it was, not as
+    # a `ComfyError`. Catching only `ComfyError` misses exactly the failure the
+    # replay exists for.
+    except (ComfyError, httpx.HTTPError) as exc:
+        key = exc.idempotency_key
+        if key is None:
+            raise  # Nothing to replay under; a fresh key would re-run and re-bill.
+        # Later — after `exc.retry_after` seconds, if the server named a pace.
+        result = client.models.run(
+            "acme/flux/dev", {"prompt": "a cat"}, idempotency_key=key
+        )
+```
+
+Against a deployment that replays a claimed key, the second call returns the
+original generation's result rather than starting a new one; while that
+generation is still running it is refused instead, with a `Retry-After` saying
+when to ask. Send the same arguments you sent the first time — a repeated key
+with a *different* body is rejected outright.
+
+Three attributes carry this:
+
+| Attribute | Value |
+|---|---|
+| `exc.idempotency_key` | the key that call was made under — the one `run` minted, or the one you passed. Present on every exception `models.run` raises, including a transport failure with no response at all (`httpx.ConnectError`, a read timeout), a cancelled `await`, and a `RouterError` whose `error_type` this SDK version does not recognise |
+| `exc.request_id` | the server's `X-Comfy-Request-Id` for the call, when the response carried one — the id to quote in a support request. `None` when there was no response, or none of that header |
+| `exc.retry_after` | seconds the server asked you to wait before asking again, from `Retry-After`. `None` when it named no pace |
+
+All three read as `None` rather than raising on any exception `models.run`
+raises, so a handler never has to guard the attribute access itself.
+
+`idempotency_key` is `None` on errors from *other* surfaces, though — it is
+`models.run` that records it, and `submit()` sends a key without stamping one.
+So `None` means "this SDK did not record a key for you", **not** "no key was
+sent, resend freely": check for it before replaying, as the snippet above does,
+rather than passing it straight back into `idempotency_key=` where `None` means
+"mint a fresh one" and starts a second billed generation.
+
 ## Sync and async
 
 `Comfy` and `AsyncComfy` expose the identical surface — swap the import and

src/comfy_low/errors.py

@@ -9,8 +9,34 @@
 
 from __future__ import annotations
 
+import re
 from typing import Any
 
+#: The leading run of characters a request id may consist of, bounded in the
+#: pattern itself. ``X-Comfy-Request-Id`` is server-controlled and the id is
+#: meant to be *displayed* — rendered in a traceback, written to a log, pasted
+#: into a support ticket — so it is reduced to something safe to display rather
+#: than kept verbatim. Matching a leading run (rather than deleting the
+#: offending bytes) also gives the right answer for the one case a well-behaved
+#: server can still produce: ``httpx.Headers.get`` joins duplicate headers with
+#: ``", "``, and a comma is not in the class, so ``"a1, a2"`` yields ``"a1"``
+#: instead of a spliced ``"a1a2"`` that identifies no call at all.
+_REQUEST_ID_RE = re.compile(r"[A-Za-z0-9._:+/=@-]{1,200}")
+
+
+def clean_request_id(raw: Any) -> str | None:
+    """``raw`` reduced to a bounded, printable request id, or ``None``.
+
+    Defined here rather than beside either reader because both error surfaces
+    parse the same header off their own response — ``comfy_low.transport`` off
+    the shared envelope, ``comfy_sdk.router_exceptions`` off the router's — and
+    an id that is safe to display on one of them has to be safe on the other.
+    """
+    if not isinstance(raw, str):
+        return None
+    match = _REQUEST_ID_RE.match(raw.strip())
+    return match.group(0) if match else None
+
 
 class ApiError(Exception):
     """Base for every error carried by the API's error envelope."""
@@ -25,6 +51,7 @@ def __init__(
         http_status: int,
         details: dict[str, Any] | None = None,
         retry_after: int | None = None,
+        request_id: str | None = None,
     ) -> None:
         super().__init__(message)
         self.message = message
@@ -33,6 +60,12 @@ def __init__(
         self.http_status = http_status
         self.details = details
         self.retry_after = retry_after
+        #: Server-minted id for the call, read off ``X-Comfy-Request-Id``.
+        #: ``None`` when the response carried no such header. Surfaced the same
+        #: way ``retry_after`` is — a response header kept on the exception,
+        #: because it is the id a user quotes in a support request and it is
+        #: unreachable once the response object is gone.
+        self.request_id = request_id
 
 
 class InvalidWorkflow(ApiError):
@@ -116,6 +149,7 @@ def error_from_envelope(
     body: dict[str, Any] | None,
     *,
     retry_after: int | None = None,
+    request_id: str | None = None,
     error_type: str | None = None,
 ) -> ApiError:
     """Build the typed exception for an error response.
@@ -169,6 +203,7 @@ def error_from_envelope(
         http_status=http_status,
         details=details if isinstance(details, dict) else None,
         retry_after=retry_after,
+        request_id=request_id,
     )
 
 

src/comfy_low/transport.py

@@ -37,7 +37,7 @@
 import httpx
 
 from . import _multipart
-from .errors import ApiError, error_from_envelope
+from .errors import ApiError, clean_request_id, error_from_envelope
 from .models import Asset, Job, JobWorkflowResponse
 from .sse import RawEvent, SSEDecoder
 
@@ -127,6 +127,25 @@ def _retry_after(resp: httpx.Response) -> int | None:
         return None
 
 
+#: Response header carrying the server-minted id for the call. Spelled here as
+#: well as in :data:`comfy_sdk.router_exceptions.REQUEST_ID_HEADER` because the
+#: two layers read it independently — the router surface off its own error
+#: response, this one off the shared error envelope — and ``comfy_low`` never
+#: imports from ``comfy_sdk``. ``tests/test_error_contract.py`` asserts the two
+#: spellings agree, so they cannot drift apart.
+REQUEST_ID_HEADER = "X-Comfy-Request-Id"
+
+
+def _request_id(resp: httpx.Response) -> str | None:
+    """``X-Comfy-Request-Id`` as a bounded, printable string, or ``None``.
+
+    Filtered rather than kept verbatim — see
+    :func:`comfy_low.errors.clean_request_id`, which both error surfaces share
+    so the id cannot be safe to display on one and not the other.
+    """
+    return clean_request_id(resp.headers.get(REQUEST_ID_HEADER))
+
+
 def origin(url: str) -> tuple[str, str, int | None]:
     """Normalized ``(scheme, host, port)`` — the parts that define same-origin.
 
@@ -212,9 +231,24 @@ def headers(self, url: str, extra: dict[str, str] | None = None) -> dict[str, st
 
     def parse_or_raise(self, resp: httpx.Response, ok: tuple[int, ...]) -> dict[str, Any]:
         if resp.status_code in ok:
-            if resp.content:
+            if not resp.content:
+                return {}
+            try:
                 return resp.json()
-            return {}
+            except ValueError as exc:
+                # A success status whose body will not decode — a proxy
+                # interstitial served as 200, a response truncated mid-stream.
+                # Raised as an ApiError rather than escaping as the raw
+                # `json.JSONDecodeError` so it lands on the surface the SDK
+                # translates and stamps: on `models.run` this is a generation
+                # that ran and was billed with the result lost, which is
+                # exactly the failure the Idempotency-Key has to ride out on.
+                raise ApiError(
+                    f"Could not decode the {resp.status_code} response body as JSON",
+                    code="invalid_response",
+                    http_status=resp.status_code,
+                    request_id=_request_id(resp),
+                ) from exc
         body: dict[str, Any] | None
         try:
             body = resp.json()
@@ -224,6 +258,7 @@ def parse_or_raise(self, resp: httpx.Response, ok: tuple[int, ...]) -> dict[str,
             resp.status_code,
             body,
             retry_after=_retry_after(resp),
+            request_id=_request_id(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