Browser streaming refactor (#4064)

skyvern/forge/sdk/routes/streaming/verify.py

@@ -1,18 +1,43 @@
+"""
+Channels (see ./channels/README.md) variously rely on the state of one of
+these database entities:
+  - browser session
+  - task
+  - workflow run
+
+That is, channels are created on the basis of one of those entities, and that
+entity must be in a valid state for the channel to continue.
+
+So, in order to continue operating a channel, we need to periodically verify
+that the underlying entity is still valid. This module provides logic to
+perform those verifications.
+"""
+
+from __future__ import annotations
+
 import asyncio
+import typing as t
 from datetime import datetime
 
 import structlog
 
-import skyvern.forge.sdk.routes.streaming.clients as sc
 from skyvern.config import settings
 from skyvern.forge import app
 from skyvern.forge.sdk.schemas.persistent_browser_sessions import AddressablePersistentBrowserSession
 from skyvern.forge.sdk.schemas.tasks import Task, TaskStatus
 from skyvern.forge.sdk.workflow.models.workflow import WorkflowRun, WorkflowRunStatus
 
+if t.TYPE_CHECKING:
+    from skyvern.forge.sdk.routes.streaming.channels.message import MessageChannel
+    from skyvern.forge.sdk.routes.streaming.channels.vnc import VncChannel
+
 LOG = structlog.get_logger()
 
 
+class Constants:
+    POLL_INTERVAL_FOR_VERIFICATION_SECONDS = 5
+
+
 async def verify_browser_session(
     browser_session_id: str,
     organization_id: str,
@@ -122,9 +147,7 @@ async def verify_task(
             **browser_session.model_dump() | {"browser_address": browser_session.browser_address},
         )
     except Exception as e:
-        LOG.error(
-            "streaming-vnc.browser-session-reify-error", task_id=task_id, organization_id=organization_id, error=e
-        )
+        LOG.error("vnc.browser-session-reify-error", task_id=task_id, organization_id=organization_id, error=e)
         return task, None
 
     return task, addressable_browser_session
@@ -238,7 +261,7 @@ async def verify_workflow_run(
     return workflow_run, addressable_browser_session
 
 
-async def loop_verify_browser_session(verifiable: sc.MessageChannel | sc.Streaming) -> None:
+async def loop_verify_browser_session(verifiable: MessageChannel | VncChannel) -> None:
     """
     Loop until the browser session is cleared or the websocket is closed.
     """
@@ -251,27 +274,27 @@ async def loop_verify_browser_session(verifiable: sc.MessageChannel | sc.Streami
 
         verifiable.browser_session = browser_session
 
-        await asyncio.sleep(2)
+        await asyncio.sleep(Constants.POLL_INTERVAL_FOR_VERIFICATION_SECONDS)
 
 
-async def loop_verify_task(streaming: sc.Streaming) -> None:
+async def loop_verify_task(vnc_channel: VncChannel) -> None:
     """
     Loop until the task is cleared or the websocket is closed.
     """
 
-    while streaming.task and streaming.is_open:
+    while vnc_channel.task and vnc_channel.is_open:
         task, browser_session = await verify_task(
-            task_id=streaming.task.task_id,
-            organization_id=streaming.organization_id,
+            task_id=vnc_channel.task.task_id,
+            organization_id=vnc_channel.organization_id,
         )
 
-        streaming.task = task
-        streaming.browser_session = browser_session
+        vnc_channel.task = task
+        vnc_channel.browser_session = browser_session
 
-        await asyncio.sleep(2)
+        await asyncio.sleep(Constants.POLL_INTERVAL_FOR_VERIFICATION_SECONDS)
 
 
-async def loop_verify_workflow_run(verifiable: sc.MessageChannel | sc.Streaming) -> None:
+async def loop_verify_workflow_run(verifiable: MessageChannel | VncChannel) -> None:
     """
     Loop until the workflow run is cleared or the websocket is closed.
     """
@@ -285,4 +308,4 @@ async def loop_verify_workflow_run(verifiable: sc.MessageChannel | sc.Streaming)
         verifiable.workflow_run = workflow_run
         verifiable.browser_session = browser_session
 
-        await asyncio.sleep(2)
+        await asyncio.sleep(Constants.POLL_INTERVAL_FOR_VERIFICATION_SECONDS)