fix: ComfyUI timeout/resume + websocket wait + multi-server + music tool

.agents/skills/comfyui/SKILL.md

@@ -1,17 +1,19 @@
 ---
 name: comfyui
-description: Use when working with ComfyUI workflows in OpenMontage, including comfyui_image/comfyui_video, custom workflow_json/workflow_path inputs, output_node selection, missing model setup, LoRAs, low-VRAM workflow choices, and community workflow imports.
+description: Use when working with ComfyUI workflows in OpenMontage, including comfyui_image/comfyui_video/comfyui_music, custom workflow_json/workflow_path inputs, output_node selection, missing model setup, LoRAs, low-VRAM workflow choices, and community workflow imports.
 ---
 
 # ComfyUI Workflows in OpenMontage
 
-Use this skill before calling `comfyui_image` or `comfyui_video`, and when converting a community ComfyUI workflow into an OpenMontage tool call.
+Use this skill before calling `comfyui_image`, `comfyui_video`, or `comfyui_music`, and when converting a community ComfyUI workflow into an OpenMontage tool call.
 
 ## Server Contract
 
 - ComfyUI must be running before the tool can generate. The default server is `http://localhost:8188`; override it with `COMFYUI_SERVER_URL`.
+- Running separate ComfyUI instances per capability (different GPU, different model set)? `COMFYUI_IMAGE_SERVER_URL` / `COMFYUI_VIDEO_SERVER_URL` / `COMFYUI_MUSIC_SERVER_URL` each override `COMFYUI_SERVER_URL` for that one tool only. Optional -- a single-server setup needs none of these.
 - Health and hardware status come from `GET /system_stats`.
 - Jobs are submitted to `POST /prompt`, completed outputs are read from `GET /history/{prompt_id}`, and artifact bytes are downloaded with `GET /view`.
+- Long waits (video, music) prefer ComfyUI's websocket feed for immediate completion/error detection and transparently fall back to REST polling if `websocket-client` isn't installed. Either way, a timeout is recoverable: pass the error's `prompt_id` back in as `resume_prompt_id` to resume waiting on the same job instead of resubmitting it.
 - Export workflows with ComfyUI's API-format JSON, not the UI layout format. If a downloaded workflow will not submit, re-export it from ComfyUI with API format enabled.
 
 ## Choosing a Workflow
@@ -52,4 +54,13 @@ Use this skill before calling `comfyui_image` or `comfyui_video`, and when conve
 - If the server is unavailable, surface the structured setup offer. Starting ComfyUI or setting `COMFYUI_SERVER_URL` is the first fix.
 - If models are missing, read `data.missing_models[]`; each item should include the file name, role, destination hint, and download URL when OpenMontage knows it.
 - If custom nodes are missing, ask the user to install them through ComfyUI Manager or the workflow author's documented install path, then restart ComfyUI.
-- If a long render times out locally, check ComfyUI history before retrying from scratch; the server may still have completed the prompt.
+- If a long render times out locally, check ComfyUI history before retrying from scratch; the server may still have completed the prompt -- or just call again with `resume_prompt_id` set to the `prompt_id` from the timeout error.
+
+## Music (`comfyui_music`)
+
+- Bundled default is ACE-Step v1 (3.5B) text-to-audio, built from ComfyUI's *native* `TextEncodeAceStepAudio`/`EmptyAceStepLatentAudio` nodes (core, not a third-party pack) -- unlike ACE-Step 1.5 or other custom node packs, v1's interface is standardized enough to bundle safely.
+- `prompt` maps to the bundled workflow's `tags` field (style/genre/mood, e.g. `"upbeat electronic pop, female vocals"`), matching the same "prompt = music description" convention `suno_music` uses. `lyrics` is a separate optional field -- leave empty for instrumental, or use `[verse]`/`[chorus]`/`[bridge]` structure tags and `[zh]`/`[ja]`/`[ko]`-style language-code prefixes for non-English lines.
+- `duration_seconds`, `steps`, `cfg`, `lyrics_strength`, and `seed` are patchable on the bundled workflow. Missing `ace_step_v1_3.5b.safetensors` surfaces through the same `data.missing_models[]` contract as image/video.
+- Need ACE-Step 1.5, a different node pack, or a non-ACE-Step audio model? Fall back to `workflow_json`/`workflow_path` + `output_node`, exactly like a custom image/video workflow -- in that mode `prompt` becomes provenance/logging only again and must already be baked into the graph.
+- `output_node` (bundled or custom) should be the node that writes the final audio -- the bundled workflow's is `SaveAudioMP3`. The client reads artifacts from that node's `"audio"` output key (parallel to `"images"` for image/video savers).
+- For custom workflows, provide `workflow_name`/`workflow_model`/`workflow_model_stack` for provenance exactly as you would for a custom image/video workflow.

docs/comfyui-adapter-plan.md

@@ -257,21 +257,36 @@ output_node:          string    # required for custom workflows
 workflow_name:        string    # optional custom workflow provenance label
 workflow_model:       string    # optional custom model/provenance label
 workflow_model_stack: []        # optional custom dependency provenance
+timeout_seconds:      integer   # optional, default 3600 (see below)
+resume_prompt_id:     string    # optional, resume a timed-out job without resubmitting
 ```
 
 **execute() flow (i2v):**
 1. Upload reference image via `client.upload_image()`
 2. Deep-copy i2v workflow template
 3. Inject prompt, uploaded image name, seed, dimensions
-4. `client.generate(workflow, output_node="108", dest=output_path, timeout=900)`
+4. `client.generate(workflow, output_node="108", dest=output_path, timeout=inputs.get("timeout_seconds", 3600), resume_prompt_id=inputs.get("resume_prompt_id"))`
 5. Return `ToolResult`
 
 **execute() flow (t2v):**
 1. Deep-copy t2v workflow template
 2. Inject prompt, seed, dimensions
-3. `client.generate(workflow, output_node="16", dest=output_path, timeout=900)`
+3. `client.generate(workflow, output_node="16", dest=output_path, timeout=inputs.get("timeout_seconds", 3600), resume_prompt_id=inputs.get("resume_prompt_id"))`
 4. Return `ToolResult`
 
+**Timeout and resume (added after real-world local-GPU testing):** the
+default client wait was raised from 900s to 3600s — non-accelerated custom
+Wan 1.3B workflows on modest local GPUs were observed taking ~1360-1630s at
+832x480/81-97 frames, and the old 900s default false-failed those jobs even
+though ComfyUI kept rendering server-side. `ComfyUIError` now carries a
+`prompt_id` on both execution errors and timeouts (`ComfyUIError.prompt_id`),
+and `ComfyUIVideo`'s `ToolResult.error`/`.data` surface it on timeout so the
+caller isn't left guessing whether the job is dead. Callers recover a
+timed-out-but-still-running job by calling `execute()` again with
+`resume_prompt_id` set to that `prompt_id` (and a longer `timeout_seconds` if
+needed) — `client.generate()` then skips `submit()` entirely and just resumes
+polling/downloading the existing job instead of queuing a duplicate.
+
 `comfyui_video` publishes `operation_statuses` in `get_info()` and implements
 `is_operation_available(operation)` for selector routing. This keeps partial
 ComfyUI installs useful for the installed mode without advertising unavailable
@@ -281,19 +296,52 @@ not promote ComfyUI for an operation whose bundled models are missing.
 
 ---
 
-### `comfyui_music` -- Music Generation (not shipped)
-
-We explored adding a `comfyui_music` tool using the ACE-Step 3.5B model.
-The model runs well in ComfyUI, but the ComfyUI node interface for
-ACE-Step is not standardized -- there are multiple custom node packs with
-different class names (`AceStepModelLoader` vs native `TextEncodeAceStepAudio`,
-etc.).  Shipping a workflow that only works with one specific custom node
-pack would break for most users.
-
-**Future path:** ACE-Step support should be revisited once OpenMontage decides
-the music-generation routing shape and a portable ComfyUI audio workflow
-contract. Current image/video workflow overrides are intentionally scoped to
-image and video artifacts, not arbitrary audio workflows.
+### `comfyui_music` -- Music Generation (shipped, with a native-node bundled workflow)
+
+`tools/audio/comfyui_music.py`. `capability="music_generation"`, `provider="comfyui"`.
+
+**Bundled default:** ACE-Step v1 (3.5B) text-to-audio, via `tools/_comfyui/workflows/ace-step-1-t2a.json`.
+The node-pack fragmentation that originally blocked this tool (`AceStepModelLoader`
+vs native `TextEncodeAceStepAudio`, etc.) turned out to be moot for ACE-Step v1:
+ComfyUI ships `TextEncodeAceStepAudio`/`EmptyAceStepLatentAudio` as **native core
+nodes** (`comfy_extras/nodes_ace.py`), not a third-party pack, and Comfy-Org's own
+[`workflow_templates`](https://github.com/Comfy-Org/workflow_templates) repo bundles
+an official ACE-Step-v1 template built entirely from those native nodes plus
+long-stable core nodes (`CheckpointLoaderSimple`, `KSampler`, `ModelSamplingSD3`,
+`VAEDecodeAudio`, `SaveAudioMP3`). Every node's `class_type` and input names in
+`ace-step-1-t2a.json` were cross-checked against ComfyUI's own source
+(`comfy_extras/nodes_ace.py`, `nodes_audio.py`, `nodes_latent.py`, `nodes.py`) --
+not guessed from the UI export -- since the UI-format template Comfy-Org ships
+isn't directly usable as the API-format JSON this client submits.
+
+`prompt` maps to ACE-Step's `tags` field (style/genre/mood description, matching
+the "prompt = description of desired music" convention `suno_music` already uses).
+`lyrics` is a separate optional field (empty for instrumental). `duration_seconds`,
+`steps`, `cfg`, `lyrics_strength`, and `seed` are all patchable; `shift` and the
+tonemap `multiplier` stay at the official template's defaults.
+
+Newer/different setups aren't locked out: `workflow_json`/`workflow_path` +
+`output_node` still works exactly like the image/video tools' override path --
+for ACE-Step 1.5, a different node pack, or a non-ACE-Step audio model entirely.
+
+**Selector integration:** no dedicated `music_selector` exists in OpenMontage
+(unlike `tts_selector`/`image_selector`/`video_selector`) -- music tools are
+already routed directly via `registry.get_by_capability("music_generation")`,
+and `comfyui_music` participates in that the same way `suno_music`/`music_gen`
+do. `fallback_tools = ["suno_music", "music_gen"]`.
+
+**Audio artifact schema:** `ToolResult.data` follows the same shape as the
+image/video tools (`provider`, `model`, `output`, `format`, `workflow_provenance`),
+plus `lyrics` and `duration_seconds` -- the latter a best-effort `ffprobe` probe
+of the downloaded file (`None` if `ffprobe` isn't on PATH), since even the bundled
+workflow doesn't report actual rendered duration back through `/history`.
+
+**Workflow/output-node contract:** identical to image/video -- `output_node`
+must be the ID of the node that writes the final artifact (the bundled workflow's
+is `SaveAudioMP3`, ComfyUI's native audio saver). `ComfyUIClient.generate()`'s
+artifact extraction now also checks the `"audio"` output key (previously only
+`"images"`/`"gifs"`), which is what `SaveAudioMP3`/`SaveAudio` write to in
+ComfyUI's `/history` response.
 
 ---
 
@@ -358,6 +406,19 @@ COMFYUI_POLL_TIMEOUT=600                    # max wait for image gen
 COMFYUI_VIDEO_TIMEOUT=900                   # max wait for video gen
 ```
 
+**Multi-server (optional):** point `comfyui_image`, `comfyui_video`, and
+`comfyui_music` at separate ComfyUI instances -- e.g. one GPU running FLUX 2,
+another running WAN 2.2, another running ACE-Step -- by setting a
+per-capability override. Each takes priority over `COMFYUI_SERVER_URL` for
+its own tool only; leave all three unset and everything talks to the single
+shared server.
+
+```bash
+COMFYUI_IMAGE_SERVER_URL=http://gpu-a:8188
+COMFYUI_VIDEO_SERVER_URL=http://gpu-b:8188
+COMFYUI_MUSIC_SERVER_URL=http://gpu-c:8188
+```
+
 **For Docker Compose setups** (ComfyUI in a container):
 
 ```bash
@@ -459,15 +520,35 @@ pipeline definition, or any schema.
    user-provided via a config directory? Bundling gives reproducibility;
    external gives flexibility.
 
-2. **Async generation:** ComfyUI supports websocket connections for real-time
-   progress. Worth implementing for long video generations, or is polling
-   sufficient?
-
-3. **Multi-server:** Should the adapter support multiple ComfyUI instances
-   (e.g., one for images, one for video) via per-capability URLs?
-
-4. **Music generation:** ACE-Step works in ComfyUI but OpenMontage needs a
-   dedicated music-generation routing contract before adding `comfyui_music`.
-   The follow-up should decide selector integration, audio artifact schemas, and
-   a portable workflow/output-node contract rather than treating music as a
-   hidden image/video workflow override.
+2. ~~**Async generation:**~~ **Resolved.** `ComfyUIClient.generate()` now
+   waits via ComfyUI's websocket feed (`wait_ws()`) by default, reacting to
+   `executing`/`execution_error` events immediately instead of sleeping
+   between REST polls — completion and errors are caught without the
+   `interval`-seconds lag, and an optional `on_progress` callback gets live
+   `progress` events (`comfyui_video` uses this to print step progress on
+   long renders). No new hard dependency: `websocket-client` is an optional
+   import, and `_wait()` transparently falls back to the original
+   `poll()` REST loop when it isn't installed or the connection fails —
+   `resume_prompt_id` recovery behaves identically either way.
+
+3. ~~**Multi-server:**~~ **Resolved.** `ComfyUIClient(capability="image"|"video"|"music")`
+   resolves its server URL from a per-capability env var first
+   (`COMFYUI_IMAGE_SERVER_URL` / `COMFYUI_VIDEO_SERVER_URL` / `COMFYUI_MUSIC_SERVER_URL`),
+   then the shared `COMFYUI_SERVER_URL`, then the `http://localhost:8188` default.
+   All three tools pass their capability at construction, so image, video, and
+   music generation can each point at different ComfyUI instances (different GPUs,
+   different model sets) with zero code changes -- single-server setups need no extra
+   configuration since all three env vars are optional. `client.capability`/
+   `client.is_default_url`/`client.unavailable_reason()` all account for the
+   override, and `COMFYUI_SETUP_OFFER.per_capability_env_var_overrides` documents
+   it for the setup-offer surfacing in `provider_menu()`.
+
+4. ~~**Music generation:**~~ **Resolved -- shipped with a bundled ACE-Step v1 workflow.**
+   `comfyui_music` is a real tool now (not a hidden image/video override), routed
+   through the existing `registry.get_by_capability("music_generation")` path
+   like `suno_music`/`music_gen`. The node-pack fragmentation that originally
+   blocked this turned out not to apply to ACE-Step v1: its ComfyUI nodes are
+   native core nodes, not a third-party pack, so `ace-step-1-t2a.json` ships as
+   the default, verified node-by-node against ComfyUI's own source. Custom
+   `workflow_json`/`workflow_path` + `output_node` remains available for other
+   versions/packs. See the `comfyui_music` section above for the full contract.