From 5145bc8a1bae57e9a644105ba8d0e9d68fecc244 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Mon, 7 Apr 2025 11:54:39 -0400
Subject: [PATCH] mark task v2 timed_out (#2113)

---
 skyvern/forge/sdk/workflow/service.py | 12 ++++++++++++
 skyvern/services/task_v2_service.py   | 17 +++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py
index 3e6ac2cb..e9694dbe 100644
--- a/skyvern/forge/sdk/workflow/service.py
+++ b/skyvern/forge/sdk/workflow/service.py
@@ -770,6 +770,18 @@ class WorkflowService:
             status=WorkflowRunStatus.canceled,
         )
 
+    async def mark_workflow_run_as_timed_out(self, workflow_run_id: str, failure_reason: str | None = None) -> None:
+        LOG.info(
+            f"Marking workflow run {workflow_run_id} as timed out",
+            workflow_run_id=workflow_run_id,
+            workflow_status="timed_out",
+        )
+        await app.DATABASE.update_workflow_run(
+            workflow_run_id=workflow_run_id,
+            status=WorkflowRunStatus.timed_out,
+            failure_reason=failure_reason,
+        )
+
     async def get_workflow_run(self, workflow_run_id: str, organization_id: str | None = None) -> WorkflowRun:
         workflow_run = await app.DATABASE.get_workflow_run(
             workflow_run_id=workflow_run_id,
diff --git a/skyvern/services/task_v2_service.py b/skyvern/services/task_v2_service.py
index c842c9f6..f853228f 100644
--- a/skyvern/services/task_v2_service.py
+++ b/skyvern/services/task_v2_service.py
@@ -1320,6 +1320,23 @@ async def mark_task_v2_as_terminated(
     return task_v2
 
 
+async def mark_task_v2_as_timed_out(
+    task_v2_id: str,
+    workflow_run_id: str | None = None,
+    organization_id: str | None = None,
+    failure_reason: str | None = None,
+) -> TaskV2:
+    task_v2 = await app.DATABASE.update_task_v2(
+        task_v2_id,
+        organization_id=organization_id,
+        status=TaskV2Status.timed_out,
+    )
+    if workflow_run_id:
+        await app.WORKFLOW_SERVICE.mark_workflow_run_as_timed_out(workflow_run_id, failure_reason)
+    await send_task_v2_webhook(task_v2)
+    return task_v2
+
+
 def _get_extracted_data_from_block_result(
     block_result: BlockResult,
     task_type: str,