From 54a3bf322e565e2238307371b767266e7d02065a Mon Sep 17 00:00:00 2001 From: Andrew Barnes Date: Sat, 18 Apr 2026 05:51:39 -0400 Subject: [PATCH] fix: add generation cancellation flow (#444) --- app/src/components/History/HistoryTable.tsx | 141 +++++++++++------- app/src/lib/api/client.ts | 6 + backend/routes/generations.py | 33 +++- backend/services/generation.py | 8 + backend/services/task_queue.py | 74 ++++++++- backend/tests/test_task_queue_cancellation.py | 54 +++++++ 6 files changed, 252 insertions(+), 64 deletions(-) create mode 100644 backend/tests/test_task_queue_cancellation.py diff --git a/app/src/components/History/HistoryTable.tsx b/app/src/components/History/HistoryTable.tsx index 9a7bdcdd..e3172e8f 100644 --- a/app/src/components/History/HistoryTable.tsx +++ b/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 (
+ ) : isGenerating ? ( + ) : ( - <> - - - - - - handlePlay(gen.id, gen.text, gen.profile_id)} - > - - Play - - handleDownloadAudio(gen.id, gen.text)} - disabled={exportGenerationAudio.isPending} - > - - Export Audio - - handleExportPackage(gen.id, gen.text)} - disabled={exportGeneration.isPending} - > - - Export Package - - handleApplyEffects(gen.id)}> - - Apply Effects - - handleRegenerate(gen.id)}> - - Regenerate - - handleDeleteClick(gen.id, gen.profile_name)} - disabled={deleteGeneration.isPending} - // className="text-destructive focus:text-destructive" - > - - Delete - - - - + + + + + + handlePlay(gen.id, gen.text, gen.profile_id)}> + + Play + + handleDownloadAudio(gen.id, gen.text)} + disabled={exportGenerationAudio.isPending} + > + + Export Audio + + handleExportPackage(gen.id, gen.text)} + disabled={exportGeneration.isPending} + > + + Export Package + + handleApplyEffects(gen.id)}> + + Apply Effects + + handleRegenerate(gen.id)}> + + Regenerate + + handleDeleteClick(gen.id, gen.profile_name)} + disabled={deleteGeneration.isPending} + // className="text-destructive focus:text-destructive" + > + + Delete + + + )}
diff --git a/app/src/lib/api/client.ts b/app/src/lib/api/client.ts index 1374f27b..dbff4e87 100644 --- a/app/src/lib/api/client.ts +++ b/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 { return this.request(`/generate/${generationId}/regenerate`, { method: 'POST', diff --git a/backend/routes/generations.py b/backend/routes/generations.py index 2af3832e..775a4e3e 100644 --- a/backend/routes/generations.py +++ b/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.""" diff --git a/backend/services/generation.py b/backend/services/generation.py index a70e633e..718fabbd 100644 --- a/backend/services/generation.py +++ b/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( diff --git a/backend/services/task_queue.py b/backend/services/task_queue.py index fbd9638e..9177a5c6 100644 --- a/backend/services/task_queue.py +++ b/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()) diff --git a/backend/tests/test_task_queue_cancellation.py b/backend/tests/test_task_queue_cancellation.py new file mode 100644 index 00000000..7ca3ecbf --- /dev/null +++ b/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)