From bea9a8103b2517881c4d3cf16d19728513660025 Mon Sep 17 00:00:00 2001
From: davelopez <46503462+davelopez@users.noreply.github.com>
Date: Thu, 14 Sep 2023 15:41:29 +0200
Subject: [PATCH] Add task status type annotation

---
 client/src/schema/schema.ts            |  8 +++++++-
 lib/galaxy/managers/tasks.py           | 27 +++++++++++++++++++++++---
 lib/galaxy/webapps/galaxy/api/tasks.py |  7 +++++--
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/client/src/schema/schema.ts b/client/src/schema/schema.ts
index 789deebb50c3..54b20f389e31 100644
--- a/client/src/schema/schema.ts
+++ b/client/src/schema/schema.ts
@@ -8688,6 +8688,12 @@ export interface components {
             | "Page"
             | "StoredWorkflow"
             | "Visualization";
+        /**
+         * TaskState
+         * @description Enum representing the possible states of a task.
+         * @enum {string}
+         */
+        TaskState: "PENDING" | "STARTED" | "RETRY" | "FAILURE" | "SUCCESS";
         /** ToolDataDetails */
         ToolDataDetails: {
             /**
@@ -16677,7 +16683,7 @@ export interface operations {
             /** @description String indicating task state. */
             200: {
                 content: {
-                    "application/json": string;
+                    "application/json": components["schemas"]["TaskState"];
                 };
             };
             /** @description Validation Error */
diff --git a/lib/galaxy/managers/tasks.py b/lib/galaxy/managers/tasks.py
index 468cf7f0863d..8509829e1ef3 100644
--- a/lib/galaxy/managers/tasks.py
+++ b/lib/galaxy/managers/tasks.py
@@ -2,11 +2,32 @@
     ABCMeta,
     abstractmethod,
 )
+from enum import Enum
+from typing import cast
 from uuid import UUID
 
 from celery.result import AsyncResult
 
 
+class TaskState(str, Enum):
+    """Enum representing the possible states of a task."""
+
+    PENDING = "PENDING"
+    """The task is waiting for execution."""
+
+    STARTED = "STARTED"
+    """The task has been started."""
+
+    RETRY = "RETRY"
+    """The task is to be retried, possibly because of failure."""
+
+    FAILURE = "FAILURE"
+    """The task raised an exception, or has exceeded the retry limit."""
+
+    SUCCESS = "SUCCESS"
+    """The task executed successfully."""
+
+
 class AsyncTasksManager(metaclass=ABCMeta):
     @abstractmethod
     def is_ready(self, task_uuid: UUID) -> bool:
@@ -28,7 +49,7 @@ def has_failed(self, task_uuid: UUID) -> bool:
         """Return `True` if the task failed."""
 
     @abstractmethod
-    def get_state(self, task_uuid: UUID) -> str:
+    def get_state(self, task_uuid: UUID) -> TaskState:
         """Returns the current state of the task as a string."""
 
 
@@ -50,9 +71,9 @@ def has_failed(self, task_uuid: UUID) -> bool:
         """Return `True` if the task failed."""
         return self._get_result(task_uuid).failed()
 
-    def get_state(self, task_uuid: UUID) -> str:
+    def get_state(self, task_uuid: UUID) -> TaskState:
         """Returns the tasks current state as a string."""
-        return str(self._get_result(task_uuid).state)
+        return cast(TaskState, str(self._get_result(task_uuid).state))
 
     def _get_result(self, task_uuid: UUID) -> AsyncResult:
         return AsyncResult(str(task_uuid))
diff --git a/lib/galaxy/webapps/galaxy/api/tasks.py b/lib/galaxy/webapps/galaxy/api/tasks.py
index f57b5847682b..158222b1bec0 100644
--- a/lib/galaxy/webapps/galaxy/api/tasks.py
+++ b/lib/galaxy/webapps/galaxy/api/tasks.py
@@ -4,7 +4,10 @@
 import logging
 from uuid import UUID
 
-from galaxy.managers.tasks import AsyncTasksManager
+from galaxy.managers.tasks import (
+    AsyncTasksManager,
+    TaskState,
+)
 from . import (
     depends,
     Router,
@@ -24,5 +27,5 @@ class FastAPITasks:
         summary="Determine state of task ID",
         response_description="String indicating task state.",
     )
-    def state(self, task_id: UUID) -> str:
+    def state(self, task_id: UUID) -> TaskState:
         return self.manager.get_state(task_id)