Labyricorn/voicebox
1
0
Fork 0
mirror of https://github.com/jamiepine/voicebox.git synced 2026-09-16 13:20:39 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: add generation cancellation flow (#444)

Browse Source
This commit is contained in:
Andrew Barnes
2026-04-18 02:51:39 -07:00
committed by GitHub
parent 476abe07fc
commit 54a3bf322e
6 changed files with 252 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/components/History/HistoryTable.tsx
+85 -56
View File
@@ -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,60 +649,71 @@ 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
variant="ghost"
size="icon"
className="h-6 w-6 text-muted-foreground/50 hover:bg-muted-foreground/20 hover:text-muted-foreground"
aria-label="Actions"
disabled={isGenerating}
>
<MoreHorizontal className="h-2 w-2" />
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="end">
<DropdownMenuItem
onClick={() => handlePlay(gen.id, gen.text, gen.profile_id)}
>
<Play className="mr-2 h-4 w-4" />
Play
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => handleDownloadAudio(gen.id, gen.text)}
disabled={exportGenerationAudio.isPending}
>
<Download className="mr-2 h-4 w-4" />
Export Audio
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => handleExportPackage(gen.id, gen.text)}
disabled={exportGeneration.isPending}
>
<FileArchive className="mr-2 h-4 w-4" />
Export Package
</DropdownMenuItem>
<DropdownMenuItem onClick={() => handleApplyEffects(gen.id)}>
<Wand2 className="mr-2 h-4 w-4" />
Apply Effects
</DropdownMenuItem>
<DropdownMenuItem onClick={() => handleRegenerate(gen.id)}>
<RotateCcw className="mr-2 h-4 w-4" />
Regenerate
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => handleDeleteClick(gen.id, gen.profile_name)}
disabled={deleteGeneration.isPending}
// className="text-destructive focus:text-destructive"
>
<Trash2 className="mr-2 h-4 w-4" />
Delete
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
</>
<DropdownMenu>
<DropdownMenuTrigger asChild>
<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="Actions"
disabled={isGenerating}
>
<MoreHorizontal className="h-2 w-2" />
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="end">
<DropdownMenuItem onClick={() => handlePlay(gen.id, gen.text, gen.profile_id)}>
<Play className="mr-2 h-4 w-4" />
Play
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => handleDownloadAudio(gen.id, gen.text)}
disabled={exportGenerationAudio.isPending}
>
<Download className="mr-2 h-4 w-4" />
Export Audio
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => handleExportPackage(gen.id, gen.text)}
disabled={exportGeneration.isPending}
>
<FileArchive className="mr-2 h-4 w-4" />
Export Package
</DropdownMenuItem>
<DropdownMenuItem onClick={() => handleApplyEffects(gen.id)}>
<Wand2 className="mr-2 h-4 w-4" />
Apply Effects
</DropdownMenuItem>
<DropdownMenuItem onClick={() => handleRegenerate(gen.id)}>
<RotateCcw className="mr-2 h-4 w-4" />
Regenerate
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => handleDeleteClick(gen.id, gen.profile_name)}
disabled={deleteGeneration.isPending}
// className="text-destructive focus:text-destructive"
>
<Trash2 className="mr-2 h-4 w-4" />
Delete
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
)}
</div>
</div>
app/src/lib/api/client.ts
+6
View File
@@ -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
+32 -1
View File
@@ -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
+8
View File
@@ -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
+67 -7
View File
@@ -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
+54
View File
@@ -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)