From df8275aa4be081b17ce8de1f983b8375edff7236 Mon Sep 17 00:00:00 2001
From: Jonathan Dobson <jon.m.dobson@gmail.com>
Date: Mon, 4 Aug 2025 15:14:05 -0400
Subject: [PATCH] set status to failed when browser session fails to initially
 connect (#3101)

---
 skyvern/forge/sdk/db/client.py                | 11 ++++-
 .../schemas/persistent_browser_sessions.py    |  6 +++
 skyvern/webeye/persistent_sessions_manager.py | 43 ++++++++++++++++++-
 3 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py
index b63f44c2..a5cb68cd 100644
--- a/skyvern/forge/sdk/db/client.py
+++ b/skyvern/forge/sdk/db/client.py
@@ -3040,7 +3040,9 @@ class AgentDB:
     async def update_persistent_browser_session(
         self,
         browser_session_id: str,
-        timeout_minutes: int,
+        *,
+        status: str | None = None,
+        timeout_minutes: int | None = None,
         organization_id: str | None = None,
     ) -> PersistentBrowserSession:
         try:
@@ -3055,7 +3057,12 @@ class AgentDB:
                 ).first()
                 if not persistent_browser_session:
                     raise NotFoundError(f"PersistentBrowserSession {browser_session_id} not found")
-                persistent_browser_session.timeout_minutes = timeout_minutes
+
+                if status:
+                    persistent_browser_session.status = status
+                if timeout_minutes:
+                    persistent_browser_session.timeout_minutes = timeout_minutes
+
                 await session.commit()
                 await session.refresh(persistent_browser_session)
                 return PersistentBrowserSession.model_validate(persistent_browser_session)
diff --git a/skyvern/forge/sdk/schemas/persistent_browser_sessions.py b/skyvern/forge/sdk/schemas/persistent_browser_sessions.py
index f0d91ca4..4b5b4f88 100644
--- a/skyvern/forge/sdk/schemas/persistent_browser_sessions.py
+++ b/skyvern/forge/sdk/schemas/persistent_browser_sessions.py
@@ -2,6 +2,12 @@ from datetime import datetime
 
 from pydantic import BaseModel, ConfigDict
 
+FINAL_STATUSES = ("completed", "failed")
+
+
+def is_final_status(status: str | None) -> bool:
+    return status in FINAL_STATUSES
+
 
 class PersistentBrowserSession(BaseModel):
     model_config = ConfigDict(from_attributes=True)
diff --git a/skyvern/webeye/persistent_sessions_manager.py b/skyvern/webeye/persistent_sessions_manager.py
index d024cc93..1cda1968 100644
--- a/skyvern/webeye/persistent_sessions_manager.py
+++ b/skyvern/webeye/persistent_sessions_manager.py
@@ -11,7 +11,7 @@ from skyvern.config import settings
 from skyvern.exceptions import BrowserSessionNotRenewable, MissingBrowserAddressError
 from skyvern.forge.sdk.db.client import AgentDB
 from skyvern.forge.sdk.db.polls import wait_on_persistent_browser_address
-from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
+from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession, is_final_status
 from skyvern.webeye.browser_factory import BrowserState
 
 LOG = structlog.get_logger()
@@ -98,7 +98,7 @@ async def renew_session(database: AgentDB, session_id: str, organization_id: str
         new_timeout_minutes = current_timeout_minutes + minutes_diff
 
         browser_session = await database.update_persistent_browser_session(
-            browser_session_id=session_id,
+            session_id,
             organization_id=organization_id,
             timeout_minutes=new_timeout_minutes,
         )
@@ -115,6 +115,40 @@ async def renew_session(database: AgentDB, session_id: str, organization_id: str
     raise BrowserSessionNotRenewable("Session has expired", session_id)
 
 
+async def update_status(
+    db: AgentDB, session_id: str, organization_id: str, status: str
+) -> PersistentBrowserSession | None:
+    persistent_browser_session = await db.get_persistent_browser_session(session_id, organization_id)
+
+    if not persistent_browser_session:
+        LOG.warning(
+            "Cannot update browser session status, browser session not found in database",
+            organization_id=organization_id,
+            session_id=session_id,
+            desired_status=status,
+        )
+        return None
+
+    if is_final_status(status):
+        if is_final_status(persistent_browser_session.status):
+            LOG.warning(
+                "Attempted to update browser session status to a final status when it is already final",
+                browser_session_id=session_id,
+                organization_id=organization_id,
+                desired_status=status,
+                current_status=persistent_browser_session.status,
+            )
+            return None
+
+    persistent_browser_session = await db.update_persistent_browser_session(
+        session_id,
+        status=status,
+        organization_id=organization_id,
+    )
+
+    return persistent_browser_session
+
+
 class PersistentSessionsManager:
     instance: PersistentSessionsManager | None = None
     _browser_sessions: dict[str, BrowserSession] = dict()
@@ -238,6 +272,11 @@ class PersistentSessionsManager:
             await self.close_session(organization_id, session_id)
             raise
 
+    async def update_status(
+        self, session_id: str, organization_id: str, status: str
+    ) -> PersistentBrowserSession | None:
+        return await update_status(self.database, session_id, organization_id, status)
+
     async def release_browser_session(self, session_id: str, organization_id: str) -> None:
         """Release a specific browser session."""
         await self.database.release_persistent_browser_session(session_id, organization_id)