add endpoint for making a new debug session (with a new browser session) SKY-5666 (#3069)
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user