diff --git a/backend/app/routers/generation.py b/backend/app/routers/generation.py index eb6bfa1..6ca8fc7 100644 --- a/backend/app/routers/generation.py +++ b/backend/app/routers/generation.py @@ -164,6 +164,11 @@ def generate_storyboard_task(task_id: str, project_id: str, user_input: str): story_blocks = [b for b in json_blocks if b.get("type") not in ["character_sheet", "comic_config"]] # --- Missing Character Check & Fix --- + session.refresh(task) + if task.status == "cancelled": + log_task_event(session, task_id, "Task execution cancelled by user.") + return + story_char_names = set() for block in story_blocks: chars = block.get("characters", []) @@ -303,6 +308,12 @@ def generate_all_images_task(task_id: str, project_id: str): total_chars = len(project.characters) log_task_event(session, task_id, f"Generating {total_chars} characters...") for i, char in enumerate(project.characters): + # Check for cancellation + session.refresh(task) + if task.status == "cancelled": + log_task_event(session, task_id, "Task execution cancelled by user.") + return + if char.image_url: log_task_event(session, task_id, f"Character {char.name} already has image, skipping.") continue @@ -353,6 +364,12 @@ def generate_all_images_task(task_id: str, project_id: str): # For "one click", let's assume we scan all items. for i, item in enumerate(items): + # Check for cancellation + session.refresh(task) + if task.status == "cancelled": + log_task_event(session, task_id, "Task execution cancelled by user.") + return + # Update progress at start of loop # task.progress = int((i / total_items) * 100) # session.add(task) @@ -479,6 +496,12 @@ def generate_all_characters_task(task_id: str, project_id: str): log_task_event(session, task_id, f"Generating {total_chars} characters...") for i, char in enumerate(project.characters): + # Check for cancellation + session.refresh(task) + if task.status == "cancelled": + log_task_event(session, task_id, "Task execution cancelled by user.") + return + # if char.image_url: # logger.info(f"Character {char.name} already has image, skipping.") # continue diff --git a/backend/app/routers/tasks.py b/backend/app/routers/tasks.py index 8494f54..cfb11a1 100644 --- a/backend/app/routers/tasks.py +++ b/backend/app/routers/tasks.py @@ -21,3 +21,19 @@ def get_project_tasks(project_id: str, session: Session = Depends(get_session)): # Filter only recent or active tasks if list is too long? # For now return all, maybe limit 20 return tasks[:20] + +@router.post("/{task_id}/cancel", response_model=TaskRead) +def cancel_task(task_id: str, session: Session = Depends(get_session)): + task = session.get(Task, task_id) + if not task: + raise HTTPException(status_code=404, detail="Task not found") + + if task.status in ["completed", "failed", "cancelled"]: + return task + + task.status = "cancelled" + task.message = "Task cancelled by user" + session.add(task) + session.commit() + session.refresh(task) + return task \ No newline at end of file diff --git a/frontend/src/views/project/TaskManager.vue b/frontend/src/views/project/TaskManager.vue index fafdb5b..5075c1e 100644 --- a/frontend/src/views/project/TaskManager.vue +++ b/frontend/src/views/project/TaskManager.vue @@ -15,6 +15,16 @@