From 9d6e30f1de0e53791669d0c341758c8f79e39a2e Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Thu, 3 Jul 2025 18:45:04 -0700
Subject: [PATCH] create task, task v2 and workflow run with browser_session_id
 (#2872)

---
 skyvern/forge/agent.py                        | 1 +
 skyvern/forge/sdk/db/client.py                | 4 ++++
 skyvern/forge/sdk/db/utils.py                 | 2 ++
 skyvern/forge/sdk/schemas/tasks.py            | 1 +
 skyvern/forge/sdk/workflow/models/workflow.py | 1 +
 skyvern/forge/sdk/workflow/service.py         | 1 +
 6 files changed, 10 insertions(+)

diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index 55eed339..18753623 100644
--- a/skyvern/forge/agent.py
+++ b/skyvern/forge/agent.py
@@ -243,6 +243,7 @@ class ForgeAgent:
             model=task_request.model,
             max_screenshot_scrolling_times=task_request.max_screenshot_scrolls,
             extra_http_headers=task_request.extra_http_headers,
+            browser_session_id=task_request.browser_session_id,
         )
         LOG.info(
             "Created new task",
diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py
index d8f3ac13..a849e357 100644
--- a/skyvern/forge/sdk/db/client.py
+++ b/skyvern/forge/sdk/db/client.py
@@ -151,6 +151,7 @@ class AgentDB:
         model: dict[str, Any] | None = None,
         max_screenshot_scrolling_times: int | None = None,
         extra_http_headers: dict[str, str] | None = None,
+        browser_session_id: str | None = None,
     ) -> Task:
         try:
             async with self.Session() as session:
@@ -180,6 +181,7 @@ class AgentDB:
                     model=model,
                     max_screenshot_scrolling_times=max_screenshot_scrolling_times,
                     extra_http_headers=extra_http_headers,
+                    browser_session_id=browser_session_id,
                 )
                 session.add(new_task)
                 await session.commit()
@@ -1553,6 +1555,7 @@ class AgentDB:
         workflow_permanent_id: str,
         workflow_id: str,
         organization_id: str,
+        browser_session_id: str | None = None,
         proxy_location: ProxyLocation | None = None,
         webhook_callback_url: str | None = None,
         totp_verification_url: str | None = None,
@@ -1567,6 +1570,7 @@ class AgentDB:
                     workflow_permanent_id=workflow_permanent_id,
                     workflow_id=workflow_id,
                     organization_id=organization_id,
+                    browser_session_id=browser_session_id,
                     proxy_location=proxy_location,
                     status="created",
                     webhook_callback_url=webhook_callback_url,
diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py
index 1cb6f65c..9ee309b4 100644
--- a/skyvern/forge/sdk/db/utils.py
+++ b/skyvern/forge/sdk/db/utils.py
@@ -144,6 +144,7 @@ def convert_to_task(task_obj: TaskModel, debug_enabled: bool = False, workflow_p
         started_at=task_obj.started_at,
         finished_at=task_obj.finished_at,
         max_screenshot_scrolls=task_obj.max_screenshot_scrolling_times,
+        browser_session_id=task_obj.browser_session_id,
     )
     return task
 
@@ -268,6 +269,7 @@ def convert_to_workflow_run(
         parent_workflow_run_id=workflow_run_model.parent_workflow_run_id,
         workflow_id=workflow_run_model.workflow_id,
         organization_id=workflow_run_model.organization_id,
+        browser_session_id=workflow_run_model.browser_session_id,
         status=WorkflowRunStatus[workflow_run_model.status],
         failure_reason=workflow_run_model.failure_reason,
         proxy_location=(
diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py
index 06fc920b..c69ed153 100644
--- a/skyvern/forge/sdk/schemas/tasks.py
+++ b/skyvern/forge/sdk/schemas/tasks.py
@@ -242,6 +242,7 @@ class Task(TaskBase):
     organization_id: str
     workflow_run_id: str | None = None
     workflow_permanent_id: str | None = None
+    browser_session_id: str | None = None
     order: int | None = None
     retry: int | None = None
     max_steps_per_run: int | None = None
diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py
index c1077c02..e959247b 100644
--- a/skyvern/forge/sdk/workflow/models/workflow.py
+++ b/skyvern/forge/sdk/workflow/models/workflow.py
@@ -111,6 +111,7 @@ class WorkflowRun(BaseModel):
     workflow_id: str
     workflow_permanent_id: str
     organization_id: str
+    browser_session_id: str | None = None
     status: WorkflowRunStatus
     extra_http_headers: dict[str, str] | None = None
     proxy_location: ProxyLocation | None = None
diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py
index 11222a7f..ebcf282a 100644
--- a/skyvern/forge/sdk/workflow/service.py
+++ b/skyvern/forge/sdk/workflow/service.py
@@ -783,6 +783,7 @@ class WorkflowService:
             workflow_permanent_id=workflow_permanent_id,
             workflow_id=workflow_id,
             organization_id=organization_id,
+            browser_session_id=workflow_request.browser_session_id,
             proxy_location=workflow_request.proxy_location,
             webhook_callback_url=workflow_request.webhook_callback_url,
             totp_verification_url=workflow_request.totp_verification_url,