feat: carry the Idempotency-Key on every exception submit() raises

CHANGELOG.md

@@ -14,6 +14,22 @@ notes for each version.
 
 ### Added
 
+- Every exception `client.submit()` — and so `client.run()`, which submits
+  through it — raises **for a failed call** now carries the `Idempotency-Key`
+  it was made under, on `.idempotency_key`, matching `client.models.run()`:
+  the mapped `ComfyError` subclasses, a `QueueFull` raised once the 429 retry
+  budget is exhausted (the same key on every retried attempt), and a transport
+  failure with no response at all (a dropped connection, a read timeout), which
+  previously escaped `submit()` untranslated and now reads `.request_id` and
+  `.retry_after` as `None` rather than raising `AttributeError`. Cancelling an
+  in-flight `AsyncComfy.submit()` yields the key too, and the cancellation
+  still propagates unchanged. The semantics differ from `models.run`'s and the
+  difference matters: `POST /jobs` **rejects** a reused key with
+  `422 idempotency_key_reuse` rather than replaying it, so the key on a
+  `submit()` error records what was sent — poll or list for the job the first
+  attempt may already have created — rather than being a replay handle to
+  resubmit under. Nothing about what is retried, what key is minted, or what
+  goes on the wire changed.
 - 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

README.md

@@ -584,12 +584,22 @@ Three attributes carry this:
 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.
+`submit()` — and so `run()`, which submits through it — stamps the key too, but
+what the key is *good for* differs, so read it with the surface in mind.
+`models.run` sends it to a surface that **replays** a claimed key, which is what
+makes it a handle on a generation you were already billed for. `POST /jobs`
+instead **rejects** a reused key with `422 idempotency_key_reuse` (see the
+[`IdempotencyKeyReuse`](#typed-errors) bullet below): keys there are single-use
+and there is no replay. So on a `submit()` failure `exc.idempotency_key` tells
+you a key *was* sent — do not resubmit blindly under it, poll or list for the
+job the first attempt may already have created.
+
+`idempotency_key` is still `None` on errors from every other surface — one that
+sends no key, and an asset upload, which mints a key per handle without
+recording it. 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
 

src/comfy_sdk/client.py

@@ -48,7 +48,7 @@
 
 from . import _core
 from .assets import AssetFactory, AsyncAssetFactory
-from .exceptions import MissingApiKey, WorkflowFormatUi, to_sdk_error
+from .exceptions import MissingApiKey, WorkflowFormatUi, to_sdk_error, translating
 from .jobs import AsyncJob, AsyncJobFactory, Job, JobFactory
 from .models import AsyncModels, Models
 from .retry import DEFAULT_RETRY, RetryPolicy
@@ -310,6 +310,14 @@ def submit(
         reused key is *rejected*, not replayed: on reuse, catch the error and
         poll/list for the job the first attempt already created.
 
+        Every exception a failed submit raises carries the key it was made
+        under on ``.idempotency_key`` — including a transport failure that
+        never reached a response (``httpx.ConnectError``, a read timeout),
+        which reads ``.request_id`` and ``.retry_after`` as ``None`` rather
+        than raising. Because a reused key is rejected rather than replayed,
+        that key is a record of what was sent, not a replay handle: use it to
+        go looking for the job, not to resubmit.
+
         ``api_key`` authenticates partner (API) nodes embedded in the workflow
         (e.g. Gemini) — unrelated to idempotency and unrelated to the bearer
         token this client was constructed with. It is never persisted or
@@ -321,17 +329,23 @@ def submit(
         key = idempotency_key or _core.new_idempotency_key()
         extra_data = _core.extra_data_for(api_key)
         deadline = _now() + _QUEUE_RETRY_BUDGET
-        while True:
-            try:
-                model = self._low.post_jobs(graph, idempotency_key=key, extra_data=extra_data)
-                return Job(self._low, model)
-            except ApiError as exc:
-                err = to_sdk_error(exc)
-                delay = _retry_delay(exc, deadline)
-                if delay is None:
-                    raise err from exc
-                time.sleep(delay)
-                continue
+        # The key is a local of this frame, so anything that propagates past
+        # here takes the caller's only record of what was sent with it. The
+        # inner handler still owns what is retried and what surfaces; this
+        # stamps the key onto whatever it lets out — including a transport
+        # failure that never reached a response to translate.
+        with translating(idempotency_key=key):
+            while True:
+                try:
+                    model = self._low.post_jobs(graph, idempotency_key=key, extra_data=extra_data)
+                    return Job(self._low, model)
+                except ApiError as exc:
+                    err = to_sdk_error(exc)
+                    delay = _retry_delay(exc, deadline)
+                    if delay is None:
+                        raise err from exc
+                    time.sleep(delay)
+                    continue
 
     def run(
         self,
@@ -426,17 +440,22 @@ async def submit(
         key = idempotency_key or _core.new_idempotency_key()
         extra_data = _core.extra_data_for(api_key)
         deadline = _now() + _QUEUE_RETRY_BUDGET
-        while True:
-            try:
-                model = await self._low.post_jobs(graph, idempotency_key=key, extra_data=extra_data)
-                return AsyncJob(self._low, model)
-            except ApiError as exc:
-                err = to_sdk_error(exc)
-                delay = _retry_delay(exc, deadline)
-                if delay is None:
-                    raise err from exc
-                await asyncio.sleep(delay)
-                continue
+        # See :meth:`Comfy.submit` — same stamp, and here it also rides out on
+        # the cancellation of an in-flight `post_jobs`.
+        with translating(idempotency_key=key):
+            while True:
+                try:
+                    model = await self._low.post_jobs(
+                        graph, idempotency_key=key, extra_data=extra_data
+                    )
+                    return AsyncJob(self._low, model)
+                except ApiError as exc:
+                    err = to_sdk_error(exc)
+                    delay = _retry_delay(exc, deadline)
+                    if delay is None:
+                        raise err from exc
+                    await asyncio.sleep(delay)
+                    continue
 
     async def run(
         self,

src/comfy_sdk/exceptions.py

@@ -23,12 +23,23 @@ class ComfyError(Exception):
     """Base for every SDK-level error."""
 
     #: The ``Idempotency-Key`` the failed call was made under. Populated by
-    #: :meth:`comfy_sdk.models.Models.run` and its async twin, which are the
-    #: operations that pass a key to :func:`translating`; ``None`` everywhere
-    #: else — including on operations that *do* send a key but do not stamp it
-    #: (``Comfy.submit()``), and on an exception constructed by hand. So
-    #: ``None`` means "this SDK did not record a key for you", never "no key
-    #: reached the server": do not infer from it that a resend is safe.
+    #: :meth:`comfy_sdk.models.Models.run` and its async twin, and by
+    #: :meth:`comfy_sdk.client.Comfy.submit` /
+    #: :meth:`comfy_sdk.client.AsyncComfy.submit` — and therefore by
+    #: ``Comfy.run`` / ``AsyncComfy.run``, which submit through them. It is
+    #: ``None`` everywhere else: on an operation that sends no key, on an
+    #: asset upload (which mints a key per handle and does not record it), and
+    #: on an exception constructed by hand. So ``None`` means "this SDK did not
+    #: record a key for you", never "no key reached the server": do not infer
+    #: from it that a resend is safe.
+    #:
+    #: What the key is *good for* differs by surface, so read it with the
+    #: operation in mind: ``models.run`` sends it to a surface that replays a
+    #: claimed key, so the key is a handle on the generation you were already
+    #: billed for. ``POST /jobs`` instead *rejects* a reused key with
+    #: ``422 idempotency_key_reuse``, so on a ``submit`` failure the key says a
+    #: key was sent — poll or list for the job the first attempt may have
+    #: created rather than resubmitting under it.
     #:
     #: Declared on the base rather than set per subclass so that a bucket this
     #: SDK version has never heard of — which arrives as a bare

tests/conftest.py

@@ -181,6 +181,11 @@ class ServerState:
     # Idempotency-Key -> job id of the first (accepted) request, so a reuse of
     # the same key can be detected and rejected (single-use, no replay).
     idempotency: dict[str, str] = field(default_factory=dict)
+    # Every Idempotency-Key seen on POST /jobs, in arrival order (`None`
+    # records a submit that arrived without the header at all). Distinct from
+    # `idempotency`, which only records the keys an *accepted* request claimed
+    # — a test about a failed submit needs the key the server actually saw.
+    jobs_idempotency_keys: list[str | None] = field(default_factory=list)
     # Raw bytes of the last POST /assets multipart body (so tests can inspect
     # the parts actually sent — e.g. how many `tags` fields were included).
     last_upload_body: bytes = b""
@@ -660,6 +665,7 @@ def _post_jobs(self) -> None:
             state.last_workflow = body.get("workflow")
             state.last_jobs_body = body
             key = self.headers.get("Idempotency-Key")
+            state.jobs_idempotency_keys.append(key)
 
             if key and key in state.idempotency:
                 # Reject-on-duplicate (single-use keys, no replay): any reuse of