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.