"""Router - handle message router (sock). Router to manage responses from a socket client. """ from __future__ import annotations from wandb.proto import wandb_internal_pb2 as pb from wandb.proto import wandb_server_pb2 as spb from wandb.sdk.lib.sock_client import SockClient, SockClientClosedError from wandb.sdk.mailbox import Mailbox from .router import MessageRouter, MessageRouterClosedError class MessageSockRouter(MessageRouter): _sock_client: SockClient _mailbox: Mailbox def __init__(self, sock_client: SockClient, mailbox: Mailbox) -> None: self._sock_client = sock_client super().__init__(mailbox=mailbox) def _read_message(self) -> spb.ServerResponse | None: try: return self._sock_client.read_server_response(timeout=1) except SockClientClosedError as e: raise MessageRouterClosedError from e def _send_message(self, record: pb.Record) -> None: self._sock_client.send_record_communicate(record)