fix: keep protocol errors off the public surface, add AsyncOutput.to_stream, widen the 429 retry

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.
+PyPI's Trusted Publishing (OIDC) — no API token is stored in this repo.

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

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,

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,12 +299,14 @@ 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:
     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))

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: