49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
"""Execution control helpers for generation harness workflows."""
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.services.generation_jobs import finish_generation_job
|
|
|
|
if TYPE_CHECKING:
|
|
from app.db.models import GenerationJob, Story
|
|
|
|
|
|
class GenerationJobCanceledError(Exception):
|
|
"""Raised when a running worker job has been canceled by the user."""
|
|
|
|
|
|
class ExecutionControl:
|
|
"""Runtime control surface for cancelable generation workflows."""
|
|
|
|
def __init__(self, db: AsyncSession):
|
|
self.db = db
|
|
|
|
async def stop_if_cancel_requested(
|
|
self,
|
|
*,
|
|
job: "GenerationJob | None",
|
|
story: "Story | None" = None,
|
|
) -> None:
|
|
"""Stop a worker-owned job at the next safe checkpoint after cancellation."""
|
|
|
|
if job is None:
|
|
return
|
|
|
|
await self.db.refresh(job)
|
|
if job.current_step != "cancel_requested":
|
|
return
|
|
|
|
await finish_generation_job(
|
|
self.db,
|
|
job=job,
|
|
story=story,
|
|
status="canceled",
|
|
current_step="generation_canceled",
|
|
error_message="Generation canceled by user.",
|
|
message="Generation job was canceled after a user request.",
|
|
)
|
|
raise GenerationJobCanceledError()
|
|
|