fix: add generation cancellation flow (#444)

app/src/components/History/HistoryTable.tsx

@@ -1,15 +1,14 @@
-import { useQueryClient } from '@tanstack/react-query';
+import { useMutation, useQueryClient } from '@tanstack/react-query';
 import { AnimatePresence, motion } from 'framer-motion';
 import {
-  AlignCenter,
   AudioLines,
-  AudioWaveform,
   Download,
   FileArchive,
   Loader2,
   MoreHorizontal,
   Play,
   RotateCcw,
+  Square,
   Star,
   Trash2,
   Wand2,
@@ -130,6 +129,23 @@ export function HistoryTable() {
   const exportGeneration = useExportGeneration();
   const exportGenerationAudio = useExportGenerationAudio();
   const importGeneration = useImportGeneration();
+  const cancelGeneration = useMutation({
+    mutationFn: (generationId: string) => apiClient.cancelGeneration(generationId),
+    onSuccess: async (data) => {
+      await queryClient.invalidateQueries({ queryKey: ['history'] });
+      toast({
+        title: 'Cancelling generation',
+        description: data.message,
+      });
+    },
+    onError: (error) => {
+      toast({
+        title: 'Cancel failed',
+        description: error instanceof Error ? error.message : 'Could not cancel generation',
+        variant: 'destructive',
+      });
+    },
+  });
   const addPendingGeneration = useGenerationStore((state) => state.addPendingGeneration);
   const setAudioWithAutoPlay = usePlayerStore((state) => state.setAudioWithAutoPlay);
   const restartCurrentAudio = usePlayerStore((state) => state.restartCurrentAudio);
@@ -483,6 +499,8 @@ export function HistoryTable() {
               const isPlayable = !isGenerating && !isFailed;
               const hasVersions = gen.versions && gen.versions.length > 1;
               const isVersionsExpanded = expandedVersionsId === gen.id;
+              const isCancelling =
+                cancelGeneration.isPending && cancelGeneration.variables === gen.id;
               return (
                 <div
                   key={gen.id}
@@ -631,8 +649,22 @@ export function HistoryTable() {
                             <Trash2 className="h-2 w-2" />
                           </Button>
                         </>
+                      ) : isGenerating ? (
+                        <Button
+                          variant="ghost"
+                          size="icon"
+                          className="h-6 w-6 text-muted-foreground/50 hover:bg-muted-foreground/20 hover:text-muted-foreground"
+                          aria-label="Cancel generation"
+                          disabled={isCancelling}
+                          onClick={() => cancelGeneration.mutate(gen.id)}
+                        >
+                          {isCancelling ? (
+                            <Loader2 className="h-2 w-2 animate-spin" />
+                          ) : (
+                            <Square className="h-2 w-2" />
+                          )}
+                        </Button>
                       ) : (
-                        <>
                         <DropdownMenu>
                           <DropdownMenuTrigger asChild>
                             <Button
@@ -646,9 +678,7 @@ export function HistoryTable() {
                             </Button>
                           </DropdownMenuTrigger>
                           <DropdownMenuContent align="end">
-                              <DropdownMenuItem
+                            <DropdownMenuItem onClick={() => handlePlay(gen.id, gen.text, gen.profile_id)}>
-                                onClick={() => handlePlay(gen.id, gen.text, gen.profile_id)}
-                              >
                               <Play className="mr-2 h-4 w-4" />
                               Play
                             </DropdownMenuItem>
@@ -684,7 +714,6 @@ export function HistoryTable() {
                             </DropdownMenuItem>
                           </DropdownMenuContent>
                         </DropdownMenu>
-                        </>
                       )}
                     </div>
                   </div>

app/src/lib/api/client.ts

@@ -234,6 +234,12 @@ class ApiClient {
     });
   }
 
+  async cancelGeneration(generationId: string): Promise<{ message: string }> {
+    return this.request<{ message: string }>(`/generate/${generationId}/cancel`, {
+      method: 'POST',
+    });
+  }
+
   async regenerateGeneration(generationId: string): Promise<GenerationResponse> {
     return this.request<GenerationResponse>(`/generate/${generationId}/regenerate`, {
       method: 'POST',

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)