Make _serializing_socket_cls used by internal WebsocketRPCEndpoint configurable from PubSubEndpoint #96

class JobFetchWebSocketProxy(SimpleWebSocket):
    data_schema = FetchExternalJobSchema

    def __init__(self, websocket: SimpleWebSocket):
        self.socket = websocket
        self.rabbitmq_connection = getattr(
            self.root.app.state, "rabbitmq_connection", None
        )
        if self.rabbitmq_connection is None:
            raise ValueError("RabbitMQ connection not found in app state")
        self.pubsub_topics = getattr(self.root.state, "pubsub_topics", None)
        self.db_session = getattr(self.root.state, "db_session", None)

    @property
    def root(self) -> WebSocket:
        root = self.socket
        while socket := (
            getattr(root, "socket", None) or getattr(root, "websocket", None)
        ):
            root = socket

        if not isinstance(root, WebSocket):
            raise ValueError("Could not find root `starlette.WebSocket` instance")
        return root

    async def connect(self, uri: str, **connect_kwargs: typing.Any):  # type: ignore
        await self.socket.connect(uri, **connect_kwargs)  # type: ignore

    def serialize(self, msg: pydantic.BaseModel) -> str:
        return pydantic_serialize(msg)

    def deserialize(self, buffer: str) -> typing.Dict[str, typing.Any]:
        data = orjson.loads(buffer)
        if not isinstance(data, dict):
            raise ValueError("Invalid data format, expected a JSON object")
        return data

    async def send(self, msg: pydantic.BaseModel):  # type: ignore
        # Convert the `RpcMessage` to a `response.Schema`
        if isinstance(msg, RpcMessage) and msg.response is not None:
            # Custom response schema applied here
            response_msg = response.Schema.model_validate(msg.response.result)
        else:
            response_msg = msg
        await self.socket.send(self.serialize(response_msg))  # type: ignore

    async def recv(self) -> typing.Dict[str, typing.Any]:  # type: ignore
        msg = await self.socket.recv()  # type: ignore
        if msg is None:
            return {"request": None}

        try:
            msg = self.deserialize(msg)
            # Data schema validated here
            data = self.data_schema.model_validate(msg)
        except pydantic.ValidationError as exc:
            logger.error(f"Failed to validate job fetch data: {exc}")
            await self.send(
                response.Schema(
                    status=response.Status.ERROR,
                    message="Invalid data",
                    detail="Failed to validate job fetch data",
                    errors=[e["msg"] for e in exc.errors()],
                )
            )
            return {"request": None}

        # Convert the data schema to a `RpcMessage`
        rpc_msg = RpcMessage(
            request=RpcRequest(
                method="fetch_job",
                arguments={
                    "db": self.db_session,
                    "job_url": str(data.job_url),
                    "metadata": data.metadata,
                    "rabbitmq_connection": self.rabbitmq_connection,
                    "pubsub_topics": self.pubsub_topics,
                },
            )
        )
        return rpc_msg.model_dump()

job_fetch_pubsub = PubSubEndpoint(
    methods_class=JobFetchRPCMethods,
    broadcaster=settings.REDIS_URL,
    on_connect=[on_job_fetch_socket_connect],  # type: ignore
    on_disconnect=[on_job_fetch_socket_disconnect],  # type: ignore
    ignore_broadcaster_disconnected=False,
)
# Manually overriding the `_serializing_socket_cls `
job_fetch_pubsub.endpoint._serializing_socket_cls = JobFetchWebSocketProxy

@router.websocket("/ws/job-fetch", dependencies=[authe.authentication_required])
async def job_fetch_socket(ws: WebSocket, user: authe.VerifiedUser):
    user_id = str(user.id)
    ws_channel_id = f"milo:core:job_fetch:ws:{user_id}"
    logger.info(
        f"User {user_id!r} is connecting to job fetch WebSocket with client ID: {ws_channel_id!r}",
        extra={"user_id": user_id, "ws_channel_id": ws_channel_id},
    )
    # Let the topics be the same as the channel ID for simplicity
    pubsub_topics = [ws_channel_id]
    ws.state.pubsub_topics = pubsub_topics
    await job_fetch_pubsub.main_loop(
        ws, channel_id=ws_channel_id, pubsub_topics=pubsub_topics
    )

job_fetch_pubsub = PubSubEndpoint(
    methods_class=JobFetchRPCMethods,
    broadcaster=settings.REDIS_URL,
    on_connect=[on_job_fetch_socket_connect],  # type: ignore
    on_disconnect=[on_job_fetch_socket_disconnect],  # type: ignore
    ignore_broadcaster_disconnected=False,
    serializing_socket_cls=JobFetchWebSocketProxy,
)