627 lines
20 KiB
Python
627 lines
20 KiB
Python
"""
|
|
Streaming VNC WebSocket connections.
|
|
|
|
NOTE(jdo:streaming-local-dev)
|
|
-----------------------------
|
|
- grep the above for local development seams
|
|
- augment those seams as indicated, then
|
|
- stand up https://github.com/jomido/whyvern
|
|
|
|
"""
|
|
|
|
import asyncio
|
|
import typing as t
|
|
from urllib.parse import urlparse
|
|
|
|
import structlog
|
|
import websockets
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
from websockets import Data
|
|
from websockets.exceptions import ConnectionClosedError
|
|
|
|
import skyvern.forge.sdk.routes.streaming_clients as sc
|
|
from skyvern.config import settings
|
|
from skyvern.forge import app
|
|
from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
|
|
from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router
|
|
from skyvern.forge.sdk.routes.streaming_agent import connected_agent
|
|
from skyvern.forge.sdk.routes.streaming_auth import auth
|
|
from skyvern.forge.sdk.routes.streaming_verify import (
|
|
loop_verify_browser_session,
|
|
loop_verify_task,
|
|
loop_verify_workflow_run,
|
|
verify_browser_session,
|
|
verify_task,
|
|
verify_workflow_run,
|
|
)
|
|
from skyvern.forge.sdk.utils.aio import collect
|
|
|
|
LOG = structlog.get_logger()
|
|
|
|
|
|
class Constants:
|
|
MissingXApiKey = "<missing-x-api-key>"
|
|
|
|
|
|
async def get_x_api_key(organization_id: str) -> str:
|
|
token = await app.DATABASE.get_valid_org_auth_token(
|
|
organization_id,
|
|
OrganizationAuthTokenType.api.value,
|
|
)
|
|
|
|
if not token:
|
|
LOG.warning(
|
|
"No valid API key found for organization when streaming.",
|
|
organization_id=organization_id,
|
|
)
|
|
x_api_key = Constants.MissingXApiKey
|
|
else:
|
|
x_api_key = token.token
|
|
|
|
return x_api_key
|
|
|
|
|
|
async def get_streaming_for_browser_session(
|
|
client_id: str,
|
|
browser_session_id: str,
|
|
organization_id: str,
|
|
websocket: WebSocket,
|
|
) -> tuple[sc.Streaming, sc.Loops] | None:
|
|
"""
|
|
Return a streaming context for a browser session, with a list of loops to run concurrently.
|
|
"""
|
|
|
|
LOG.info("Getting streaming context for browser session.", browser_session_id=browser_session_id)
|
|
|
|
browser_session = await verify_browser_session(
|
|
browser_session_id=browser_session_id,
|
|
organization_id=organization_id,
|
|
)
|
|
|
|
if not browser_session:
|
|
LOG.info(
|
|
"No initial browser session found.", browser_session_id=browser_session_id, organization_id=organization_id
|
|
)
|
|
return None
|
|
|
|
x_api_key = await get_x_api_key(organization_id)
|
|
|
|
streaming = sc.Streaming(
|
|
client_id=client_id,
|
|
interactor="agent",
|
|
organization_id=organization_id,
|
|
vnc_port=settings.SKYVERN_BROWSER_VNC_PORT,
|
|
browser_session=browser_session,
|
|
x_api_key=x_api_key,
|
|
websocket=websocket,
|
|
)
|
|
|
|
LOG.info("Got streaming context for browser session.", streaming=streaming)
|
|
|
|
loops = [
|
|
asyncio.create_task(loop_verify_browser_session(streaming)),
|
|
asyncio.create_task(loop_stream_vnc(streaming)),
|
|
]
|
|
|
|
return streaming, loops
|
|
|
|
|
|
async def get_streaming_for_task(
|
|
client_id: str,
|
|
task_id: str,
|
|
organization_id: str,
|
|
websocket: WebSocket,
|
|
) -> tuple[sc.Streaming, sc.Loops] | None:
|
|
"""
|
|
Return a streaming context for a task, with a list of loops to run concurrently.
|
|
"""
|
|
|
|
task, browser_session = await verify_task(task_id=task_id, organization_id=organization_id)
|
|
|
|
if not task:
|
|
LOG.info("No initial task found.", task_id=task_id, organization_id=organization_id)
|
|
return None
|
|
|
|
if not browser_session:
|
|
LOG.info("No initial browser session found for task.", task_id=task_id, organization_id=organization_id)
|
|
return None
|
|
|
|
x_api_key = await get_x_api_key(organization_id)
|
|
|
|
streaming = sc.Streaming(
|
|
client_id=client_id,
|
|
interactor="agent",
|
|
organization_id=organization_id,
|
|
vnc_port=settings.SKYVERN_BROWSER_VNC_PORT,
|
|
x_api_key=x_api_key,
|
|
websocket=websocket,
|
|
browser_session=browser_session,
|
|
task=task,
|
|
)
|
|
|
|
loops = [
|
|
asyncio.create_task(loop_verify_task(streaming)),
|
|
asyncio.create_task(loop_stream_vnc(streaming)),
|
|
]
|
|
|
|
return streaming, loops
|
|
|
|
|
|
async def get_streaming_for_workflow_run(
|
|
client_id: str,
|
|
workflow_run_id: str,
|
|
organization_id: str,
|
|
websocket: WebSocket,
|
|
) -> tuple[sc.Streaming, sc.Loops] | None:
|
|
"""
|
|
Return a streaming context for a workflow run, with a list of loops to run concurrently.
|
|
"""
|
|
|
|
LOG.info("Getting streaming context for workflow run.", workflow_run_id=workflow_run_id)
|
|
|
|
workflow_run, browser_session = await verify_workflow_run(
|
|
workflow_run_id=workflow_run_id,
|
|
organization_id=organization_id,
|
|
)
|
|
|
|
if not workflow_run:
|
|
LOG.info("No initial workflow run found.", workflow_run_id=workflow_run_id, organization_id=organization_id)
|
|
return None
|
|
|
|
if not browser_session:
|
|
LOG.info(
|
|
"No initial browser session found for workflow run.",
|
|
workflow_run_id=workflow_run_id,
|
|
organization_id=organization_id,
|
|
)
|
|
return None
|
|
|
|
x_api_key = await get_x_api_key(organization_id)
|
|
|
|
streaming = sc.Streaming(
|
|
client_id=client_id,
|
|
interactor="agent",
|
|
organization_id=organization_id,
|
|
vnc_port=settings.SKYVERN_BROWSER_VNC_PORT,
|
|
browser_session=browser_session,
|
|
workflow_run=workflow_run,
|
|
x_api_key=x_api_key,
|
|
websocket=websocket,
|
|
)
|
|
|
|
LOG.info("Got streaming context for workflow run.", streaming=streaming)
|
|
|
|
loops = [
|
|
asyncio.create_task(loop_verify_workflow_run(streaming)),
|
|
asyncio.create_task(loop_stream_vnc(streaming)),
|
|
]
|
|
|
|
return streaming, loops
|
|
|
|
|
|
def verify_message_channel(
|
|
message_channel: sc.MessageChannel | None, streaming: sc.Streaming
|
|
) -> sc.MessageChannel | t.Literal[False]:
|
|
if message_channel and message_channel.is_open:
|
|
return message_channel
|
|
|
|
LOG.warning(
|
|
"No message channel found for client, or it is not open",
|
|
message_channel=message_channel,
|
|
client_id=streaming.client_id,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
|
|
return False
|
|
|
|
|
|
async def copy_text(streaming: sc.Streaming) -> None:
|
|
try:
|
|
async with connected_agent(streaming) as agent:
|
|
copied_text = await agent.get_selected_text()
|
|
|
|
LOG.info(
|
|
"Retrieved selected text via CDP",
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
|
|
message_channel = sc.get_message_client(streaming.client_id)
|
|
|
|
if cc := verify_message_channel(message_channel, streaming):
|
|
await cc.send_copied_text(copied_text, streaming)
|
|
else:
|
|
LOG.warning(
|
|
"No message channel found for client, or it is not open",
|
|
message_channel=message_channel,
|
|
client_id=streaming.client_id,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
except Exception:
|
|
LOG.exception(
|
|
"Failed to retrieve selected text via CDP",
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
|
|
|
|
async def ask_for_clipboard(streaming: sc.Streaming) -> None:
|
|
try:
|
|
LOG.info(
|
|
"Asking for clipboard data via CDP",
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
|
|
message_channel = sc.get_message_client(streaming.client_id)
|
|
|
|
if cc := verify_message_channel(message_channel, streaming):
|
|
await cc.ask_for_clipboard(streaming)
|
|
except Exception:
|
|
LOG.exception(
|
|
"Failed to ask for clipboard via CDP",
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
|
|
|
|
async def loop_stream_vnc(streaming: sc.Streaming) -> None:
|
|
"""
|
|
Actually stream the VNC session data between a frontend and a browser
|
|
session.
|
|
|
|
Loops until the task is cleared or the websocket is closed.
|
|
"""
|
|
|
|
if not streaming.browser_session:
|
|
LOG.info("No browser session found for task.", task=streaming.task, organization_id=streaming.organization_id)
|
|
return
|
|
|
|
vnc_url: str = ""
|
|
if streaming.browser_session.ip_address:
|
|
if ":" in streaming.browser_session.ip_address:
|
|
ip, _ = streaming.browser_session.ip_address.split(":")
|
|
vnc_url = f"ws://{ip}:{streaming.vnc_port}"
|
|
else:
|
|
vnc_url = f"ws://{streaming.browser_session.ip_address}:{streaming.vnc_port}"
|
|
else:
|
|
browser_address = streaming.browser_session.browser_address
|
|
|
|
parsed_browser_address = urlparse(browser_address)
|
|
host = parsed_browser_address.hostname
|
|
vnc_url = f"ws://{host}:{streaming.vnc_port}"
|
|
|
|
# NOTE(jdo:streaming-local-dev)
|
|
# vnc_url = "ws://localhost:9001/ws/novnc"
|
|
|
|
LOG.info(
|
|
"Connecting to VNC URL.",
|
|
vnc_url=vnc_url,
|
|
task=streaming.task,
|
|
workflow_run=streaming.workflow_run,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
|
|
async with websockets.connect(vnc_url) as novnc_ws:
|
|
|
|
async def frontend_to_browser() -> None:
|
|
LOG.info("Starting frontend-to-browser data transfer.", streaming=streaming)
|
|
data: Data | None = None
|
|
|
|
while streaming.is_open:
|
|
try:
|
|
data = await streaming.websocket.receive_bytes()
|
|
|
|
if data:
|
|
message_type = data[0]
|
|
|
|
if message_type == sc.MessageType.Keyboard.value:
|
|
streaming.update_key_state(data)
|
|
|
|
if streaming.key_state.is_copy(data):
|
|
await copy_text(streaming)
|
|
|
|
if streaming.key_state.is_paste(data):
|
|
await ask_for_clipboard(streaming)
|
|
|
|
if streaming.key_state.is_forbidden(data):
|
|
continue
|
|
|
|
if message_type == sc.MessageType.Mouse.value:
|
|
if sc.Mouse.Up.Right(data):
|
|
continue
|
|
|
|
if not streaming.interactor == "user" and message_type in (
|
|
sc.MessageType.Keyboard.value,
|
|
sc.MessageType.Mouse.value,
|
|
):
|
|
LOG.info(
|
|
"Blocking user message.", task=streaming.task, organization_id=streaming.organization_id
|
|
)
|
|
continue
|
|
|
|
except WebSocketDisconnect:
|
|
LOG.info("Frontend disconnected.", task=streaming.task, organization_id=streaming.organization_id)
|
|
raise
|
|
except ConnectionClosedError:
|
|
LOG.info(
|
|
"Frontend closed the streaming session.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
raise
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception:
|
|
LOG.exception(
|
|
"An unexpected exception occurred.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
raise
|
|
|
|
if not data:
|
|
continue
|
|
|
|
try:
|
|
await novnc_ws.send(data)
|
|
except WebSocketDisconnect:
|
|
LOG.info(
|
|
"Browser disconnected from the streaming session.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
raise
|
|
except ConnectionClosedError:
|
|
LOG.info(
|
|
"Browser closed the streaming session.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
raise
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception:
|
|
LOG.exception(
|
|
"An unexpected exception occurred in frontend-to-browser loop.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
raise
|
|
|
|
async def browser_to_frontend() -> None:
|
|
LOG.info("Starting browser-to-frontend data transfer.", streaming=streaming)
|
|
data: Data | None = None
|
|
|
|
while streaming.is_open:
|
|
try:
|
|
data = await novnc_ws.recv()
|
|
|
|
except WebSocketDisconnect:
|
|
LOG.info(
|
|
"Browser disconnected from the streaming session.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
await streaming.close(reason="browser-disconnected")
|
|
except ConnectionClosedError:
|
|
LOG.info(
|
|
"Browser closed the streaming session.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
await streaming.close(reason="browser-closed")
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception:
|
|
LOG.exception(
|
|
"An unexpected exception occurred in browser-to-frontend loop.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
raise
|
|
|
|
if not data:
|
|
continue
|
|
|
|
try:
|
|
await streaming.websocket.send_bytes(data)
|
|
except WebSocketDisconnect:
|
|
LOG.info(
|
|
"Frontend disconnected from the streaming session.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
await streaming.close(reason="frontend-disconnected")
|
|
except ConnectionClosedError:
|
|
LOG.info(
|
|
"Frontend closed the streaming session.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
await streaming.close(reason="frontend-closed")
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception:
|
|
LOG.exception(
|
|
"An unexpected exception occurred.",
|
|
task=streaming.task,
|
|
organization_id=streaming.organization_id,
|
|
)
|
|
raise
|
|
|
|
loops = [
|
|
asyncio.create_task(frontend_to_browser()),
|
|
asyncio.create_task(browser_to_frontend()),
|
|
]
|
|
|
|
try:
|
|
await collect(loops)
|
|
except Exception:
|
|
LOG.exception(
|
|
"An exception occurred in loop stream.", task=streaming.task, organization_id=streaming.organization_id
|
|
)
|
|
finally:
|
|
LOG.info("Closing the loop stream.", task=streaming.task, organization_id=streaming.organization_id)
|
|
await streaming.close(reason="loop-stream-vnc-closed")
|
|
|
|
|
|
@base_router.websocket("/stream/vnc/browser_session/{browser_session_id}")
|
|
async def browser_session_stream(
|
|
websocket: WebSocket,
|
|
browser_session_id: str,
|
|
apikey: str | None = None,
|
|
client_id: str | None = None,
|
|
token: str | None = None,
|
|
) -> None:
|
|
await stream(websocket, apikey=apikey, client_id=client_id, browser_session_id=browser_session_id, token=token)
|
|
|
|
|
|
@legacy_base_router.websocket("/stream/vnc/task/{task_id}")
|
|
async def task_stream(
|
|
websocket: WebSocket,
|
|
task_id: str,
|
|
apikey: str | None = None,
|
|
client_id: str | None = None,
|
|
token: str | None = None,
|
|
) -> None:
|
|
await stream(websocket, apikey=apikey, client_id=client_id, task_id=task_id, token=token)
|
|
|
|
|
|
@legacy_base_router.websocket("/stream/vnc/workflow_run/{workflow_run_id}")
|
|
async def workflow_run_stream(
|
|
websocket: WebSocket,
|
|
workflow_run_id: str,
|
|
apikey: str | None = None,
|
|
client_id: str | None = None,
|
|
token: str | None = None,
|
|
) -> None:
|
|
await stream(websocket, apikey=apikey, client_id=client_id, workflow_run_id=workflow_run_id, token=token)
|
|
|
|
|
|
async def stream(
|
|
websocket: WebSocket,
|
|
*,
|
|
apikey: str | None = None,
|
|
browser_session_id: str | None = None,
|
|
client_id: str | None = None,
|
|
task_id: str | None = None,
|
|
token: str | None = None,
|
|
workflow_run_id: str | None = None,
|
|
) -> None:
|
|
if not client_id:
|
|
LOG.error(
|
|
"Client ID not provided for VNC stream.",
|
|
browser_session_id=browser_session_id,
|
|
task_id=task_id,
|
|
workflow_run_id=workflow_run_id,
|
|
)
|
|
return
|
|
|
|
LOG.info(
|
|
"Starting VNC stream.",
|
|
browser_session_id=browser_session_id,
|
|
client_id=client_id,
|
|
task_id=task_id,
|
|
workflow_run_id=workflow_run_id,
|
|
)
|
|
|
|
organization_id = await auth(apikey=apikey, token=token, websocket=websocket)
|
|
|
|
if not organization_id:
|
|
LOG.error("Authentication failed.", task_id=task_id, workflow_run_id=workflow_run_id)
|
|
return
|
|
|
|
streaming: sc.Streaming
|
|
loops: list[asyncio.Task] = []
|
|
|
|
if browser_session_id:
|
|
result = await get_streaming_for_browser_session(
|
|
client_id=client_id,
|
|
browser_session_id=browser_session_id,
|
|
organization_id=organization_id,
|
|
websocket=websocket,
|
|
)
|
|
|
|
if not result:
|
|
LOG.error(
|
|
"No streaming context found for the browser session.",
|
|
browser_session_id=browser_session_id,
|
|
organization_id=organization_id,
|
|
)
|
|
await websocket.close(code=1013)
|
|
return
|
|
|
|
streaming, loops = result
|
|
|
|
LOG.info(
|
|
"Starting streaming for browser session.",
|
|
browser_session_id=browser_session_id,
|
|
organization_id=organization_id,
|
|
)
|
|
elif task_id:
|
|
result = await get_streaming_for_task(
|
|
client_id=client_id,
|
|
task_id=task_id,
|
|
organization_id=organization_id,
|
|
websocket=websocket,
|
|
)
|
|
|
|
if not result:
|
|
LOG.error("No streaming context found for the task.", task_id=task_id, organization_id=organization_id)
|
|
await websocket.close(code=1013)
|
|
return
|
|
|
|
streaming, loops = result
|
|
|
|
LOG.info("Starting streaming for task.", task_id=task_id, organization_id=organization_id)
|
|
|
|
elif workflow_run_id:
|
|
LOG.info(
|
|
"Starting streaming for workflow run.", workflow_run_id=workflow_run_id, organization_id=organization_id
|
|
)
|
|
result = await get_streaming_for_workflow_run(
|
|
client_id=client_id,
|
|
workflow_run_id=workflow_run_id,
|
|
organization_id=organization_id,
|
|
websocket=websocket,
|
|
)
|
|
|
|
if not result:
|
|
LOG.error(
|
|
"No streaming context found for the workflow run.",
|
|
workflow_run_id=workflow_run_id,
|
|
organization_id=organization_id,
|
|
)
|
|
await websocket.close(code=1013)
|
|
return
|
|
|
|
streaming, loops = result
|
|
|
|
LOG.info(
|
|
"Starting streaming for workflow run.", workflow_run_id=workflow_run_id, organization_id=organization_id
|
|
)
|
|
else:
|
|
LOG.error("Neither task ID nor workflow run ID was provided.")
|
|
return
|
|
|
|
try:
|
|
LOG.info(
|
|
"Starting streaming loops.",
|
|
task_id=task_id,
|
|
workflow_run_id=workflow_run_id,
|
|
organization_id=organization_id,
|
|
)
|
|
await collect(loops)
|
|
except Exception:
|
|
LOG.exception(
|
|
"An exception occurred in the stream function.",
|
|
task_id=task_id,
|
|
workflow_run_id=workflow_run_id,
|
|
organization_id=organization_id,
|
|
)
|
|
finally:
|
|
LOG.info(
|
|
"Closing the streaming session.",
|
|
task_id=task_id,
|
|
workflow_run_id=workflow_run_id,
|
|
organization_id=organization_id,
|
|
)
|
|
await streaming.close(reason="stream-closed")
|