From 521abe160d6cbd5aa3e67888c297dec7f220a14b Mon Sep 17 00:00:00 2001
From: Muhammed Salih Altun <muhammedsalihaltun@gmail.com>
Date: Wed, 24 Apr 2024 22:08:45 +0300
Subject: [PATCH] Allow querying multiple statuses

---
 skyvern-frontend/src/routes/tasks/list/TaskList.tsx | 11 +++++++----
 skyvern/forge/sdk/db/client.py                      |  4 ++--
 skyvern/forge/sdk/routes/agent_protocol.py          |  2 +-
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/skyvern-frontend/src/routes/tasks/list/TaskList.tsx b/skyvern-frontend/src/routes/tasks/list/TaskList.tsx
index f397d3385..6da6fb42e 100644
--- 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
index 051a83065..f26c50934 100644
--- 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
index 4855e9e4f..f61d6d30f 100644
--- 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:
     """