Settle generation scratch in the worker that owns it

Four-reviewer simplify pass over the Stop change. A stopped run leaked
its scratch tree: Gradio closes the generator before the sentinel stops
the worker, so the worker_finished cleanup guard never fired. Releasing
in the worker's own finally covers return, failure, and stop alike, and
deletes the flag plus the conditional cleanup. The stop test now proves
a mid-run stop deterministically instead of an abort-at-entry, and the
session fixture takes data_dir instead of being duplicated.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

fdanyone_app.py

@@ -518,8 +518,6 @@ class Session:
     """Decoded source stills for the diffusion pane, keyed by canonical frame."""
     events: queue.Queue[str | None] = field(default_factory=queue.Queue)
     """Stage labels the pipeline hooks push from the worker thread."""
-    worker_finished: bool = False
-    """Whether the last ``_pump`` call's pipeline thread returned."""
 
 
 class RunCancelled(FourDAnyoneError):
@@ -557,7 +555,6 @@ def _pump(session: Session, work: Callable[[], T]) -> Iterator[str]:
 
     outcome: list[T] = []
     failure: list[BaseException] = []
-    session.worker_finished = False
 
     def target() -> None:
         try:
@@ -575,7 +572,6 @@ def _pump(session: Session, work: Callable[[], T]) -> Iterator[str]:
             break
         yield label
     worker.join()
-    session.worker_finished = True
     if failure:
         raise failure[0]
     return outcome[0]
@@ -803,31 +799,29 @@ def generate_phase(link: Link, session: Session) -> Iterator[str]:
 
     if session.prepared is None:
         raise FourDAnyoneError("The motion phase did not finish.")
+    prepared: PreparedRun = session.prepared
     session.source_frames = _decode_source_stills(session.spec)
     link.recording.send_blueprint(diffusion_blueprint(), make_active=True)
     yield "Generating six views."
 
     def work() -> dict:
-        return generate_run(
-            session.prepared, seed=session.spec.seed, on_denoise_step=_denoise_hook(link, session)
-        )
+        # The worker owns the scratch, so it must also settle it: a Stop closes
+        # the outer generator, which then never sees the pipeline call end.
+        try:
+            return generate_run(
+                prepared, seed=session.spec.seed, on_denoise_step=_denoise_hook(link, session)
+            )
+        finally:
+            release_run(prepared)
 
-    try:
-        steps: Iterator[str] = _pump(session, work)
-        while True:
-            try:
-                label: str = next(steps)
-            except StopIteration as done:
-                session.summary = done.value
-                break
-            yield f"Denoising: {label} of {SETTINGS.num_inference_steps}."
-    finally:
-        # Stop closes this generator while the worker is still inside
-        # generate_run, and its scratch directory must not be deleted out from
-        # under it. Only a run whose pipeline call returned is safe to settle;
-        # a cancelled one leaks its scratch into the Space's ephemeral disk.
-        if session.worker_finished and session.prepared is not None:
-            release_run(session.prepared)
+    steps: Iterator[str] = _pump(session, work)
+    while True:
+        try:
+            label: str = next(steps)
+        except StopIteration as done:
+            session.summary = done.value
+            break
+        yield f"Denoising: {label} of {SETTINGS.num_inference_steps}."
 
     _status(link.recording, "generate: six views published")
     yield "Publishing the result."

tests/test_app_helpers.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import os
 import pickle
+import threading
 import traceback
 import uuid
 from collections.abc import Callable
@@ -75,7 +76,7 @@ def test_probe_clip_rejects_a_missing_file(tmp_path: Path) -> None:
         fdanyone_app.probe_clip(tmp_path / "absent.mp4", 0.0)
 
 
-def _pump_session() -> fdanyone_app.Session:
+def _pump_session(data_dir: Path = Path("/nonexistent")) -> fdanyone_app.Session:
     """A Session with only the fields `_pump` touches."""
 
     return fdanyone_app.Session(
@@ -85,7 +86,7 @@ def _pump_session() -> fdanyone_app.Session:
             start_time=0.0,
             seed=0,
             fps=Fraction(25, 1),
-            data_dir=Path("/nonexistent"),
+            data_dir=data_dir,
         )
     )
 
@@ -107,45 +108,38 @@ def test_pump_yields_hook_labels_and_returns_the_work_result() -> None:
             assert done.value == "prepared"
             break
     assert seen == ["bboxes", "smplx"]
-    assert session.worker_finished is True
 
 
-def test_stop_sentinel_unwinds_the_pump_worker(tmp_path: Path) -> None:
+def test_stop_sentinel_unwinds_a_running_pump_worker(tmp_path: Path) -> None:
     """Stop must reach the pipeline thread itself, not just Gradio's generator.
 
     ``cancels`` closes the request generator, but the worker thread (or the
     forked ZeroGPU child) never hears it; the sentinel checked by the hooks is
-    the only thing that actually frees the GPU.
+    the only thing that actually frees the GPU. The barrier makes the order
+    deterministic: the worker is provably mid-run when Stop arrives.
     """
 
-    session: fdanyone_app.Session = _pump_session()
-    spec: fdanyone_app.RunSpec = fdanyone_app.RunSpec(
-        token=session.spec.token,
-        video_path=session.spec.video_path,
-        start_time=0.0,
-        seed=0,
-        fps=Fraction(25, 1),
-        data_dir=tmp_path,
-    )
-    session = fdanyone_app.Session(spec)
     assert fdanyone_app.request_stop(None) == "Nothing is running."
-    fdanyone_app.request_stop(spec)
-    assert (tmp_path / "stop-requested").exists()
 
-    ticks: list[int] = []
+    session: fdanyone_app.Session = _pump_session(tmp_path)
+    stop_requested: threading.Event = threading.Event()
 
     def work() -> str:
-        # What every pipeline hook does on entry, once per stage or step.
-        for index in range(1000):
-            fdanyone_app._check_stop(spec)
-            ticks.append(index)
+        # A pipeline hook checks the sentinel once per stage; model the stage
+        # boundary the run is inside when the user presses Stop.
+        fdanyone_app._check_stop(session.spec)
+        session.events.put("running")
+        assert stop_requested.wait(timeout=5.0)
+        fdanyone_app._check_stop(session.spec)
         return "finished"
 
+    pump = fdanyone_app._pump(session, work)
+    assert next(pump) == "running"
+    fdanyone_app.request_stop(session.spec)
+    assert (tmp_path / "stop-requested").exists()
+    stop_requested.set()
     with pytest.raises(fdanyone_app.RunCancelled):
-        list(fdanyone_app._pump(session, work))
-    assert ticks == []
-    # The worker joined before re-raising, so the thread is truly gone.
-    assert session.worker_finished is True
+        list(pump)
 
 
 def test_pump_reraises_a_worker_failure_on_the_caller_thread() -> None:
@@ -158,8 +152,6 @@ def test_pump_reraises_a_worker_failure_on_the_caller_thread() -> None:
 
     with pytest.raises(FourDAnyoneError, match="boom"):
         list(fdanyone_app._pump(session, work))
-    # The worker joined before re-raising, so releasing its scratch is safe.
-    assert session.worker_finished is True
 
 
 def test_blueprints_build_for_every_phase() -> None: