Repository: litestar-org/litestar

Issue:
Feature: Websockets: best way to accept() custom headers with `WebsocketListener` class
### Discussed in https://github.com/orgs/litestar-org/discussions/1552

<div type='discussions-op-text'>

<sup>Originally posted by **hadware** April 20, 2023</sup>
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, :

```python
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](https://docs.litestar.dev/2/usage/websockets.html#class-based-websocket-handling), 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`?</div>


Discussion:

Resolving patch:
diff --git a/docs/examples/websockets/setting_custom_connection_headers.py b/docs/examples/websockets/setting_custom_connection_headers.py
new file mode 100644
--- /dev/null
+++ b/docs/examples/websockets/setting_custom_connection_headers.py
@@ -0,0 +1,13 @@
+from litestar import Litestar, WebSocket, websocket_listener
+
+
+async def accept_connection(socket: WebSocket) -> None:
+    await socket.accept(headers={"Cookie": "custom-cookie"})
+
+
+@websocket_listener("/", connection_accept_handler=accept_connection)
+def handler(data: str) -> str:
+    return data
+
+
+app = Litestar([handler])
diff --git a/litestar/handlers/websocket_handlers/_utils.py b/litestar/handlers/websocket_handlers/_utils.py
--- a/litestar/handlers/websocket_handlers/_utils.py
+++ b/litestar/handlers/websocket_handlers/_utils.py
@@ -113,9 +113,10 @@ def create_handler_function(
     listener_context: ListenerContext,
     on_accept: AsyncCallable | None,
     on_disconnect: AsyncCallable | None,
+    accept_connection_handler: Callable[[WebSocket], Coroutine[Any, Any, None]],
 ) -> Callable[..., Coroutine[None, None, None]]:
     async def handler_fn(socket: WebSocket, **kwargs: Any) -> None:
-        await socket.accept()
+        await accept_connection_handler(socket)
 
         listener_callback = AsyncCallable(listener_context.listener_callback)
         ctx = ConnectionContext.from_connection(socket)
diff --git a/litestar/handlers/websocket_handlers/listener.py b/litestar/handlers/websocket_handlers/listener.py
--- a/litestar/handlers/websocket_handlers/listener.py
+++ b/litestar/handlers/websocket_handlers/listener.py
@@ -7,6 +7,7 @@
 from msgspec.json import Encoder as JsonEncoder
 
 from litestar._signature import create_signature_model
+from litestar.connection import WebSocket
 from litestar.dto.interface import HandlerContext
 from litestar.exceptions import ImproperlyConfiguredException
 from litestar.serialization import default_serializer
@@ -29,7 +30,9 @@
 from .route_handler import WebsocketRouteHandler
 
 if TYPE_CHECKING:
-    from litestar import Litestar, WebSocket
+    from typing import Coroutine
+
+    from litestar import Litestar
     from litestar.dto.interface import DTOInterface
     from litestar.types.asgi_types import WebSocketMode
 
@@ -52,12 +55,14 @@ class websocket_listener(WebsocketRouteHandler):
         "_receive_mode",
         "_send_mode",
         "_listener_context",
+        "accept_connection_handler",
     )
 
     def __init__(
         self,
         path: str | None | list[str] | None = None,
         *,
+        connection_accept_handler: Callable[[WebSocket], Coroutine[Any, Any, None]] = WebSocket.accept,
         dependencies: Dependencies | None = None,
         dto: type[DTOInterface] | None | EmptyType = Empty,
         exception_handlers: dict[int | type[Exception], ExceptionHandler] | None = None,
@@ -79,6 +84,8 @@ def __init__(
         Args:
             path: A path fragment for the route handler function or a sequence of path fragments. If not given defaults
                 to ``/``
+            connection_accept_handler: A callable that accepts a :class:`WebSocket <.connection.WebSocket>` instance
+                and returns a coroutine that when awaited, will accept the connection. Defaults to ``WebSocket.accept``.
             dependencies: A string keyed mapping of dependency :class:`Provider <.di.Provide>` instances.
             dto: :class:`DTOInterface <.dto.interface.DTOInterface>` to use for (de)serializing and
                 validation of request data.
@@ -107,6 +114,7 @@ def __init__(
         self._send_mode: WebSocketMode = send_mode
         self._on_accept = AsyncCallable(on_accept) if on_accept else None
         self._on_disconnect = AsyncCallable(on_disconnect) if on_disconnect else None
+        self.accept_connection_handler = connection_accept_handler
         self.type_encoders = type_encoders
 
         super().__init__(
@@ -146,6 +154,7 @@ def __call__(self, listener_callback: AnyCallable) -> websocket_listener:
             listener_context=self._listener_context,
             on_accept=self._on_accept,
             on_disconnect=self._on_disconnect,
+            accept_connection_handler=self.accept_connection_handler,
         )
         return super().__call__(handler_function)
 


Test patch:
diff --git a/tests/handlers/websocket/test_listeners.py b/tests/handlers/websocket/test_listeners.py
--- a/tests/handlers/websocket/test_listeners.py
+++ b/tests/handlers/websocket/test_listeners.py
@@ -262,3 +262,16 @@ def handler_body(data: str, body: bytes) -> None:
             ...
 
         handler_body.on_registration(Litestar())
+
+
+def test_listener_accept_connection_callback() -> None:
+    async def accept_connection(socket: WebSocket) -> None:
+        await socket.accept(headers={"Cookie": "custom-cookie"})
+
+    @websocket_listener("/", connection_accept_handler=accept_connection)
+    def handler(data: bytes) -> None:
+        return None
+
+    client = create_test_client([handler])
+    with client.websocket_connect("/") as ws:
+        assert ws.extra_headers == [(b"cookie", b"custom-cookie")]