From 0742a448caec4eded5ae373f68bdacf6a8b1653e Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 26 Aug 2026 20:12:48 -0700 Subject: [PATCH 1/3] feat: carry the Idempotency-Key and request id on every exception models.run raises MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `models.run` mints an `Idempotency-Key` per call and sends it on every attempt, but the key was a local of that frame: when the call raised, the key went with it. That is the one value a caller needs to collect a generation they were already billed for after a lost response — a 504 where the server stopped holding the connection at its own deadline, a connection dropped mid-generation — since recovering it means asking again under the *same* key. Only callers who had chosen and stored their own `idempotency_key=` could do that; an auto-minted key left the paid-for generation uncollectable, and the server echoes the key on no response header, so there was nowhere else to read it from. `translating()` grows an optional `idempotency_key=` and stamps it onto whatever leaves the block, and both `run` loops pass their key. Stamping at that boundary rather than in each exception's constructor is what makes the base `RouterError` — what an `error_type` this version has never heard of falls through to — carry it too, along with a transport failure that has no response to translate. `request_id` comes along for the same reason: `X-Comfy-Request-Id` was read off a router error response but dropped from the shared error envelope, so the id a user quotes in a support request was unreachable once the response was gone. It is now on `ApiError` and on the `ComfyError` base, `None` where the response named none. Both attributes are declared on `ComfyError` with a `None` default, so `exc.idempotency_key` is always safe to read rather than a `getattr` dance, and `tests/test_error_contract.py` pins the base error's whole attribute surface as a written-out list so adding one is a deliberate edit and removing one is a failing test. Additive throughout: `translating()` with no key behaves exactly as before, nothing changes about what is retried, and nothing changes on the wire. --- CHANGELOG.md | 18 +++ README.md | 41 +++++++ src/comfy_low/errors.py | 9 ++ src/comfy_low/transport.py | 24 +++- src/comfy_sdk/exceptions.py | 70 ++++++++++- src/comfy_sdk/models.py | 27 +++- tests/conftest.py | 17 ++- tests/test_error_contract.py | 90 +++++++++++++- tests/test_models_run.py | 232 ++++++++++++++++++++++++++++++++++- 9 files changed, 511 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8662687..648b650 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,24 @@ notes for each version. ### Added +- Every exception `client.models.run()` raises now carries the + `Idempotency-Key` the call 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). `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. - 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 diff --git a/README.md b/README.md index c33814a..3d3f617 100644 --- a/README.md +++ b/README.md @@ -468,6 +468,47 @@ you just asked for. `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 +from comfy_sdk import Comfy, ComfyError + +with Comfy() as client: + try: + result = client.models.run("acme/flux/dev", {"prompt": "a cat"}) + except ComfyError as exc: + # Later — after the Retry-After the server named, if it named one. + result = client.models.run( + "acme/flux/dev", {"prompt": "a cat"}, idempotency_key=exc.idempotency_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. + +Two 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) 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 | + +Both are `None` rather than absent on an error from a surface that sends no +key, so `exc.idempotency_key` is always safe to read on a `ComfyError`. + ## Sync and async `Comfy` and `AsyncComfy` expose the identical surface — swap the import and diff --git a/src/comfy_low/errors.py b/src/comfy_low/errors.py index 8a3a12f..6182981 100644 --- a/src/comfy_low/errors.py +++ b/src/comfy_low/errors.py @@ -25,6 +25,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 +34,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): @@ -103,6 +110,7 @@ def error_from_envelope( body: dict[str, Any] | None, *, retry_after: int | None = None, + request_id: str | None = None, ) -> ApiError: """Build the typed exception for an error response. @@ -127,6 +135,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, ) diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index beb06e9..4bea998 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -127,6 +127,23 @@ 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 non-empty string, or ``None``.""" + raw = resp.headers.get(REQUEST_ID_HEADER) + if raw is None: + return None + return raw.strip() or None + + def origin(url: str) -> tuple[str, str, int | None]: """Normalized ``(scheme, host, port)`` — the parts that define same-origin. @@ -220,7 +237,12 @@ 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), + request_id=_request_id(resp), + ) def parse_expiry(url: str) -> datetime | None: diff --git a/src/comfy_sdk/exceptions.py b/src/comfy_sdk/exceptions.py index c4ca138..e65c0d5 100644 --- a/src/comfy_sdk/exceptions.py +++ b/src/comfy_sdk/exceptions.py @@ -10,7 +10,9 @@ from collections.abc import Iterator from contextlib import contextmanager -from typing import Any +from typing import Any, TypeVar + +import httpx from comfy_low.errors import ApiError from comfy_low.models import JobError @@ -19,6 +21,21 @@ class ComfyError(Exception): """Base for every SDK-level error.""" + #: The ``Idempotency-Key`` the failed call was made under, when the + #: operation that raised this sends one — see :func:`translating`. It is + #: ``None`` on every error from an operation that sends no key, and on one + #: constructed by hand. + #: + #: 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 + #: :class:`~comfy_sdk.router_exceptions.RouterError` — still carries it. + idempotency_key: str | None = None + + #: Server-minted id for the call, from ``X-Comfy-Request-Id``, or ``None`` + #: when the response carried no such header (and on a failure with no + #: response at all). The id a user quotes in a support request. + request_id: str | None = None + def __init__( self, message: str, @@ -26,12 +43,14 @@ def __init__( code: str | None = None, http_status: int | None = None, details: dict[str, Any] | None = None, + request_id: str | None = None, ) -> None: super().__init__(message) self.message = message self.code = code self.http_status = http_status self.details = details + self.request_id = request_id class MissingApiKey(ComfyError): @@ -141,6 +160,7 @@ def to_sdk_error(exc: ApiError) -> ComfyError: code=exc.code, http_status=exc.http_status, details=exc.details, + request_id=exc.request_id, ) cls = _BY_CODE.get(exc.code, ComfyError) return cls( @@ -148,18 +168,62 @@ def to_sdk_error(exc: ApiError) -> ComfyError: code=exc.code, http_status=exc.http_status, details=exc.details, + request_id=exc.request_id, ) +#: What an ``idempotency_key=`` gets stamped onto. Deliberately not +#: "everything": these are the failures of the *call* — the SDK's own errors +#: (which includes every ``RouterError``) and an httpx failure that means the +#: request did not complete. Anything else leaving the block is a bug in the +#: SDK or the caller's own code, where a key is noise, and a ``KeyboardInterrupt`` +#: must not be touched at all. +_STAMPABLE: tuple[type[BaseException], ...] = (ComfyError, httpx.HTTPError) + +_E = TypeVar("_E", bound=BaseException) + + +def _stamp(exc: _E, idempotency_key: str | None) -> _E: + """Attach ``idempotency_key`` to ``exc`` in place and hand it back. + + ``setattr`` rather than a constructor argument because the transport-level + members of :data:`_STAMPABLE` are httpx's classes, which this SDK does not + build. A ``None`` key writes nothing, so an operation that sends no key + leaves ``ComfyError.idempotency_key`` at its class default. + """ + if idempotency_key is not None: + exc.idempotency_key = idempotency_key # type: ignore[attr-defined] + return exc + + @contextmanager -def translating() -> Iterator[None]: +def translating(*, idempotency_key: str | None = None) -> Iterator[None]: """Re-raise any protocol ``ApiError`` as its idiomatic SDK exception. Wrap the SDK-level operations that call ``comfy_low`` with this so integrators only ever catch ``comfy_sdk`` exceptions (``MissingAsset``, ``HashMismatch``, ``NotFound``, ...), never the raw protocol error. + + ``idempotency_key`` is the key the wrapped call was made under. Give it and + every failure that leaves this block carries it as ``.idempotency_key`` — + the one place that can, because the key is a local of the caller's frame and + is otherwise lost the moment the exception propagates past it. That matters + on :meth:`comfy_sdk.models.Models.run`, where the router's replay contract + lets a caller who lost the response collect the generation they were already + billed for by resending under the *same* key. Stamping here rather than in + each exception's constructor is what makes an unrecognised ``error_type``, + which falls through to the base ``RouterError``, carry it too. Omit it and + this behaves exactly as it did before the parameter existed. """ try: yield except ApiError as exc: - raise to_sdk_error(exc) from exc + raise _stamp(to_sdk_error(exc), idempotency_key) from exc + except _STAMPABLE as exc: + # Already on a surface a caller catches — a RouterError the transport + # raised directly, or an httpx failure with no response to translate. + # Nothing to convert, but the key still has to ride out with it. A bare + # `raise` keeps the original traceback, so with no key given this branch + # is indistinguishable from not catching at all. + _stamp(exc, idempotency_key) + raise diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index 116db98..b09dadb 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -147,6 +147,18 @@ def run( with it the key claimed, so those need ``RetryPolicy(retry_possibly_in_flight=True)``. See :mod:`comfy_sdk.retry`, and ``Comfy(retry=NO_RETRY)`` to switch it off. + + **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`` + ``504``, a dropped connection — can still collect a generation they were + already billed for: ``client.models.run(model, arguments, + idempotency_key=exc.idempotency_key)`` returns the original result + (``200`` + ``Idempotent-Replayed``) against a deployment that replays a + claimed key, or is refused ``409`` "still in progress" with a + ``Retry-After`` saying when to ask again. Without that attribute an + auto-minted key died with the call and the paid-for generation was + uncollectable. """ low = cast(ComfyLow, self._low) # Minted once, outside the loop: reusing this exact value on every @@ -158,7 +170,10 @@ def run( # same-key-different-body case the contract rejects outright. payload = dict(arguments) retrier = Retrier(self._retry, now=_now) - with translating(): + # 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. + with translating(idempotency_key=key): while True: try: return low.post_model_run(model, payload, idempotency_key=key, timeout=timeout) @@ -187,8 +202,9 @@ async def run( """Awaitable :meth:`Models.run` — same arguments, same result shape. This *is* the async form of ``run``: awaiting it on ``AsyncComfy`` is - the whole difference from the sync client — including the retry policy - and the one-key-per-call rule. See :meth:`Models.run`. + the whole difference from the sync client — including the retry policy, + the one-key-per-call rule, and the ``.idempotency_key`` every exception + it raises carries for the replay. See :meth:`Models.run`. """ low = cast(AsyncComfyLow, self._low) key = idempotency_key or new_idempotency_key() @@ -197,7 +213,10 @@ async def run( # while the retry sleeps. payload = dict(arguments) retrier = Retrier(self._retry, now=_now) - with translating(): + # 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. + with translating(idempotency_key=key): while True: try: return await low.post_model_run( diff --git a/tests/conftest.py b/tests/conftest.py index 1a393e2..cb45449 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -118,6 +118,9 @@ class ServerState: # is the 429 the policy is allowed to pace itself against. `None` sends no # header at all, which is the 429 the policy must *not* retry. model_run_retry_after: str | None = None + # Sent as X-Comfy-Request-Id alongside a failed run. `None` sends no header, + # which is the response an intermediary that never reached the router gives. + model_run_request_id: str | None = None # --- counters the tests assert on --- upload_count: int = 0 @@ -474,12 +477,16 @@ def claim_if_outcome_unknown(status: int) -> None: def fail(status: int, code: str, message: str) -> None: claim_if_outcome_unknown(status) - headers = ( - {"Retry-After": state.model_run_retry_after} - if state.model_run_retry_after is not None - else None + headers: dict[str, str] = {} + if state.model_run_retry_after is not None: + headers["Retry-After"] = state.model_run_retry_after + if state.model_run_request_id is not None: + headers["X-Comfy-Request-Id"] = state.model_run_request_id + self._json( + status, + {"error": {"code": code, "message": message}}, + headers=headers or None, ) - self._json(status, {"error": {"code": code, "message": message}}, headers=headers) if state.model_run_fail_times > 0: state.model_run_fail_times -= 1 diff --git a/tests/test_error_contract.py b/tests/test_error_contract.py index d9f9787..4ac1463 100644 --- a/tests/test_error_contract.py +++ b/tests/test_error_contract.py @@ -5,6 +5,13 @@ Regression guard for the 10 entry points that used to skip it: the four ``outputs.py`` download methods (sync + async), both asset/job factories' ``get()``, and the non-501 raise in ``events()``. + +The second half of the file pins the *shape* of that surface: the attributes +every SDK exception is guaranteed to answer to. They are asserted here rather +than left to whichever raise site happens to set them, because an attribute a +caller reads inside an ``except`` block is a published contract the moment the +package ships — a later release cannot take one away, and one that is present +on some errors and missing on others forces every caller into ``getattr``. """ from __future__ import annotations @@ -15,9 +22,11 @@ import pytest +import comfy_low.transport as low_transport from comfy_low.errors import ApiError as LowApiError from comfy_sdk import AsyncComfy, Comfy, Forbidden, NotFound -from comfy_sdk.exceptions import ComfyError +from comfy_sdk.exceptions import ComfyError, translating +from comfy_sdk.router_exceptions import REQUEST_ID_HEADER, ROUTER_EXCEPTIONS, RouterError def _wf(client: Comfy | AsyncComfy): @@ -141,3 +150,82 @@ async def _drain() -> None: pass await _assert_no_leak_async("AsyncJob.events", _drain(), Forbidden) + + +# -- the attributes every SDK exception answers to --------------------------- + + +#: The whole published attribute surface of the base error. Spelled out rather +#: than introspected so that *adding* one is a deliberate edit to this list — +#: which is what makes it reviewable — and *removing* one is a failing test +#: rather than a silent break of somebody's ``except`` block. +_BASE_ATTRIBUTES = ( + "message", + "code", + "http_status", + "details", + "request_id", + "idempotency_key", +) + + +@pytest.mark.parametrize("name", _BASE_ATTRIBUTES) +def test_the_base_error_answers_to_its_whole_attribute_surface(name: str) -> None: + assert hasattr(ComfyError("boom"), name) + + +@pytest.mark.parametrize("name", ("request_id", "idempotency_key")) +def test_those_attributes_default_to_none_rather_than_being_absent(name: str) -> None: + # An operation that sends no Idempotency-Key, and a response that named no + # request id, both leave the attribute readable as `None`. A caller writes + # `exc.idempotency_key` unconditionally; there is no surface on which that + # is an AttributeError. + assert getattr(ComfyError("boom"), name) is None + + +@pytest.mark.parametrize("cls", ROUTER_EXCEPTIONS + (RouterError,), ids=lambda c: c.__name__) +def test_every_router_bucket_answers_to_the_key_including_the_base(cls: type) -> None: + # Including `RouterError` itself, which is what an `error_type` this SDK + # version has never heard of falls through to. + assert cls("detail", http_status=500).idempotency_key is None + + +def test_translating_stamps_the_key_onto_a_translated_protocol_error() -> None: + with pytest.raises(NotFound) as excinfo: + with translating(idempotency_key="k-01"): + raise LowApiError("gone", code="not_found", http_status=404) + assert excinfo.value.idempotency_key == "k-01" + + +def test_translating_stamps_the_key_onto_an_sdk_error_it_did_not_build() -> None: + with pytest.raises(RouterError) as excinfo: + with translating(idempotency_key="k-02"): + raise RouterError("nope", error_type="unheard_of", http_status=503) + assert excinfo.value.idempotency_key == "k-02" + + +def test_translating_without_a_key_leaves_the_attribute_alone() -> None: + # The parameter is additive: every existing caller passes nothing and gets + # exactly the behaviour it had before the parameter existed. + with pytest.raises(NotFound) as excinfo: + with translating(): + raise LowApiError("gone", code="not_found", http_status=404) + assert excinfo.value.idempotency_key is None + + +def test_a_bug_in_the_sdk_is_not_stamped_and_not_swallowed() -> None: + # `_STAMPABLE` is deliberately not "everything": a programming error is not + # a failed call, a key means nothing on it, and it must reach the caller + # unaltered. Same reasoning as `models._CANDIDATE_FAILURES`. + with pytest.raises(ZeroDivisionError) as excinfo: + with translating(idempotency_key="k-03"): + raise ZeroDivisionError("a bug, not a failed request") + assert not hasattr(excinfo.value, "idempotency_key") + + +def test_the_two_layers_spell_the_request_id_header_the_same_way() -> None: + # `comfy_low` reads it off the shared error envelope and `comfy_sdk` reads + # it off a router error response, so the name is written out in both — and + # neither may import the other's copy (comfy_low never imports comfy_sdk). + # A drift here would silently drop the id on one of the two surfaces. + assert low_transport.REQUEST_ID_HEADER == REQUEST_ID_HEADER diff --git a/tests/test_models_run.py b/tests/test_models_run.py index 53fa4b3..70ae5d0 100644 --- a/tests/test_models_run.py +++ b/tests/test_models_run.py @@ -4,8 +4,9 @@ call blocks and returns the finished result, the async call awaits to the same shape, the awaitable form is the *async client* rather than a suffixed method (asserted, not merely absent), the wait is sized for a server that polls -upstream inside the call, an ``Idempotency-Key`` is plumbed onto the wire, and -the result handed back is the provider's own payload rather than a wrapper. +upstream inside the call, an ``Idempotency-Key`` is plumbed onto the wire and rides +out on whatever the call raises, and the result handed back is the provider's +own payload rather than a wrapper. Everything here runs against the stubbed server in ``conftest.py``. """ @@ -14,14 +15,22 @@ import inspect import re +from collections.abc import Mapping +from typing import Any, cast import httpx import pytest from comfy_low.transport import MODEL_RUN_TIMEOUT, AsyncComfyLow, ComfyLow, model_run_request -from comfy_sdk import NO_RETRY, AsyncComfy, Comfy +from comfy_sdk import NO_RETRY, AsyncComfy, Comfy, RetryPolicy from comfy_sdk.exceptions import ComfyError, NotFound, Unauthorized from comfy_sdk.models import AsyncModels, Models +from comfy_sdk.router_exceptions import ( + ERROR_TYPE_HEADER, + DeadlineExceeded, + RouterError, + error_from_response, +) MODEL = "acme/flux/dev" ARGS = {"prompt": "a cat", "steps": 4} @@ -243,3 +252,220 @@ def test_an_unmapped_failure_still_lands_as_a_comfy_error(server) -> None: client.models.run(MODEL, ARGS) assert excinfo.value.code == "model_unavailable" assert excinfo.value.http_status == 503 + + +# --- the key survives the failure ---------------------------------------- +# +# The router's replay contract lets a caller who lost a response resend the +# same request under the same `Idempotency-Key` and collect the generation +# they were already billed for. `run` mints that key itself, so unless the +# exception carries it the key dies with the call and a paid-for generation is +# uncollectable — the only callers who could replay were the ones who had +# passed `idempotency_key=` and stored it themselves. + + +class _RaisingLow: + """A transport whose every attempt raises ``exc``, recording the key sent. + + A stub HTTP server that is listening cannot produce a connect failure, and + ``post_model_run`` does not raise ``RouterError`` today — so the two cases + that never reach the wire are asserted against a transport that can make + them. + """ + + def __init__(self, exc: BaseException) -> None: + self._exc = exc + self.keys: list[str | None] = [] + + def post_model_run( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + timeout: Any = None, + ) -> dict[str, Any]: + self.keys.append(idempotency_key) + raise self._exc + + +class _AsyncRaisingLow(_RaisingLow): + async def post_model_run( # type: ignore[override] + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + timeout: Any = None, + ) -> dict[str, Any]: + self.keys.append(idempotency_key) + raise self._exc + + +def _models_over(exc: BaseException) -> tuple[Models, _RaisingLow]: + low = _RaisingLow(exc) + return Models(cast(Any, low), NO_RETRY), low + + +def _async_models_over(exc: BaseException) -> tuple[AsyncModels, _RaisingLow]: + low = _AsyncRaisingLow(exc) + return AsyncModels(cast(Any, low), NO_RETRY), low + + +def test_a_deadline_exceeded_504_exposes_the_exact_auto_minted_key(server) -> None: + # The headline case: Comfy stopped holding the connection at its own bound, + # the generation may well have completed and been billed, and this + # attribute is the caller's only route back to it. + server.state.model_run_error = (504, "deadline_exceeded") + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + (sent,) = server.state.model_run_idempotency_keys + assert sent is not None + # The exact key, not merely "a key": asserting truthiness would pass on a + # freshly minted one, which is precisely the value that cannot collect the + # generation. + assert excinfo.value.idempotency_key == sent + + +async def test_an_async_deadline_exceeded_504_exposes_the_key_too(server) -> None: + server.state.model_run_error = (504, "deadline_exceeded") + async with AsyncComfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + await client.models.run(MODEL, ARGS) + (sent,) = server.state.model_run_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.model_run_error = (504, "deadline_exceeded") + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS, idempotency_key="caller-chosen-03") + assert excinfo.value.idempotency_key == "caller-chosen-03" + assert server.state.model_run_idempotency_keys == ["caller-chosen-03"] + + +async def test_a_caller_supplied_key_round_trips_on_the_async_client(server) -> None: + server.state.model_run_error = (504, "deadline_exceeded") + async with AsyncComfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + await client.models.run(MODEL, ARGS, idempotency_key="caller-chosen-04") + assert excinfo.value.idempotency_key == "caller-chosen-04" + + +def test_the_replay_idiom_from_the_docstring_collects_the_generation(server) -> None: + # End to end, against a deployment that replays a claimed key: the 504 + # raises, the caller reads the key off the exception and re-runs under it, + # and the result of the generation they were already billed for comes back. + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_replays_idempotency_key = True + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + server.state.model_run_error = None + replayed = client.models.run(MODEL, ARGS, idempotency_key=excinfo.value.idempotency_key) + assert replayed == server.state.model_run_result + first, second = server.state.model_run_idempotency_keys + assert first == second + + +def test_a_transport_level_failure_carries_the_key() -> None: + # No response at all, so nothing to translate — and the case the replay + # contract names alongside the 504, since a connection dropped mid-run says + # nothing about whether the generation ran. + models, low = _models_over(httpx.ConnectError("connection refused")) + with pytest.raises(httpx.ConnectError) as excinfo: + models.run(MODEL, ARGS) + assert low.keys == [excinfo.value.idempotency_key] # type: ignore[attr-defined] + assert excinfo.value.idempotency_key is not None # type: ignore[attr-defined] + + +async def test_an_async_transport_level_failure_carries_the_key() -> None: + models, low = _async_models_over(httpx.ReadTimeout("no answer")) + with pytest.raises(httpx.ReadTimeout) as excinfo: + await models.run(MODEL, ARGS) + assert low.keys == [excinfo.value.idempotency_key] # type: ignore[attr-defined] + + +def test_a_client_side_timeout_against_a_live_server_carries_the_key(server) -> None: + # The read timeout the retry module calls the important member of the + # unknown-outcome class: the server is most likely still generating. + server.state.model_run_delay = 0.4 + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(httpx.TimeoutException) as excinfo: + client.models.run(MODEL, ARGS, timeout=0.05) + (sent,) = server.state.model_run_idempotency_keys + assert excinfo.value.idempotency_key == sent # type: ignore[attr-defined] + + +def test_an_unknown_error_type_lands_on_the_base_router_error_with_the_key() -> None: + # The reason the stamp lives at the translation boundary rather than in + # each subclass's constructor: a bucket this SDK version has never heard of + # falls through to `RouterError` itself, and it has to carry the key too. + unknown = error_from_response( + 503, {ERROR_TYPE_HEADER: "a_bucket_from_the_future"}, {"detail": "nope"} + ) + assert type(unknown) is RouterError + models, low = _models_over(unknown) + with pytest.raises(RouterError) as excinfo: + models.run(MODEL, ARGS) + assert excinfo.value.error_type == "a_bucket_from_the_future" + assert low.keys == [excinfo.value.idempotency_key] + assert excinfo.value.idempotency_key is not None + + +def test_a_typed_router_subclass_carries_the_key_as_well() -> None: + models, low = _models_over(DeadlineExceeded("we stopped waiting", http_status=504)) + with pytest.raises(DeadlineExceeded) as excinfo: + models.run(MODEL, ARGS) + assert low.keys == [excinfo.value.idempotency_key] + + +def test_the_key_on_the_exception_is_the_one_every_retry_reused(server) -> None: + # One key across the attempts, and that same key on the exception the + # exhausted retry finally raises — not the first attempt's, not a fresh one. + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_replays_idempotency_key = True + policy = RetryPolicy( + max_elapsed=0.5, initial_backoff=0.01, max_backoff=0.02, retry_possibly_in_flight=True + ) + with Comfy(retry=policy) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + sent = server.state.model_run_idempotency_keys + assert len(sent) > 1 + assert set(sent) == {excinfo.value.idempotency_key} + + +# --- request_id, from the server's own header ---------------------------- + + +def test_a_failed_run_carries_the_servers_request_id(server) -> None: + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_request_id = "req_abc123" + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert excinfo.value.request_id == "req_abc123" + + +async def test_an_async_failed_run_carries_the_servers_request_id(server) -> None: + server.state.model_run_error = (500, "internal_error") + server.state.model_run_request_id = "req_def456" + async with AsyncComfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + await client.models.run(MODEL, ARGS) + assert excinfo.value.request_id == "req_def456" + + +def test_request_id_is_none_when_the_response_named_none(server) -> None: + server.state.model_run_error = (504, "deadline_exceeded") + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert excinfo.value.request_id is None + # The key is still there — the two are independent, and the one that + # enables the replay is minted client-side. + assert excinfo.value.idempotency_key is not None From 7c2b0bfe6913f74ce02a4f62ef7e632c1f09891a Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 26 Aug 2026 21:09:17 -0700 Subject: [PATCH 2/3] fix: address the review panel on the exception idempotency-key surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ten unresolved cursor-review findings on the replay contract this PR documents. The high-severity one is that the README's recovery snippet — the whole point of the feature — could not reach the replay for the failure it exists for. - The recovery example caught only `ComfyError`, but a dropped connection or read timeout never reaches a response to translate and arrives as the `httpx` error it was. Catch both, and check the key for `None` before passing it back: `idempotency_key=None` means "mint one", so replaying with a key that was never recorded starts a second billed generation instead of collecting the first. - `_stamp` wrote only `idempotency_key`, so `exc.request_id` on an httpx failure was an `AttributeError` rather than the documented `None` — on exactly the no-response failures the pair matters most on. Both it and the new `retry_after` are now defaulted onto the exceptions this SDK does not build, and a test pins that list against the base's readable attribute surface so a third one cannot be added without it. - `to_sdk_error` dropped `Retry-After` for every code but `queue_full`, so a `deadline_exceeded` 504 that named a pace reached a caller who had been told to wait with nothing to read the wait off. `ComfyError` now carries `retry_after` for every code. - `_Prepared.parse_or_raise` called `resp.json()` unguarded on the ok path, so a 200 with a non-JSON body escaped as `json.JSONDecodeError` — outside the translated surface, and so unstamped. On `models.run` that is a generation that ran and was billed with the result lost, which is precisely the failure the key has to ride out on. - `asyncio.CancelledError` is a `BaseException` and passed unstamped, so the `asyncio.wait_for` a caller wraps a ten-minute run in abandoned a possibly-dispatched generation and took the key with it. Stamped and re-raised bare, so the cancellation still propagates; `KeyboardInterrupt` stays untouched and has a test saying so. - `X-Comfy-Request-Id` was stored verbatim from a server-controlled header and is documented as something to display and paste into support tickets. Bounded and filtered through one function both error surfaces share, so it cannot be safe to display on one and not the other. - The replay test asserted only that the payloads matched, which holds just as well for a *second* generation returning an equal payload — the double charge the feature prevents. The stub now models a lost response properly (records the result against the key, serves it back without re-running) and the test asserts one generation across two requests, with the negative case alongside it. - The `models.run` docstring named a `409` "still in progress" refusal, but a body-less 409 maps to `hash_mismatch`; it now describes the refusal without promising a class the mapping does not deliver. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 27 +++++- README.md | 31 +++++-- src/comfy_low/errors.py | 26 ++++++ src/comfy_low/transport.py | 33 ++++++-- src/comfy_sdk/exceptions.py | 75 ++++++++++++++--- src/comfy_sdk/models.py | 20 ++++- src/comfy_sdk/router_exceptions.py | 7 +- tests/conftest.py | 56 +++++++++++++ tests/test_error_contract.py | 105 ++++++++++++++++++++++- tests/test_models_run.py | 129 ++++++++++++++++++++++++++++- 10 files changed, 471 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 648b650..2dd0756 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,7 +18,10 @@ notes for each version. `Idempotency-Key` the call 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). `run` mints that key + 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. `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 @@ -31,7 +34,27 @@ notes for each version. 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. + 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 diff --git a/README.md b/README.md index 3d3f617..2ac60b4 100644 --- a/README.md +++ b/README.md @@ -481,15 +481,24 @@ 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"}) - except ComfyError as exc: - # Later — after the Retry-After the server named, if it named one. + # 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=exc.idempotency_key + "acme/flux/dev", {"prompt": "a cat"}, idempotency_key=key ) ``` @@ -499,15 +508,23 @@ 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. -Two attributes carry this: +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) and a `RouterError` whose `error_type` this SDK version does not recognise | +| `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 | -Both are `None` rather than absent on an error from a surface that sends no -key, so `exc.idempotency_key` is always safe to read on a `ComfyError`. +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 diff --git a/src/comfy_low/errors.py b/src/comfy_low/errors.py index 6182981..d77a400 100644 --- a/src/comfy_low/errors.py +++ b/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.""" diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index 4bea998..ccb0290 100644 --- a/src/comfy_low/transport.py +++ b/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 @@ -137,11 +137,13 @@ def _retry_after(resp: httpx.Response) -> int | None: def _request_id(resp: httpx.Response) -> str | None: - """``X-Comfy-Request-Id`` as a non-empty string, or ``None``.""" - raw = resp.headers.get(REQUEST_ID_HEADER) - if raw is None: - return None - return raw.strip() or 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]: @@ -229,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() diff --git a/src/comfy_sdk/exceptions.py b/src/comfy_sdk/exceptions.py index e65c0d5..a964194 100644 --- a/src/comfy_sdk/exceptions.py +++ b/src/comfy_sdk/exceptions.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio from collections.abc import Iterator from contextlib import contextmanager from typing import Any, TypeVar @@ -21,10 +22,13 @@ class ComfyError(Exception): """Base for every SDK-level error.""" - #: The ``Idempotency-Key`` the failed call was made under, when the - #: operation that raised this sends one — see :func:`translating`. It is - #: ``None`` on every error from an operation that sends no key, and on one - #: constructed by hand. + #: 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. #: #: 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 @@ -36,6 +40,15 @@ class ComfyError(Exception): #: response at all). The id a user quotes in a support request. request_id: str | None = None + #: Seconds the server asked the caller to wait before asking again, from + #: ``Retry-After``, or ``None`` when it named no pace. Carried on the base + #: because the header is not the throttled buckets' alone — a + #: ``deadline_exceeded`` ``504`` names the pace at which a replay of the + #: same ``Idempotency-Key`` may be attempted, and a caller told to wait for + #: it needs somewhere to read it. :class:`QueueFull` narrows it to a + #: required ``int``. + retry_after: int | None = None + def __init__( self, message: str, @@ -44,6 +57,7 @@ def __init__( http_status: int | None = None, details: dict[str, Any] | None = None, request_id: str | None = None, + retry_after: int | None = None, ) -> None: super().__init__(message) self.message = message @@ -51,6 +65,7 @@ def __init__( self.http_status = http_status self.details = details self.request_id = request_id + self.retry_after = retry_after class MissingApiKey(ComfyError): @@ -169,6 +184,11 @@ def to_sdk_error(exc: ApiError) -> ComfyError: http_status=exc.http_status, details=exc.details, request_id=exc.request_id, + # Forwarded for every code, not just the throttled ones: `Retry-After` + # is how a `deadline_exceeded` 504 paces the same-key replay that + # collects an already-billed generation, and dropping it here left the + # caller told to wait with nothing to wait on. + retry_after=exc.retry_after, ) @@ -178,7 +198,25 @@ def to_sdk_error(exc: ApiError) -> ComfyError: #: request did not complete. Anything else leaving the block is a bug in the #: SDK or the caller's own code, where a key is noise, and a ``KeyboardInterrupt`` #: must not be touched at all. -_STAMPABLE: tuple[type[BaseException], ...] = (ComfyError, httpx.HTTPError) +#: +#: ``asyncio.CancelledError`` is the one ``BaseException`` here, and it earns +#: the place: cancelling an in-flight ``AsyncModels.run`` — which is what the +#: ``asyncio.wait_for`` a caller wraps a ten-minute call in does — abandons a +#: generation that may already be dispatched and billed, and that is precisely +#: the case the key exists to collect. It is re-raised bare like the rest of +#: this branch, so nothing is swallowed and the cancellation still propagates. +_STAMPABLE: tuple[type[BaseException], ...] = ( + ComfyError, + httpx.HTTPError, + asyncio.CancelledError, +) + +#: Attributes a stamped exception is guaranteed to answer to, defaulted to +#: ``None`` on the ones that do not declare them. Kept beside +#: :data:`_STAMPABLE` so an attribute added to :class:`ComfyError` for the +#: caller to read inside an ``except`` block is added here too — +#: ``tests/test_error_contract.py`` pins the pairing. +_STAMPED_ATTRIBUTES = ("request_id", "retry_after") _E = TypeVar("_E", bound=BaseException) @@ -190,9 +228,22 @@ def _stamp(exc: _E, idempotency_key: str | None) -> _E: members of :data:`_STAMPABLE` are httpx's classes, which this SDK does not build. A ``None`` key writes nothing, so an operation that sends no key leaves ``ComfyError.idempotency_key`` at its class default. + + The rest of :data:`_STAMPED_ATTRIBUTES` is defaulted alongside it, for the + same reason and onto exactly the same exceptions: :class:`ComfyError` + declares them on the class, but ``httpx.ConnectError`` and + ``asyncio.CancelledError`` do not — so without this the documented surface + would be uniform on everything *except* the no-response failures it is most + needed on, where ``exc.request_id`` would raise ``AttributeError`` instead + of reading ``None``. Never overwritten: a stamped exception that already + carries one of them keeps its own value. """ - if idempotency_key is not None: - exc.idempotency_key = idempotency_key # type: ignore[attr-defined] + if idempotency_key is None: + return exc + exc.idempotency_key = idempotency_key # type: ignore[attr-defined] + for name in _STAMPED_ATTRIBUTES: + if not hasattr(exc, name): + setattr(exc, name, None) return exc @@ -221,9 +272,11 @@ def translating(*, idempotency_key: str | None = None) -> Iterator[None]: raise _stamp(to_sdk_error(exc), idempotency_key) from exc except _STAMPABLE as exc: # Already on a surface a caller catches — a RouterError the transport - # raised directly, or an httpx failure with no response to translate. - # Nothing to convert, but the key still has to ride out with it. A bare - # `raise` keeps the original traceback, so with no key given this branch - # is indistinguishable from not catching at all. + # raised directly, an httpx failure with no response to translate, or a + # cancellation of a call that may already be dispatched. Nothing to + # convert, but the key still has to ride out with it. A bare `raise` + # keeps the original traceback and the cancellation semantics, so with + # no key given this branch is indistinguishable from not catching at + # all. _stamp(exc, idempotency_key) raise diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index b09dadb..0ef0207 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -155,10 +155,22 @@ def run( already billed for: ``client.models.run(model, arguments, idempotency_key=exc.idempotency_key)`` returns the original result (``200`` + ``Idempotent-Replayed``) against a deployment that replays a - claimed key, or is refused ``409`` "still in progress" with a - ``Retry-After`` saying when to ask again. Without that attribute an - auto-minted key died with the call and the paid-for generation was - uncollectable. + claimed key, or is refused while that generation is still running, with + ``exc.retry_after`` naming when to ask again when the server sent a + ``Retry-After``. Without that attribute an auto-minted key died with the + call and the paid-for generation was uncollectable. + + Note that a dropped connection surfaces as an ``httpx`` error rather + than a :class:`~comfy_sdk.ComfyError` — it never reached a response to + translate — so a handler written for the replay has to catch both; see + the README's "Collecting a generation after a lost response". + + Pass ``exc.idempotency_key`` back only after checking it is not + ``None``. This parameter treats ``None`` as "mint one", so replaying + with a key that was never recorded silently starts a *second* billed + generation instead of collecting the first. Every exception *this* + method raises carries a real key, but an ``except ComfyError`` that also + catches errors from other surfaces can hand you one that does not. """ low = cast(ComfyLow, self._low) # Minted once, outside the loop: reusing this exact value on every diff --git a/src/comfy_sdk/router_exceptions.py b/src/comfy_sdk/router_exceptions.py index 26941aa..f377d28 100644 --- a/src/comfy_sdk/router_exceptions.py +++ b/src/comfy_sdk/router_exceptions.py @@ -75,6 +75,8 @@ class docstrings below reproduce. ``tests/test_router_spec_contract.py`` reads from dataclasses import dataclass from typing import Any +from comfy_low.errors import clean_request_id + from .exceptions import ComfyError #: Response header carrying the coarse failure bucket. Set on every router error @@ -460,7 +462,10 @@ def error_from_response( response would replace a diagnosable failure with an undiagnosable one. """ lowered = _lowercase_headers(headers) - request_id = _clean(lowered.get(REQUEST_ID_HEADER.lower())) + # Bounded and filtered rather than merely stripped: the id is displayed + # and pasted into support tickets, and this header is server-controlled. + # Shared with `comfy_low.transport` so both surfaces clean it identically. + request_id = clean_request_id(lowered.get(REQUEST_ID_HEADER.lower())) error_type = _clean(lowered.get(ERROR_TYPE_HEADER.lower())) retry_after = _retry_after(lowered.get(RETRY_AFTER_HEADER.lower())) diff --git a/tests/conftest.py b/tests/conftest.py index cb45449..1cb12ad 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -108,12 +108,25 @@ class ServerState: model_run_transient_error: tuple[int, str] = (503, "internal_error") # Status code for a successful run (201 exercises the created-shaped path). model_run_status: int = 200 + # Answer a successful run with a body that is not JSON at all — a proxy + # interstitial served under a 200, a response truncated mid-stream. The + # generation ran and was billed; only the result is unreadable. + model_run_undecodable_body: bool = False # Model the deployment `retry_possibly_in_flight` exists for: one that # *replays* a repeated Idempotency-Key rather than rejecting it, so a key # is released rather than claimed when a request fails 5xx. Default False # is the vendored contract (reject-on-duplicate, no replay), under which a # same-key retry after a 5xx can only come back 422. model_run_replays_idempotency_key: bool = False + # The stronger property the flag above does *not* model: the generation ran + # to completion server side and only the *response* was lost (the + # `deadline_exceeded` 504 the replay contract is written for). With this on, + # a failed run whose outcome is unknown records its result against the key, + # and a later request presenting that key is answered with the recorded + # result — without running the model again. Kept separate because the flag + # above only releases the key: on its own it lets a same-key resend *re-run* + # the model, which is the double charge, not the replay. + model_run_replays_lost_result: bool = False # Seconds sent as Retry-After alongside `model_run_error`, when that error # is the 429 the policy is allowed to pace itself against. `None` sends no # header at all, which is the 429 the policy must *not* retry. @@ -154,6 +167,15 @@ class ServerState: # as POST /jobs rejects one. Kept apart from `idempotency` only so a model # test cannot perturb a workflow test's bookkeeping. model_run_idempotency: dict[str, str] = field(default_factory=dict) + # Idempotency-Key -> the result recorded for it under + # `model_run_replays_lost_result`, served verbatim to a later request + # presenting the same key. + model_run_replay_store: dict[str, dict[str, Any]] = field(default_factory=dict) + # How many times the model actually *ran*, as distinct from how many + # requests arrived (`model_run_count`). A replay serves a recorded result + # and does not increment this, which is what lets a test tell a real replay + # apart from a second generation that merely returns an equal payload. + model_run_generations: int = 0 def _asset_json(asset_id: str, hash_: str, created_new: bool, size: int) -> dict: @@ -238,6 +260,15 @@ def _json(self, status: int, payload: dict, headers: dict | None = None) -> None self.end_headers() self.wfile.write(body) + def _raw(self, status: int, body: bytes, content_type: str) -> None: + """A response whose body is *not* JSON — the case a client that + calls ``.json()`` unguarded on a success status falls over on.""" + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + def _err(self, status: int, code: str, message: str = "err") -> None: self._json(status, {"error": {"code": code, "message": message}}) @@ -456,6 +487,18 @@ def _post_model_run(self) -> None: key = self.headers.get("Idempotency-Key") state.model_run_idempotency_keys.append(key) + # Checked before the reject-on-duplicate rule below: a deployment + # that replays a claimed key answers with the recorded result + # rather than rejecting the resend, and the model does not run + # again — which is the whole point of asking under the same key. + if key and key in state.model_run_replay_store: + self._json( + 200, + state.model_run_replay_store[key], + headers={"Idempotent-Replayed": "true"}, + ) + return + # The same reject-on-duplicate rule `_post_jobs` implements, for # the same reason: a stub more permissive than the contract would # let a retry design that the real server rejects pass its tests. @@ -477,6 +520,11 @@ def claim_if_outcome_unknown(status: int) -> None: def fail(status: int, code: str, message: str) -> None: claim_if_outcome_unknown(status) + if state.model_run_replays_lost_result and key and status >= 500: + # The generation completed; only the answer was lost. Bill + # it once and record it, so the same key collects it. + state.model_run_generations += 1 + state.model_run_replay_store[key] = state.model_run_result headers: dict[str, str] = {} if state.model_run_retry_after is not None: headers["Retry-After"] = state.model_run_retry_after @@ -499,6 +547,14 @@ def fail(status: int, code: str, message: str) -> None: return if key: state.model_run_idempotency[key] = "done" + state.model_run_generations += 1 + if state.model_run_undecodable_body: + self._raw( + state.model_run_status, + b"502 from an intermediary", + "text/html", + ) + return self._json(state.model_run_status, state.model_run_result) def _post_jobs(self) -> None: diff --git a/tests/test_error_contract.py b/tests/test_error_contract.py index 4ac1463..0ded0dc 100644 --- a/tests/test_error_contract.py +++ b/tests/test_error_contract.py @@ -16,17 +16,25 @@ from __future__ import annotations +import asyncio import io from collections.abc import Awaitable, Callable from typing import Any +import httpx import pytest import comfy_low.transport as low_transport from comfy_low.errors import ApiError as LowApiError +from comfy_low.errors import clean_request_id from comfy_sdk import AsyncComfy, Comfy, Forbidden, NotFound -from comfy_sdk.exceptions import ComfyError, translating -from comfy_sdk.router_exceptions import REQUEST_ID_HEADER, ROUTER_EXCEPTIONS, RouterError +from comfy_sdk.exceptions import _STAMPED_ATTRIBUTES, ComfyError, translating +from comfy_sdk.router_exceptions import ( + REQUEST_ID_HEADER, + ROUTER_EXCEPTIONS, + RouterError, + error_from_response, +) def _wf(client: Comfy | AsyncComfy): @@ -166,6 +174,7 @@ async def _drain() -> None: "details", "request_id", "idempotency_key", + "retry_after", ) @@ -174,7 +183,7 @@ def test_the_base_error_answers_to_its_whole_attribute_surface(name: str) -> Non assert hasattr(ComfyError("boom"), name) -@pytest.mark.parametrize("name", ("request_id", "idempotency_key")) +@pytest.mark.parametrize("name", ("request_id", "idempotency_key", "retry_after")) def test_those_attributes_default_to_none_rather_than_being_absent(name: str) -> None: # An operation that sends no Idempotency-Key, and a response that named no # request id, both leave the attribute readable as `None`. A caller writes @@ -223,6 +232,96 @@ def test_a_bug_in_the_sdk_is_not_stamped_and_not_swallowed() -> None: assert not hasattr(excinfo.value, "idempotency_key") +def test_a_keyboard_interrupt_is_never_touched() -> None: + # `_STAMPABLE` gained one BaseException (`asyncio.CancelledError`), which + # makes this the assertion that the widening stopped there: an interrupt is + # the user asking the process to stop, not a failed call. + with pytest.raises(KeyboardInterrupt) as excinfo: + with translating(idempotency_key="k-04"): + raise KeyboardInterrupt + assert not hasattr(excinfo.value, "idempotency_key") + + +def test_a_cancelled_call_is_stamped_and_still_cancelled() -> None: + # Cancelling an in-flight run abandons a generation that may already be + # dispatched and billed, so the key rides out on it — but the cancellation + # itself is re-raised bare, never converted into an ordinary error. + with pytest.raises(asyncio.CancelledError) as excinfo: + with translating(idempotency_key="k-05"): + raise asyncio.CancelledError + assert excinfo.value.idempotency_key == "k-05" # type: ignore[attr-defined] + assert excinfo.value.request_id is None # type: ignore[attr-defined] + + +@pytest.mark.parametrize("name", _STAMPED_ATTRIBUTES) +def test_a_stamped_transport_error_reads_every_attribute_as_none(name: str) -> None: + # httpx's classes declare none of these; the stamp defaults them so the + # documented surface is uniform on the no-response failures it matters most + # on, rather than uniform everywhere except there. + with pytest.raises(httpx.ConnectError) as excinfo: + with translating(idempotency_key="k-06"): + raise httpx.ConnectError("refused") + assert getattr(excinfo.value, name) is None + + +def test_every_readable_base_attribute_is_defaulted_onto_a_stamped_error() -> None: + # The pairing that keeps the two lists honest: an attribute added to + # `ComfyError` for a caller to read inside an `except` block has to be + # defaulted onto the exceptions this SDK does not build, or it is readable + # on some errors and an AttributeError on others. + readable = set(_BASE_ATTRIBUTES) - {"message", "code", "http_status", "details"} + assert readable == {"idempotency_key", *_STAMPED_ATTRIBUTES} + + +def test_the_stamp_never_overwrites_an_attribute_that_is_already_set() -> None: + err = LowApiError("slow down", code="queue_full", http_status=429, retry_after=7) + with pytest.raises(ComfyError) as excinfo: + with translating(idempotency_key="k-07"): + raise err + assert excinfo.value.retry_after == 7 + + +@pytest.mark.parametrize( + ("raw", "expected"), + ( + ("req_abc123", "req_abc123"), + (" req_abc123 ", "req_abc123"), + # `httpx.Headers.get` joins duplicate headers with ", " — take the + # first id rather than splicing two into one that identifies no call. + ("req_a, req_b", "req_a"), + # A terminal escape truncates at the first character outside the class + # rather than being written into somebody's log verbatim. + ("req_ared", "req_a"), + ("req_a", None), + ("", None), + (" ", None), + (None, None), + (b"req_bytes", None), + ), +) +def test_a_request_id_is_bounded_and_printable_before_it_is_stored( + raw: Any, expected: str | None +) -> None: + assert clean_request_id(raw) == expected + + +def test_a_request_id_is_length_bounded() -> None: + # No length bound at all meant a server (or anything between) could put an + # arbitrarily long string into every traceback and log line. + assert clean_request_id("x" * 5000) == "x" * 200 + + +def test_both_layers_clean_the_request_id_through_the_same_function() -> None: + # Not merely "both sanitise": the same function, so an id that is safe to + # display on one surface cannot be unsafe on the other. + hostile = "req_ok" + "x" * 500 + from_envelope = low_transport._request_id( + httpx.Response(500, headers={REQUEST_ID_HEADER: hostile}) + ) + from_router = error_from_response(500, {REQUEST_ID_HEADER: hostile}).request_id + assert from_envelope == from_router == "req_ok" + + def test_the_two_layers_spell_the_request_id_header_the_same_way() -> None: # `comfy_low` reads it off the shared error envelope and `comfy_sdk` reads # it off a router error response, so the name is written out in both — and diff --git a/tests/test_models_run.py b/tests/test_models_run.py index 70ae5d0..973488b 100644 --- a/tests/test_models_run.py +++ b/tests/test_models_run.py @@ -13,6 +13,7 @@ from __future__ import annotations +import asyncio import inspect import re from collections.abc import Mapping @@ -360,15 +361,42 @@ def test_the_replay_idiom_from_the_docstring_collects_the_generation(server) -> # raises, the caller reads the key off the exception and re-runs under it, # and the result of the generation they were already billed for comes back. server.state.model_run_error = (504, "deadline_exceeded") - server.state.model_run_replays_idempotency_key = True + # The generation completed server side and only the answer was lost, so the + # stub records the result against the key and serves it back without + # running the model again. Asserting on the payload alone would not + # distinguish that from a *second* generation returning an equal payload — + # which is the double charge this feature exists to avoid — so the run + # counter below is the assertion that actually holds the contract. + server.state.model_run_replays_lost_result = True with Comfy(retry=NO_RETRY) as client: with pytest.raises(ComfyError) as excinfo: client.models.run(MODEL, ARGS) - server.state.model_run_error = None replayed = client.models.run(MODEL, ARGS, idempotency_key=excinfo.value.idempotency_key) assert replayed == server.state.model_run_result first, second = server.state.model_run_idempotency_keys assert first == second + # Two requests arrived, one generation happened. The `model_run_error` knob + # is deliberately left set: the second call succeeds because the key + # collected the recorded result, not because the failure was turned off. + assert server.state.model_run_count == 2 + assert server.state.model_run_generations == 1 + + +def test_replaying_without_the_key_would_start_a_second_generation(server) -> None: + # The negative of the test above, and the reason both the README snippet + # and the docstring tell a caller to check `idempotency_key` for None: a + # resend that does not carry the key is a new call, mints a new one, and + # bills a second generation. Asserted so the guard is not quietly dropped. + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_replays_lost_result = True + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError): + client.models.run(MODEL, ARGS) + server.state.model_run_error = None + client.models.run(MODEL, ARGS) + first, second = server.state.model_run_idempotency_keys + assert first != second + assert server.state.model_run_generations == 2 def test_a_transport_level_failure_carries_the_key() -> None: @@ -469,3 +497,100 @@ def test_request_id_is_none_when_the_response_named_none(server) -> None: # The key is still there — the two are independent, and the one that # enables the replay is minted client-side. assert excinfo.value.idempotency_key is not None + + +def test_request_id_reads_as_none_on_a_failure_with_no_response_at_all() -> None: + # The documented pair has to be uniform on exactly the failures it is most + # needed on. `httpx.ConnectError` is not one of this SDK's classes and + # declares no `request_id`, so without the stamp defaulting it, the + # attribute access the docs invite would raise AttributeError here. + models, _ = _models_over(httpx.ConnectError("connection refused")) + with pytest.raises(httpx.ConnectError) as excinfo: + models.run(MODEL, ARGS) + assert excinfo.value.request_id is None # type: ignore[attr-defined] + assert excinfo.value.idempotency_key is not None # type: ignore[attr-defined] + + +async def test_an_async_transport_failure_reads_request_id_as_none() -> None: + models, _ = _async_models_over(httpx.ReadTimeout("no answer")) + with pytest.raises(httpx.ReadTimeout) as excinfo: + await models.run(MODEL, ARGS) + assert excinfo.value.request_id is None # type: ignore[attr-defined] + + +def test_a_stamped_error_that_already_has_a_request_id_keeps_it(server) -> None: + # The default must never overwrite a real id — that would erase the one + # thing a user quotes in a support request. + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_request_id = "req_keepme" + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert excinfo.value.request_id == "req_keepme" + + +# --- retry_after, the pace the replay is meant to wait for ---------------- + + +def test_a_504_forwards_the_retry_after_the_server_named(server) -> None: + # The docs tell a caller to replay "after the Retry-After the server + # named". `deadline_exceeded` is not a throttled bucket, so before this was + # forwarded for every code the caller was told to wait with nothing to read + # the wait off. + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "2" + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert excinfo.value.retry_after == 2 + + +def test_retry_after_is_none_when_the_server_named_no_pace(server) -> None: + server.state.model_run_error = (504, "deadline_exceeded") + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert excinfo.value.retry_after is None + + +def test_retry_after_reads_as_none_on_a_failure_with_no_response() -> None: + # The README's table promises all three attributes read rather than raise + # on anything `models.run` raises, and a transport failure is the case that + # is neither one of this SDK's classes nor a response. + models, _ = _models_over(httpx.ConnectError("connection refused")) + with pytest.raises(httpx.ConnectError) as excinfo: + models.run(MODEL, ARGS) + assert excinfo.value.retry_after is None # type: ignore[attr-defined] + + +# --- cancellation, the one BaseException the key rides out on ------------- + + +async def test_cancelling_an_in_flight_run_still_yields_the_key() -> None: + # `asyncio.wait_for` around a ten-minute run is the ordinary way this + # happens: the request may already be dispatched and billed, so the key is + # the caller's route back to it. The cancellation itself must still + # propagate — it is re-raised bare, not converted. + models, low = _async_models_over(asyncio.CancelledError()) + with pytest.raises(asyncio.CancelledError) as excinfo: + await models.run(MODEL, ARGS) + assert low.keys == [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] + + +# --- a success status whose body will not decode -------------------------- + + +def test_an_undecodable_success_body_is_a_stamped_sdk_error(server) -> None: + # A 200 whose body is not JSON — a proxy interstitial, a truncated + # response. The generation ran and was billed with the result lost, which + # is exactly the case the key has to ride out on, so it must not escape as + # the raw json.JSONDecodeError from outside the translated surface. + server.state.model_run_undecodable_body = True + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert excinfo.value.idempotency_key is not None + assert excinfo.value.http_status == 200 + assert excinfo.value.code == "invalid_response" From c4d757eb348073929f55fcc7921a0ae8bd5afaea Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 26 Aug 2026 21:14:14 -0700 Subject: [PATCH 3/3] docs: scope the changelog claim to failed calls rather than every exception MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit is right that 'every exception' is literally overbroad: a programming error escaping the call, and KeyboardInterrupt, are deliberately left unstamped — a key means nothing on either, and there are tests asserting both. The enumeration that followed already said so implicitly; the leading claim now says it outright. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2dd0756..4662bda 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,22 +14,25 @@ notes for each version. ### Added -- Every exception `client.models.run()` raises now carries the - `Idempotency-Key` the call 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 +- 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. `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. + 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