diff --git a/CHANGELOG.md b/CHANGELOG.md index 8662687..4662bda 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,50 @@ notes for each version. ### Added +- Every exception `client.models.run()` raises **for a failed call** now + carries the `Idempotency-Key` it was made under, on `.idempotency_key` — the + typed `RouterError` buckets, a `RouterError` whose `error_type` this version + does not recognise, any other `ComfyError`, and a transport failure with no + response at all (a dropped connection, a read timeout), and a cancelled + `await` of `AsyncModels.run` — the `asyncio.wait_for` a caller wraps a + ten-minute call in abandons a generation that may already be dispatched and + billed, and the cancellation still propagates unchanged. "Failed call" is the + boundary, not "every exception": a programming error escaping the call is not + a failed request, a key means nothing on it, and it reaches you untouched — + as does `KeyboardInterrupt`. `run` mints that key itself unless you pass + `idempotency_key=`, and it used to be a local of the call: when the call + raised, the key went with it. Since collecting a generation you were already + billed for after a lost response means asking again under the *same* key, + that made the auto-minted case uncollectable — only callers who chose and + stored their own key could recover. The recovery idiom is now + `client.models.run(model, arguments, idempotency_key=exc.idempotency_key)`; + see the README. Nothing about what is retried, or what goes on the wire, + changed. +- `ComfyError.request_id` — the server's `X-Comfy-Request-Id` for the failed + call, when the response carried that header, on every SDK exception rather + than only on `RouterError`. It is the id to quote in a support request, and it + was previously unreachable once the response object was gone. `None` when the + response named none, or when there was no response — including on a transport + failure, where the attribute now reads as `None` rather than being absent, so + a handler never has to guard the access. The header is bounded and filtered + before it is stored (it is server-controlled and the id is meant to be + displayed and pasted into support tickets), identically on both error + surfaces. +- `ComfyError.retry_after` — seconds the server asked the caller to wait, from + `Retry-After`, now forwarded for every error code rather than only for + `queue_full`. The replay documented above tells a caller to ask again "after + the `Retry-After` the server named", and a `deadline_exceeded` `504` that + carried one had nowhere to surface it, so the caller had nothing to wait on. + `None` when the server named no pace. + +### Fixed + +- A success status whose body will not decode (a proxy interstitial served + under a `200`, a response truncated mid-stream) now raises a translated SDK + error instead of letting `json.JSONDecodeError` escape from outside the + translated surface. On `models.run` that is a generation that ran and was + billed with the result lost — precisely the failure the `Idempotency-Key` + has to ride out on, and it previously carried no key. - Automatic retry for `client.models.run`, on by default, with the `Idempotency-Key` sent unconditionally on **every** attempt of one logical call — a new call mints a new key. That is what keeps a retry from being diff --git a/README.md b/README.md index dba5df1..6b59c01 100644 --- a/README.md +++ b/README.md @@ -493,6 +493,64 @@ is the collect class's alone. `retry` governs `client.models` only. `submit()`/`run()` on the client keep their own 429 handling, which follows the server's `Retry-After`. +### Collecting a generation after a lost response + +Everything above is about the retry the SDK makes *for* you. When it gives up — +or when you turned it off — the failure that reaches you may still be a +generation that ran and was billed: a `504` where the server stopped holding the +connection at its own deadline, or a connection that dropped while it was +generating. Recovering that generation means asking again under the **same** +`Idempotency-Key`, and if `run` minted the key for you then that key is the one +thing you need and the one thing you never saw. + +So every exception `models.run` raises carries it: + +```python +import httpx +from comfy_sdk import Comfy, ComfyError + +with Comfy() as client: + try: + result = client.models.run("acme/flux/dev", {"prompt": "a cat"}) + # Both, and the second is not optional: a dropped connection or a read + # timeout — one of the two cases this section is about — never reached a + # response to translate, so it arrives as the `httpx` error it was, not as + # a `ComfyError`. Catching only `ComfyError` misses exactly the failure the + # replay exists for. + except (ComfyError, httpx.HTTPError) as exc: + key = exc.idempotency_key + if key is None: + raise # Nothing to replay under; a fresh key would re-run and re-bill. + # Later — after `exc.retry_after` seconds, if the server named a pace. + result = client.models.run( + "acme/flux/dev", {"prompt": "a cat"}, idempotency_key=key + ) +``` + +Against a deployment that replays a claimed key, the second call returns the +original generation's result rather than starting a new one; while that +generation is still running it is refused instead, with a `Retry-After` saying +when to ask. Send the same arguments you sent the first time — a repeated key +with a *different* body is rejected outright. + +Three attributes carry this: + +| Attribute | Value | +|---|---| +| `exc.idempotency_key` | the key that call was made under — the one `run` minted, or the one you passed. Present on every exception `models.run` raises, including a transport failure with no response at all (`httpx.ConnectError`, a read timeout), a cancelled `await`, and a `RouterError` whose `error_type` this SDK version does not recognise | +| `exc.request_id` | the server's `X-Comfy-Request-Id` for the call, when the response carried one — the id to quote in a support request. `None` when there was no response, or none of that header | +| `exc.retry_after` | seconds the server asked you to wait before asking again, from `Retry-After`. `None` when it named no pace | + +All three read as `None` rather than raising on any exception `models.run` +raises, so a handler never has to guard the attribute access itself. + +`idempotency_key` is `None` on errors from *other* surfaces, though — it is +`models.run` that records it, and `submit()` sends a key without stamping one. +So `None` means "this SDK did not record a key for you", **not** "no key was +sent, resend freely": check for it before replaying, as the snippet above does, +rather than passing it straight back into `idempotency_key=` where `None` means +"mint a fresh one" and starts a second billed generation. + ## Sync and async `Comfy` and `AsyncComfy` expose the identical surface — swap the import and diff --git a/src/comfy_low/errors.py b/src/comfy_low/errors.py index 1736dcb..3e4f194 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.""" @@ -25,6 +51,7 @@ def __init__( http_status: int, details: dict[str, Any] | None = None, retry_after: int | None = None, + request_id: str | None = None, ) -> None: super().__init__(message) self.message = message @@ -33,6 +60,12 @@ def __init__( self.http_status = http_status self.details = details self.retry_after = retry_after + #: Server-minted id for the call, read off ``X-Comfy-Request-Id``. + #: ``None`` when the response carried no such header. Surfaced the same + #: way ``retry_after`` is — a response header kept on the exception, + #: because it is the id a user quotes in a support request and it is + #: unreachable once the response object is gone. + self.request_id = request_id class InvalidWorkflow(ApiError): @@ -116,6 +149,7 @@ def error_from_envelope( body: dict[str, Any] | None, *, retry_after: int | None = None, + request_id: str | None = None, error_type: str | None = None, ) -> ApiError: """Build the typed exception for an error response. @@ -169,6 +203,7 @@ def error_from_envelope( http_status=http_status, details=details if isinstance(details, dict) else None, retry_after=retry_after, + request_id=request_id, ) diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index 8a5e56c..bb1fc09 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 @@ -127,6 +127,25 @@ def _retry_after(resp: httpx.Response) -> int | None: return None +#: Response header carrying the server-minted id for the call. Spelled here as +#: well as in :data:`comfy_sdk.router_exceptions.REQUEST_ID_HEADER` because the +#: two layers read it independently — the router surface off its own error +#: response, this one off the shared error envelope — and ``comfy_low`` never +#: imports from ``comfy_sdk``. ``tests/test_error_contract.py`` asserts the two +#: spellings agree, so they cannot drift apart. +REQUEST_ID_HEADER = "X-Comfy-Request-Id" + + +def _request_id(resp: httpx.Response) -> str | None: + """``X-Comfy-Request-Id`` as a bounded, printable string, or ``None``. + + Filtered rather than kept verbatim — see + :func:`comfy_low.errors.clean_request_id`, which both error surfaces share + so the id cannot be safe to display on one and not the other. + """ + return clean_request_id(resp.headers.get(REQUEST_ID_HEADER)) + + def origin(url: str) -> tuple[str, str, int | None]: """Normalized ``(scheme, host, port)`` — the parts that define same-origin. @@ -212,9 +231,24 @@ def headers(self, url: str, extra: dict[str, str] | None = None) -> dict[str, st def parse_or_raise(self, resp: httpx.Response, ok: tuple[int, ...]) -> dict[str, Any]: if resp.status_code in ok: - if resp.content: + if not resp.content: + return {} + try: return resp.json() - return {} + except ValueError as exc: + # A success status whose body will not decode — a proxy + # interstitial served as 200, a response truncated mid-stream. + # Raised as an ApiError rather than escaping as the raw + # `json.JSONDecodeError` so it lands on the surface the SDK + # translates and stamps: on `models.run` this is a generation + # that ran and was billed with the result lost, which is + # exactly the failure the Idempotency-Key has to ride out on. + raise ApiError( + f"Could not decode the {resp.status_code} response body as JSON", + code="invalid_response", + http_status=resp.status_code, + request_id=_request_id(resp), + ) from exc body: dict[str, Any] | None try: body = resp.json() @@ -224,6 +258,7 @@ def parse_or_raise(self, resp: httpx.Response, ok: tuple[int, ...]) -> dict[str, resp.status_code, body, retry_after=_retry_after(resp), + request_id=_request_id(resp), # Router repeats its coarse bucket on this header, and the model-run # route answers in Router's body shape rather than the envelope's. # Passing it here is what keeps the bucket alive across the layer diff --git a/src/comfy_sdk/exceptions.py b/src/comfy_sdk/exceptions.py index c4ca138..a964194 100644 --- a/src/comfy_sdk/exceptions.py +++ b/src/comfy_sdk/exceptions.py @@ -8,9 +8,12 @@ from __future__ import annotations +import asyncio 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 +22,33 @@ 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. + #: + #: 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 + + #: 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, @@ -26,12 +56,16 @@ def __init__( code: str | None = None, 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 self.code = code self.http_status = http_status self.details = details + self.request_id = request_id + self.retry_after = retry_after class MissingApiKey(ComfyError): @@ -141,6 +175,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 +183,100 @@ def to_sdk_error(exc: ApiError) -> ComfyError: code=exc.code, 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, ) +#: 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. +#: +#: ``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) + + +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. + + 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 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 + + @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, 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 59f3f3b..9e5def3 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -160,6 +160,30 @@ def run( outcome genuinely unknown and 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 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 @@ -177,7 +201,10 @@ def run( # the wire, so everything legal in it is deep-copyable. payload = deepcopy(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) @@ -206,8 +233,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() @@ -217,7 +245,10 @@ async def run( # nested values included. payload = deepcopy(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/src/comfy_sdk/router_exceptions.py b/src/comfy_sdk/router_exceptions.py index 1453647..e1ceec1 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 @@ -465,7 +467,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 8889f90..6710100 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 # The narrower carry the router contract describes for `deadline_exceeded`: # the reservation survives the 504 with a handle to the generation stamped # on it, so the SAME key presented again *collects* that generation instead @@ -128,6 +141,9 @@ class ServerState: # 409). `None` sends no header at all, which is the same failure 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 # Answer model-run failures in Router's own error shape -- the coarse bucket # on `X-Comfy-Error-Type` plus a `{detail, error_type}` body -- instead of # the v2 `{error: {code, message}}` envelope. `POST /api/v2/models/run` is @@ -167,6 +183,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: @@ -251,6 +276,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}}) @@ -469,6 +503,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. @@ -510,9 +556,16 @@ def claim_if_outcome_unknown(status: int, code: str) -> None: def fail(status: int, code: str, message: str) -> None: claim_if_outcome_unknown(status, code) - headers = {} + 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 + if state.model_run_request_id is not None: + headers["X-Comfy-Request-Id"] = state.model_run_request_id if state.model_run_router_error_shape: # What a real Router failure looks like: the coarse bucket # on `X-Comfy-Error-Type` and a `{detail, error_type}` body, @@ -541,6 +594,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 d9f9787..0ded0dc 100644 --- a/tests/test_error_contract.py +++ b/tests/test_error_contract.py @@ -5,19 +5,36 @@ 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 +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 +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): @@ -141,3 +158,173 @@ 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", + "retry_after", +) + + +@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", "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 + # `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_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 + # 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..973488b 100644 --- a/tests/test_models_run.py +++ b/tests/test_models_run.py @@ -4,24 +4,34 @@ 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``. """ from __future__ import annotations +import asyncio 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 +253,344 @@ 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") + # 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) + 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: + # 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 + + +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"