From ee2155d984d9cc946e54382bf5ae252ad22e9d2d Mon Sep 17 00:00:00 2001 From: Wei Hai Date: Wed, 12 Aug 2026 20:13:02 -0700 Subject: [PATCH] feat: add get_workflow() to fetch the graph that produced a job MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A caller holding a job — including one rehydrated by id — had no way to see the workflow behind it. The SDK holds the graph only when it submitted the job in the same process; otherwise it is gone. Calls GET /api/v2/jobs/{id}/workflow and returns both the graph and the format discriminator. "api" is the executed graph, with editor-only constructs already resolved away; "save" is the authoring workflow at the version the job ran. Returning the discriminator matters: a caller must be able to tell which shape it holds. Written by hand rather than through the generated client, because the vendored spec does not describe this endpoint yet. It moves onto the generated client once the spec re-syncs. Named to match the existing get_download_url(), and mirrored as getWorkflow() in the TypeScript SDK. Co-Authored-By: Claude Opus 5 (1M context) --- src/comfy_low/transport.py | 26 +++++++++++++ src/comfy_sdk/__init__.py | 3 +- src/comfy_sdk/jobs.py | 45 ++++++++++++++++++++- tests/conftest.py | 20 ++++++++++ tests/test_job_workflow.py | 80 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 172 insertions(+), 2 deletions(-) create mode 100644 tests/test_job_workflow.py diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index 207f6bf..fae271d 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -462,6 +462,22 @@ def cancel_job(self, job_id_or_url: str, *, timeout: Any = _UNSET) -> Job: resp = self.raw_request("POST", path, timeout=timeout) return Job.model_validate(self._p.parse_or_raise(resp, (200,))) + def get_job_workflow(self, job_id_or_url: str, *, timeout: Any = _UNSET) -> dict[str, Any]: + """GET /api/v2/jobs/{id}/workflow. + + Hand-written: this operation is not yet in spec/openapi.yaml — the + server endpoint is still in review, landing separately. There is no + generated model to validate against yet, so this returns the raw, + parsed envelope (``{"workflow": {...}, "format": "save" | "api"}``) + instead of a typed model; move this onto the generated client once the + spec re-syncs. + """ + path = ( + job_id_or_url if _looks_like_path(job_id_or_url) else f"/jobs/{job_id_or_url}/workflow" + ) + resp = self.raw_request("GET", path, timeout=timeout) + return self._p.parse_or_raise(resp, (200,)) + class AsyncComfyLow: """Asynchronous protocol bindings — mirrors :class:`ComfyLow`.""" @@ -709,6 +725,16 @@ async def cancel_job(self, job_id_or_url: str, *, timeout: Any = _UNSET) -> Job: resp = await self.raw_request("POST", path, timeout=timeout) return Job.model_validate(self._p.parse_or_raise(resp, (200,))) + async def get_job_workflow( + self, job_id_or_url: str, *, timeout: Any = _UNSET + ) -> dict[str, Any]: + """See the sync ``get_job_workflow`` for why this returns a raw dict.""" + path = ( + job_id_or_url if _looks_like_path(job_id_or_url) else f"/jobs/{job_id_or_url}/workflow" + ) + resp = await self.raw_request("GET", path, timeout=timeout) + return self._p.parse_or_raise(resp, (200,)) + def _looks_like_path(s: str) -> bool: return s.startswith("http") or s.startswith("/") diff --git a/src/comfy_sdk/__init__.py b/src/comfy_sdk/__init__.py index 399fe99..69fbd51 100644 --- a/src/comfy_sdk/__init__.py +++ b/src/comfy_sdk/__init__.py @@ -55,7 +55,7 @@ Unauthorized, WorkflowFormatUi, ) -from .jobs import AsyncJob, Job +from .jobs import AsyncJob, Job, JobWorkflow from .outputs import AsyncOutput, DownloadUrl, Output from .workflows import Workflow, WorkflowFactory @@ -82,6 +82,7 @@ "WorkflowFactory", "Job", "AsyncJob", + "JobWorkflow", "Output", "AsyncOutput", "DownloadUrl", diff --git a/src/comfy_sdk/jobs.py b/src/comfy_sdk/jobs.py index 59872d9..2535c5d 100644 --- a/src/comfy_sdk/jobs.py +++ b/src/comfy_sdk/jobs.py @@ -12,7 +12,8 @@ import time from collections.abc import AsyncIterator, Iterator -from typing import Any +from dataclasses import dataclass +from typing import Any, Literal import httpx @@ -29,6 +30,24 @@ _RECONNECT_PAUSE = 0.1 +@dataclass(frozen=True) +class JobWorkflow: + """The workflow that produced a job — see :meth:`Job.get_workflow`. + + ``format`` discriminates the shape of ``graph``, so a caller can tell + which one it holds instead of the two shapes being silently collapsed: + + * ``"api"`` — the executed graph. API-format, so frontend-only constructs + (Note nodes, Get/Set) are already resolved away. + * ``"save"`` — the authoring workflow at the version the job ran, + un-mangled. Only occurs for a job that pins a workflow version; a job + submitted through this SDK today always gets ``"api"``. + """ + + graph: dict[str, Any] + format: Literal["save", "api"] + + class Job: """Synchronous job handle.""" @@ -108,6 +127,24 @@ def cancel(self) -> Job: self._model = self._low.cancel_job(self._model.urls.cancel or self._model.id) return self + def get_workflow(self) -> JobWorkflow: + """Fetch the workflow that produced this job. + + Needed for a job rehydrated purely by id (e.g. via + ``client.jobs.get``) — the SDK only holds the workflow it submitted + for as long as the same process's :class:`Job` handle is alive, so + this is the only way to see the graph otherwise. + + Calls ``GET /api/v2/jobs/{id}/workflow`` via + :meth:`comfy_low.transport.ComfyLow.get_job_workflow`, which is + hand-written because the operation is not yet in the vendored spec — + the server endpoint is still in review. A missing job raises the + SDK's normal :class:`~comfy_sdk.exceptions.NotFound`. + """ + with translating(): + data = self._low.get_job_workflow(self._model.id) + return JobWorkflow(graph=data["workflow"], format=data["format"]) + # -- live events (best-effort, reconnecting) -------------------------- def events(self) -> Iterator[Event]: """Typed live event iterator. Auto-reconnects with no replay; falls back @@ -217,6 +254,12 @@ async def cancel(self) -> AsyncJob: self._model = await self._low.cancel_job(self._model.urls.cancel or self._model.id) return self + async def get_workflow(self) -> JobWorkflow: + """Async :meth:`Job.get_workflow`.""" + with translating(): + data = await self._low.get_job_workflow(self._model.id) + return JobWorkflow(graph=data["workflow"], format=data["format"]) + async def events(self) -> AsyncIterator[Event]: """Async :meth:`Job.events` — typed live stream, auto-reconnecting with no replay and the poll path as its backstop. diff --git a/tests/conftest.py b/tests/conftest.py index 4f547e8..95f5de5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -55,6 +55,13 @@ class ServerState: # Set to a list of output dicts (see `_output_json`) to test a job with # multiple outputs, each backed by a distinct asset id. job_outputs: list[dict] | None = None + # GET /jobs/{id}/workflow response. `job_workflow_not_found=True` answers + # 404 job_not_found instead, for the missing-job path. + job_workflow_graph: dict[str, Any] = field( + default_factory=lambda: {"3": {"class_type": "KSampler", "inputs": {}}} + ) + job_workflow_format: str = "api" + job_workflow_not_found: bool = False # --- counters the tests assert on --- upload_count: int = 0 @@ -197,6 +204,10 @@ def do_GET(self) -> None: if m: self._serve_events(m.group(1)) return + m = re.match(r"/api/v2/jobs/([^/]+)/workflow$", self.path) + if m: + self._serve_job_workflow(m.group(1)) + return m = re.match(r"/api/v2/jobs/([^/]+)$", self.path) if m: self._serve_job(m.group(1)) @@ -243,6 +254,15 @@ def _serve_job(self, job_id: str) -> None: outputs = [] self._json(200, _job_json(job_id, status, outputs)) + def _serve_job_workflow(self, job_id: str) -> None: + if state.job_workflow_not_found: + self._err(404, "job_not_found", "no such job") + return + self._json( + 200, + {"workflow": state.job_workflow_graph, "format": state.job_workflow_format}, + ) + def _serve_events(self, job_id: str) -> None: state.events_connect_count += 1 if state.events_not_implemented: diff --git a/tests/test_job_workflow.py b/tests/test_job_workflow.py new file mode 100644 index 0000000..bd9d33c --- /dev/null +++ b/tests/test_job_workflow.py @@ -0,0 +1,80 @@ +"""Job.get_workflow() / AsyncJob.get_workflow() — GET /api/v2/jobs/{id}/workflow. + +This endpoint is not yet in spec/openapi.yaml (the server side is still in +review), so the stub server in conftest.py stands in for it directly rather +than the SDK talking to a generated model. +""" + +from __future__ import annotations + +import pytest + +from comfy_sdk import AsyncComfy, Comfy, JobWorkflow, NotFound + + +def _wf(client: Comfy | AsyncComfy): + return client.workflows.from_json({"3": {"class_type": "KSampler", "inputs": {}}}) + + +def test_get_workflow_returns_api_format(server) -> None: + server.state.job_workflow_format = "api" + server.state.job_workflow_graph = {"3": {"class_type": "KSampler", "inputs": {}}} + with Comfy() as client: + job = client.submit(_wf(client)) + wf = job.get_workflow() + assert isinstance(wf, JobWorkflow) + assert wf.format == "api" + assert wf.graph == {"3": {"class_type": "KSampler", "inputs": {}}} + + +def test_get_workflow_returns_save_format(server) -> None: + # "save" is the authoring workflow at the pinned version, un-mangled — + # a different shape than the executed "api" graph, so both must be + # reachable, not collapsed into one. + server.state.job_workflow_format = "save" + server.state.job_workflow_graph = { + "nodes": [{"id": 3, "type": "KSampler"}], + "links": [], + "last_node_id": 3, + } + with Comfy() as client: + job = client.submit(_wf(client)) + wf = job.get_workflow() + assert wf.format == "save" + assert wf.graph["nodes"][0]["type"] == "KSampler" + + +def test_get_workflow_on_job_rehydrated_by_id(server) -> None: + # The motivating case: a job the SDK did not submit in this process (e.g. + # rehydrated purely by id via client.jobs.get) still exposes its workflow. + with Comfy() as client: + submitted = client.submit(_wf(client)) + rehydrated = client.jobs.get(submitted.id) + wf = rehydrated.get_workflow() + assert wf.format == "api" + + +def test_get_workflow_not_found_raises_sdk_not_found(server) -> None: + server.state.job_workflow_not_found = True + with Comfy() as client: + job = client.submit(_wf(client)) + with pytest.raises(NotFound): + job.get_workflow() + + +async def test_async_get_workflow_mirrors_sync(server) -> None: + server.state.job_workflow_format = "save" + server.state.job_workflow_graph = {"nodes": [], "links": [], "last_node_id": 0} + async with AsyncComfy() as client: + job = await client.submit(_wf(client)) + wf = await job.get_workflow() + assert wf.format == "save" + assert wf.graph == {"nodes": [], "links": [], "last_node_id": 0} + + +async def test_async_get_workflow_not_found_raises_sdk_not_found(server) -> None: + server.state.job_workflow_not_found = True + async with AsyncComfy() as client: + job = await client.submit(_wf(client)) + with pytest.raises(NotFound): + await job.get_workflow()