[litestarV2] Websockets: best way to accept() custom headers with the new WebsocketListener class

[hadware]

Hello,

I'm wondering if it's possible (or even relevant) to be using the WebsocketListener class if I want to send custom headers back to the client when accepting the connection.

The "old" way of doing this would be, IIRC, :

from starlite import websocket, WebSocket
from starlite.exceptions import WebSocketDisconnect


@websocket("/")
async def handler(socket: WebSocket) -> str:
  # for instance, here, sending a cookie to be set on the client
  await socket.accept({'Cookie': 'custom-cookie'})
  while True:
    try:
      data = await socket.receive_text()
      await socket.send_text(data)
    except WebSocketDisconnect:
      break

In the new WebsocketListener handler, which seems to fit much better in the OOP logic encouraged by litestar, as per the documentation, there's an on_accept(...) method, but it seems that once this method is called by litestar, it's "already too late" to set any custom parameters for websocket.accept().

Is there any way to do this? Is it even a good idea to use WebsocketListener in my case? Should I maybe start an issue/PR to suggest the addition of an optionally overloadable method such as on_connect (the autobahn way) for WebsocketListener?

[provinzkraut]

Hi, thanks for opening this.

You are indeed correct that there currently is not a way to do what you want. I actually considered this but decided to address this at a later point. I guess that point is now.

I think the best way to handle this is to simply add an accept callback, which by default just calls socket.accept.

Alternatively, we could change the behaviour of on_accept so that it calls socket.accept itself, but that would mean it would have to be async, breaking the paradigm I was trying to establish with this interface.

@peterschutt wdyt?

[peterschutt]

I think the best way to handle this is to simply add an accept callback, which by default just calls socket.accept.

Something like this?

class websocket_listener(WebsocketRouteHandler):
    def __init__(
        ...,
        acceptor: Callable[[WebSocket], Coroutine[Any, Any, None]] = WebSocket.accept,
        ...,
    ) -> None:
        ...

[provinzkraut]

Something like that, yes. Although "Acceptor" sounds like a very odd super hero 🙃

[Goldziher]

why not inject a helper called send_headers? its going to be much simpler and easier to comprehend.

[provinzkraut]

That would only solve this particular use case and not allow further customization. Also send_headers I would assume sends headers with every message, not just on connect.

I also feel like having a separate callback for this is more in-line with the rest of the interface.

[Goldziher]

you cant send headers with every message. The ASGI protocol only allows on connection

[provinzkraut]

Yeah, I know. That's part of the reason I think it might be misleading.

[Goldziher]

Hmm... So maybe instead simply call it accept_connection_handler?

[peterschutt]

Created #1571