fix: add generation cancellation flow (#444)

backend/routes/generations.py

@@ -14,7 +14,7 @@ from .. import models
 from ..services import history, profiles, tts
 from ..database import Generation as DBGeneration, VoiceProfile as DBVoiceProfile, get_db
 from ..services.generation import run_generation
-from ..services.task_queue import enqueue_generation
+from ..services.task_queue import cancel_generation as cancel_generation_job, enqueue_generation
 from ..utils.tasks import get_task_manager
 
 router = APIRouter()
@@ -82,6 +82,7 @@ async def generate_speech(
                 pass
 
     enqueue_generation(
+        generation_id,
         run_generation(
             generation_id=generation_id,
             profile_id=data.profile_id,
@@ -127,6 +128,7 @@ async def retry_generation(generation_id: str, db: Session = Depends(get_db)):
     )
 
     enqueue_generation(
+        generation_id,
         run_generation(
             generation_id=generation_id,
             profile_id=gen.profile_id,
@@ -170,6 +172,7 @@ async def regenerate_generation(generation_id: str, db: Session = Depends(get_db
     version_id = str(uuid.uuid4())
 
     enqueue_generation(
+        generation_id,
         run_generation(
             generation_id=generation_id,
             profile_id=gen.profile_id,
@@ -187,6 +190,34 @@ async def regenerate_generation(generation_id: str, db: Session = Depends(get_db
     return models.GenerationResponse.model_validate(gen)
 
 
+@router.post("/generate/{generation_id}/cancel")
+async def cancel_generation(generation_id: str, db: Session = Depends(get_db)):
+    """Cancel a queued or running generation."""
+    gen = db.query(DBGeneration).filter_by(id=generation_id).first()
+    if not gen:
+        raise HTTPException(status_code=404, detail="Generation not found")
+
+    if (gen.status or "completed") not in ("loading_model", "generating"):
+        raise HTTPException(status_code=400, detail="Only active generations can be cancelled")
+
+    cancellation_state = cancel_generation_job(generation_id)
+    if cancellation_state is None:
+        raise HTTPException(status_code=409, detail="Generation is no longer cancellable")
+
+    if cancellation_state == "queued":
+        task_manager = get_task_manager()
+        task_manager.complete_generation(generation_id)
+        await history.update_generation_status(
+            generation_id=generation_id,
+            status="failed",
+            db=db,
+            error="Generation cancelled",
+        )
+        return {"message": "Queued generation cancelled"}
+
+    return {"message": "Generation cancellation requested"}
+
+
 @router.get("/generate/{generation_id}/status")
 async def get_generation_status(generation_id: str, db: Session = Depends(get_db)):
     """SSE endpoint that streams generation status updates."""

backend/services/generation.py

@@ -16,6 +16,7 @@ Mode differences:
 
 from __future__ import annotations
 
+import asyncio
 import traceback
 from typing import Literal, Optional
 
@@ -126,6 +127,13 @@ async def run_generation(
             duration=duration,
         )
 
+    except asyncio.CancelledError:
+        await history.update_generation_status(
+            generation_id=generation_id,
+            status="failed",
+            db=bg_db,
+            error="Generation cancelled",
+        )
     except Exception as e:
         traceback.print_exc()
         await history.update_generation_status(

backend/services/task_queue.py

@@ -5,12 +5,27 @@ to avoid GPU contention.
 
 import asyncio
 import traceback
+from dataclasses import dataclass
+from typing import Coroutine, Literal
 
 # Keep references to fire-and-forget background tasks to prevent GC
 _background_tasks: set = set()
 
+
+@dataclass
+class GenerationJob:
+    """Queued generation work plus the generation ID it belongs to."""
+
+    generation_id: str
+    coro: Coroutine
+
+
 # Generation queue — serializes TTS inference to avoid GPU contention
 _generation_queue: asyncio.Queue = None  # type: ignore  # initialized at startup
+_generation_worker_task: asyncio.Task | None = None
+_queued_generation_ids: set[str] = set()
+_running_generation_tasks: dict[str, asyncio.Task] = {}
+_cancelled_generation_ids: set[str] = set()
 
 
 def create_background_task(coro) -> asyncio.Task:
@@ -24,25 +39,70 @@ def create_background_task(coro) -> asyncio.Task:
 async def _generation_worker():
     """Worker that processes generation tasks one at a time."""
     while True:
-        coro = await _generation_queue.get()
+        job = await _generation_queue.get()
         try:
-            await coro
+            if job.generation_id in _cancelled_generation_ids:
+                _cancelled_generation_ids.discard(job.generation_id)
+                job.coro.close()
+                continue
+
+            task = asyncio.create_task(job.coro)
+            _running_generation_tasks[job.generation_id] = task
+            _queued_generation_ids.discard(job.generation_id)
+            try:
+                await task
+            except asyncio.CancelledError:
+                if not task.cancelled():
+                    raise
         except Exception:
             traceback.print_exc()
         finally:
+            _running_generation_tasks.pop(job.generation_id, None)
+            _queued_generation_ids.discard(job.generation_id)
             _generation_queue.task_done()
 
 
-def enqueue_generation(coro):
+def enqueue_generation(generation_id: str, coro):
     """Add a generation coroutine to the serial queue."""
-    _generation_queue.put_nowait(coro)
+    if _generation_queue is None:
+        raise RuntimeError("Generation queue has not been initialized")
+
+    _queued_generation_ids.add(generation_id)
+    _generation_queue.put_nowait(GenerationJob(generation_id=generation_id, coro=coro))
 
 
-def init_queue():
+def cancel_generation(generation_id: str) -> Literal["queued", "running"] | None:
+    """Cancel a queued or running generation if it is still active."""
+    running_task = _running_generation_tasks.get(generation_id)
+    if running_task is not None:
+        running_task.cancel()
+        return "running"
+
+    if generation_id in _queued_generation_ids:
+        _queued_generation_ids.discard(generation_id)
+        _cancelled_generation_ids.add(generation_id)
+        return "queued"
+
+    return None
+
+
+def init_queue(force: bool = False):
     """Initialize the generation queue and start the worker.
 
     Must be called once during application startup (inside a running event loop).
     """
-    global _generation_queue
+    global _generation_queue, _generation_worker_task
+    global _queued_generation_ids, _running_generation_tasks, _cancelled_generation_ids
+
+    if _generation_worker_task is not None and not _generation_worker_task.done():
+        if not force:
+            return
+        _generation_worker_task.cancel()
+        for task in list(_running_generation_tasks.values()):
+            task.cancel()
+
     _generation_queue = asyncio.Queue()
-    create_background_task(_generation_worker())
+    _queued_generation_ids = set()
+    _running_generation_tasks = {}
+    _cancelled_generation_ids = set()
+    _generation_worker_task = create_background_task(_generation_worker())

backend/tests/test_task_queue_cancellation.py

@@ -0,0 +1,54 @@
+import asyncio
+
+import pytest
+
+from backend.services import task_queue
+
+
+@pytest.mark.asyncio
+async def test_cancel_queued_generation_skips_execution():
+    task_queue.init_queue(force=True)
+
+    running_started = asyncio.Event()
+    release_running = asyncio.Event()
+    queued_ran = asyncio.Event()
+
+    async def running_job():
+        running_started.set()
+        await release_running.wait()
+
+    async def queued_job():
+        queued_ran.set()
+
+    task_queue.enqueue_generation("gen-running", running_job())
+    await asyncio.wait_for(running_started.wait(), timeout=1)
+
+    task_queue.enqueue_generation("gen-queued", queued_job())
+    assert task_queue.cancel_generation("gen-queued") == "queued"
+
+    release_running.set()
+    await asyncio.sleep(0.1)
+
+    assert not queued_ran.is_set()
+
+
+@pytest.mark.asyncio
+async def test_cancel_running_generation_cancels_task():
+    task_queue.init_queue(force=True)
+
+    running_started = asyncio.Event()
+    running_cancelled = asyncio.Event()
+
+    async def running_job():
+        running_started.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            running_cancelled.set()
+            raise
+
+    task_queue.enqueue_generation("gen-running", running_job())
+    await asyncio.wait_for(running_started.wait(), timeout=1)
+
+    assert task_queue.cancel_generation("gen-running") == "running"
+    await asyncio.wait_for(running_cancelled.wait(), timeout=1)