mirror of
https://github.com/Skyvern-AI/skyvern.git
synced 2026-04-26 10:41:14 +00:00
328bce3cdd
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Shuchang Zheng <wintonzheng0325@gmail.com>
92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
"""
|
|
Contains registries for coordinating active WS connections (aka "channels", see
|
|
`./channels/README.md`).
|
|
|
|
NOTE: in AWS we had to turn on what amounts to sticky sessions for frontend apps,
|
|
so that an individual frontend app instance is guaranteed to always connect to
|
|
the same backend api instance. This is beccause the two registries here are
|
|
tied together via a `client_id` string.
|
|
|
|
The tale-of-the-tape is this:
|
|
- frontend app requires two different channels (WS connections) to the backend api
|
|
- one dedicated to streaming VNC's RFB protocol
|
|
- the other dedicated to messaging (JSON)
|
|
- both of these channels are stateful and need to coordinate with one another
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import typing as t
|
|
|
|
import structlog
|
|
|
|
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()
|
|
|
|
|
|
# a registry for VNC channels, keyed by `client_id`
|
|
vnc_channels: dict[str, VncChannel] = {}
|
|
|
|
|
|
def add_vnc_channel(vnc_channel: VncChannel) -> None:
|
|
vnc_channels[vnc_channel.client_id] = vnc_channel
|
|
|
|
|
|
def get_vnc_channel(client_id: str) -> VncChannel | None:
|
|
return vnc_channels.get(client_id)
|
|
|
|
|
|
def del_vnc_channel(client_id: str, *, expected: VncChannel | None = None) -> None:
|
|
candidate = vnc_channels.get(client_id)
|
|
|
|
if candidate is None:
|
|
return
|
|
|
|
# Prevent stale channel shutdown from deleting a newer channel that reused
|
|
# the same client_id during route transitions/reconnects.
|
|
if expected is not None and candidate is not expected:
|
|
return
|
|
|
|
del vnc_channels[client_id]
|
|
|
|
|
|
# a registry for message channels, keyed by `client_id`
|
|
message_channels: dict[str, MessageChannel] = {}
|
|
|
|
|
|
def add_message_channel(message_channel: MessageChannel) -> None:
|
|
message_channels[message_channel.client_id] = message_channel
|
|
|
|
|
|
def get_message_channel(client_id: str) -> MessageChannel | None:
|
|
candidate = message_channels.get(client_id)
|
|
|
|
if candidate is None:
|
|
return None
|
|
|
|
if candidate.is_open:
|
|
return candidate
|
|
|
|
LOG.info(
|
|
"MessageChannel: message channel is not open; deleting it",
|
|
client_id=candidate.client_id,
|
|
)
|
|
del_message_channel(candidate.client_id, expected=candidate)
|
|
return None
|
|
|
|
|
|
def del_message_channel(client_id: str, *, expected: MessageChannel | None = None) -> None:
|
|
candidate = message_channels.get(client_id)
|
|
|
|
if candidate is None:
|
|
return
|
|
|
|
# Prevent stale channel shutdown from deleting a newer channel that reused
|
|
# the same client_id during route transitions/reconnects.
|
|
if expected is not None and candidate is not expected:
|
|
return
|
|
|
|
del message_channels[client_id]
|