From c846d3871b8ccbe32958ed26daab2f560ccaba50 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Mon, 17 Feb 2025 02:40:47 +0800
Subject: [PATCH] do not clean up workflow in a task v2 block (#1771)

---
 skyvern/forge/agent.py                         |  6 +++++-
 skyvern/forge/sdk/services/observer_service.py | 10 +---------
 skyvern/webeye/browser_manager.py              | 14 ++++++++++++--
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index 4c8b5da3..9b5fcd15 100644
--- a/skyvern/forge/agent.py
+++ b/skyvern/forge/agent.py
@@ -352,6 +352,7 @@ class ForgeAgent:
                 detailed_output,
             ) = await self.initialize_execution_state(task, step, workflow_run, browser_session_id)
 
+            # mark step as completed and mark task as completed
             if (
                 not task.navigation_goal
                 and not task.data_extraction_goal
@@ -359,7 +360,10 @@ class ForgeAgent:
                 and not task.terminate_criterion
             ):
                 # most likely a GOTO_URL task block
-                # mark step as completed and mark task as completed
+                page = await browser_state.must_get_working_page()
+                current_url = page.url
+                if current_url.rstrip("/") != task.url.rstrip("/"):
+                    await page.goto(task.url)
                 step = await self.update_step(
                     step, status=StepStatus.completed, is_last=True, output=AgentStepOutput(action_results=[])
                 )
diff --git a/skyvern/forge/sdk/services/observer_service.py b/skyvern/forge/sdk/services/observer_service.py
index 59d3e564..5fdf064c 100644
--- a/skyvern/forge/sdk/services/observer_service.py
+++ b/skyvern/forge/sdk/services/observer_service.py
@@ -262,7 +262,7 @@ async def run_observer_task(
             organization_id=organization_id,
         )
     finally:
-        if workflow and workflow_run:
+        if workflow and workflow_run and workflow_run.parent_workflow_run_id is None:
             await app.WORKFLOW_SERVICE.clean_up_workflow(
                 workflow=workflow,
                 workflow_run=workflow_run,
@@ -699,14 +699,6 @@ async def handle_block_result(
         )
         await app.WORKFLOW_SERVICE.mark_workflow_run_as_canceled(workflow_run_id=workflow_run.workflow_run_id)
 
-        # TODO: we can also support webhook by adding api_key to the function signature
-        await app.WORKFLOW_SERVICE.clean_up_workflow(
-            workflow=workflow,
-            workflow_run=workflow_run,
-            need_call_webhook=False,
-            close_browser_on_completion=browser_session_id is None,
-            browser_session_id=browser_session_id,
-        )
     elif block_result.status == BlockStatus.failed:
         LOG.error(
             f"Block with type {block.block_type} failed for workflow run {workflow_run_id}",
diff --git a/skyvern/webeye/browser_manager.py b/skyvern/webeye/browser_manager.py
index 4d273b9f..065a9134 100644
--- a/skyvern/webeye/browser_manager.py
+++ b/skyvern/webeye/browser_manager.py
@@ -130,11 +130,19 @@ class BrowserManager:
         url: str | None = None,
         browser_session_id: str | None = None,
     ) -> BrowserState:
-        browser_state = self.get_for_workflow_run(workflow_run_id=workflow_run.workflow_run_id)
+        parent_workflow_run_id = workflow_run.parent_workflow_run_id
+        workflow_run_id = workflow_run.workflow_run_id
+        browser_state = self.get_for_workflow_run(workflow_run_id=workflow_run_id)
+        if parent_workflow_run_id:
+            browser_state = self.get_for_workflow_run(workflow_run_id=parent_workflow_run_id)
+            if browser_state:
+                self.pages[workflow_run_id] = browser_state
+
         if browser_state is not None:
             return browser_state
 
         if browser_session_id:
+            # TODO: what if there's a parent workflow run?
             LOG.info(
                 "Getting browser state for workflow run from persistent sessions manager",
                 browser_session_id=browser_session_id,
@@ -171,7 +179,9 @@ class BrowserManager:
                 organization_id=workflow_run.organization_id,
             )
 
-        self.pages[workflow_run.workflow_run_id] = browser_state
+        self.pages[workflow_run_id] = browser_state
+        if parent_workflow_run_id:
+            self.pages[parent_workflow_run_id] = browser_state
 
         # The URL here is only used when creating a new page, and not when using an existing page.
         # This will make sure browser_state.page is not None.