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)