Remove frontend hack for requesting persistent browser sessions, part ii (backend) (#3052)

skyvern/forge/sdk/routes/agent_protocol.py

@@ -1,5 +1,7 @@
 import asyncio
+from datetime import datetime, timedelta, timezone
 from enum import Enum
+from math import floor
 from typing import Annotated, Any
 
 import structlog
@@ -37,6 +39,7 @@ from skyvern.forge.sdk.routes.code_samples import (
 )
 from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router, legacy_v2_router
 from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestionBase, AISuggestionRequest
+from skyvern.forge.sdk.schemas.debug_sessions import DebugSession
 from skyvern.forge.sdk.schemas.organizations import (
     GetOrganizationAPIKeysResponse,
     GetOrganizationsResponse,
@@ -2005,3 +2008,110 @@ async def _flatten_workflow_run_timeline(organization_id: str, workflow_run_id:
         final_workflow_run_block_timeline.extend(thought_timeline)
     final_workflow_run_block_timeline.sort(key=lambda x: x.created_at, reverse=True)
     return final_workflow_run_block_timeline
+
+
+@base_router.get(
+    "/debug-session/{workflow_permanent_id}",
+    include_in_schema=False,
+)
+async def get_or_create_debug_session_by_user_and_workflow_permanent_id(
+    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:
+    """
+    `current_user_id` is a unique identifier for a user, but does not map to an
+    entity in the database (at time of writing)
+
+    If the debug session does not exist, a new one will be created.
+
+    If the debug session exists, the timeout will be extended to 4 hours from
+    the time of the request.
+    """
+
+    debug_session = await app.DATABASE.get_debug_session(
+        organization_id=current_org.organization_id,
+        user_id=current_user_id,
+        workflow_permanent_id=workflow_permanent_id,
+    )
+
+    if not debug_session:
+        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,
+        )
+
+        return debug_session
+
+    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:
+        if browser_session.started_at is None or browser_session.timeout_minutes is None:
+            LOG.warning(
+                "Persistent browser session started_at or timeout_minutes is None; assumption == this is normal, and they will become non-None shortly",
+                debug_session_id=debug_session.debug_session_id,
+                organization_id=current_org.organization_id,
+                workflow_permanent_id=workflow_permanent_id,
+                user_id=current_user_id,
+            )
+            return debug_session
+        right_now = datetime.now(timezone.utc)
+        current_timeout_minutes = browser_session.timeout_minutes
+        started_at_utc = (
+            browser_session.started_at.replace(tzinfo=timezone.utc)
+            if browser_session.started_at.tzinfo is None
+            else browser_session.started_at
+        )
+        current_timeout_datetime = started_at_utc + timedelta(minutes=float(current_timeout_minutes))
+        minutes_left = (current_timeout_datetime - right_now).total_seconds() / 60
+
+        if minutes_left >= settings.DEBUG_SESSION_TIMEOUT_THRESHOLD_MINUTES:
+            new_timeout_datetime = right_now + timedelta(minutes=settings.DEBUG_SESSION_TIMEOUT_MINUTES)
+            minutes_diff = floor((new_timeout_datetime - current_timeout_datetime).total_seconds() / 60)
+            new_timeout_minutes = current_timeout_minutes + minutes_diff
+
+            LOG.info(
+                f"Extended persistent browser session (for debugging) by {minutes_diff} minute(s)",
+                minutes_diff=minutes_diff,
+                debug_session_id=debug_session.debug_session_id,
+                organization_id=current_org.organization_id,
+                workflow_permanent_id=workflow_permanent_id,
+                user_id=current_user_id,
+            )
+
+            await app.DATABASE.update_persistent_browser_session(
+                browser_session_id=debug_session.browser_session_id,
+                organization_id=current_org.organization_id,
+                timeout_minutes=new_timeout_minutes,
+            )
+
+            return debug_session
+        else:
+            LOG.info(
+                "pbs for debug session has expired",
+                debug_session_id=debug_session.debug_session_id,
+                organization_id=current_org.organization_id,
+                workflow_permanent_id=workflow_permanent_id,
+                user_id=current_user_id,
+            )
+
+    browser_session = await app.PERSISTENT_SESSIONS_MANAGER.create_session(
+        organization_id=current_org.organization_id,
+        timeout_minutes=settings.DEBUG_SESSION_TIMEOUT_MINUTES,
+    )
+
+    await app.DATABASE.update_debug_session(
+        debug_session_id=debug_session.debug_session_id,
+        browser_session_id=browser_session.persistent_browser_session_id,
+    )
+
+    return debug_session