From 75ef44dacc4656a73d53c9a60cd71b0ae261173d Mon Sep 17 00:00:00 2001 From: Wei Hai Date: Thu, 13 Aug 2026 11:06:46 -0700 Subject: [PATCH] fix: expose job_id and expires_at on the public wrappers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recent spec sync added Asset.job_id, Asset.expires_at and Output.job_id to the generated models, but the hand-written wrapper classes never surfaced them — so a caller could not reach the field the change exists to provide without touching a private attribute. Found by end-to-end testing against Comfy Cloud staging and a self-hosted proxy: the value was correct on the wire from both, and unreachable through the public API. Existing tests missed it because they assert on the request and response bodies, never on the wrapper's surface; the new tests assert on the wrapper. Co-Authored-By: Claude Opus 5 (1M context) --- src/comfy_sdk/assets.py | 17 +++++++++++++++++ src/comfy_sdk/outputs.py | 12 ++++++++++++ tests/conftest.py | 7 ++++++- tests/test_assets.py | 23 +++++++++++++++++++++++ tests/test_jobs.py | 9 +++++++++ 5 files changed, 67 insertions(+), 1 deletion(-) diff --git a/src/comfy_sdk/assets.py b/src/comfy_sdk/assets.py index 23d092d..7d2bcf8 100644 --- a/src/comfy_sdk/assets.py +++ b/src/comfy_sdk/assets.py @@ -14,6 +14,7 @@ import mimetypes from collections.abc import Callable from dataclasses import dataclass +from datetime import datetime from io import BytesIO from os import PathLike from os.path import basename, getsize @@ -63,6 +64,8 @@ def __init__(self, source: _Source) -> None: self._id: str | None = None self._created_new: bool | None = None self._url: str | None = None + self._job_id: str | None = None + self._expires_at: datetime | None = None self._idempotency_key = _core.new_idempotency_key() @property @@ -84,12 +87,26 @@ def hash(self) -> str: def created_new(self) -> bool | None: return self._created_new + @property + def job_id(self) -> str | None: + """The id of the job that produced this asset, or ``None`` for an + asset with no producing job (e.g. a plain upload).""" + return self._job_id + + @property + def expires_at(self) -> datetime | None: + """Retention deadline for this asset, or ``None`` if it doesn't + expire.""" + return self._expires_at + def _apply(self, asset: LowAsset) -> None: self._id = asset.id if asset.hash: self._hash = asset.hash self._created_new = asset.created_new self._url = str(asset.url) + self._job_id = asset.job_id + self._expires_at = asset.expires_at def __repr__(self) -> str: state = self._id or "uncommitted" diff --git a/src/comfy_sdk/outputs.py b/src/comfy_sdk/outputs.py index bec243a..d666992 100644 --- a/src/comfy_sdk/outputs.py +++ b/src/comfy_sdk/outputs.py @@ -65,6 +65,12 @@ def size_bytes(self) -> int: def content_type(self) -> str: return self._model.content_type + @property + def job_id(self) -> str | None: + """The id of the job that produced this output, or ``None`` if the + backend didn't report one.""" + return self._model.job_id + def to_file(self, path: str | PathLike[str], *, range: tuple[int, int] | None = None) -> Path: """Stream this output to ``path`` and return the path written. @@ -154,6 +160,12 @@ def size_bytes(self) -> int: def content_type(self) -> str: return self._model.content_type + @property + def job_id(self) -> str | None: + """The id of the job that produced this output, or ``None`` if the + backend didn't report one.""" + return self._model.job_id + async def to_file( self, path: str | PathLike[str], *, range: tuple[int, int] | None = None ) -> Path: diff --git a/tests/conftest.py b/tests/conftest.py index 072bdd2..65b2793 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -99,6 +99,9 @@ def _asset_json(asset_id: str, hash_: str, created_new: bool, size: int) -> dict "created_at": "2026-07-10T18:00:00Z", "url": "http://example.invalid/blob", "url_expires_at": "2026-07-10T19:00:00Z", + # Retention deadline for the asset itself — distinct from url_expires_at + # above. Omitted (-> None) for a plain upload with no producing job. + "expires_at": "2026-08-09T18:00:00Z", } @@ -268,7 +271,9 @@ def _serve_job(self, job_id: str) -> None: if state.job_poll_count >= state.polls_to_succeed: status = state.terminal_status if status == "succeeded": - outputs = state.job_outputs if state.job_outputs is not None else [_OUTPUT] + raw = state.job_outputs if state.job_outputs is not None else [_OUTPUT] + # Stamp the producing job id, same as a real server would. + outputs = [{**o, "job_id": job_id} for o in raw] else: outputs = [] else: diff --git a/tests/test_assets.py b/tests/test_assets.py index d90946b..f975bd9 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -3,6 +3,7 @@ from __future__ import annotations import io +from datetime import datetime, timezone import pytest @@ -30,6 +31,28 @@ def test_dedup_fast_path_skips_upload(server, tmp_path) -> None: assert server.state.upload_count == 0 +def test_uploaded_asset_has_no_producing_job(server, tmp_path) -> None: + # A plain upload (never a job output) must report job_id as None on the + # public wrapper, not just on the private generated model. + p = tmp_path / "photo.png" + p.write_bytes(b"plain-upload-bytes") + + with Comfy() as client: + asset = client.assets.from_file(p) + asset.commit() + assert asset.job_id is None + + +def test_committed_asset_exposes_expiry(server, tmp_path) -> None: + p = tmp_path / "photo.png" + p.write_bytes(b"expiring-bytes") + + with Comfy() as client: + asset = client.assets.from_file(p) + asset.commit() + assert asset.expires_at == datetime(2026, 8, 9, 18, 0, tzinfo=timezone.utc) + + class _ReadRecorder(io.BytesIO): """A file object that records the largest single read() it served.""" diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 6c7c567..600e73f 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -37,6 +37,15 @@ def test_run_completes_via_polling_when_sse_absent(server, tmp_path) -> None: assert server.state.events_connect_count == 0 # SSE never used +def test_output_exposes_producing_job_id(server) -> None: + # The public Output wrapper must surface job_id, not just the private + # generated model — this is the field the whole feature exists to provide. + with Comfy() as client: + job = client.run(_wf(client)) + outs = job.get_outputs("13") + assert outs[0].job_id == job.id + + def test_idempotent_submit_rejects_reused_key(server) -> None: # Keys are single-use (reject-on-duplicate, no replay): reusing the same # explicit key raises IdempotencyKeyReuse rather than replaying the job.