Stop button also stops Upscale and Interpolate jobs

Jobs end as stopped (kept for the poll, no half file, no restarted message); Stop shows while one runs.
Python and JS tests for the cancel path.

app/ui/features/generate/index.js

@@ -2,6 +2,7 @@
 import { composer as c } from "../../composer/composer.js";
 import { isGenerative, isVideoClip, genUnitId, unitRoot, effFrames, effFps } from "../../shell/scenes.js";
 import { QUEUED, RUNNING } from "../../shell/run.js";
+import { jobsRunning, onJobs, stopJobs } from "../../shell/renders.js";
 import { inPlace } from "../../shell/place.js";
 import { buildInputs, rolesAt } from "./inputs.js";
 import { offer } from "../../shell/actions.js";
@@ -201,15 +202,18 @@ export default {
     let drawn = "";
     function draw() {
       const open = p.project;
-      const picked = open ? pickedScenes().length : 0, key = `${busy}|${held}|${g.state().phase}|${Boolean(open)}|${picked}`;
+      const picked = open ? pickedScenes().length : 0, key = `${busy}|${held}|${g.state().phase}|${Boolean(open)}|${picked}|${jobsRunning()}`;
       if (key === drawn) return;          // progress ticks must not rebuild the buttons
       drawn = key;
       const working = busy || held || g.state().phase === QUEUED || g.state().phase === RUNNING;     // also a run this page did not start
       const go = c.button.sm({ label: working ? "Generating…" : "▶ Generate", tone: "primary", disabled: working || !open, title: "Generate every scene on the timeline, except left-out scenes, video clips and Auto Montage pieces. An earlier render of a scene stays as a take.",
         onClick: () => runUnits(unitsOf(p.scenes)) });
       const one = c.button.sm({ label: picked > 1 ? `Selected (${picked})` : "Selected", title: "Generate only the picked scenes", disabled: working || !picked, onClick: () => runUnits(unitsOf(pickedScenes())) });
-      const stop = c.button.sm({ label: "■ Stop", tone: "danger", title: "Stop the run that is going", onClick: () => { stopped = true; g.cancel(); } });
-      stop.node.hidden = !working;
+      const stop = c.button.sm({ label: "■ Stop", tone: "danger", title: "Stop the run that is going, and any Upscale or Interpolate in progress", onClick: () => {
+        stopJobs(app.api.stopJob);
+        if (working) { stopped = true; g.cancel(); }
+      } });
+      stop.node.hidden = !working && !jobsRunning();
       put(go.node, one.node, stop.node);
     }
     // A run found after a reload belongs to this page's last session: its result still goes on its clips.
@@ -244,11 +248,11 @@ export default {
       });
     });
     draw();
-    const offRun = g.subscribe(draw);
+    const offRun = g.subscribe(draw), offJobs = onJobs(draw);          // the Stop button also shows while an Upscale or Interpolate runs
     const offApp = app.on((what) => { if (what === "generate.scene" && p.selected) runUnits(unitsOf([p.selected]));
       else if (what === "generate.selected" && pickedScenes().length) runUnits(unitsOf(pickedScenes())); else draw(); });
     const offers = [offer(app, { id: "generate-all", label: "Generate", icon: "▶", run: () => { if (p.project) runUnits(unitsOf(p.scenes)); } }),
       offer(app, { id: "generate-selected", label: "Generate selected", icon: "▶", run: () => { if (pickedScenes().length) runUnits(unitsOf(pickedScenes())); } })];
-    return () => { offRun(); offApp(); offers.forEach((f) => f()); };
+    return () => { offRun(); offJobs(); offApp(); offers.forEach((f) => f()); };
   },
 };

app/ui/features/interpolate/index.js

@@ -48,7 +48,7 @@ export default {
         const swapped = await p.editFor(pid, (pr) => swap(pr, media, out) > 0);
         if (!swapped) c.toast.warn({ text: "Interpolated, but that render was replaced meanwhile, so nothing was swapped." });
         else c.toast.good({ text: `Interpolated: the render now plays at ${opts().multiplier} times the frames.` });
-      } catch (err) { c.toast.warn({ text: `Interpolate failed: ${err.message}` }); }
+      } catch (err) { c.toast[err.stopped ? "info" : "warn"]({ text: err.stopped ? "Interpolate stopped. The render stays as it was." : `Interpolate failed: ${err.message}` }); }
       finally { busy.delete(key); draw(); }
     }
 

app/ui/features/interpolate/job.test.js

@@ -1,7 +1,7 @@
 import test from "node:test";
 import assert from "node:assert/strict";
 import { interpolate, modeLabel } from "./job.js";
-import { derived, gaveUp, onGaveUp, swap, wentThrough } from "../../shell/renders.js";
+import { derived, gaveUp, onGaveUp, swap, wentThrough, waitJob, jobsRunning, onJobs, stopJobs, JobStopped } from "../../shell/renders.js";
 
 const m = (f, sub = "") => ({ filename: f, subfolder: sub, type: "output" });
 
@@ -46,3 +46,20 @@ test("a step that gave up is announced, so a step waiting behind it can go on wi
   gaveUp(m("r.mp4"), "upscale");
   assert.deepEqual(heard, [["r.mp4", "upscale"]]);
 });
+
+test("Stop: a job being waited for can be stopped, ends as stopped (not failed, not restarted), and the Stop button hears about it", async () => {
+  const stopped = [], heard = [];
+  const off = onJobs(() => heard.push(jobsRunning()));
+  let polls = 0;
+  const read = async () => (++polls < 2 ? { state: "running" } : { state: "stopped" });
+  const waiting = waitJob(read, "j1", async () => { if (polls === 0) await stopJobs((id) => stopped.push(id)); });
+  await assert.rejects(waiting, (err) => err instanceof JobStopped && err.stopped === true && err.message === "Stopped.");
+  assert.deepEqual(stopped, ["j1"]);
+  assert.deepEqual(heard, [1, 0]);                 // shown while it waits, hidden when it ends
+  assert.equal(jobsRunning(), 0);
+  off();
+});
+
+test("Stop: a stopped job that is read once more is still there (no 404 turned into 'ComfyUI restarted')", async () => {
+  await assert.rejects(waitJob(async () => ({ state: "stopped" }), "j2", async () => {}), (err) => err.stopped === true);
+});

app/ui/features/upscale/index.js

@@ -53,7 +53,10 @@ export default {
         const swapped = await p.editFor(pid, (pr) => swap(pr, media, out) > 0);
         if (!swapped) c.toast.warn({ text: "Upscaled, but that render was replaced meanwhile, so nothing was swapped." });
         else c.toast.good({ text: "Upscaled: the render now plays the larger video." });
-      } catch (err) { gaveUp(media, "upscale"); c.toast.warn({ text: `Upscale failed: ${err.message}` }); }
+      } catch (err) {
+        if (err.stopped) c.toast.info({ text: "Upscale stopped. The render stays as it was." });
+        else { gaveUp(media, "upscale"); c.toast.warn({ text: `Upscale failed: ${err.message}` }); }
+      }
       finally { busy.delete(key); draw(); }
     }
 

app/ui/shell/api.js

@@ -138,6 +138,8 @@ export const api = {
   /** The upscale is a background job of FunUI, not a ComfyUI prompt (that would make ComfyUI drop the workflow's loaded models). -> {job_id}. */
   upscale: (media, model, rtx) => call("POST", "/api/m/render/upscale", { media: { filename: media.filename, subfolder: media.subfolder || "", type: media.type === "temp" ? "temp" : "output" }, model, ...(rtx ? { rtx_options: rtx } : {}) }),
   upscaleJob: (id) => call("GET", `/api/m/render/upscale/${encodeURIComponent(id)}`),
+  /** The Stop button for Upscale and Interpolate jobs (one job list serves both). -> {stopped} */
+  stopJob: (id) => call("POST", `/api/m/render/jobs/${encodeURIComponent(id)}/stop`),
   frameGeneration: () => call("GET", "/api/m/interpolate/frame_generation"),
   /** Interpolation is a background job of FunUI, not a ComfyUI prompt (that would make ComfyUI drop the workflow's loaded models). -> {job_id}. */
   interpolate: (media, options) => call("POST", "/api/m/interpolate/run", { media: { filename: media.filename, subfolder: media.subfolder || "", type: media.type === "temp" ? "temp" : "output" }, options }),

app/ui/shell/renders.js

@@ -2,21 +2,37 @@
 // scene that plays it, and remembering which steps a file has been through so two "always" modes can never feed each other forever. Pure of the page.
 export const keyOf = (m) => `${m.type || "output"}/${m.subfolder || ""}/${m.filename}`;
 
-/** Wait for a FunUI background job. `read(id)` -> {state: running | done | error, media, detail}. -> the media, or throws a sentence. `sleep` is injectable so a test need not wait. */
+const active = new Set();                                      // ids of the Upscale / Interpolate jobs this page is waiting for
+const listeners = [];
+const changed = () => { for (const fn of listeners) { try { fn(); } catch { /* one listener must not stop the others */ } } };
+/** How many jobs the Stop button can stop right now. */
+export const jobsRunning = () => active.size;
+/** -> a function that stops listening. */
+export const onJobs = (fn) => { listeners.push(fn); return () => { const i = listeners.indexOf(fn); if (i >= 0) listeners.splice(i, 1); }; };
+/** The Stop button: ask every job this page is waiting for to stop. `stop(id)` is the API call. */
+export const stopJobs = (stop) => Promise.all([...active].map((id) => Promise.resolve(stop(id)).catch(() => {})));
+/** What waitJob throws for a job that was stopped: not a failure, so callers say "stopped" and not "failed". */
+export class JobStopped extends Error { constructor() { super("Stopped."); this.stopped = true; } }
+
+/** Wait for a FunUI background job. `read(id)` -> {state: running | done | error | stopped, media, detail}. -> the media, or throws a sentence (a JobStopped when it was stopped). `sleep` is injectable so a test need not wait. */
 export async function waitJob(read, id, sleep = (ms) => new Promise((r) => setTimeout(r, ms))) {
   let lost = 0;
-  for (;;) {
-    await sleep(1000);
-    let r;
-    try { r = await read(id); } catch (err) {
-      if (err.status === 404) throw new Error("ComfyUI restarted while it was working");
-      if (++lost < 30) continue;                                    // a dropped reply is not a failed job
-      throw new Error("lost contact with ComfyUI");
+  active.add(id); changed();
+  try {
+    for (;;) {
+      await sleep(1000);
+      let r;
+      try { r = await read(id); } catch (err) {
+        if (err.status === 404) throw new Error("ComfyUI restarted while it was working");
+        if (++lost < 30) continue;                                    // a dropped reply is not a failed job
+        throw new Error("lost contact with ComfyUI");
+      }
+      lost = 0;
+      if (r.state === "stopped") throw new JobStopped();
+      if (r.state === "error") throw new Error(r.detail || "failed inside ComfyUI");
+      if (r.state === "done") return r.media;
     }
-    lost = 0;
-    if (r.state === "error") throw new Error(r.detail || "failed inside ComfyUI");
-    if (r.state === "done") return r.media;
-  }
+  } finally { active.delete(id); changed(); }
 }
 
 /** One file backs a whole chain run, so every scene render and ghost that plays `from` gets `to`. -> how many were swapped. */

core/jobs.py

@@ -8,12 +8,32 @@ from collections import OrderedDict
 MAX_JOBS = 50                         # finished jobs are kept this long for the poll that reads them
 
 _jobs: "OrderedDict[str, dict]" = OrderedDict()
+_stop: set = set()                    # ids asked to stop; the job notices between chunks
+
+
+class Stopped(Exception):
+    """Raised inside a job that was asked to stop; its partial output is removed by whoever was writing it."""
+
+
+def stop(job: str) -> bool:
+    """Ask a queued or running job to stop. False when there is no such job or it already ended."""
+    entry = _jobs.get(job)
+    if not entry or entry.get("state") not in ("queued", "running"):
+        return False
+    _stop.add(job)
+    return True
+
+
+def check(job: str) -> None:
+    """Raises Stopped once `stop(job)` was called: the hook a long job calls between its chunks."""
+    if job in _stop:
+        raise Stopped()
 
 
 def new() -> str:
     job = uuid.uuid4().hex
     _jobs[job] = {"state": "queued"}
-    for old in [k for k, v in _jobs.items() if v.get("state") in ("done", "error")]:
+    for old in [k for k, v in _jobs.items() if v.get("state") in ("done", "error", "stopped")]:
         if len(_jobs) <= MAX_JOBS:
             break
         del _jobs[old]                                 # only finished ones: a running job's poll must keep finding it
@@ -24,14 +44,22 @@ def get(job: str) -> dict | None:
     return _jobs.get(job)
 
 
-async def run(job: str, fn, *args, error) -> None:
+async def run(job: str, fn, *args, error, stoppable: bool = False) -> None:
     """Run the blocking `fn(*args)` off the event loop and record how it ended. `error` is the
     exception class whose message is meant for the person; anything else is reported plainly."""
+    if job in _stop:                                        # stopped while it was still queued
+        _stop.discard(job)
+        _jobs[job] = {"state": "stopped"}
+        return
     _jobs[job] = {"state": "running"}
     try:
-        result = await asyncio.to_thread(fn, *args)
+        result = await asyncio.to_thread(fn, *args, **({"interrupted": lambda: check(job)} if stoppable else {}))
         _jobs[job] = {"state": "done", **result}
+    except Stopped:
+        _jobs[job] = {"state": "stopped"}
     except error as exc:
         _jobs[job] = {"state": "error", "detail": str(exc)}
     except Exception as exc:                          # noqa: BLE001
         _jobs[job] = {"state": "error", "detail": f"{type(exc).__name__}: {exc}"}
+    finally:
+        _stop.discard(job)

modules/output/interpolate/__init__.py

@@ -43,7 +43,7 @@ def routes(table, base, web):
         except ValueError as exc:
             return bad(str(exc))
         job = jobs.new()
-        asyncio.create_task(jobs.run(job, nodes.interpolate, media, body.get("options") or {}, error=nodes.InterpolateError))
+        asyncio.create_task(jobs.run(job, nodes.interpolate, media, body.get("options") or {}, error=nodes.InterpolateError, stoppable=True))
         return web.json_response({"job_id": job})
 
     @table.get(base + "/jobs/{job}")

modules/output/render/__init__.py

@@ -102,9 +102,14 @@ def routes(table, base, web):
         if not isinstance(media, dict) or not isinstance(media.get("filename"), str) or not media["filename"] or not isinstance(model, str) or not model:
             return bad("Send {media: {filename, subfolder, type}, model, rtx_options}.")
         job = jobs.new()
-        asyncio.create_task(jobs.run(job, upscale.run, media, model, body.get("rtx_options") or "{}", error=(RuntimeError, ValueError, files.ClipError)))
+        asyncio.create_task(jobs.run(job, upscale.run, media, model, body.get("rtx_options") or "{}", error=(RuntimeError, ValueError, files.ClipError), stoppable=True))
         return web.json_response({"job_id": job})
 
+    @table.post(base + "/jobs/{job}/stop")
+    async def _stop_job(req):
+        """Stop an Upscale or Interpolate job (the page's Stop button). Both ride this module's job list."""
+        return web.json_response({"stopped": jobs.stop(req.match_info["job"])})
+
     @table.get(base + "/upscale/{job}")
     async def _upscale_job(req):
         job = jobs.get(req.match_info["job"])

modules/output/render/tests/test_jobs_routes.py

@@ -149,3 +149,63 @@ def test_render_keeps_the_real_frame_rate_of_interpolated_clips_only_when_asked(
         job = finish(server, f"{BASE}/projects/{proj.id}/render", queued["job_id"])
         assert job["state"] == "done"
         assert files.frame_rate(str(comfy.temp / job["media"]["filename"])) == expected
+
+
+def test_a_running_job_asked_to_stop_ends_stopped_and_stays_pollable(monkeypatch):
+    import asyncio
+    import threading
+    monkeypatch.setattr(jobs, "MAX_JOBS", 3)
+    started, seen = threading.Event(), []
+
+    def long_job(interrupted):
+        started.set()
+        while True:
+            interrupted()                                   # between chunks, as stream_file does
+            seen.append(1)
+            threading.Event().wait(0.01)
+
+    async def go():
+        job = jobs.new()
+        task = asyncio.create_task(jobs.run(job, long_job, error=RuntimeError, stoppable=True))
+        while not started.is_set():
+            await asyncio.sleep(0.01)
+        assert jobs.stop(job) is True
+        await task
+        return job
+
+    job = asyncio.run(go())
+    assert jobs.get(job) == {"state": "stopped"}              # a poll finds it: no 404, so no "ComfyUI restarted"
+    assert jobs.stop(job) is False and jobs.stop("nope") is False
+
+
+def test_a_job_stopped_before_it_started_never_runs():
+    import asyncio
+    ran = []
+    job = jobs.new()
+    assert jobs.stop(job) is True
+    asyncio.run(jobs.run(job, lambda interrupted: ran.append(1), error=RuntimeError, stoppable=True))
+    assert not ran and jobs.get(job) == {"state": "stopped"}
+
+
+@needs_ffmpeg
+def test_a_stopped_stream_leaves_no_half_written_file(comfy):
+    from core import video_pipe
+    src = comfy.output / "a.mp4"
+    make_clip(src, seconds=1)
+    dst = comfy.output / "out.mp4"
+    calls = []
+
+    def stop_after_first_chunk():
+        calls.append(1)
+        if len(calls) > 1:
+            raise jobs.Stopped()
+    with pytest.raises(jobs.Stopped):
+        video_pipe.stream_file(str(src), str(dst), lambda x: x, stop_after_first_chunk)
+    assert not dst.exists()
+
+
+def test_the_stop_route_answers_whether_it_stopped_anything(comfy, server):
+    job = jobs.new()
+    status, body = call(server, "POST", f"{BASE}/jobs/{job}/stop")
+    assert status == 200 and body == {"stopped": True}
+    assert call(server, "POST", f"{BASE}/jobs/nope/stop")[1] == {"stopped": False}