From 239ea184bca54f1237b6e63a910fe7b963ef01fd Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Sat, 5 Sep 2026 14:41:24 -0700 Subject: [PATCH] feat: carry the Idempotency-Key on every exception submit() raises MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Comfy.submit` / `AsyncComfy.submit` mint an `Idempotency-Key` per call and send it on `POST /jobs`, but the key was a local of the submitting frame: any exception the call raised took the caller's only record of what was sent with it. A transport failure escaped `submit()` untranslated as well, so the documented `.request_id` / `.retry_after` pair raised `AttributeError` there rather than reading `None`. Wrap each retry loop in `translating(idempotency_key=key)` — the same boundary `models.run` already uses. The inner `except ApiError` handler still owns what is retried and what surfaces; the outer stamp only attaches the key to whatever it lets out, so nothing about the retry policy, the key minting or the wire changed. The semantics differ from `models.run`'s and the docs say so: `POST /jobs` rejects a reused key (`422 idempotency_key_reuse`) rather than replaying it, so the key on a `submit()` error is a record of what was sent — poll or list for the job the first attempt may have created — not a replay handle. --- CHANGELOG.md | 16 ++ README.md | 22 ++- src/comfy_sdk/client.py | 65 +++--- src/comfy_sdk/exceptions.py | 23 ++- tests/conftest.py | 6 + tests/test_client_submit_idempotency.py | 252 ++++++++++++++++++++++++ 6 files changed, 349 insertions(+), 35 deletions(-) create mode 100644 tests/test_client_submit_idempotency.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f4968d..38d70c8 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index ac04659..bca503e 100644 --- a/README.md +++ b/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 diff --git a/src/comfy_sdk/client.py b/src/comfy_sdk/client.py index a7aa403..9c469bc 100644 --- a/src/comfy_sdk/client.py +++ b/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, diff --git a/src/comfy_sdk/exceptions.py b/src/comfy_sdk/exceptions.py index be2986a..d37c610 100644 --- a/src/comfy_sdk/exceptions.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index 799351a..ffd073e 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_client_submit_idempotency.py b/tests/test_client_submit_idempotency.py new file mode 100644 index 0000000..cb9db27 --- /dev/null +++ b/tests/test_client_submit_idempotency.py @@ -0,0 +1,252 @@ +"""``submit()`` stamps its ``Idempotency-Key`` onto whatever it raises. + +``Comfy.submit`` / ``AsyncComfy.submit`` mint a key per call and send it on +``POST /jobs``. Until now that key was a local of the submitting frame, so it +died with any exception the call raised — and the caller of an auto-minted +submit had no record of what had been sent. These tests pin the same contract +``tests/test_models_run.py`` pins for ``models.run``: every failure of the call +carries ``.idempotency_key``, including a transport failure that never reached +a response and a cancellation of an in-flight submit. + +The semantics differ from ``models.run``'s and the difference is deliberate: +``POST /jobs`` *rejects* a reused key (``422 idempotency_key_reuse``) rather +than replaying it, so this key is a record of what was sent — poll or list for +the job the first attempt may have created — not a replay handle. + +Everything here runs against the stubbed server in ``conftest.py``, except the +two cases a listening server cannot produce (a connect failure, a cancellation). +""" + +from __future__ import annotations + +import asyncio +import socket +from typing import Any + +import httpx +import pytest + +import comfy_sdk.client as _client_module +from comfy_sdk import AsyncComfy, Comfy, ComfyError, IdempotencyKeyReuse, QueueFull +from comfy_sdk.client import BASE_URL_ENV_VAR + +_GRAPH = {"3": {"class_type": "KSampler", "inputs": {}}} + + +def _wf(client: Comfy | AsyncComfy): + return client.workflows.from_json(_GRAPH) + + +def _closed_port() -> int: + """A port nothing is listening on — bound, read, then released.""" + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +class _StepClock: + """A monotonic clock that advances a fixed step per read. + + The queue-full budget is spent by *reads* rather than by wall time, so + "the budget ran out after N attempts" is exact instead of a race against a + loaded machine. + """ + + def __init__(self, step: float = 1.0) -> None: + self._t = 0.0 + self._step = step + + def __call__(self) -> float: + now = self._t + self._t += self._step + return now + + +# --- the auto-minted key, which the caller never saw ---------------------- + + +def test_a_failed_submit_exposes_the_exact_auto_minted_key(server) -> None: + # An unmapped code, so this lands on the base `ComfyError` — the stamp has + # to reach it, not only the typed subclasses. + server.state.job_error = (422, "validation_error") + with Comfy() as client: + with pytest.raises(ComfyError) as excinfo: + client.submit(_wf(client)) + (sent,) = server.state.jobs_idempotency_keys + assert sent is not None + # The exact key the server received, not merely "a key": truthiness would + # pass on a freshly minted one, which records nothing about what was sent. + assert excinfo.value.idempotency_key == sent + + +async def test_a_failed_async_submit_exposes_the_exact_auto_minted_key(server) -> None: + server.state.job_error = (422, "validation_error") + async with AsyncComfy() as client: + with pytest.raises(ComfyError) as excinfo: + await client.submit(_wf(client)) + (sent,) = server.state.jobs_idempotency_keys + assert sent is not None + assert excinfo.value.idempotency_key == sent + + +def test_a_caller_supplied_key_round_trips_onto_the_exception(server) -> None: + server.state.job_error = (422, "validation_error") + with Comfy() as client: + with pytest.raises(ComfyError) as excinfo: + client.submit(_wf(client), idempotency_key="abc") + assert excinfo.value.idempotency_key == "abc" + assert server.state.jobs_idempotency_keys == ["abc"] + + +async def test_a_caller_supplied_key_round_trips_on_the_async_client(server) -> None: + server.state.job_error = (422, "validation_error") + async with AsyncComfy() as client: + with pytest.raises(ComfyError) as excinfo: + await client.submit(_wf(client), idempotency_key="abc-async") + assert excinfo.value.idempotency_key == "abc-async" + assert server.state.jobs_idempotency_keys == ["abc-async"] + + +# --- the retried paths: one key across every attempt ---------------------- + + +def test_an_exhausted_queue_full_budget_carries_the_retried_key(server, monkeypatch) -> None: + # The server never clears backpressure, so the retry loop runs to the end + # of its budget and surfaces `QueueFull`. That exception has to carry the + # key — and it must be the *same* key every attempt was made under, not a + # fresh one minted somewhere along the way. + server.state.queue_full_times = 1_000_000 + monkeypatch.setattr(_client_module, "_now", _StepClock()) + monkeypatch.setattr(_client_module, "_QUEUE_RETRY_BUDGET", 3.0) + with Comfy() as client: + with pytest.raises(QueueFull) as excinfo: + client.submit(_wf(client)) + sent = server.state.jobs_idempotency_keys + assert len(sent) > 1 + assert set(sent) == {excinfo.value.idempotency_key} + assert excinfo.value.idempotency_key is not None + + +async def test_an_exhausted_async_queue_full_budget_carries_the_retried_key( + server, monkeypatch +) -> None: + server.state.queue_full_times = 1_000_000 + monkeypatch.setattr(_client_module, "_now", _StepClock()) + monkeypatch.setattr(_client_module, "_QUEUE_RETRY_BUDGET", 3.0) + async with AsyncComfy() as client: + with pytest.raises(QueueFull) as excinfo: + await client.submit(_wf(client)) + sent = server.state.jobs_idempotency_keys + assert len(sent) > 1 + assert set(sent) == {excinfo.value.idempotency_key} + + +# --- reject-not-replay: the key that was refused -------------------------- + + +def test_a_rejected_reused_key_is_the_key_on_the_exception(server) -> None: + # `POST /jobs` keys are single-use: the second call under the same key is + # refused rather than replayed. The exception names the key that was + # rejected, which is what a caller polls or lists for the first job by. + with Comfy() as client: + client.submit(_wf(client), idempotency_key="reused-01") + with pytest.raises(IdempotencyKeyReuse) as excinfo: + client.submit(_wf(client), idempotency_key="reused-01") + assert excinfo.value.idempotency_key == "reused-01" + assert server.state.jobs_idempotency_keys == ["reused-01", "reused-01"] + + +async def test_a_rejected_reused_key_is_the_key_on_the_async_exception(server) -> None: + async with AsyncComfy() as client: + await client.submit(_wf(client), idempotency_key="reused-02") + with pytest.raises(IdempotencyKeyReuse) as excinfo: + await client.submit(_wf(client), idempotency_key="reused-02") + assert excinfo.value.idempotency_key == "reused-02" + + +# --- no response at all: the case the key is most needed on --------------- + + +def test_a_transport_level_failure_carries_the_key(monkeypatch) -> None: + # Nothing to translate — no response reached the client — so before this + # the httpx error escaped `submit()` untouched and the caller was left with + # a request that may or may not have created a job and no key to look for + # it by. + monkeypatch.setenv(BASE_URL_ENV_VAR, f"http://127.0.0.1:{_closed_port()}") + with Comfy() as client: + with pytest.raises(httpx.ConnectError) as excinfo: + client.submit(_wf(client)) + assert excinfo.value.idempotency_key is not None # type: ignore[attr-defined] + # The documented trio has to *read* rather than raise on exactly the + # failures it is most needed on: httpx's classes declare neither attribute. + assert excinfo.value.request_id is None # type: ignore[attr-defined] + assert excinfo.value.retry_after is None # type: ignore[attr-defined] + + +async def test_an_async_transport_level_failure_carries_the_key(monkeypatch) -> None: + monkeypatch.setenv(BASE_URL_ENV_VAR, f"http://127.0.0.1:{_closed_port()}") + async with AsyncComfy() as client: + with pytest.raises(httpx.ConnectError) as excinfo: + await client.submit(_wf(client)) + assert excinfo.value.idempotency_key is not None # type: ignore[attr-defined] + assert excinfo.value.request_id is None # type: ignore[attr-defined] + assert excinfo.value.retry_after is None # type: ignore[attr-defined] + + +def test_a_caller_supplied_key_survives_a_transport_failure(monkeypatch) -> None: + monkeypatch.setenv(BASE_URL_ENV_VAR, f"http://127.0.0.1:{_closed_port()}") + with Comfy() as client: + with pytest.raises(httpx.ConnectError) as excinfo: + client.submit(_wf(client), idempotency_key="abc-transport") + assert excinfo.value.idempotency_key == "abc-transport" # type: ignore[attr-defined] + + +# --- run(), which submits through submit() -------------------------------- + + +def test_run_surfaces_the_key_of_a_failed_submit_phase(server) -> None: + # `run` mints no key of its own — it calls `submit`, so the stamp has to + # reach the caller through it. + server.state.job_error = (422, "invalid_workflow") + with Comfy() as client: + with pytest.raises(ComfyError) as excinfo: + client.run(_wf(client)) + (sent,) = server.state.jobs_idempotency_keys + assert sent is not None + assert excinfo.value.idempotency_key == sent + + +async def test_async_run_surfaces_the_key_of_a_failed_submit_phase(server) -> None: + server.state.job_error = (422, "invalid_workflow") + async with AsyncComfy() as client: + with pytest.raises(ComfyError) as excinfo: + await client.run(_wf(client)) + (sent,) = server.state.jobs_idempotency_keys + assert sent is not None + assert excinfo.value.idempotency_key == sent + + +# --- cancellation, the one BaseException the key rides out on ------------- + + +async def test_cancelling_an_in_flight_submit_still_yields_the_key(server) -> None: + # A submit cancelled mid-flight — what `asyncio.wait_for` around it does — + # may already have reached the server and created a job, so the key is the + # caller's only record of what to look for. The cancellation itself must + # still propagate: it is re-raised bare, never converted. + seen: list[str | None] = [] + + async def _cancelled_post_jobs( + graph: dict[str, Any], *, idempotency_key: str | None = None, **kw: Any + ) -> Any: + seen.append(idempotency_key) + raise asyncio.CancelledError + + async with AsyncComfy() as client: + wf = _wf(client) + client._low.post_jobs = _cancelled_post_jobs # type: ignore[method-assign] + with pytest.raises(asyncio.CancelledError) as excinfo: + await client.submit(wf) + assert seen == [excinfo.value.idempotency_key] # type: ignore[attr-defined] + assert excinfo.value.idempotency_key is not None # type: ignore[attr-defined] + assert excinfo.value.request_id is None # type: ignore[attr-defined]