diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 27982d7a..681d5359 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/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) diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index b9d80aa4..6b2be75b 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/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 diff --git a/skyvern/forge/sdk/schemas/debug_sessions.py b/skyvern/forge/sdk/schemas/debug_sessions.py index d588cc3d..48aa658d 100644 --- a/skyvern/forge/sdk/schemas/debug_sessions.py +++ b/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