Repository navigation
Fix generation cancellation flow - #444
Conversation
📝 WalkthroughWalkthroughThis change implements a full-stack task cancellation feature allowing users to stop in-progress generation tasks. The frontend UI now displays a cancel button for active tasks, which triggers an API endpoint that coordinates cancellation across the backend's route layer, task queue manager, and generation service. Changes
Sequence DiagramsequenceDiagram
participant UI as User / HistoryTable
participant API as ApiClient
participant Endpoint as POST /generate/{id}/cancel
participant DB as Database
participant TQ as TaskQueue Manager
participant Worker as Worker Task
participant Service as Generation Service
UI->>API: cancelGeneration(generationId)
API->>Endpoint: POST /generate/{id}/cancel
Endpoint->>DB: Fetch generation record
Endpoint->>TQ: cancel_generation(generation_id)
alt Generation is Queued
TQ->>TQ: Mark ID as cancelled
TQ->>Worker: Skip queued job execution
Endpoint->>DB: Update status → failed, error: "Generation cancelled"
Endpoint->>API: 200 {"message": "Queued generation cancelled"}
else Generation is Running
TQ->>Worker: Close coroutine (asyncio.CancelledError)
Worker->>Service: Catch CancelledError
Service->>DB: Update status → failed, error: "Generation cancelled"
Endpoint->>API: 200 {"message": "Generation cancellation requested"}
end
API->>UI: Show success toast
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (4)
backend/services/task_queue.py (1)
9-9: ImportCoroutinefromcollections.abc.
typing.Coroutinehas been a deprecated alias since Python 3.9 (flagged by Ruff UP035).♻️ Proposed fix
-from dataclasses import dataclass -from typing import Coroutine, Literal +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Literal🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/services/task_queue.py` at line 9, Replace the deprecated typing.Coroutine import by importing Coroutine from collections.abc: update the import statement that currently references "Coroutine" (in backend/services/task_queue.py) to import Coroutine from collections.abc instead of typing, keeping other imports like Literal unchanged so functions/types using Coroutine continue to work with the non-deprecated symbol.backend/tests/test_task_queue_cancellation.py (2)
1-54: Add teardown to stop the worker between tests.Both tests call
init_queue(force=True)at the top, which cancels the previous worker, but they never stop the worker after the test completes. Test 2 in particular leaves a cancelled-task coroutine and a live_generation_worker_taskbackground task in the module. Withpytest-asyncio's function-scoped event loop, this background task is torn down with the loop and can produceTask was destroyed but it is pending!/ unawaited-coroutine warnings that mask real issues.Consider a fixture (module- or function-scoped) that calls
init_queue(force=True)in setup and cancels_generation_worker_task+ awaits it in teardown.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/tests/test_task_queue_cancellation.py` around lines 1 - 54, Tests start new background worker with task_queue.init_queue(force=True) but never stop it, causing pending background tasks and warnings; add a pytest async fixture (function- or module-scoped) that calls task_queue.init_queue(force=True) in setup and in teardown checks for task_queue._generation_worker_task, calls .cancel() on it and awaits it (e.g., inside an asyncio.wait_for to avoid hanging) to ensure the worker is fully stopped between tests; update both tests to use this fixture instead of calling init_queue directly or remove their init_queue calls if the fixture performs setup.
29-32: Replace the fixedasyncio.sleep(0.1)with deterministic synchronization.
await asyncio.sleep(0.1)assumes the worker can drain the queue within 100 ms, which is fine locally but can flake on a loaded CI runner. Prefer waiting on an observable signal, e.g.await task_queue._generation_queue.join()(with the test already guaranteeing both jobs were enqueued), or poll on_queued_generation_ids/_cancelled_generation_idsbeing empty viaasyncio.wait_for.♻️ Proposed fix
release_running.set() - await asyncio.sleep(0.1) + await asyncio.wait_for(task_queue._generation_queue.join(), timeout=1) assert not queued_ran.is_set()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/tests/test_task_queue_cancellation.py` around lines 29 - 32, The test uses a fixed await asyncio.sleep(0.1) which can flake; replace it with deterministic synchronization by waiting for the queue to drain or for observable internal state changes. Specifically, after release_running.set(), await task_queue._generation_queue.join() (since the test already enqueued both jobs) or use asyncio.wait_for to poll until task_queue._queued_generation_ids and/or task_queue._cancelled_generation_ids reflect the expected empty/state, then assert not queued_ran.is_set(); update the assertion block around release_running, queued_ran and the task_queue internals (e.g., _generation_queue.join(), _queued_generation_ids, _cancelled_generation_ids) to remove the sleep and use these deterministic waits.app/src/components/History/HistoryTable.tsx (1)
652-666: LGTM on the cancel button, minor a11y nit on the icon size.Behavior and
aria-label="Cancel generation"are correct. Theh-2 w-2on bothLoader2andSquarematches the surrounding icons in this row, but those are unusually small (~8 px) and below WCAG 2.5.5 minimum touch-target guidance even inside a 24 px button. Considerh-3 w-3for the cancel icon specifically so it's visually discoverable as a "stop" affordance. Non-blocking.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@app/src/components/History/HistoryTable.tsx` around lines 652 - 666, Increase the icon size for the cancel button to improve discoverability: in the isGenerating branch where the Button with aria-label="Cancel generation" (and props disabled={isCancelling}, onClick={() => cancelGeneration.mutate(gen.id)}) renders Loader2 and Square, change their className from "h-2 w-2" to "h-3 w-3" so the cancel icon is slightly larger while keeping the surrounding Button styling unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@app/src/components/History/HistoryTable.tsx`:
- Around line 132-148: The onSuccess of cancelGeneration currently immediately
calls queryClient.invalidateQueries and refetches history while backend
cancellation is still being processed, causing a race where the row remains in
in-progress state and the cancel button stays enabled; fix this by tracking the
cancelled generation id locally (e.g. add a useState holding a Set<string> of
pendingCancels), push the generationId into that set inside
cancelGeneration.onSuccess, and update the History row/button render logic to
disable or hide the cancel button when the id is present until the refetch shows
a non-in-progress status (then remove the id from the set); alternatively, if
you prefer a lighter change, defer the invalidateQueries call in
cancelGeneration.onSuccess with a short timeout (e.g. 500–1000ms) so the
backend’s CancelledError handler has time to flip the status before refetch.
In `@backend/routes/generations.py`:
- Around line 193-218: The queued-cancellation branch can leave the DB row in an
orphaned active state if history.update_generation_status raises after
cancel_generation_job and task_manager.complete_generation have already modified
the queue; wrap the call to history.update_generation_status(...) in a
try/except that catches database errors (e.g. OperationalError/Exception), logs
the failure with logger.info or logger.error including generation_id and
cancellation_state, and returns a partial-success response (e.g. message
indicating queue removed but DB update failed) instead of letting the exception
propagate; alternatively persist the new status (call
history.update_generation_status) before calling
get_task_manager().complete_generation to ensure DB is updated first. Also add a
minimal response_model (e.g. {"message": str}) to the cancel_generation route
decorator and add a logger.info("Cancellation requested for generation %s
(state=%s)", generation_id, cancellation_state) before performing queue changes
to improve visibility.
- Around line 200-201: The current status check is racy; replace the expression
(gen.status or "completed") not in ("loading_model", "generating") with a strict
"never-was-active" check: if gen.status is None: raise
HTTPException(status_code=400, detail="Generation was never active and cannot be
cancelled"). Leave cancellation and the subsequent conflict handling to
cancel_generation_job(generation_id) so transitions after this check are handled
consistently; update the error message to refer only to the never-active case
and reference gen.status and cancel_generation_job for locating the code.
In `@backend/services/generation.py`:
- Around line 130-136: The except block that catches asyncio.CancelledError
around the generation flow (catching asyncio.CancelledError and calling
history.update_generation_status with generation_id and db=bg_db) must re-raise
the CancelledError after persisting the failure state so cancellation propagates
correctly; modify the handler to first persist the "failed" status via
history.update_generation_status (use
asyncio.shield(history.update_generation_status(...)) to protect the DB update
from being cancelled mid-write) and then re-raise the original
asyncio.CancelledError.
In `@backend/services/task_queue.py`:
- Around line 39-62: The worker _generation_worker relies on run_generation to
propagate asyncio.CancelledError so the worker can detect cancelled jobs; update
backend/services/generation.py so that any except asyncio.CancelledError handler
in run_generation (or any coroutine invoked by it) does not swallow the
exception — either remove the CancelledError catch or re-raise it (raise) after
any cleanup, ensuring CancelledError reaches the awaiting code in
_generation_worker.
- Around line 89-108: init_queue(force=True) cancels the existing
_generation_worker_task and running tasks but immediately resets the module
state, risking race conditions where cancelled tasks' finally blocks access
cleared structures; instead, after calling .cancel() on _generation_worker_task
and each task from _running_generation_tasks, await their termination before
clearing/overwriting _generation_queue, _queued_generation_ids,
_running_generation_tasks, and _cancelled_generation_ids — e.g., collect the
cancelled asyncio.Task objects (including _generation_worker_task and values()
of _running_generation_tasks), await them using asyncio.gather(...,
return_exceptions=True) or asyncio.wait with appropriate shielding
(asyncio.shield) to ensure they run their cleanup, then only after they complete
reset the module-level tracking and start a new worker with
create_background_task(_generation_worker()).
---
Nitpick comments:
In `@app/src/components/History/HistoryTable.tsx`:
- Around line 652-666: Increase the icon size for the cancel button to improve
discoverability: in the isGenerating branch where the Button with
aria-label="Cancel generation" (and props disabled={isCancelling}, onClick={()
=> cancelGeneration.mutate(gen.id)}) renders Loader2 and Square, change their
className from "h-2 w-2" to "h-3 w-3" so the cancel icon is slightly larger
while keeping the surrounding Button styling unchanged.
In `@backend/services/task_queue.py`:
- Line 9: Replace the deprecated typing.Coroutine import by importing Coroutine
from collections.abc: update the import statement that currently references
"Coroutine" (in backend/services/task_queue.py) to import Coroutine from
collections.abc instead of typing, keeping other imports like Literal unchanged
so functions/types using Coroutine continue to work with the non-deprecated
symbol.
In `@backend/tests/test_task_queue_cancellation.py`:
- Around line 1-54: Tests start new background worker with
task_queue.init_queue(force=True) but never stop it, causing pending background
tasks and warnings; add a pytest async fixture (function- or module-scoped) that
calls task_queue.init_queue(force=True) in setup and in teardown checks for
task_queue._generation_worker_task, calls .cancel() on it and awaits it (e.g.,
inside an asyncio.wait_for to avoid hanging) to ensure the worker is fully
stopped between tests; update both tests to use this fixture instead of calling
init_queue directly or remove their init_queue calls if the fixture performs
setup.
- Around line 29-32: The test uses a fixed await asyncio.sleep(0.1) which can
flake; replace it with deterministic synchronization by waiting for the queue to
drain or for observable internal state changes. Specifically, after
release_running.set(), await task_queue._generation_queue.join() (since the test
already enqueued both jobs) or use asyncio.wait_for to poll until
task_queue._queued_generation_ids and/or task_queue._cancelled_generation_ids
reflect the expected empty/state, then assert not queued_ran.is_set(); update
the assertion block around release_running, queued_ran and the task_queue
internals (e.g., _generation_queue.join(), _queued_generation_ids,
_cancelled_generation_ids) to remove the sleep and use these deterministic
waits.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1b3c03b5-8c95-446c-890c-25b28a8762c6
📒 Files selected for processing (6)
app/src/components/History/HistoryTable.tsxapp/src/lib/api/client.tsbackend/routes/generations.pybackend/services/generation.pybackend/services/task_queue.pybackend/tests/test_task_queue_cancellation.py
| 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', | ||
| }); | ||
| }, | ||
| }); |
There was a problem hiding this comment.
onSuccess invalidation races the actual status transition for running cancels.
For the "running" branch, the backend only requests cancellation — the DB row remains in generating/loading_model until run_generation's CancelledError handler flushes status="failed". The mutation's onSuccess fires immediately and invalidates ['history'], so the refetch will almost always return the row still in progress, leaving the cancel button visible and enabled. A quick second click then hits the backend during the narrow window and gets a 409 "Generation is no longer cancellable" toast even though cancellation was already in flight.
Two light-touch options:
- After cancelling, locally track the generation id (e.g.
useState<Set<string>>) and keep the button disabled (or hide it) until the row'sstatusleaves the in-progress states. - Defer the
invalidateQueriesslightly (e.g.setTimeout(..., 1000)) — less clean, but avoids the immediate stale refetch.
The first option also makes the "Cancelling…" UI accurate end-to-end.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@app/src/components/History/HistoryTable.tsx` around lines 132 - 148, The
onSuccess of cancelGeneration currently immediately calls
queryClient.invalidateQueries and refetches history while backend cancellation
is still being processed, causing a race where the row remains in in-progress
state and the cancel button stays enabled; fix this by tracking the cancelled
generation id locally (e.g. add a useState holding a Set<string> of
pendingCancels), push the generationId into that set inside
cancelGeneration.onSuccess, and update the History row/button render logic to
disable or hide the cancel button when the id is present until the refetch shows
a non-in-progress status (then remove the id from the set); alternatively, if
you prefer a lighter change, defer the invalidateQueries call in
cancelGeneration.onSuccess with a short timeout (e.g. 500–1000ms) so the
backend’s CancelledError handler has time to flip the status before refetch.
| @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"} |
There was a problem hiding this comment.
Inconsistent state possible if the DB update fails after the queue has been marked cancelled.
For the "queued" branch, the sequence is:
cancel_generation_job— marks the ID as cancelled intask_queue(side effect committed).task_manager.complete_generation— removes it from active tasks (side effect committed).await history.update_generation_status(..., status="failed", ...)— may raise.
If step 3 raises (e.g. OperationalError), the client sees a 500 but the queue has already discarded the job, and the DB row stays at generating forever (orphaned — exactly the problem referenced in issue #435). Wrap step 3 in a try/except that at minimum logs and returns a partial-success response, or persist the status before touching the queue.
Additionally:
- Consider adding a minimal
response_model(e.g.{"message": str}) for schema consistency with the rest of the module. - A
logger.info("Cancellation requested for generation %s (state=%s)", generation_id, cancellation_state)would meaningfully improve the "logging/visibility for task status endpoints" item from issue#435.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/routes/generations.py` around lines 193 - 218, The
queued-cancellation branch can leave the DB row in an orphaned active state if
history.update_generation_status raises after cancel_generation_job and
task_manager.complete_generation have already modified the queue; wrap the call
to history.update_generation_status(...) in a try/except that catches database
errors (e.g. OperationalError/Exception), logs the failure with logger.info or
logger.error including generation_id and cancellation_state, and returns a
partial-success response (e.g. message indicating queue removed but DB update
failed) instead of letting the exception propagate; alternatively persist the
new status (call history.update_generation_status) before calling
get_task_manager().complete_generation to ensure DB is updated first. Also add a
minimal response_model (e.g. {"message": str}) to the cancel_generation route
decorator and add a logger.info("Cancellation requested for generation %s
(state=%s)", generation_id, cancellation_state) before performing queue changes
to improve visibility.
| if (gen.status or "completed") not in ("loading_model", "generating"): | ||
| raise HTTPException(status_code=400, detail="Only active generations can be cancelled") |
There was a problem hiding this comment.
Status check is susceptible to a benign TOCTOU.
Between this check and cancel_generation_job(generation_id) below, the generation can transition to completed/failed and the 400 message ("Only active generations can be cancelled") will no longer be accurate — the request will instead correctly bounce to 409 from the next check. Not a correctness bug (the 409 path is sound), but worth noting that the 400 branch is racy and primarily serves the clear "never was active" case. No change required.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/routes/generations.py` around lines 200 - 201, The current status
check is racy; replace the expression (gen.status or "completed") not in
("loading_model", "generating") with a strict "never-was-active" check: if
gen.status is None: raise HTTPException(status_code=400, detail="Generation was
never active and cannot be cancelled"). Leave cancellation and the subsequent
conflict handling to cancel_generation_job(generation_id) so transitions after
this check are handled consistently; update the error message to refer only to
the never-active case and reference gen.status and cancel_generation_job for
locating the code.
| except asyncio.CancelledError: | ||
| await history.update_generation_status( | ||
| generation_id=generation_id, | ||
| status="failed", | ||
| db=bg_db, | ||
| error="Generation cancelled", | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, find and read the generation.py file around the specified lines
fd generation.py --type fRepository: jamiepine/voicebox
Length of output: 92
🏁 Script executed:
# Read the generation.py file to see the context around lines 130-136
cat -n backend/services/generation.py | head -180Repository: jamiepine/voicebox
Length of output: 7032
🏁 Script executed:
# Find and read task_queue.py to understand worker behavior
fd task_queue.py --type fRepository: jamiepine/voicebox
Length of output: 92
🏁 Script executed:
# Read task_queue.py to see how the worker handles task completion
cat -n backend/services/task_queue.pyRepository: jamiepine/voicebox
Length of output: 4499
🏁 Script executed:
# Search for history.update_generation_status to understand what it does
fd history.py --type f | grep -E "services|backend"Repository: jamiepine/voicebox
Length of output: 115
🏁 Script executed:
# Read history.py to see the update_generation_status function
cat -n backend/services/history.pyRepository: jamiepine/voicebox
Length of output: 13039
🏁 Script executed:
# Let's verify the exact behavior by checking if there are any tests that demonstrate the cancellation behavior
fd -e py test | head -20Repository: jamiepine/voicebox
Length of output: 362
🏁 Script executed:
# Read the task queue cancellation test to understand expected behavior
cat -n backend/tests/test_task_queue_cancellation.pyRepository: jamiepine/voicebox
Length of output: 1903
🏁 Script executed:
# Let's check if there are any other places where run_generation is called or cancelled
rg "run_generation|cancel_generation" --type py -B 2 -A 2 | head -80Repository: jamiepine/voicebox
Length of output: 4094
🏁 Script executed:
# Let's verify the task.cancelled() behavior by checking Python's asyncio documentation
# and the actual implementation of how task.cancelled() works
python3 << 'EOF'
import asyncio
# Test 1: Task that re-raises CancelledError
async def task_with_reraise():
try:
await asyncio.sleep(10)
except asyncio.CancelledError:
print("Caught, re-raising")
raise
# Test 2: Task that swallows CancelledError
async def task_swallowing():
try:
await asyncio.sleep(10)
except asyncio.CancelledError:
print("Caught and swallowed")
async def test():
# Test re-raise
t1 = asyncio.create_task(task_with_reraise())
await asyncio.sleep(0.01)
t1.cancel()
try:
await t1
except asyncio.CancelledError:
print(f"Test 1 - task.cancelled() = {t1.cancelled()}, task was re-raised")
# Test swallow
t2 = asyncio.create_task(task_swallowing())
await asyncio.sleep(0.01)
t2.cancel()
try:
await t2
except asyncio.CancelledError:
print(f"Test 2 - Should not reach here")
else:
print(f"Test 2 - task.cancelled() = {t2.cancelled()}, task was swallowed (no exception)")
asyncio.run(test())
EOFRepository: jamiepine/voicebox
Length of output: 223
🏁 Script executed:
# Let's verify the double-cancellation scenario by checking the init_queue behavior more carefully
cat -n backend/services/task_queue.py | sed -n '89,103p'Repository: jamiepine/voicebox
Length of output: 752
Re-raise CancelledError after persisting the failure state.
Swallowing asyncio.CancelledError is an asyncio anti-pattern:
- Task appears successful to the worker. The task in
task_queue.pycompletes "normally," sotask.cancelled()isFalseand the worker'sif not task.cancelled(): raiseguard becomes dead code. Observability (e.g.Task.result()/ logging) will show a successful completion of what was actually a cancellation. - Breaks structured concurrency. Any caller that cancels this coroutine and awaits it (today: the worker; tomorrow: potentially a supervisor during shutdown) will see it complete successfully and may proceed assuming the task ran to completion.
- Double-cancellation hazard. If the owning task is cancelled a second time (e.g. during worker shutdown via
init_queue(force=True)), theawait history.update_generation_status(...)inside the handler is itself a cancellation point and can raiseCancelledErrormidway, leaving the DB row stuck ingenerating.
🐛 Proposed fix — persist, then re-raise
except asyncio.CancelledError:
await history.update_generation_status(
generation_id=generation_id,
status="failed",
db=bg_db,
error="Generation cancelled",
)
+ raise
except Exception as e:For extra robustness against the double-cancel case, consider shielding the status update:
await asyncio.shield(history.update_generation_status(...))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| except asyncio.CancelledError: | |
| await history.update_generation_status( | |
| generation_id=generation_id, | |
| status="failed", | |
| db=bg_db, | |
| error="Generation cancelled", | |
| ) | |
| except asyncio.CancelledError: | |
| await history.update_generation_status( | |
| generation_id=generation_id, | |
| status="failed", | |
| db=bg_db, | |
| error="Generation cancelled", | |
| ) | |
| raise |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/services/generation.py` around lines 130 - 136, The except block that
catches asyncio.CancelledError around the generation flow (catching
asyncio.CancelledError and calling history.update_generation_status with
generation_id and db=bg_db) must re-raise the CancelledError after persisting
the failure state so cancellation propagates correctly; modify the handler to
first persist the "failed" status via history.update_generation_status (use
asyncio.shield(history.update_generation_status(...)) to protect the DB update
from being cancelled mid-write) and then re-raise the original
asyncio.CancelledError.
| 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() |
There was a problem hiding this comment.
Worker cancellation handling looks correct, but depends on run_generation re-raising CancelledError.
The except asyncio.CancelledError / if not task.cancelled(): raise pattern only works as intended if the inner coroutine propagates CancelledError. In this PR, backend/services/generation.py catches CancelledError and does not re-raise (see the separate comment there), so await task here will return normally after a user-initiated cancel and task.cancelled() will be False. The branch is effectively dead and the worker loses the ability to distinguish cancelled-runs from completed-runs.
Once run_generation is fixed to re-raise, this block will function correctly. No change needed here beyond the cross-file fix.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/services/task_queue.py` around lines 39 - 62, The worker
_generation_worker relies on run_generation to propagate asyncio.CancelledError
so the worker can detect cancelled jobs; update backend/services/generation.py
so that any except asyncio.CancelledError handler in run_generation (or any
coroutine invoked by it) does not swallow the exception — either remove the
CancelledError catch or re-raise it (raise) after any cleanup, ensuring
CancelledError reaches the awaiting code in _generation_worker.
| 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()) |
There was a problem hiding this comment.
init_queue(force=True) leaves the old worker in a weird state before re-init.
When cancelling _generation_worker_task and the running tasks, the cancels are fire-and-forget — you immediately overwrite the tracking structures on lines 104–108 without awaiting the prior worker/task teardown. If a cancelled in-flight task is partway through run_generation (e.g., awaiting DB I/O in the new CancelledError handler), its finally may still try to call into stale module state after you've reset _running_generation_tasks, etc.
This is primarily exercised in tests (init_queue(force=True) at the start of each test) and at process shutdown, so it's a minor concern, but consider awaiting shielded teardown of the old worker before resetting, or document that force=True is only safe during shutdown.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/services/task_queue.py` around lines 89 - 108, init_queue(force=True)
cancels the existing _generation_worker_task and running tasks but immediately
resets the module state, risking race conditions where cancelled tasks' finally
blocks access cleared structures; instead, after calling .cancel() on
_generation_worker_task and each task from _running_generation_tasks, await
their termination before clearing/overwriting _generation_queue,
_queued_generation_ids, _running_generation_tasks, and _cancelled_generation_ids
— e.g., collect the cancelled asyncio.Task objects (including
_generation_worker_task and values() of _running_generation_tasks), await them
using asyncio.gather(..., return_exceptions=True) or asyncio.wait with
appropriate shielding (asyncio.shield) to ensure they run their cleanup, then
only after they complete reset the module-level tracking and start a new worker
with create_background_task(_generation_worker()).
Closes #435
What changed
Verification
Summary by CodeRabbit
New Features
Tests