diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 259c235f..051a8306 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -390,7 +390,13 @@ class AgentDB: LOG.error("UnexpectedError", exc_info=True) raise - async def get_tasks(self, page: int = 1, page_size: int = 10, organization_id: str | None = None) -> list[Task]: + async def get_tasks( + self, + page: int = 1, + page_size: int = 10, + task_status: TaskStatus | None = None, + organization_id: str | None = None, + ) -> list[Task]: """ Get all tasks. :param page: Starts at 1 @@ -403,15 +409,11 @@ class AgentDB: try: async with self.Session() as session: db_page = page - 1 # offset logic is 0 based - tasks = ( - await session.scalars( - select(TaskModel) - .filter_by(organization_id=organization_id) - .order_by(TaskModel.created_at.desc()) - .limit(page_size) - .offset(db_page * page_size) - ) - ).all() + query = select(TaskModel).filter_by(organization_id=organization_id) + if task_status: + query = query.filter_by(status=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] except SQLAlchemyError: LOG.error("SQLAlchemyError", exc_info=True) diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index 3c839ea4..4855e9e4 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -313,6 +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, current_org: Organization = Depends(org_auth_service.get_current_org), ) -> Response: """ @@ -323,7 +324,9 @@ async def get_agent_tasks( get_agent_task endpoint. """ analytics.capture("skyvern-oss-agent-tasks-get") - tasks = await app.DATABASE.get_tasks(page, page_size, organization_id=current_org.organization_id) + tasks = await app.DATABASE.get_tasks( + page, page_size, task_status=task_status, organization_id=current_org.organization_id + ) return ORJSONResponse([task.to_task_response().model_dump() for task in tasks])