79 lines
2.2 KiB
Python
79 lines
2.2 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) -> t.Union[VncChannel, None]:
|
|
return vnc_channels.get(client_id, None)
|
|
|
|
|
|
def del_vnc_channel(client_id: str) -> None:
|
|
try:
|
|
del vnc_channels[client_id]
|
|
except KeyError:
|
|
pass
|
|
|
|
|
|
# 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) -> t.Union[MessageChannel, None]:
|
|
candidate = message_channels.get(client_id, None)
|
|
|
|
if candidate and candidate.is_open:
|
|
return candidate
|
|
|
|
if candidate:
|
|
LOG.info(
|
|
"MessageChannel: message channel is not open; deleting it",
|
|
client_id=candidate.client_id,
|
|
)
|
|
|
|
del_message_channel(candidate.client_id)
|
|
|
|
return None
|
|
|
|
|
|
def del_message_channel(client_id: str) -> None:
|
|
try:
|
|
del message_channels[client_id]
|
|
except KeyError:
|
|
pass
|