If you're using (or interested) in the WebSocket support in Mangum, please read

[jordaneremieff]

Hello.

It has been about 6 months since WebSocket API support was re-introduced in version 0.12.0. Since that time I've not received any issues or seen mention elsewhere at all related to the WebSocket support, good or bad, which is surprising to me.

This lack of activity seems to indicate that either the design of the WebSocket event flow is flawed in some way (either in how it handles multiple events or the steps required to configure it), or it may just be an unpopular feature and most users of ASGI frameworks prefer traditional server solutions. I can only speculate at this point.

If anyone has thoughts on this topic then please share here. I want to get a better idea of how valuable this feature is to consider if it is worth continuing to maintain in this project, or potentially a niche enough use-case that it should be extracted into a different package.

[jordaneremieff]

@koxudaxi you're the only WebSockets user that I have personally been able to confirm previously. Are you still using it, or have any feedback/thoughts on it becoming a separate project?

[koxudaxi]

@jordaneremieff

Are you still using it,

I used the feature in a previous company.
I don't use it now 😅

or have any feedback/thoughts on it becoming a separate project?

I think good. The use case is separated completely.

[jordaneremieff]

Thanks @koxudaxi. I think I will separate these at some point then to make it easier to maintain the primary use-case.

[radialsolutions]

WebSocket API support

I've been looking into websockets with Mangum, looks like this link is broken @jordaneremieff. I was wondering if there were any examples elsewhere? I saw some mention of adding handler = Mangum(app, dsn=db_connection), but wasn't sure where to go from there.

[jordaneremieff]

@radialsolutions you can look at https://github.com/jordaneremieff/mangum/releases/tag/0.12.0 to see the source code. At this point I don't have the time/interest in supporting this feature myself, but anyone else is welcome to look at the removed implementation and support this in a separate project.

[natflow]

One use case we had was a Lambda accepting both SQS events and API Gateway WebSocket events, using FastAPI with the later. Our devs didn't see a public API for accessing the WebSocket connection when handling an SQS event and so chose to build their own implementation. What we'd need would be to look up the WebSocket by user_id.

[khamaileon]

I'm currently working on a big real time application using django-channels. Now it's time to deploy it and I would like to do it serverless.
Unfortunately the support of Websocket is no longer available in Mangum. It's a bummer because for me it's one of the main specification of ASGI. I need to make it work and I'm willing to spend some time on it. How can I help?

[nitsujri]

Hrmmm, what part of Channels is being used at that point?

On the surface it appears agnostic because API Gateway is action based a opposed to URL and then if boto3 is sending the message via post_to_connection not sure what situation we would want to go through the Channels class WebsocketConsumer or its async sibling?

We're quite interested in making it work with Channels because Channels would be used for local dev while obviously API Gateway for anything server based, but ideally the developer wouldn't need to worry about duplicating business logic for the two paths.

[khamaileon]

It's not "agnostic" per se, as messages are sent directly via boto (we may be able to improve on this). Here's some code:

# projects/handlers.py

from mangum import Mangum

from project.asgi import application
from utils.mangum.handlers import WebSocketHandler

# FIXME: https://github.com/jordaneremieff/mangum/issues/143
def main(event, context):
    if event.get("body") is not None:
        content_length = str(len(event["body"]))
        if "multiValueHeaders" in event:
            event["multiValueHeaders"]["Content-Length"] = [content_length]
        if "headers" in event:
            event["headers"]["Content-Length"] = content_length

    asgi_handler = Mangum(
        application,
        lifespan="off",
        custom_handlers=[WebSocketHandler],
        text_mime_types=[
            "text/",
            "application/json",
            "application/javascript",
            "application/xml",
            "application/vnd.api+json",
            "application/vnd.oai.openapi",
            "application/geo+json",
        ],
    )
    return asgi_handler(event, context)

# utils/mangum/handlers.py

from typing import Type

from mangum.handlers.utils import maybe_encode_body
from mangum.types import (
    Cycle,
    LambdaConfig,
    LambdaContext,
    LambdaEvent,
    Response,
    Scope,
)

from utils.mangum.protocols import WebSocketCycle


class WebSocketHandler:
    @classmethod
    def infer(
        cls, event: LambdaEvent, context: LambdaContext, config: LambdaConfig
    ) -> bool:
        return "requestContext" in event and "connectionId" in event["requestContext"]

    def __init__(
        self, event: LambdaEvent, context: LambdaContext, config: LambdaConfig
    ) -> None:
        self.event = event
        self.context = context
        self.config = config

    @property
    def cycle_cls(self) -> Type[Cycle]:
        return WebSocketCycle

    @property
    def headers(self) -> dict:
        headers_dict = {
            k.lower().encode(): v.encode()
            for k, v in self.event.get("headers", {}).items()
        }
        if b"origin" not in headers_dict:
            domain_name = self.event["requestContext"]["domainName"]
            headers_dict[b"origin"] = f"http://{domain_name}".encode()

        return headers_dict

    @property
    def body(self) -> bytes:
        return maybe_encode_body(
            self.event.get("body", b""),
            is_base64=self.event.get("isBase64Encoded", False),
        )

    @property
    def connection_id(self) -> str:
        return self.event["requestContext"]["connectionId"]

    @property
    def event_type(self) -> str:
        return self.event["requestContext"]["eventType"]

    @property
    def query_string(self) -> str:
        return ";".join(
            [
                f"{key}={value}"
                for key, value in self.event.get("queryStringParameters", {}).items()
            ]
        ).encode()

    @property
    def source_ip(self) -> str:
        return self.event["requestContext"]["identity"]["sourceIp"]

    @property
    def scope(self) -> Scope:
        host = self.headers.get("host", "mangum")
        port = self.headers.get("x-forwarded-for", 80)
        return {
            "type": "websocket",
            "path": "/",
            "raw_path": b"/",
            "headers": list(self.headers.items()),
            "query_string": self.query_string,
            "client": [self.source_ip, 1000],
            "server": [host, port],
            "subprotocols": [],
            "asgi": {"version": "3.0"},
            "apigateway_websocket_connection_id": self.connection_id,
            "apigateway_event_type": self.event_type,
        }

    def __call__(self, response: Response) -> dict:
        return {"statusCode": response["status"]}

# project/consumers.py

import logging

import orjson
from channels.db import database_sync_to_async
from channels.exceptions import DenyConnection, StopConsumer
from channels.generic.websocket import AsyncJsonWebsocketConsumer

from account.models import Session

logger = logging.getLogger("my-super-app")


class WebsocketConsumer(AsyncJsonWebsocketConsumer):
    @property
    def connection_id(self):
        return self.scope.get("apigateway_websocket_connection_id")

    @database_sync_to_async
    def store_connection_id(self, session_key):
        logger.info(
            f"Store connection id {self.connection_id} in session {session_key}"
        )
        Session.objects.filter(session_key=session_key).update(
            apigateway_websocket_connection_id=self.connection_id
        )

    @database_sync_to_async
    def get_session_key(self):
        logger.info(f"Get session key with connection id {self.connection_id}")
        return Session.objects.get(
            apigateway_websocket_connection_id=self.connection_id
        ).session_key

    @database_sync_to_async
    def remove_connection_id(self):
        logger.info(f"Remove connection id {self.connection_id}")
        Session.objects.filter(
            apigateway_websocket_connection_id=self.connection_id
        ).update(apigateway_websocket_connection_id=None)

    async def websocket_connect(self, message):
        logger.info("Websocket connect")
        user = self.scope["user"]
        session_key = self.scope["session"].session_key
        group = session_key

        if not user.is_authenticated:
            logger.info("Websocket connection denied: user must be authenticated")
            raise DenyConnection

        if self.connection_id:
            await self.store_connection_id(session_key)

        await self.channel_layer.group_add(group, self.channel_name)
        await self.accept()

    async def websocket_disconnect(self, message):
        logger.info("Websocket disconnect")

        if self.connection_id:
            group = await self.get_session_key()
            await self.remove_connection_id()

        else:
            group = self.scope["session"].session_key

        await self.channel_layer.group_discard(group, self.channel_name)
        raise StopConsumer()

    async def send_data(self, event):
        await self.send_json(event["payload"])

    @classmethod
    async def decode_json(cls, text_data):
        return orjson.loads(text_data)

    @classmethod
    async def encode_json(cls, content):
        return orjson.dumps(content).decode()

# project/asgi.py

import os

from channels.auth import AuthMiddlewareStack
from channels.routing import ProtocolTypeRouter, URLRouter
from channels.security.websocket import AllowedHostsOriginValidator
from django.core.asgi import get_asgi_application

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "project.settings")
# Initialize Django ASGI application early to ensure the AppRegistry
# is populated before importing code that may import ORM models.
django_asgi_app = get_asgi_application()

from project.urls import websocket_urlpatterns  # noqa: E402

application = ProtocolTypeRouter(
    {
        "http": django_asgi_app,
        "websocket": AllowedHostsOriginValidator(
            AuthMiddlewareStack(URLRouter(websocket_urlpatterns))
        ),
    }
)

# project/urls.py

from project.consumers import WebsocketConsumer

...

websocket_urlpatterns = [
    path(
        "",
        WebsocketConsumer.as_asgi(),
    ),
]

# account/models.py

from django.contrib.sessions.base_session import AbstractBaseSession

...

class Session(AbstractBaseSession):
    user = models.ForeignKey(User, on_delete=models.CASCADE, null=True)
    apigateway_websocket_connection_id = models.CharField(
        max_length=16, blank=True, null=True
    )

    @classmethod
    def get_session_store_class(cls):
        from account.sessions import SessionStore

        return SessionStore

# account/sessions.py

from django.contrib.sessions.backends.cached_db import SessionStore as CachedDBStore


class SessionStore(CachedDBStore):
    @classmethod
    def get_model_class(cls):
        from account.models import Session

        return Session

    def create_model_instance(self, data):
        session = super().create_model_instance(data)
        user_id = data.get("_auth_user_id", "")
        if user_id:
            session.user_id = user_id

        return session

# common/channels.py

import boto3
import orjson
from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer
from django.conf import settings
from hashid_field.hashid import Hashid

from account.models import Session


def default(obj):
    if isinstance(obj, Hashid):
        return str(obj)
    raise TypeError


def notify_users(user_ids, payload):
    session_qs = Session.objects.filter(user_id__in=user_ids)
    if settings.ENV == "prod":
        session_qs = session_qs.exclude(apigateway_websocket_connection_id=None)
        connection_ids = session_qs.values_list(
            "apigateway_websocket_connection_id", flat=True
        )
        client = boto3.client(
            "apigatewaymanagementapi",
            endpoint_url=settings.API_GATEWAY_WEBSOCKET_ENDPOINT_URL,
        )
        for connection_id in connection_ids:
            try:
                data = orjson.dumps(payload, default=default)
                client.post_to_connection(Data=data, ConnectionId=connection_id)

            except client.exceptions.GoneException:
                Session.objects.filter(
                    apigateway_websocket_connection_id=connection_id
                ).update(apigateway_websocket_connection_id=None)

    else:
        session_keys = session_qs.values_list("session_key", flat=True)
        channel_layer = get_channel_layer()
        for session_key in session_keys:
            message = {
                "type": "send.data",
                "payload": payload,
            }
            group = str(session_key)
            async_to_sync(channel_layer.group_send)(group, message)

Don't hesitate to give me feedback if you see any possible improvements (for example, if the API Gateway interface has evolved).

[nitsujri]

@khamaileon thanks so much for the code! it's super helpful.

Could you share WebSocketCycle or is the exact same one as in 0.14.3?

[khamaileon]

Oh you're right, I forgot that one 😅

# utils/mangum/protocols.py

import asyncio
import logging

from mangum.exceptions import UnexpectedMessage
from mangum.types import ASGI, Message, Response, Scope

logger = logging.getLogger("my-super-app")


class WebSocketAccepted(Exception):
    pass


class WebSocketClosed(Exception):
    pass


class WebSocketCycle:
    def __init__(self, scope: Scope, body: bytes) -> None:
        self.scope = scope
        self.body = body
        self.response_status = 200
        self.app_queue: asyncio.Queue[Message] = asyncio.Queue()

    def __call__(self, app: ASGI) -> Response:
        event_type = self.scope["apigateway_event_type"]
        logger.info(f"WebSocketCycle called with '{event_type}' event")

        if event_type == "CONNECT":
            self.app_queue.put_nowait({"type": "websocket.connect"})

        if event_type == "DISCONNECT":
            self.app_queue.put_nowait({"type": "websocket.disconnect", "code": 1001})

        asgi_instance = self.run(app)
        loop = asyncio.get_event_loop()
        asgi_task = loop.create_task(asgi_instance)
        loop.run_until_complete(asgi_task)
        return {
            "status": self.response_status,
            "headers": [],
            "body": b"",
        }

    async def run(self, app: ASGI) -> None:
        try:
            await app(self.scope, self.receive, self.send)

        except WebSocketAccepted:
            self.response_status = 200
            logger.info("WebSocket accepted")

        except WebSocketClosed:
            self.response_status = 200
            logger.info("WebSocket closed")

        except UnexpectedMessage as e:
            self.response_status = 500
            logger.error(e)

    async def receive(self) -> Message:
        return await self.app_queue.get()

    async def send(self, message: Message) -> None:
        if message["type"] == "websocket.accept":
            raise WebSocketAccepted

        elif message["type"] == "websocket.close":
            raise WebSocketClosed

        else:
            raise UnexpectedMessage(f"Unexpected {message['type']}")

[khamaileon]

And be aware that I follow a specific logic. In my case a channel layer = a WS session, because I want to target each user individually. It's possible to keep the “channel” logic, but in my case a user can be on several "channels" and I don't want to open several WS per user. There is a way to do it by multiplexing, I had even found a lib for that but it became too complex for my use case.

[dbramwell]

I've recently thought of a hobby project that I want to use websockets for and wondered if mangum supported it, leading me here. I've had a look through the code base at v0.12.4 but haven't had chance to try it out yet. One thing that struck me is that I think the implementation may have been trying to do too much, particularly around the backend/storage stuff.

I'm not familiar with many ASGI frameworks, but the fastapi websocket docs end with an example of storing connections in a simple list (obviously not suitable for serverless architecture, but that's not my point here), where as it looks like if I ran mangum 0.12.4 the connect and disconnect messages would never reach fastapi but instead be handled by which ever of the storage mechanisms I used? That feels like it should be the responsibility of my fastapi code and mangum should just pass the connect and disconnect messages through, am I missing something?

[dbramwell]

Further question. If I could get something reasonable working with the backend stuff stripped out, would you accept a PR for it and add it back to mangum rather than me maintaining it in a different repo?

[jordaneremieff]

the connect and disconnect messages would never reach fastapi but instead be handled by which ever of the storage mechanisms I used?

This was because connect, message, and disconnect events all represented a separate route and invocation of the lambda. If something has since changed in AWS WebSocket support or you are able to discover some way that doesn't require a separate storage backend to reliably support WebSocket events, then I would consider reintroducing it.

The primary reason I dropped support was because I agree with this:

One thing that struck me is that I think the implementation may have been trying to do too much, particularly around the backend/storage stuff.

However, to my knowledge there isn't a way to get around this because the initial connection event is not available in the message route but contains the necessary information to form the ASGI scope.

[jrobbins-LiveData]

Latecomer to this discussion, but we are happy mangum users and are now investigating ASGI websocket support.

Some background context: we have a Lambda-based app that uses FastAPI for REST APIs. We also have an API Gateway websocket API, that we manage "on the side", by inspecting the Lambda handler's incoming event and dealing with the websocket CONNECT, MESSAGE, and DISCONNECT with our own code.

I am investigating what it would take to fit the websocket event traffic into the ASGI flow, so that our FastAPI routes could handle the websocket messages. The main motivation for this is to make the FastAPI code simple, facilitating hosting it in, e.g. uvicorn also, to facilitate development. I tried picking up closed PR 190 and making my own Mangum-based class like so

class WsMangum(Mangum):
    def __call__(self, event: LambdaEvent, context: LambdaContext) -> dict:
        handler = self.infer(event, context)
        logger.debug(f'{handler=}')

        with ExitStack() as stack:
            if self.lifespan in ("auto", "on"):
                lifespan_cycle = LifespanCycle(self.app, self.lifespan)
                stack.enter_context(lifespan_cycle)

            if isinstance(handler, WsAPIGateway):
                logger.debug('WebSocketCycle')
                ws_storage = WebSocketStorage()
                ws_cycle = WebSocketCycle(handler.request, handler.message_type, handler.connection_id, ws_storage)
                resp = ws_cycle(self.app, handler.body)
            else:
                logger.debug('HTTPCycle')
                http_cycle = HTTPCycle(handler.scope, handler.body)
                resp = http_cycle(self.app)

            logger.debug(f'{resp=}')
            return handler(resp)

        assert False, "unreachable"  # pragma: no cover

which then allows this FastAPI route to work

@router.websocket_route('/')
async def websocket_endpoint(websocket: WebSocket):
    if websocket.client_state is not WebSocketState.CONNECTED:
        await websocket.accept()
    data = await websocket.receive_text()
    await websocket.send_text(f'Message text was: {data}')

Even though I have not (yet) implemented a storage handler to reconstitute the scope on each new websocket MESSAGE, I had the fantasy that, assuming I hit the same Lambda instance (by luck) on subsequent MESSAGEs, I'd see the same websocket instance of starlette\websockets.py's WebSocket class. But I see a fresh websocket on each call, even if the Lambda function instance is the same.

My hope for an interface between API Gateway websocket and ASGI would be that, from ASGI's point of view, an incoming MESSAGE event would be treated as yet another message on a persistent websocket, and somehow the FastAPI/Starlette code could be "told" that this is the case, avoiding all the setup work on each and every MESSAGE.

It seems a shame to go to all the work to persist the scope, only to have the app framework have to start over on its view of the websocket state. In reading about Starlette's WebSocketEndpoint, I was hoping that somehow (again, indirectly via FastAPI) I could get the ASGI app framework to maintain its state between Lambda invocations.

I realize that a real app can't rely on "luck". Assuming I can see the behavior I want, I'd then look at a cache in each Lambda function so it could know if a given MESSAGE on a given connectionId was already initialized in this particular Lambda instance, or rather it would need to somehow explicitly initialize the app's state on behalf of this particular websocket. The storage backend would only be needed if the connectionId wasn't already cached in this Lambda instance.

[nitsujri]

@jrobbins-LiveData i've also been interested in this for a while now and our team is getting to the point of putting some effort into this. So thanks so much for trying to resurrect it.

Apologies for the naiveness of my questions as I'm just beginning to scratch the surface.

My hope for an interface between API Gateway websocket and ASGI would be that, from ASGI's point of view, an incoming MESSAGE event would be treated as yet another message on a persistent websocket, and somehow the FastAPI/Starlette code could be "told" that this is the case, avoiding all the setup work on each and every MESSAGE.

What is the primary drawback of re-setting up on every MESSAGE?

From my limited understanding, the Lambda function should be stateless in regards to the websocket because, as you've mentioned, the Lambda invocation, at any point, be new. The state of the connection is held "open" by the API Gateway itself and the connectionId held in either pg or redis?

The application shouldn't care whether it holds a live connection or not? As long it knows which connectionId to work with and post_to_connection will send the message along. Especially true if the message is coming from an SQS or direct invoke.

[okigan]

So, I am interested in this, but cannot find a succinct example how to use/configure it.

The referenced page (https://mangum.io/websockets/) returns a 404....please advise.

[okigan]

I was able to get web sockets working with mangum/fastapi/custom code. You can check out results [here](https://qnamod.pragmavant.com) .. create event and share/open it in several browsers, then post some questions.

[mrmuskrat]

Any chance you have this in a repo we can review?