From fa0fce490427b05b0b4b56772a0ea675a52eda11 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Thu, 12 Jun 2025 13:39:16 -0700
Subject: [PATCH] return pbs id on workflow run endpoint (#2699)

---
 skyvern/client/types/workflow_run_response.py |  5 +++++
 skyvern/forge/sdk/routes/agent_protocol.py    | 19 +++++++++++++++++++
 skyvern/forge/sdk/workflow/models/workflow.py |  1 +
 skyvern/schemas/runs.py                       |  5 +++++
 4 files changed, 30 insertions(+)

diff --git a/skyvern/client/types/workflow_run_response.py b/skyvern/client/types/workflow_run_response.py
index e8cd407d..98c0407f 100644
--- a/skyvern/client/types/workflow_run_response.py
+++ b/skyvern/client/types/workflow_run_response.py
@@ -67,6 +67,11 @@ class WorkflowRunResponse(UniversalBaseModel):
     The original request parameters used to start this workflow run
     """
 
+    browser_session_id: typing.Optional[str] = pydantic.Field(default=None)
+    """
+    ID of the Skyvern persistent browser session used for this run
+    """
+
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
     else:
diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
index 8a3fdfb8..fe3bcc12 100644
--- a/skyvern/forge/sdk/routes/agent_protocol.py
+++ b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -1357,12 +1357,31 @@ async def get_workflow_run_with_workflow_id(
         include_cost=True,
     )
     return_dict = workflow_run_status_response.model_dump()
+
+    browser_session = await app.DATABASE.get_persistent_browser_session_by_runnable_id(
+        runnable_id=workflow_run_id,
+        organization_id=current_org.organization_id,
+    )
+
+    browser_session_id = browser_session.persistent_browser_session_id if browser_session else None
+
+    LOG.info(
+        "workflow-run.assign-browser-session-id",
+        runnable_id=workflow_run_id,
+        browser_session_id=browser_session_id,
+        organization_id=current_org.organization_id,
+    )
+
+    return_dict["browser_session_id"] = browser_session_id
+
     task_v2 = await app.DATABASE.get_task_v2_by_workflow_run_id(
         workflow_run_id=workflow_run_id,
         organization_id=current_org.organization_id,
     )
+
     if task_v2:
         return_dict["task_v2"] = task_v2.model_dump(by_alias=True)
+
     return return_dict
 
 
diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py
index ca555337..1880446c 100644
--- a/skyvern/forge/sdk/workflow/models/workflow.py
+++ b/skyvern/forge/sdk/workflow/models/workflow.py
@@ -161,3 +161,4 @@ class WorkflowRunResponseBase(BaseModel):
     total_cost: float | None = None
     task_v2: TaskV2 | None = None
     workflow_title: str | None = None
+    browser_session_id: str | None = None
diff --git a/skyvern/schemas/runs.py b/skyvern/schemas/runs.py
index 4ca0afee..f1e64ac2 100644
--- a/skyvern/schemas/runs.py
+++ b/skyvern/schemas/runs.py
@@ -365,6 +365,11 @@ class BaseRunResponse(BaseModel):
         description="URL to the application UI where the run can be viewed",
         examples=["https://app.skyvern.com/tasks/tsk_123", "https://app.skyvern.com/workflows/wpid_123/wr_123"],
     )
+    browser_session_id: str | None = Field(
+        default=None,
+        description="ID of the Skyvern persistent browser session used for this run",
+        examples=["pbs_123"],
+    )
 
 
 class TaskRunResponse(BaseRunResponse):