add endpoint for making a new debug session (with a new browser session) SKY-5666 (#3069)

skyvern/forge/sdk/db/client.py

@@ -3417,6 +3417,9 @@ class AgentDB:
                     .filter_by(organization_id=organization_id)
                     .filter_by(workflow_permanent_id=workflow_permanent_id)
                     .filter_by(user_id=user_id)
+                    .filter_by(deleted_at=None)
+                    .filter_by(status="created")
+                    .order_by(DebugSessionModel.created_at.desc())
                 )
             ).first()
 
@@ -3425,6 +3428,37 @@ class AgentDB:
 
             return DebugSession.model_validate(debug_session)
 
+    async def complete_debug_sessions(
+        self,
+        *,
+        organization_id: str,
+        user_id: str | None = None,
+        workflow_permanent_id: str | None = None,
+    ) -> list[DebugSession]:
+        async with self.Session() as session:
+            query = (
+                select(DebugSessionModel)
+                .filter_by(organization_id=organization_id)
+                .filter_by(deleted_at=None)
+                .filter_by(status="created")
+            )
+
+            if user_id:
+                query = query.filter_by(user_id=user_id)
+            if workflow_permanent_id:
+                query = query.filter_by(workflow_permanent_id=workflow_permanent_id)
+
+            models = (await session.scalars(query)).all()
+
+            for model in models:
+                model.status = "completed"
+
+            debug_sessions = [DebugSession.model_validate(model) for model in models]
+
+            await session.commit()
+
+            return debug_sessions
+
     async def create_debug_session(
         self,
         *,
@@ -3439,6 +3473,7 @@ class AgentDB:
                 workflow_permanent_id=workflow_permanent_id,
                 user_id=user_id,
                 browser_session_id=browser_session_id,
+                status="created",
             )
 
             session.add(debug_session)

skyvern/forge/sdk/routes/agent_protocol.py

@@ -1,6 +1,7 @@
 import asyncio
 from datetime import datetime, timedelta, timezone
 from enum import Enum
+from functools import partial
 from math import floor
 from typing import Annotated, Any
 
@@ -46,6 +47,7 @@ from skyvern.forge.sdk.schemas.organizations import (
     Organization,
     OrganizationUpdate,
 )
+from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
 from skyvern.forge.sdk.schemas.task_generations import GenerateTaskRequest, TaskGeneration
 from skyvern.forge.sdk.schemas.task_v2 import TaskV2Request
 from skyvern.forge.sdk.schemas.tasks import (
@@ -2159,3 +2161,112 @@ async def get_or_create_debug_session_by_user_and_workflow_permanent_id(
     )
 
     return debug_session
+
+
+@base_router.post(
+    "/debug-session/{workflow_permanent_id}/new",
+    include_in_schema=False,
+)
+async def new_browser_session_for_debug_session(
+    workflow_permanent_id: str,
+    current_org: Organization = Depends(org_auth_service.get_current_org),
+    current_user_id: str = Depends(org_auth_service.get_current_user_id),
+) -> DebugSession:
+    """
+    Create a new debug session, along with a new browser session. If any
+    existing debug session are found, close the browser sessions associated
+    with them.
+    """
+
+    completed_debug_sessions = await app.DATABASE.complete_debug_sessions(
+        organization_id=current_org.organization_id,
+        user_id=current_user_id,
+        workflow_permanent_id=workflow_permanent_id,
+    )
+
+    LOG.info(
+        f"Completed {len(completed_debug_sessions)} pre-existing debug session(s)",
+        num_completed_debug_sessions=len(completed_debug_sessions),
+        organization_id=current_org.organization_id,
+        user_id=current_user_id,
+        workflow_permanent_id=workflow_permanent_id,
+    )
+
+    if completed_debug_sessions:
+        closeable_browser_sessions: list[PersistentBrowserSession] = []
+
+        for debug_session in completed_debug_sessions:
+            browser_session = await app.DATABASE.get_persistent_browser_session(
+                debug_session.browser_session_id,
+                current_org.organization_id,
+            )
+
+            if browser_session and browser_session.completed_at is None:
+                closeable_browser_sessions.append(browser_session)
+
+        LOG.info(
+            f"Closing browser {len(closeable_browser_sessions)} browser session(s)",
+            organization_id=current_org.organization_id,
+            user_id=current_user_id,
+            workflow_permanent_id=workflow_permanent_id,
+        )
+
+        def handle_close_browser_session_error(
+            browser_session_id: str,
+            organization_id: str,
+            task: asyncio.Task,
+        ) -> None:
+            if task.exception():
+                LOG.error(
+                    f"Failed to close session: {task.exception()}",
+                    browser_session_id=browser_session_id,
+                    organization_id=organization_id,
+                )
+
+        for browser_session in closeable_browser_sessions:
+            LOG.info(
+                "Closing existing browser session for debug session",
+                browser_session_id=browser_session.persistent_browser_session_id,
+                organization_id=current_org.organization_id,
+            )
+
+            # NOTE(jdo): these may fail to actually close on infra, but the user
+            # wants (and should get) a new session regardless - so we will just
+            # log the error and continue
+            task = asyncio.create_task(
+                app.PERSISTENT_SESSIONS_MANAGER.close_session(
+                    current_org.organization_id,
+                    browser_session.persistent_browser_session_id,
+                )
+            )
+
+            task.add_done_callback(
+                partial(
+                    handle_close_browser_session_error,
+                    browser_session.persistent_browser_session_id,
+                    current_org.organization_id,
+                )
+            )
+
+    new_browser_session = await app.PERSISTENT_SESSIONS_MANAGER.create_session(
+        organization_id=current_org.organization_id,
+        timeout_minutes=settings.DEBUG_SESSION_TIMEOUT_MINUTES,
+    )
+
+    debug_session = await app.DATABASE.create_debug_session(
+        browser_session_id=new_browser_session.persistent_browser_session_id,
+        organization_id=current_org.organization_id,
+        user_id=current_user_id,
+        workflow_permanent_id=workflow_permanent_id,
+    )
+
+    LOG.info(
+        "Created new debug session",
+        debug_session_id=debug_session.debug_session_id,
+        browser_session_id=new_browser_session.persistent_browser_session_id,
+        organization_id=current_org.organization_id,
+        user_id=current_user_id,
+        workflow_permanent_id=workflow_permanent_id,
+    )
+
+    return debug_session

skyvern/forge/sdk/schemas/debug_sessions.py

@@ -1,7 +1,10 @@
+import typing as t
 from datetime import datetime
 
 from pydantic import BaseModel, ConfigDict
 
+DebugSessionStatus = t.Literal["created", "completed"]
+
 
 class DebugSession(BaseModel):
     model_config = ConfigDict(from_attributes=True)
@@ -11,3 +14,5 @@ class DebugSession(BaseModel):
     workflow_permanent_id: str | None = None
     created_at: datetime
     modified_at: datetime
+    deleted_at: datetime | None = None
+    status: DebugSessionStatus