Allow querying multiple statuses (#232)

Skyvern-AI · Apr 24, 2024 · 95f8ad1 · 95f8ad1
1 parent 1d32953
commit 95f8ad1
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/skyvern-frontend/src/routes/tasks/list/TaskList.tsx b/skyvern-frontend/src/routes/tasks/list/TaskList.tsx
@@ -46,12 +46,15 @@ function TaskList() {
   } = useQuery<Array<TaskApiResponse>>({
     queryKey: ["tasks", "all", page],
     queryFn: async () => {
+      const params = new URLSearchParams();
+      params.append("page", String(page));
+      params.append("page_size", String(PAGE_SIZE));
+      params.append("task_status", "completed");
+      params.append("task_status", "failed");
+      params.append("task_status", "terminated");
       return client
         .get("/tasks", {
-          params: {
-            page,
-            page_size: PAGE_SIZE,
-          },
+          params,
         })
         .then((response) => response.data);
     },

diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py
@@ -394,7 +394,7 @@ async def get_tasks(
         self,
         page: int = 1,
         page_size: int = 10,
-        task_status: TaskStatus | None = None,
+        task_status: list[TaskStatus] | None = None,
         organization_id: str | None = None,
     ) -> list[Task]:
         """
@@ -411,7 +411,7 @@ async def get_tasks(
                 db_page = page - 1  # offset logic is 0 based
                 query = select(TaskModel).filter_by(organization_id=organization_id)
                 if task_status:
-                    query = query.filter_by(status=task_status)
+                    query = query.filter(TaskModel.status.in_(task_status))
                 query = query.order_by(TaskModel.created_at.desc()).limit(page_size).offset(db_page * page_size)
                 tasks = (await session.scalars(query)).all()
                 return [convert_to_task(task, debug_enabled=self.debug_enabled) for task in tasks]

diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -313,7 +313,7 @@ async def get_task_internal(
 async def get_agent_tasks(
     page: int = Query(1, ge=1),
     page_size: int = Query(10, ge=1),
-    task_status: TaskStatus | None = None,
+    task_status: Annotated[list[TaskStatus] | None, Query()] = None,
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> Response:
     """