From 19fcd4065fad9ae1432622ab7d52c722f1e3475e Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Fri, 15 Nov 2024 20:07:01 -0800
Subject: [PATCH] handle navigation failure and fail a task in workflow run
 (#1205)

---
 skyvern/forge/sdk/workflow/models/block.py | 37 ++++++++++++++++++----
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py
index 34ea7b5c..63e026dc 100644
--- a/skyvern/forge/sdk/workflow/models/block.py
+++ b/skyvern/forge/sdk/workflow/models/block.py
@@ -298,14 +298,37 @@ class TaskBlock(Block):
                 raise Exception(f"Organization is missing organization_id={workflow.organization_id}")
 
             browser_state: BrowserState | None = None
-            if is_first_task:
-                browser_state = await app.BROWSER_MANAGER.get_or_create_for_workflow_run(
-                    workflow_run=workflow_run, url=self.url
+            try:
+                if is_first_task:
+                    browser_state = await app.BROWSER_MANAGER.get_or_create_for_workflow_run(
+                        workflow_run=workflow_run, url=self.url
+                    )
+                else:
+                    browser_state = app.BROWSER_MANAGER.get_for_workflow_run(workflow_run_id=workflow_run_id)
+                    if browser_state is None:
+                        raise MissingBrowserState(task_id=task.task_id, workflow_run_id=workflow_run_id)
+            except FailedToNavigateToUrl as e:
+                # Make sure the task is marked as failed in the database before raising the exception
+                await app.DATABASE.update_task(
+                    task.task_id,
+                    status=TaskStatus.failed,
+                    organization_id=workflow.organization_id,
+                    failure_reason=str(e),
                 )
-            else:
-                browser_state = app.BROWSER_MANAGER.get_for_workflow_run(workflow_run_id=workflow_run_id)
-                if browser_state is None:
-                    raise MissingBrowserState(task_id=task.task_id, workflow_run_id=workflow_run_id)
+                raise e
+            except Exception as e:
+                await app.DATABASE.update_task(
+                    task.task_id,
+                    status=TaskStatus.failed,
+                    organization_id=workflow.organization_id,
+                    failure_reason=str(e),
+                )
+                LOG.exception(
+                    "Failed to get browser state for task",
+                    task_id=task.task_id,
+                    workflow_run_id=workflow_run_id,
+                )
+                raise e
 
             working_page = await browser_state.get_working_page()
             if not working_page: