feat: add get_workflow() to fetch the graph that produced a job

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("/")

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",

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.

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:

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()