fix: translate errors on every public entry point, add AsyncOutput.to_stream, widen 429 retry

Three defects behind one untested guarantee: nothing asserted that the
public surface only ever raises comfy_sdk exceptions.

- 10 public entry points skipped translating() and leaked the raw
  comfy_low.errors.ApiError: Output/AsyncOutput's download methods, both
  AssetFactory.get, both JobFactory.get, and the non-501 raise in
  Job.events()/AsyncJob.events(). Because comfy_low.errors and
  comfy_sdk.exceptions both export a NotFound and they are unrelated, the
  documented `except NotFound` around a download silently never fired.

- AsyncOutput was missing to_stream entirely, breaking the documented
  "swap the import and add await" parity.

- submit() gated its 429 retry on isinstance(err, QueueFull), i.e. on the
  server's error code. The contract disambiguates a retryable 429 by
  status + Retry-After, so a 429 named deployment_not_ready -- a
  serverless cold start -- hard-failed on the first attempt.

The retry delay is now clamped to what remains of the retry budget and
floored at zero. Retry-After is untrusted server input: it was previously
passed to sleep() unbounded, so one response could park a caller far past
_QUEUE_RETRY_BUDGET, and a negative value raised an uncaught ValueError on
the sync path while spinning the async path with no pause.

Also pins an incidental fix: the old `retry_after or _DEFAULT_RETRY_AFTER`
treated Retry-After: 0 as absent and slept the full default.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VaAqCLd2mvvwsYkDTHMyBX

Author: wei-hai

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
 
@@ -140,7 +140,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``)
@@ -166,11 +166,24 @@ def submit(
                 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)
+                # Disambiguated by status + Retry-After, not `code` alone
+                # (e.g. `deployment_not_ready`); a bare `queue_full` 429 (no
+                # header) must still retry on the default pause — dropping
+                # that fallback is the regression this predicate already hit.
+                retryable = exc.http_status == 429 and (
+                    exc.retry_after is not None or exc.code == "queue_full"
+                )
+                remaining = deadline - time.monotonic()
+                if retryable and remaining > 0:
+                    raw_delay = (
+                        exc.retry_after if exc.retry_after is not None else _DEFAULT_RETRY_AFTER
+                    )
+                    # Clamp: a server-supplied Retry-After is untrusted input —
+                    # never sleep past the retry budget, and never negative.
+                    delay = max(0.0, min(raw_delay, remaining))
+                    time.sleep(delay)
                     continue
-                raise err from exc
+                raise to_sdk_error(exc) from exc
 
     def run(
         self,
@@ -249,11 +262,19 @@ async def submit(
                 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)
+                # See the sync `submit` above for the predicate and clamp.
+                retryable = exc.http_status == 429 and (
+                    exc.retry_after is not None or exc.code == "queue_full"
+                )
+                remaining = deadline - time.monotonic()
+                if retryable and remaining > 0:
+                    raw_delay = (
+                        exc.retry_after if exc.retry_after is not None else _DEFAULT_RETRY_AFTER
+                    )
+                    delay = max(0.0, min(raw_delay, remaining))
+                    await asyncio.sleep(delay)
                     continue
-                raise err from exc
+                raise to_sdk_error(exc) from exc
 
     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,6 +17,8 @@
 from comfy_low.models import Output as LowOutput
 from comfy_low.transport import AsyncComfyLow, ComfyLow
 
+from .exceptions import translating
+
 _CHUNK = 64 * 1024
 
 
@@ -80,7 +82,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,7 +95,7 @@ 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)
@@ -107,7 +109,7 @@ 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)
@@ -122,7 +124,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 +174,36 @@ 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):
+                    stream.write(chunk)
+                    written += len(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:

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}")

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:

tests/test_async.py

@@ -2,10 +2,12 @@
 
 from __future__ import annotations
 
+import asyncio
+
 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):
@@ -150,6 +152,45 @@ 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)
+    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 len(sleeps) == 1
+    assert 0 <= sleeps[0] <= 5.0
+
+
+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")