diff --git a/CHANGELOG.md b/CHANGELOG.md index b40e60c..ddbf57c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,7 +12,17 @@ notes for each version. ## [Unreleased] -_Nothing yet — add an entry here when your change lands._ +### Added + +- `client.models.run(model, arguments)` on both `Comfy` and `AsyncComfy` — one + call that returns the completed generation. Where the platform has to + submit-and-poll an upstream provider, that polling happens server side inside + the call, so the client contract stays a single request. The result is the + provider's native payload, returned as-is rather than wrapped. The awaitable + form is `AsyncComfy`, not a `run_async()` suffix — there is deliberately no + suffixed variant, and a test asserts its absence. Runs use their own + 10-minute timeout (the client's default is sized for ordinary API calls) and + send an `Idempotency-Key` on every call. ## [0.1.8] - 2026-08-13 diff --git a/README.md b/README.md index 02ad2c1..e4adebc 100644 --- a/README.md +++ b/README.md @@ -276,6 +276,36 @@ namespace, and nothing extra is imported or constructed for it: `base_url` and `timeout` are a read-only view of that shared configuration; model operations are added to this namespace as they land. +### `models.run` — one call, one result + +```python +result = client.models.run("acme/flux/dev", {"prompt": "a cat", "steps": 4}) +result["images"][0]["url"] +``` + +`run` returns when the generation is **complete**. There is no submit step and +nothing to poll: where the platform has to submit-and-poll an upstream +provider, that happens server side inside this one call. The value you get back +is the provider's own payload — decoded JSON, handed over as-is, with no +wrapper class between you and the fields the provider documented. + +The awaitable form is the **async client**, not a differently-named method: + +```python +async with AsyncComfy(api_key="comfyui-...") as client: + result = await client.models.run("acme/flux/dev", {"prompt": "a cat"}) +``` + +There is no `run_async()`, and there will not be one — one operation, one name, +and `await` is what makes it asynchronous. + +Because the server may legitimately hold the connection for minutes, `run` uses +its own 10-minute timeout rather than the client's (which is sized for ordinary +API calls). Pass `timeout=` seconds, an `httpx.Timeout`, or `None` to wait +indefinitely. Each call also sends a fresh `Idempotency-Key`, so an accidental +exact resend is rejected by the server instead of billing a second generation; +pass `idempotency_key=` to choose the value yourself. + ## Sync and async `Comfy` and `AsyncComfy` expose the identical surface — swap the import and diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index f57fcc8..2efa512 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -10,6 +10,12 @@ * **per-request timeout / abort** — every method takes ``timeout`` and the raw httpx cancellation applies. +One binding is *not* backed by an ``operationId``: ``post_model_run``. The +vendored contract declares no model routes yet, so it is hand-written against +the agreed wire shape, kept out of ``comfy_low.OPERATION_IDS``, and confined to +``model_run_request`` / ``_MODEL_RUN_PATH`` so vendoring the real route later is +a one-place change. + This layer contains no orchestration, retries, hashing, or reconnection — those live in ``comfy_sdk``. """ @@ -20,7 +26,7 @@ import platform import secrets import sys -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timedelta, timezone from importlib.metadata import PackageNotFoundError @@ -46,9 +52,47 @@ # blocking iter_lines() forever. (Pass timeout=None to opt out of the timeout.) _SSE_IDLE_TIMEOUT = httpx.Timeout(10.0, read=45.0) +#: Default timeout for a model run. A run is *awaited server-side*: the server +#: holds the connection until the generation is complete, polling the upstream +#: provider itself when that provider is submit/poll. So the client has to be +#: willing to wait minutes, not the tens of seconds a normal API call gets — +#: the client's own default (30s) would abort a perfectly healthy generation. +#: ``connect`` stays short: an unreachable host is not a slow generation. +MODEL_RUN_TIMEOUT = httpx.Timeout(600.0, connect=10.0) + +#: Route for a model run. NOT an ``operationId`` from ``spec/openapi.yaml`` — +#: the vendored v2 contract declares no model routes, so this binding is +#: hand-written and is deliberately absent from ``comfy_low.OPERATION_IDS`` +#: (the spec-coverage test asserts that set equals the spec's, exactly). +#: Everything about the wire shape is confined to this constant and +#: :func:`model_run_request` so it is one place to reconcile when the route is +#: vendored into the spec and the models are regenerated from it. +_MODEL_RUN_PATH = "/models/run" + _DEFAULT_PORTS = {"http": 80, "https": 443} +def model_run_request( + model: str, + arguments: Mapping[str, Any], + idempotency_key: str | None, +) -> tuple[str, dict[str, Any], dict[str, str]]: + """Sans-IO ``(path, json_body, headers)`` for one model run. + + The model id travels in the *body*, not the path: ids are commonly + provider-namespaced and contain ``/`` (``vendor/family/variant``), which in + a path segment needs percent-encoding that intermediaries normalize + inconsistently. A named body field also leaves room for sibling fields + later without moving the route. + + ``arguments`` is copied into a plain dict so any ``Mapping`` is accepted and + the caller's object is never handed to the JSON encoder directly. + """ + body: dict[str, Any] = {"model": model, "arguments": dict(arguments)} + headers = {"Idempotency-Key": idempotency_key} if idempotency_key else {} + return _MODEL_RUN_PATH, body, headers + + def _build_user_agent(client_info: str | None) -> str: """SDK identity sent on every request. This is request metadata (not telemetry — no phone-home), so adoption is measurable server-side from @@ -490,6 +534,31 @@ def get_job_workflow(self, job_id_or_url: str, *, timeout: Any = _UNSET) -> JobW resp = self.raw_request("GET", path, timeout=timeout) return JobWorkflowResponse.model_validate(self._p.parse_or_raise(resp, (200,))) + # -- models ----------------------------------------------------------- + def post_model_run( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + timeout: Any = MODEL_RUN_TIMEOUT, + ) -> dict[str, Any]: + """POST /api/v2/models/run — run a model, awaited server-side. + + One request, one response: the server does not answer until the + generation is complete, so the decoded body *is* the finished result. + That holds for a submit/poll provider too — the polling happens server + side, inside this call, which is why the default ``timeout`` is + :data:`MODEL_RUN_TIMEOUT` rather than the client's own. + + The body is returned verbatim — the provider's native payload, with no + model class layered over it. This is not a spec operation; see + :data:`_MODEL_RUN_PATH`. + """ + path, body, headers = model_run_request(model, arguments, idempotency_key) + resp = self.raw_request("POST", path, headers=headers, json=body, timeout=timeout) + return self._p.parse_or_raise(resp, (200, 201)) + class AsyncComfyLow: """Asynchronous protocol bindings — mirrors :class:`ComfyLow`.""" @@ -763,6 +832,20 @@ async def get_job_workflow( resp = await self.raw_request("GET", path, timeout=timeout) return JobWorkflowResponse.model_validate(self._p.parse_or_raise(resp, (200,))) + # -- models ----------------------------------------------------------- + async def post_model_run( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + timeout: Any = MODEL_RUN_TIMEOUT, + ) -> dict[str, Any]: + """Async :meth:`ComfyLow.post_model_run`.""" + path, body, headers = model_run_request(model, arguments, idempotency_key) + resp = await self.raw_request("POST", path, headers=headers, json=body, timeout=timeout) + return self._p.parse_or_raise(resp, (200, 201)) + def _looks_like_path(s: str) -> bool: return s.startswith("http") or s.startswith("/") diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index 0f3ac31..1c40863 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -12,15 +12,29 @@ namespace plus a read-only view of that shared configuration; model operations are added to this object as they land, never to a parallel client. +``run`` is the one model operation today, and it exists in exactly one form per +client: ``Comfy().models.run(...)`` blocks, ``AsyncComfy().models.run(...)`` is +awaited. **The awaitable form is the async client, not a suffixed method** — +there is deliberately no ``run_async``, and there never will be: two names for +one operation is a published signature that cannot be withdrawn once released. +``tests/test_models_run.py`` asserts the suffix's absence rather than leaving it +to convention. + Callers do not import anything for this: ``from comfy_sdk import Comfy`` stays the only entry point, and ``client.models`` is the whole surface. """ from __future__ import annotations +from collections.abc import Mapping +from typing import Any, cast + import httpx -from comfy_low.transport import AsyncComfyLow, ComfyLow +from comfy_low.transport import MODEL_RUN_TIMEOUT, AsyncComfyLow, ComfyLow + +from ._core import new_idempotency_key +from .exceptions import translating class _ModelsBase: @@ -52,9 +66,71 @@ class Models(_ModelsBase): def __init__(self, low: ComfyLow) -> None: self._low = low + def run( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + timeout: float | httpx.Timeout | None = MODEL_RUN_TIMEOUT, + ) -> dict[str, Any]: + """Run ``model`` with ``arguments`` and return the completed result. + + One call, one result. It blocks until the generation is finished — + including for a provider the platform has to submit-and-poll, where the + polling happens server side inside this call, invisible to the caller. + There is no separate submit/await step and no ``run_async`` variant: the + awaitable form of this method is :meth:`AsyncModels.run` on + ``AsyncComfy``. + + The return value is the provider's own payload, decoded from JSON and + handed back as-is — no wrapper class stands between the caller and the + fields the provider documented. + + Because the server may legitimately hold the connection for minutes, + ``timeout`` defaults to :data:`~comfy_low.transport.MODEL_RUN_TIMEOUT` + (10 minutes) rather than the client's own default, which is sized for + ordinary API calls and would abort a healthy run. Pass a number of + seconds, an ``httpx.Timeout``, or ``None`` to wait indefinitely. + + An ``Idempotency-Key`` is sent on every run; a fresh one is minted per + call unless ``idempotency_key`` is given, so an accidental exact resend + is the server's to reject rather than a second charged generation. + """ + low = cast(ComfyLow, self._low) + with translating(): + return low.post_model_run( + model, + arguments, + idempotency_key=idempotency_key or new_idempotency_key(), + timeout=timeout, + ) + class AsyncModels(_ModelsBase): """``client.models`` on :class:`~comfy_sdk.client.AsyncComfy` — mirrors :class:`Models`.""" def __init__(self, low: AsyncComfyLow) -> None: self._low = low + + async def run( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + timeout: float | httpx.Timeout | None = MODEL_RUN_TIMEOUT, + ) -> dict[str, Any]: + """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. See :meth:`Models.run`. + """ + low = cast(AsyncComfyLow, self._low) + with translating(): + return await low.post_model_run( + model, + arguments, + idempotency_key=idempotency_key or new_idempotency_key(), + timeout=timeout, + ) diff --git a/tests/conftest.py b/tests/conftest.py index ff4e143..5dbad9f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -83,6 +83,25 @@ class ServerState: job_workflow_format: str = "api" job_workflow_not_found: bool = False + # --- POST /models/run (the awaited model run) --- + # The provider's native payload the run resolves to. Deliberately not a + # Comfy-shaped envelope: the SDK must hand it back untouched. + model_run_result: dict[str, Any] = field( + default_factory=lambda: { + "images": [{"url": "http://example.invalid/gen.png", "width": 1024, "height": 1024}], + "seed": 42, + "timings": {"inference": 3.5}, + } + ) + # Seconds the run holds the connection before answering — stands in for a + # generation the server polls internally, so a client whose timeout is too + # short aborts a healthy run. + model_run_delay: float = 0.0 + # (status, code) answered instead of the result. + model_run_error: tuple[int, str] | None = None + # Status code for a successful run (201 exercises the created-shaped path). + model_run_status: int = 200 + # --- counters the tests assert on --- upload_count: int = 0 from_hash_count: int = 0 @@ -106,6 +125,11 @@ class ServerState: # `None` before any request; `""` if a request arrived without one. last_auth_header: str | None = None last_user_agent: str | None = None + model_run_count: int = 0 + last_model_run_body: dict[str, Any] | None = None + # Every Idempotency-Key seen on POST /models/run, in arrival order (`None` + # records a run that arrived without the header at all). + model_run_idempotency_keys: list[str | None] = field(default_factory=list) def _asset_json(asset_id: str, hash_: str, created_new: bool, size: int) -> dict: @@ -168,6 +192,17 @@ class Handler(BaseHTTPRequestHandler): def log_message(self, *a: Any) -> None: pass + def handle_one_request(self) -> None: + # A test that deliberately aborts a slow request (an over-short + # client timeout against `model_run_delay`) closes the socket while + # this thread is still writing. That is the scenario under test, + # not a server fault — so don't dump a traceback for it. Only these + # two exception types are swallowed; anything else still surfaces. + try: + super().handle_one_request() + except (BrokenPipeError, ConnectionResetError): + self.close_connection = True + # -- helpers -- def _json(self, status: int, payload: dict, headers: dict | None = None) -> None: body = json.dumps(payload).encode() @@ -364,6 +399,9 @@ def do_POST(self) -> None: if self.path == "/api/v2/jobs": self._post_jobs() return + if self.path == "/api/v2/models/run": + self._post_model_run() + return m = re.match(r"/api/v2/jobs/([^/]+)/cancel$", self.path) if m: self._json(200, _job_json(m.group(1), "canceling")) @@ -388,6 +426,19 @@ def _post_from_hash(self) -> None: else: self._err(404, "blob_not_found", "no such blob") + def _post_model_run(self) -> None: + state.model_run_count += 1 + state.last_model_run_body = json.loads(self._read_body() or b"{}") + state.model_run_idempotency_keys.append(self.headers.get("Idempotency-Key")) + if state.model_run_delay: + # The server holding the connection while it polls upstream. + time.sleep(state.model_run_delay) + if state.model_run_error is not None: + status, code = state.model_run_error + self._err(status, code, f"model run error {code}") + return + self._json(state.model_run_status, state.model_run_result) + def _post_jobs(self) -> None: state.submit_count += 1 body = json.loads(self._read_body() or b"{}") diff --git a/tests/test_models_run.py b/tests/test_models_run.py new file mode 100644 index 0000000..8c3eab7 --- /dev/null +++ b/tests/test_models_run.py @@ -0,0 +1,241 @@ +"""``client.models.run(model, arguments)`` — one call, awaits completion. + +Covers the whole contract of the headline model API on both clients: the sync +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. + +Everything here runs against the stubbed server in ``conftest.py``. +""" + +from __future__ import annotations + +import inspect +import re + +import httpx +import pytest + +from comfy_low.transport import MODEL_RUN_TIMEOUT, AsyncComfyLow, ComfyLow, model_run_request +from comfy_sdk import AsyncComfy, Comfy +from comfy_sdk.exceptions import ComfyError, NotFound, Unauthorized +from comfy_sdk.models import AsyncModels, Models + +MODEL = "acme/flux/dev" +ARGS = {"prompt": "a cat", "steps": 4} + + +# --- the result of a completed generation ------------------------------- + + +def test_run_returns_the_completed_result(server) -> None: + with Comfy() as client: + result = client.models.run(MODEL, ARGS) + assert result == server.state.model_run_result + assert server.state.model_run_count == 1 + + +async def test_async_run_is_awaitable_and_returns_the_same_shape(server) -> None: + async with AsyncComfy() as client: + result = await client.models.run(MODEL, ARGS) + assert result == server.state.model_run_result + + +async def test_both_clients_return_an_identical_result(server) -> None: + with Comfy() as client: + sync_result = client.models.run(MODEL, ARGS) + async with AsyncComfy() as client: + async_result = await client.models.run(MODEL, ARGS) + assert sync_result == async_result + + +def test_the_result_is_the_providers_native_payload_not_a_wrapper(server) -> None: + # The provider's own field names reach the caller untouched — no SDK class + # in between, nothing renamed, nothing dropped. + payload = {"video": {"url": "http://example.invalid/v.mp4"}, "nsfw": False} + server.state.model_run_result = payload + with Comfy() as client: + result = client.models.run(MODEL, ARGS) + assert type(result) is dict + assert result == payload + + +def test_a_created_shaped_success_is_also_a_result(server) -> None: + server.state.model_run_status = 201 + with Comfy() as client: + assert client.models.run(MODEL, ARGS) == server.state.model_run_result + + +def test_run_sends_the_model_and_arguments(server) -> None: + with Comfy() as client: + client.models.run(MODEL, ARGS) + assert server.state.last_model_run_body == {"model": MODEL, "arguments": ARGS} + + +def test_run_accepts_any_mapping_and_does_not_alias_the_callers_object(server) -> None: + from types import MappingProxyType + + caller_args = {"prompt": "a dog"} + with Comfy() as client: + client.models.run(MODEL, MappingProxyType(caller_args)) + assert server.state.last_model_run_body == {"model": MODEL, "arguments": {"prompt": "a dog"}} + _path, body, _headers = model_run_request(MODEL, caller_args, None) + body["arguments"]["prompt"] = "mutated" + assert caller_args == {"prompt": "a dog"} + + +# --- the awaitable form is AsyncClient, not a run_async() suffix --------- + +_SUFFIXED = re.compile(r"(^async_|_async$|_sync$)") +_SURFACE = [Comfy, AsyncComfy, Models, AsyncModels, ComfyLow, AsyncComfyLow] + + +@pytest.mark.parametrize("cls", _SURFACE, ids=lambda c: c.__name__) +def test_no_run_async_method_exists(cls: type) -> None: + # Asserted rather than left to absence: one async mechanism (the async + # client) is a decided, published contract, and a second name for the same + # operation cannot be withdrawn once it ships. + assert not hasattr(cls, "run_async") + assert not hasattr(cls, "arun") + + +@pytest.mark.parametrize("cls", _SURFACE, ids=lambda c: c.__name__) +def test_no_suffixed_async_variant_on_the_public_surface(cls: type) -> None: + # Generalizes the rule past `run`: no public method on either client may + # signal sync-vs-async in its *name*. (`aclose` is the one deliberate + # rename and does not match — see tests/test_sync_async_parity.py.) + offenders = [n for n in dir(cls) if not n.startswith("_") and _SUFFIXED.search(n)] + assert offenders == [], f"{cls.__name__} exposes suffixed async variants: {offenders}" + + +def test_run_is_blocking_on_the_sync_client_and_awaitable_on_the_async_one() -> None: + assert not inspect.iscoroutinefunction(Models.run) + assert inspect.iscoroutinefunction(AsyncModels.run) + + +def test_both_clients_spell_the_operation_the_same_way() -> None: + # Not an exact-contents assertion — later model operations are expected to + # land here. What must hold is that the two namespaces name the *same* + # operations, which is the whole point of there being no suffixed variant. + def names(cls: type) -> set[str]: + return {n for n, v in vars(cls).items() if not n.startswith("_") and callable(v)} + + assert names(Models) == names(AsyncModels) + assert "run" in names(Models) + + +# --- the wait is sized for in-call polling ------------------------------ + + +def test_the_default_timeout_is_minutes_not_tens_of_seconds() -> None: + assert isinstance(MODEL_RUN_TIMEOUT, httpx.Timeout) + assert MODEL_RUN_TIMEOUT.read is not None and MODEL_RUN_TIMEOUT.read >= 120 + # Connecting is not generating: an unreachable host still fails fast. + assert MODEL_RUN_TIMEOUT.connect is not None and MODEL_RUN_TIMEOUT.connect <= 30 + + +def test_a_run_outlives_the_clients_own_timeout(server) -> None: + # The client is configured for ordinary API calls; the run holds the + # connection past that and must still complete. + server.state.model_run_delay = 0.75 + with Comfy(timeout=0.2) as client: + assert client.models.run(MODEL, ARGS) == server.state.model_run_result + + +async def test_an_async_run_outlives_the_clients_own_timeout(server) -> None: + server.state.model_run_delay = 0.75 + async with AsyncComfy(timeout=0.2) as client: + assert await client.models.run(MODEL, ARGS) == server.state.model_run_result + + +def test_the_clients_own_timeout_would_have_aborted_that_run(server) -> None: + # Control for the two tests above: without the run-sized default, a 0.2s + # client really does abort at 0.75s — so they are proving the override, + # not a stub that answers instantly. + server.state.model_run_delay = 0.75 + with Comfy(timeout=0.2) as client: + with pytest.raises(httpx.TimeoutException): + client.models.run(MODEL, ARGS, timeout=0.2) + + +def test_an_explicit_timeout_overrides_the_default(server) -> None: + server.state.model_run_delay = 0.5 + with Comfy() as client: + with pytest.raises(httpx.TimeoutException): + client.models.run(MODEL, ARGS, timeout=0.05) + + +# --- Idempotency-Key ----------------------------------------------------- + + +def test_run_sends_an_idempotency_key(server) -> None: + with Comfy() as client: + client.models.run(MODEL, ARGS) + (key,) = server.state.model_run_idempotency_keys + assert key + + +async def test_async_run_sends_an_idempotency_key(server) -> None: + async with AsyncComfy() as client: + await client.models.run(MODEL, ARGS) + (key,) = server.state.model_run_idempotency_keys + assert key + + +def test_each_run_mints_a_fresh_key(server) -> None: + # A second run is a second generation, not a retry of the first. + with Comfy() as client: + client.models.run(MODEL, ARGS) + client.models.run(MODEL, ARGS) + first, second = server.state.model_run_idempotency_keys + assert first and second and first != second + + +def test_an_explicit_key_is_sent_verbatim(server) -> None: + with Comfy() as client: + client.models.run(MODEL, ARGS, idempotency_key="caller-chosen-01") + assert server.state.model_run_idempotency_keys == ["caller-chosen-01"] + + +async def test_an_explicit_key_is_sent_verbatim_on_the_async_client(server) -> None: + async with AsyncComfy() as client: + await client.models.run(MODEL, ARGS, idempotency_key="caller-chosen-02") + assert server.state.model_run_idempotency_keys == ["caller-chosen-02"] + + +# --- the shared client configuration still applies ----------------------- + + +def test_run_carries_the_host_clients_credentials(server) -> None: + server.state.require_auth = True + with Comfy(api_key="k-run") as client: + client.models.run(MODEL, ARGS) + assert server.state.last_auth_header == "Bearer k-run" + + +# --- errors stay on the SDK's own surface -------------------------------- + + +def test_a_failing_run_raises_the_sdk_exception_not_the_protocol_one(server) -> None: + server.state.model_run_error = (404, "not_found") + with Comfy() as client: + with pytest.raises(NotFound): + client.models.run(MODEL, ARGS) + + +async def test_a_failing_async_run_raises_the_sdk_exception(server) -> None: + server.state.model_run_error = (401, "unauthorized") + async with AsyncComfy() as client: + with pytest.raises(Unauthorized): + await client.models.run(MODEL, ARGS) + + +def test_an_unmapped_failure_still_lands_as_a_comfy_error(server) -> None: + server.state.model_run_error = (503, "model_unavailable") + with Comfy() as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert excinfo.value.code == "model_unavailable" + assert excinfo.value.http_status == 503