feat: add embed_workflow to embed the graph in output metadata

[wei-hai]

Recovering the workflow from an image it produced isn't possible today, which makes debugging harder. Adds an opt-in flag.

flowchart LR
  W["workflow"] --> M["materialize<br/>(assets → core/ASSET)"]
  M --> G["graph"]
  G -- "workflow" --> R["POST /jobs"]
  G -- "extra_data.extra_pnginfo.workflow<br/>(only if embed_workflow)" --> R
  R --> P["PNG: workflow chunk"]

Loading

submit() and run() — sync and async — take a keyword-only embed_workflow: bool = False. When set, the SDK sends the materialized graph as extra_data.extra_pnginfo.workflow, the key local ComfyUI's SaveImage writes into output PNG metadata.

Behavior

• Off by default. With no api_key and the flag off, the wire format is byte-identical to today: no extra_data at all, never {}. Existing tests pinning this are untouched.
• Merges with the partner-node key when both are given.
• Embeds the materialized graph — the same object sent as workflow — not the caller's input. A test asserts the core/ASSET substitution actually happened, so it can't pass by embedding the wrong object.

Reviewer notes

• Needs the matching server-side contract change, landing separately — don't merge this first, or the flag is a no-op.
• spec/openapi.yaml is untouched by design (synced one-way); it picks the field up on the next sync.
• The TypeScript SDK ships the same feature as embedWorkflow, producing identical bytes.

ruff, mypy, pytest (133 passed, 4 pre-existing skips), spec-drift and hygiene checks all clean.

🤖 Generated with Claude Code

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 543e0786-9bc6-4497-938a-b6bf87fff80c

📥 Commits

Reviewing files that changed from the base of the PR and between 8bfbd3c and 8983b26.

📒 Files selected for processing (5)

• README.md
• src/comfy_sdk/_core.py
• src/comfy_sdk/client.py
• tests/test_async.py
• tests/test_jobs.py

---

📝 Walkthrough

Walkthrough

submit() and run() now support optional embedding of materialized workflows in extra_data.extra_pnginfo.workflow for synchronous and asynchronous clients. API-key metadata remains supported, and the default remains disabled.

Changes

Workflow embedding

Layer / File(s)	Summary
Extra data metadata contract src/comfy_sdk/_core.py	extra_data_for() accepts an optional workflow graph and stores it under extra_pnginfo.workflow. It returns None when no API key or workflow is provided.
Synchronous and asynchronous submission flow src/comfy_sdk/client.py	submit() and run() accept embed_workflow, default it to False, and pass the materialized workflow to extra_data_for() when enabled.
Behavior validation and documentation tests/test_jobs.py, tests/test_async.py, README.md	Tests cover embedding, API-key merging, materialized asset handles, disabled behavior, and run() forwarding. The README documents the metadata format and option behavior.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant extra_data_for
  participant ComfyAPI
  Client->>extra_data_for: Build metadata with materialized workflow
  extra_data_for->>ComfyAPI: Submit extra_pnginfo.workflow and optional API key
  ComfyAPI-->>Client: Return submitted job

Loading

Suggested reviewers: alexisrolland

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the main change: adding opt-in workflow embedding in output metadata.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch wei/be-7128-embed-workflow-metadata

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[wei-hai]

Found one blocking workflow-format issue.

[wei-hai]

Blocking: graph is API-format, but extra_pnginfo.workflow is consumed by ComfyUI as the editor/save-format graph. A generated PNG already gets the API graph in its prompt chunk from SaveImage; adding this second, API-shaped workflow chunk makes the frontend prioritize it, call loadGraphData() (which expects version/nodes), and return before falling back to prompt, yielding an empty/invalid canvas. Please rely on the existing prompt metadata for API-format recovery, or accept a real UI-format workflow separately / add explicit frontend handling instead of putting graph in this slot.

[wei-hai]

Closing — this approach turns out to be wrong, and the reason is worth recording.

extra_pnginfo.workflow is the save-format slot, not the API-format one. ComfyUI's SaveImage already writes the API-format graph into a separate prompt chunk on every execution, and on import the frontend takes the workflow chunk if present, calls loadGraphData() (which expects version/nodes), and returns before reaching the prompt fallback. So writing an API-format graph there loads a broken graph and shadows the one that already works — a regression rather than a no-op.

Beyond the mechanics: user feedback reports the workflow recovered from an image is "too mangled to be useful" — Note nodes dropped, Get/Set expanded — because those are editor-only constructs already resolved away by the time the SDK holds the graph. No embedding option recovers them. The real ask is an API that links an output file back to the full original workflow, and one that works for audio, video, 3D and text outputs, which image metadata cannot serve at all.

Replacement work will go in that direction. Nothing here is lost — the review findings stay in this thread.