Skip to content

Commit

Permalink
Add retry endpoint for failed tasks
Browse files Browse the repository at this point in the history
  • Loading branch information
Aleksandr Movchan committed Feb 4, 2025
1 parent e95212a commit b4823f5
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 0 deletions.
26 changes: 26 additions & 0 deletions aana/routers/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,32 @@ async def delete_task(
return TaskInfo.from_entity(task)


@router.post(
"/tasks/{task_id}/retry",
summary="Retry Failed Task",
description="Retry a failed task by resetting its status to CREATED.",
)
async def retry_task(
task_id: str, db: GetDbDependency, user_id: UserIdDependency
) -> TaskInfo:
"""Retry a failed task by resetting its status."""
task_repo = TaskRepository(db)
task = task_repo.read(task_id, check=False)
if not task or task.user_id != user_id:
raise HTTPException(
status_code=404,
detail="Task not found",
)
if task.status != TaskStatus.FAILED:
raise HTTPException(
status_code=400,
detail="Only failed tasks can be retried",
)

updated_task = task_repo.retry_task(task.id)
return TaskInfo.from_entity(updated_task)


# Legacy endpoints (to be removed in the future)


Expand Down
19 changes: 19 additions & 0 deletions aana/storage/repository/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,25 @@ def update_status(
self.session.commit()
return task

def retry_task(self, task_id: str) -> TaskEntity:
"""Retry a task. The task will reset to CREATED status.
Args:
task_id (str): The ID of the task.
Returns:
TaskEntity: The updated task.
"""
task = self.read(task_id)
task.status = TaskStatus.CREATED
task.progress = 0
task.result = None
task.num_retries = 0
task.assigned_at = None
task.completed_at = None
self.session.commit()
return task

def get_active_tasks(self) -> list[TaskEntity]:
"""Fetches all active tasks.
Expand Down

0 comments on commit b4823f5

Please sign in to comment.