diff --git a/README.md b/README.md index a6b8e98..de8ed63 100644 --- a/README.md +++ b/README.md @@ -231,6 +231,8 @@ can pull down whichever way suits the caller: ```python out = job.get_outputs("13")[0] out.to_file("result.png") # stream to disk in chunks +with open("result.bin", "wb") as stream: + written = out.to_stream(stream) # write to an already-open binary stream data = out.to_bytes() # buffer into memory out.to_file("head.png", range=(0, 1023)) # range-aware: first 1 KiB only ``` @@ -265,6 +267,8 @@ async def main() -> None: wf = client.workflows.from_file("workflow_api.json") job = await client.run(wf) await job.outputs[0].to_file("out.png") + with open("out.bin", "wb") as stream: + await job.outputs[0].to_stream(stream) ``` ## Typed errors @@ -273,6 +277,10 @@ async def main() -> None: exceptions, all importable from the top-level package and all subclasses of `ComfyError`: +Catch these SDK-level exceptions around `Comfy`/`AsyncComfy` methods. Public +asset, job, event, and output helpers translate protocol errors, so catches of +`comfy_low.errors.*` belong only around direct low-level transport calls. + - `Unauthorized`, `Forbidden`, `NotFound` — auth and lookup failures. - `InvalidWorkflow`, `WorkflowFormatUi` — the graph itself was rejected; `WorkflowFormatUi` specifically means a UI-export (`nodes`/`links`/ @@ -289,8 +297,8 @@ exceptions, all importable from the top-level package and all subclasses of with the same key. - `InsufficientCredits` — the account can't afford the job. - `QueueFull` — backpressure; carries `.retry_after` seconds. `client.submit` - already retries this automatically for a bounded budget before giving up - and raising it. + retries 429 responses with `Retry-After` for a bounded budget (including + deployment warm-up), then raises the translated error if backpressure remains. - `JobFailed` — a job reached a non-`succeeded` terminal state; `.error` carries node-level detail when the platform provided one. @@ -356,4 +364,4 @@ python scripts/check_drift.py # same check CI runs; fails if committed models Releases are published to PyPI from a GitHub Release (tag `vX.Y.Z`) by [`.github/workflows/publish.yml`](.github/workflows/publish.yml), using -PyPI's Trusted Publishing (OIDC) — no API token is stored in this repo. \ No newline at end of file +PyPI's Trusted Publishing (OIDC) — no API token is stored in this repo. diff --git a/src/comfy_sdk/assets.py b/src/comfy_sdk/assets.py index 7d2bcf8..ce13764 100644 --- a/src/comfy_sdk/assets.py +++ b/src/comfy_sdk/assets.py @@ -270,7 +270,8 @@ def from_url(self, url: str) -> Asset: def get(self, asset_id: str) -> Asset: """Rehydrate an already-committed asset by UUID.""" - model = self._low.get_asset(asset_id) + with translating(): + model = self._low.get_asset(asset_id) asset = Asset(self._low, _rehydrated_source(model, asset_id)) asset._apply(model) return asset @@ -319,7 +320,8 @@ async def from_url(self, url: str) -> AsyncAsset: return AsyncAsset(self._low, _bytes_source(content, filename, ct)) async def get(self, asset_id: str) -> AsyncAsset: - model = await self._low.get_asset(asset_id) + with translating(): + model = await self._low.get_asset(asset_id) asset = AsyncAsset(self._low, _rehydrated_source(model, asset_id)) asset._apply(model) return asset diff --git a/src/comfy_sdk/client.py b/src/comfy_sdk/client.py index 65544b3..4787af6 100644 --- a/src/comfy_sdk/client.py +++ b/src/comfy_sdk/client.py @@ -30,7 +30,7 @@ from . import _core from .assets import AssetFactory, AsyncAssetFactory -from .exceptions import QueueFull, WorkflowFormatUi, to_sdk_error +from .exceptions import WorkflowFormatUi, to_sdk_error from .jobs import AsyncJob, AsyncJobFactory, Job, JobFactory from .workflows import Workflow, WorkflowFactory @@ -42,6 +42,25 @@ BASE_URL_ENV_VAR = "COMFY_BASE_URL" _DEFAULT_RETRY_AFTER = 2 +_now = time.monotonic + + +def _retry_delay(exc: ApiError, deadline: float) -> float | None: + """Return a bounded 429 retry delay, or ``None`` when the error should surface.""" + if exc.http_status != 429: + return None + if exc.retry_after is None: + # The spec requires Retry-After on deployment_not_ready. Keep the + # fallback only for legacy queue_full responses that omit it. + if exc.code != "queue_full": + return None + raw_delay = _DEFAULT_RETRY_AFTER + else: + raw_delay = exc.retry_after + remaining = deadline - _now() + if remaining <= 0: + return None + return max(0.0, min(raw_delay, remaining)) def _resolve_base_url() -> str: @@ -140,7 +159,7 @@ def submit( api_key: str | None = None, idempotency_key: str | None = None, ) -> Job: - """Submit a workflow. Retries ``queue_full`` with ``Retry-After``. + """Submit a workflow. Retries any 429 that carries ``Retry-After``. Sends an auto-generated ``Idempotency-Key`` so the server rejects an accidental exact resend of *this* request (``422 idempotency_key_reuse``) @@ -160,17 +179,18 @@ def submit( graph = self._materialize(workflow) key = idempotency_key or _core.new_idempotency_key() extra_data = _core.extra_data_for(api_key) - deadline = time.monotonic() + _QUEUE_RETRY_BUDGET + deadline = _now() + _QUEUE_RETRY_BUDGET while True: try: model = self._low.post_jobs(graph, idempotency_key=key, extra_data=extra_data) return Job(self._low, model) except ApiError as exc: err = to_sdk_error(exc) - if isinstance(err, QueueFull) and time.monotonic() < deadline: - time.sleep(err.retry_after or _DEFAULT_RETRY_AFTER) - continue - raise err from exc + delay = _retry_delay(exc, deadline) + if delay is None: + raise err from exc + time.sleep(delay) + continue def run( self, @@ -243,17 +263,18 @@ async def submit( graph = await self._materialize(workflow) key = idempotency_key or _core.new_idempotency_key() extra_data = _core.extra_data_for(api_key) - deadline = time.monotonic() + _QUEUE_RETRY_BUDGET + deadline = _now() + _QUEUE_RETRY_BUDGET while True: try: model = await self._low.post_jobs(graph, idempotency_key=key, extra_data=extra_data) return AsyncJob(self._low, model) except ApiError as exc: err = to_sdk_error(exc) - if isinstance(err, QueueFull) and time.monotonic() < deadline: - await asyncio.sleep(err.retry_after or _DEFAULT_RETRY_AFTER) - continue - raise err from exc + delay = _retry_delay(exc, deadline) + if delay is None: + raise err from exc + await asyncio.sleep(delay) + continue async def run( self, diff --git a/src/comfy_sdk/jobs.py b/src/comfy_sdk/jobs.py index 0e18a58..fb61cb1 100644 --- a/src/comfy_sdk/jobs.py +++ b/src/comfy_sdk/jobs.py @@ -24,7 +24,7 @@ from . import _core from .events import Event, StatusChange, event_from_raw -from .exceptions import JobFailed, translating +from .exceptions import JobFailed, to_sdk_error, translating from .outputs import AsyncOutput, Output _RECONNECT_PAUSE = 0.1 @@ -165,7 +165,7 @@ def events(self) -> Iterator[Event]: except ApiError as exc: if exc.http_status == 501: return # surface has no SSE — poll paths remain authoritative - raise + raise to_sdk_error(exc) from exc except (httpx.HTTPError, httpx.StreamError): pass # connection dropped mid-stream — reconnect below if terminal_seen: @@ -277,7 +277,7 @@ async def events(self) -> AsyncIterator[Event]: except ApiError as exc: if exc.http_status == 501: return # surface has no SSE — poll paths remain authoritative - raise + raise to_sdk_error(exc) from exc except (httpx.HTTPError, httpx.StreamError): pass if terminal_seen: @@ -299,7 +299,8 @@ def __init__(self, low: ComfyLow) -> None: self._low = low def get(self, job_id: str) -> Job: - return Job(self._low, self._low.get_job(job_id)) + with translating(): + return Job(self._low, self._low.get_job(job_id)) class AsyncJobFactory: @@ -307,4 +308,5 @@ def __init__(self, low: AsyncComfyLow) -> None: self._low = low async def get(self, job_id: str) -> AsyncJob: - return AsyncJob(self._low, await self._low.get_job(job_id)) + with translating(): + return AsyncJob(self._low, await self._low.get_job(job_id)) diff --git a/src/comfy_sdk/outputs.py b/src/comfy_sdk/outputs.py index d666992..56cf43d 100644 --- a/src/comfy_sdk/outputs.py +++ b/src/comfy_sdk/outputs.py @@ -17,9 +17,24 @@ from comfy_low.models import Output as LowOutput from comfy_low.transport import AsyncComfyLow, ComfyLow +from .exceptions import translating + _CHUNK = 64 * 1024 +def _write_all(stream: BinaryIO, chunk: bytes) -> int: + """Write a complete chunk, including through partial writes.""" + written = 0 + while written < len(chunk): + count = stream.write(chunk[written:]) + if count is None or count <= 0: + raise OSError("stream.write() made no progress") + if count > len(chunk) - written: + raise OSError("stream.write() returned an invalid byte count") + written += count + return written + + @dataclass(frozen=True) class DownloadUrl: """A directly-fetchable URL for one output, e.g. to hand to a downstream @@ -80,7 +95,7 @@ def to_file(self, path: str | PathLike[str], *, range: tuple[int, int] | None = ``range=(0, 4)`` yields the first five bytes. """ dest = Path(path) - with self._low.get_asset_content(self._model.id, range=range) as resp: + with translating(), self._low.get_asset_content(self._model.id, range=range) as resp: with open(dest, "wb") as fh: for chunk in resp.iter_bytes(_CHUNK): fh.write(chunk) @@ -93,10 +108,9 @@ def to_stream(self, stream: BinaryIO, *, range: tuple[int, int] | None = None) - closed — that stays the caller's. See :meth:`to_file` for ``range``. """ written = 0 - with self._low.get_asset_content(self._model.id, range=range) as resp: + with translating(), self._low.get_asset_content(self._model.id, range=range) as resp: for chunk in resp.iter_bytes(_CHUNK): - stream.write(chunk) - written += len(chunk) + written += _write_all(stream, chunk) return written def to_bytes(self, *, range: tuple[int, int] | None = None) -> bytes: @@ -107,13 +121,13 @@ def to_bytes(self, *, range: tuple[int, int] | None = None) -> bytes: See :meth:`to_file` for ``range``. """ buf = bytearray() - with self._low.get_asset_content(self._model.id, range=range) as resp: + with translating(), self._low.get_asset_content(self._model.id, range=range) as resp: for chunk in resp.iter_bytes(_CHUNK): buf.extend(chunk) return bytes(buf) def get_download_url(self) -> DownloadUrl: - """A directly-fetchable URL for this output — never throws. + """A directly-fetchable URL for this output. On a Cloud/serverless backend this is a short-lived, self-authorizing signed URL for object storage: anyone holding it can read the bytes @@ -122,7 +136,8 @@ def get_download_url(self) -> DownloadUrl: ``expires_at`` is ``None``. (A genuine failure — e.g. an unknown output id — still raises the same typed error as any other call.) """ - url, expires_at = self._low.get_asset_content_url(self._model.id) + with translating(): + url, expires_at = self._low.get_asset_content_url(self._model.id) return DownloadUrl(url=url, expires_at=expires_at) def __repr__(self) -> str: @@ -171,23 +186,35 @@ async def to_file( ) -> Path: """Async :meth:`Output.to_file` — same chunked write and inclusive ``range``.""" dest = Path(path) - async with self._low.get_asset_content(self._model.id, range=range) as resp: - with open(dest, "wb") as fh: - async for chunk in resp.aiter_bytes(_CHUNK): - fh.write(chunk) + with translating(): + async with self._low.get_asset_content(self._model.id, range=range) as resp: + with open(dest, "wb") as fh: + async for chunk in resp.aiter_bytes(_CHUNK): + fh.write(chunk) return dest + async def to_stream(self, stream: BinaryIO, *, range: tuple[int, int] | None = None) -> int: + """Async :meth:`Output.to_stream` — same write-only semantics.""" + written = 0 + with translating(): + async with self._low.get_asset_content(self._model.id, range=range) as resp: + async for chunk in resp.aiter_bytes(_CHUNK): + written += _write_all(stream, chunk) + return written + async def to_bytes(self, *, range: tuple[int, int] | None = None) -> bytes: """Async :meth:`Output.to_bytes` — buffers the whole body in memory.""" buf = bytearray() - async with self._low.get_asset_content(self._model.id, range=range) as resp: - async for chunk in resp.aiter_bytes(_CHUNK): - buf.extend(chunk) + with translating(): + async with self._low.get_asset_content(self._model.id, range=range) as resp: + async for chunk in resp.aiter_bytes(_CHUNK): + buf.extend(chunk) return bytes(buf) async def get_download_url(self) -> DownloadUrl: """See the sync ``Output.get_download_url`` for the redirect/inline split.""" - url, expires_at = await self._low.get_asset_content_url(self._model.id) + with translating(): + url, expires_at = await self._low.get_asset_content_url(self._model.id) return DownloadUrl(url=url, expires_at=expires_at) def __repr__(self) -> str: diff --git a/tests/conftest.py b/tests/conftest.py index 65b2793..ff4e143 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -35,8 +35,28 @@ class ServerState: require_auth: bool = False # POST /jobs returns 429 queue_full this many times before succeeding. queue_full_times: int = 0 + # Like `queue_full_times`, but the 429 carries no Retry-After header at + # all — the bare-`queue_full` path, which retries using the client's + # default pause rather than a server-given delay. + queue_full_times_no_retry_after: int = 0 + # POST /jobs answers 429 queue_full with this literal Retry-After header + # value once, then succeeds — for a test to send an out-of-range value + # (e.g. a huge number, to prove the client clamps to its retry budget, + # or a negative one, to prove a malformed header doesn't crash the + # sync loop or busy-loop the async one). + queue_full_retry_after_header: str | None = None + # POST /jobs returns a 429 naming a code OTHER than `queue_full` (with a + # Retry-After header) this many times before succeeding — the contract + # disambiguates a retryable 429 by status + Retry-After, not by `code` + # (e.g. `deployment_not_ready` on a serverless cold start). + retryable_429_times: int = 0 + retryable_429_code: str = "deployment_not_ready" # POST /jobs returns this error envelope (status, code) instead of 201. job_error: tuple[int, str] | None = None + # GET /jobs/{id} answers 404 job_not_found instead of the job. + job_not_found: bool = False + # GET /jobs/{id}/events answers this (status, code) instead of connecting. + events_error: tuple[int, str] | None = None # Number of GET /jobs/{id} polls before the job reports succeeded. polls_to_succeed: int = 1 # Terminal status the job reaches. @@ -267,6 +287,9 @@ def _serve_content(self) -> None: self.wfile.write(data) def _serve_job(self, job_id: str) -> None: + if state.job_not_found: + self._err(404, "job_not_found", "no such job") + return state.job_poll_count += 1 if state.job_poll_count >= state.polls_to_succeed: status = state.terminal_status @@ -292,6 +315,10 @@ def _serve_job_workflow(self, job_id: str) -> None: def _serve_events(self, job_id: str) -> None: state.events_connect_count += 1 + if state.events_error is not None: + status, code = state.events_error + self._err(status, code, f"events error {code}") + return if state.events_not_implemented: self._err(501, "not_implemented", "SSE is not supported on this surface") return @@ -383,6 +410,30 @@ def _post_jobs(self) -> None: ) return + if state.queue_full_times_no_retry_after > 0: + state.queue_full_times_no_retry_after -= 1 + self._json(429, {"error": {"code": "queue_full", "message": "full"}}) + return + + if state.queue_full_retry_after_header is not None: + header = state.queue_full_retry_after_header + state.queue_full_retry_after_header = None + self._json( + 429, + {"error": {"code": "queue_full", "message": "full"}}, + headers={"Retry-After": header}, + ) + return + + if state.retryable_429_times > 0: + state.retryable_429_times -= 1 + self._json( + 429, + {"error": {"code": state.retryable_429_code, "message": "warming up"}}, + headers={"Retry-After": "0"}, + ) + return + if state.job_error is not None: status, code = state.job_error self._err(status, code, f"job error {code}") diff --git a/tests/test_assets.py b/tests/test_assets.py index f975bd9..0589f32 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -7,8 +7,7 @@ import pytest -from comfy_low.errors import NotFound -from comfy_sdk import Comfy, HashMismatch +from comfy_sdk import Comfy, HashMismatch, NotFound def test_dedup_fast_path_skips_upload(server, tmp_path) -> None: diff --git a/tests/test_async.py b/tests/test_async.py index f630fd7..34f91d7 100644 --- a/tests/test_async.py +++ b/tests/test_async.py @@ -2,16 +2,24 @@ from __future__ import annotations +import asyncio +import io + import pytest -from comfy_low.errors import NotFound -from comfy_sdk import AsyncComfy, MissingAsset, Progress, StatusChange +import comfy_sdk.client as _client_module +from comfy_sdk import AsyncComfy, MissingAsset, NotFound, Progress, StatusChange def _wf(client: AsyncComfy): return client.workflows.from_json({"3": {"class_type": "KSampler", "inputs": {}}}) +class _ShortWriter(io.BytesIO): + def write(self, data) -> int: + return super().write(data[:3]) + + async def test_async_run_and_download(server, tmp_path) -> None: server.state.polls_to_succeed = 2 async with AsyncComfy() as client: @@ -22,6 +30,15 @@ async def test_async_run_and_download(server, tmp_path) -> None: assert data == server.state.content_bytes +async def test_async_output_to_stream_writes_all_bytes(server) -> None: + stream = _ShortWriter() + async with AsyncComfy() as client: + job = await client.run(_wf(client)) + written = await job.get_outputs("13")[0].to_stream(stream) + assert written == len(server.state.content_bytes) + assert stream.getvalue() == server.state.content_bytes + + async def test_async_events_stream_to_terminal(server) -> None: async with AsyncComfy() as client: job = await client.submit(_wf(client)) @@ -150,6 +167,46 @@ async def test_async_queue_full_retries_with_retry_after(server) -> None: assert server.state.submit_count == 3 +async def test_async_submit_clamps_huge_retry_after_to_remaining_budget( + server, monkeypatch +) -> None: + # Async counterpart of the sync clamp test — same predicate, same clamp, + # both loops must bound a single sleep to what's left of the budget. + monkeypatch.setattr(_client_module, "_QUEUE_RETRY_BUDGET", 5.0) + times = iter([100.0, 104.25]) + monkeypatch.setattr(_client_module, "_now", lambda: next(times)) + sleeps: list[float] = [] + + async def _fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr(asyncio, "sleep", _fake_sleep) + server.state.queue_full_retry_after_header = "10000000" + async with AsyncComfy() as client: + job = await client.submit(_wf(client)) + assert job.id.startswith("job_") + assert server.state.submit_count == 2 + assert sleeps == [0.75] + + +async def test_async_submit_negative_retry_after_does_not_storm(server, monkeypatch) -> None: + # Where the sync loop crashes on `time.sleep(-5)` (ValueError), the async + # loop's `asyncio.sleep(-5)` returns instantly and would busy-loop the + # server for the whole retry budget with no pause. The clamp floors the + # delay at 0 either way; this proves the async side specifically. + sleeps: list[float] = [] + + async def _fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr(asyncio, "sleep", _fake_sleep) + server.state.queue_full_retry_after_header = "-5" + async with AsyncComfy() as client: + job = await client.submit(_wf(client)) + assert job.id.startswith("job_") + assert sleeps == [0.0] + + async def test_async_delete_asset_by_id(server) -> None: async with AsyncComfy() as client: await client.assets.delete("asset_uuid_01") diff --git a/tests/test_download_and_workflows.py b/tests/test_download_and_workflows.py index 605e82e..d05ce3b 100644 --- a/tests/test_download_and_workflows.py +++ b/tests/test_download_and_workflows.py @@ -2,6 +2,10 @@ from __future__ import annotations +import io + +import pytest + from comfy_sdk import Comfy from comfy_sdk._core import find_asset_handles, substitute_asset_handles @@ -10,6 +14,15 @@ def _wf(client: Comfy): return client.workflows.from_json({"3": {"class_type": "KSampler", "inputs": {}}}) +class _ShortWriter(io.BytesIO): + def __init__(self, limit: int) -> None: + super().__init__() + self.limit = limit + + def write(self, data) -> int: + return super().write(data[: self.limit]) + + def test_range_download_returns_partial(server) -> None: server.state.content_bytes = b"0123456789abcdef" with Comfy() as client: @@ -39,6 +52,23 @@ def test_full_download(server, tmp_path) -> None: assert dest.read_bytes() == server.state.content_bytes +def test_to_stream_retries_partial_writes(server) -> None: + stream = _ShortWriter(limit=3) + with Comfy() as client: + job = client.run(_wf(client)) + written = job.get_outputs("13")[0].to_stream(stream) + assert written == len(server.state.content_bytes) + assert stream.getvalue() == server.state.content_bytes + + +def test_to_stream_rejects_zero_byte_writes(server) -> None: + stream = _ShortWriter(limit=0) + with Comfy() as client: + job = client.run(_wf(client)) + with pytest.raises(OSError, match="made no progress"): + job.get_outputs("13")[0].to_stream(stream) + + def test_core_asset_substitution(server, tmp_path) -> None: p = tmp_path / "photo.png" p.write_bytes(b"pixels") diff --git a/tests/test_error_contract.py b/tests/test_error_contract.py new file mode 100644 index 0000000..d9f9787 --- /dev/null +++ b/tests/test_error_contract.py @@ -0,0 +1,143 @@ +"""The ``translating()`` contract: no public entry point may leak the +protocol-level ``comfy_low.errors.ApiError`` — every error must surface as its +``comfy_sdk`` typed equivalent (see ``comfy_sdk.exceptions.translating``). + +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()``. +""" + +from __future__ import annotations + +import io +from collections.abc import Awaitable, Callable +from typing import Any + +import pytest + +from comfy_low.errors import ApiError as LowApiError +from comfy_sdk import AsyncComfy, Comfy, Forbidden, NotFound +from comfy_sdk.exceptions import ComfyError + + +def _wf(client: Comfy | AsyncComfy): + return client.workflows.from_json({"3": {"class_type": "KSampler", "inputs": {}}}) + + +def _assert_no_leak(entry_point: str, fn: Callable[[], Any], expected: type[ComfyError]) -> None: + try: + fn() + except ComfyError as exc: + assert isinstance(exc, expected), ( + f"{entry_point} raised {type(exc).__name__}, expected {expected.__name__}" + ) + except LowApiError as exc: + pytest.fail( + f"{entry_point} leaked comfy_low.errors.{type(exc).__name__} " + "instead of translating it to a comfy_sdk ComfyError" + ) + else: + pytest.fail(f"{entry_point} did not raise; the scenario should force an error") + + +async def _assert_no_leak_async( + entry_point: str, coro: Awaitable[Any], expected: type[ComfyError] +) -> None: + try: + await coro + except ComfyError as exc: + assert isinstance(exc, expected), ( + f"{entry_point} raised {type(exc).__name__}, expected {expected.__name__}" + ) + except LowApiError as exc: + pytest.fail( + f"{entry_point} leaked comfy_low.errors.{type(exc).__name__} " + "instead of translating it to a comfy_sdk ComfyError" + ) + else: + pytest.fail(f"{entry_point} did not raise; the scenario should force an error") + + +# -- outputs.py: to_file / to_stream / to_bytes / get_download_url ---------- + + +def test_output_download_methods_translate_on_deleted_asset(server, tmp_path) -> None: + with Comfy() as client: + job = client.run(_wf(client)) + out = job.get_outputs("13")[0] + client.assets.delete(out.id) # server now 404s this asset's content + + _assert_no_leak("Output.to_bytes", out.to_bytes, NotFound) + _assert_no_leak("Output.to_file", lambda: out.to_file(tmp_path / "o.bin"), NotFound) + _assert_no_leak("Output.to_stream", lambda: out.to_stream(io.BytesIO()), NotFound) + _assert_no_leak("Output.get_download_url", out.get_download_url, NotFound) + + +async def test_async_output_download_methods_translate_on_deleted_asset(server, tmp_path) -> None: + async with AsyncComfy() as client: + job = await client.run(_wf(client)) + out = job.get_outputs("13")[0] + await client.assets.delete(out.id) + + await _assert_no_leak_async("AsyncOutput.to_bytes", out.to_bytes(), NotFound) + await _assert_no_leak_async( + "AsyncOutput.to_file", out.to_file(tmp_path / "o.bin"), NotFound + ) + await _assert_no_leak_async("AsyncOutput.to_stream", out.to_stream(io.BytesIO()), NotFound) + await _assert_no_leak_async( + "AsyncOutput.get_download_url", out.get_download_url(), NotFound + ) + + +# -- AssetFactory.get / AsyncAssetFactory.get -------------------------------- + + +def test_asset_factory_get_translates_on_deleted_asset(server) -> None: + with Comfy() as client: + client.assets.delete("asset_out_01") + _assert_no_leak("AssetFactory.get", lambda: client.assets.get("asset_out_01"), NotFound) + + +async def test_async_asset_factory_get_translates_on_deleted_asset(server) -> None: + async with AsyncComfy() as client: + await client.assets.delete("asset_out_01") + await _assert_no_leak_async( + "AsyncAssetFactory.get", client.assets.get("asset_out_01"), NotFound + ) + + +# -- JobFactory.get / AsyncJobFactory.get ------------------------------------ + + +def test_job_factory_get_translates_on_missing_job(server) -> None: + server.state.job_not_found = True + with Comfy() as client: + _assert_no_leak("JobFactory.get", lambda: client.jobs.get("no_such_job"), NotFound) + + +async def test_async_job_factory_get_translates_on_missing_job(server) -> None: + server.state.job_not_found = True + async with AsyncComfy() as client: + await _assert_no_leak_async("AsyncJobFactory.get", client.jobs.get("no_such_job"), NotFound) + + +# -- Job.events() / AsyncJob.events() non-501 raise -------------------------- + + +def test_job_events_translates_non_501_error(server) -> None: + with Comfy() as client: + job = client.run(_wf(client)) + server.state.events_error = (403, "forbidden") + _assert_no_leak("Job.events", lambda: list(job.events()), Forbidden) + + +async def test_async_job_events_translates_non_501_error(server) -> None: + async with AsyncComfy() as client: + job = await client.run(_wf(client)) + server.state.events_error = (403, "forbidden") + + async def _drain() -> None: + async for _ in job.events(): + pass + + await _assert_no_leak_async("AsyncJob.events", _drain(), Forbidden) diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 600e73f..429cb64 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -4,12 +4,15 @@ from __future__ import annotations +import time + import pytest import comfy_sdk.client as _client_module from comfy_low.errors import IdempotencyKeyReuse as LowIdempotencyKeyReuse from comfy_sdk import ( Comfy, + ComfyError, IdempotencyKeyReuse, InvalidWorkflow, JobFailed, @@ -77,6 +80,88 @@ def test_queue_full_retries_with_retry_after(server) -> None: assert server.state.submit_count == 3 # two rejections + one success +def test_submit_retries_429_with_a_code_other_than_queue_full(server) -> None: + # The contract disambiguates a retryable 429 by status + Retry-After, not + # by `code` — `deployment_not_ready` (a serverless cold start) must retry + # exactly like `queue_full` does, even though `to_sdk_error` doesn't map it + # to `QueueFull`. + server.state.retryable_429_times = 2 + with Comfy() as client: + job = client.submit(_wf(client)) + assert job.id.startswith("job_") + assert server.state.submit_count == 3 # two rejections + one success + + +def test_queue_full_retries_with_no_retry_after_header(server, monkeypatch) -> None: + # A bare `queue_full` 429 (no Retry-After header at all) must still + # retry, using the client's default pause — this was the one 429 path + # that already worked before the status+Retry-After predicate landed, + # and it regressed when that predicate first shipped without an + # explicit `code == "queue_full"` fallback. + monkeypatch.setattr(_client_module, "_DEFAULT_RETRY_AFTER", 0) + server.state.queue_full_times_no_retry_after = 2 + with Comfy() as client: + job = client.submit(_wf(client)) + assert job.id.startswith("job_") + assert server.state.submit_count == 3 # two rejections + one success + + +def test_submit_clamps_huge_retry_after_to_remaining_budget(server, monkeypatch) -> None: + # Security regression: the deadline used to be checked before the sleep + # but never against the sleep's own duration, so a server-supplied + # Retry-After of e.g. 10_000_000 committed the caller to a ~115-day + # sleep in one shot. The delay must be clamped to what's left of the + # retry budget, not to the header's raw value. + monkeypatch.setattr(_client_module, "_QUEUE_RETRY_BUDGET", 5.0) + times = iter([100.0, 104.25]) + monkeypatch.setattr(_client_module, "_now", lambda: next(times)) + sleeps: list[float] = [] + monkeypatch.setattr(time, "sleep", lambda s: sleeps.append(s)) + server.state.queue_full_retry_after_header = "10000000" + with Comfy() as client: + job = client.submit(_wf(client)) + assert job.id.startswith("job_") + assert server.state.submit_count == 2 + assert sleeps == [0.75] + + +def test_submit_negative_retry_after_does_not_crash(server, monkeypatch) -> None: + # comfy_low's header parsing is a bare `int(raw)`, so "-5" parses to -5 + # rather than None. Unclamped, `time.sleep(-5)` raises ValueError — the + # caller would get a raw ValueError instead of the documented ComfyError + # contract. The clamp must floor the delay at 0. + sleeps: list[float] = [] + monkeypatch.setattr(time, "sleep", lambda s: sleeps.append(s)) + server.state.queue_full_retry_after_header = "-5" + with Comfy() as client: + job = client.submit(_wf(client)) # must not raise ValueError + assert job.id.startswith("job_") + assert sleeps == [0.0] + + +def test_submit_retry_after_zero_sleeps_for_zero_not_default(server, monkeypatch) -> None: + # Pins the incidental fix from the explicit `is not None` check: the old + # `retry_after or _DEFAULT_RETRY_AFTER` treated a literal `Retry-After: 0` + # as absent and slept the full default. Without this assertion a + # regression back to that pattern only shows up as CI getting slower. + sleeps: list[float] = [] + monkeypatch.setattr(time, "sleep", lambda s: sleeps.append(s)) + server.state.queue_full_times = 1 # stub sends Retry-After: 0 + with Comfy() as client: + client.submit(_wf(client)) + assert sleeps == [0.0] + + +def test_submit_does_not_retry_429_without_retry_after_and_non_queue_full_code(server) -> None: + # The OpenAPI contract requires deployment_not_ready to carry + # Retry-After. A malformed response without it surfaces immediately. + server.state.job_error = (429, "deployment_not_ready") + with Comfy() as client: + with pytest.raises(ComfyError): + client.submit(_wf(client)) + assert server.state.submit_count == 1 # no retry + + def test_queue_full_gives_up_once_retry_budget_elapses(server, monkeypatch) -> None: # The server never clears backpressure. Before this test, only the # "eventually succeeds" retry path was covered — nothing proved the client diff --git a/tests/test_sync_async_parity.py b/tests/test_sync_async_parity.py new file mode 100644 index 0000000..04765f5 --- /dev/null +++ b/tests/test_sync_async_parity.py @@ -0,0 +1,47 @@ +"""Sync/async public-surface parity across the 7 mirrored class pairs. + +The README promises "swap the import and add ``await``" as the only +difference; this asserts the public method names actually match. Regression +guard for ``AsyncOutput`` shipping without ``to_stream``. + +Dunders (``__enter__``/``__aenter__``, ...) are excluded — they're a mechanical +consequence of sync vs async, not part of the call surface a swap has to +match. The intentional ``close``/``aclose`` rename is normalised away. +""" + +from __future__ import annotations + +import pytest + +from comfy_low.transport import AsyncComfyLow, ComfyLow +from comfy_sdk.assets import Asset, AssetFactory, AsyncAsset, AsyncAssetFactory +from comfy_sdk.client import AsyncComfy, Comfy +from comfy_sdk.jobs import AsyncJob, AsyncJobFactory, Job, JobFactory +from comfy_sdk.outputs import AsyncOutput, Output + +_PAIRS: list[tuple[str, type, type]] = [ + ("Comfy", Comfy, AsyncComfy), + ("Asset", Asset, AsyncAsset), + ("AssetFactory", AssetFactory, AsyncAssetFactory), + ("Job", Job, AsyncJob), + ("JobFactory", JobFactory, AsyncJobFactory), + ("Output", Output, AsyncOutput), + ("ComfyLow", ComfyLow, AsyncComfyLow), +] + + +def _public_names(cls: type) -> set[str]: + names = {n for n in dir(cls) if not n.startswith("_")} + return {"close" if n == "aclose" else n for n in names} + + +@pytest.mark.parametrize("label,sync_cls,async_cls", _PAIRS, ids=[p[0] for p in _PAIRS]) +def test_sync_async_public_surface_matches(label: str, sync_cls: type, async_cls: type) -> None: + sync_names = _public_names(sync_cls) + async_names = _public_names(async_cls) + sync_only = sync_names - async_names + async_only = async_names - sync_names + assert not sync_only and not async_only, ( + f"{label}/Async{label} public surface diverges: " + f"sync-only={sorted(sync_only)} async-only={sorted(async_only)}" + )