AsyncAPI Channels, WebSockets & Server-Sent Events¶

"""AsyncAPI 3.0 document normalization and realtime WebSocket channel typegen."""

from __future__ import annotations

from pathlib import Path
from typing import Any

from litestar_vite import TypeGenConfig
from litestar_vite.codegen import normalize_asyncapi_document


def build_sample_asyncapi_document() -> dict[str, Any]:
    """Construct a sample AsyncAPI 3.0 document for realtime notifications."""
    raw_doc: dict[str, Any] = {
        "asyncapi": "3.0.0",
        "info": {"title": "Realtime Notifications", "version": "1.0.0"},
        "channels": {
            "notifications": {
                "address": "/ws/notifications",
                "messages": {
                    "OrderCreated": {
                        "payload": {
                            "type": "object",
                            "required": ["orderId", "amountCents"],
                            "properties": {"orderId": {"type": "string"}, "amountCents": {"type": "integer"}},
                        }
                    }
                },
            }
        },
        "operations": {"receiveNotifications": {"action": "receive", "channel": {"$ref": "#/channels/notifications"}}},
    }
    return normalize_asyncapi_document(raw_doc)


def build_channel_typegen_config() -> TypeGenConfig:
    """Configure ``TypeGenConfig`` to export ``asyncapi.json`` and ``channels.ts``."""
    return TypeGenConfig(output=Path("src/generated"), generate_channels=True)


def main() -> None:
    """Verify AsyncAPI 3.0 document normalization and ``TypeGenConfig`` settings."""
    normalized = build_sample_asyncapi_document()
    cfg = build_channel_typegen_config()
    assert normalized["asyncapi"] == "3.0.0"
    assert "ws__notifications" in normalized["channels"]
    assert cfg.generate_channels is True
    assert cfg.channels_ts_path == Path("src/generated/channels.ts")


__all__ = ("build_channel_typegen_config", "build_sample_asyncapi_document", "main")

from litestar import Litestar, post
from litestar.channels import ChannelsPlugin
from litestar.channels.backends.memory import MemoryChannelsBackend


@post("/api/rooms/{room_id:str}/broadcast", sync_to_thread=False)
def broadcast_to_room(room_id: str, message: str, channels: ChannelsPlugin) -> dict[str, str]:
    """Publish a payload to a dynamic channel subscribed by WebSocket clients."""
    channels.publish({"room_id": room_id, "message": message}, channels=[f"chat_{room_id}"])
    return {"status": "published"}


app = Litestar(
    route_handlers=[broadcast_to_room],
    plugins=[
        ChannelsPlugin(
            backend=MemoryChannelsBackend(history=100),
            arbitrary_channels_allowed=True,
            create_ws_route_handlers=True,
            ws_handler_base_path="/ws",
        ),
    ],
)

from litestar import WebSocket, websocket, websocket_listener
from msgspec import Struct


class ChatInbound(Struct):
    text: str


class ChatOutbound(Struct):
    room_id: str
    sender: str
    text: str


@websocket_listener("/ws/chat/{room_id:str}", name="chat:room")
def handle_chat_message(data: ChatInbound, room_id: str) -> ChatOutbound:
    """Receive typed JSON and broadcast a typed response."""
    return ChatOutbound(room_id=room_id, sender="user", text=data.text)


@websocket("/ws/raw/{stream_id:str}", name="raw:stream")
async def raw_socket_handler(socket: WebSocket, stream_id: str) -> None:
    """Low-level WebSocket lifecycle handler."""
    await socket.accept()
    await socket.send_json({"stream_id": stream_id, "status": "connected"})
    await socket.close()

import asyncio
from collections.abc import AsyncGenerator
from litestar import get
from litestar.response import ServerSentEvent, ServerSentEventMessage
from litestar.serialization import encode_json
from msgspec import Struct


class OrderStatusEvent(Struct):
    order_id: str
    status: str


async def _order_event_generator(order_id: str) -> AsyncGenerator[ServerSentEventMessage, None]:
    for step, status in enumerate(("queued", "processing", "completed"), start=1):
        payload = encode_json(OrderStatusEvent(order_id=order_id, status=status)).decode("utf-8")
        yield ServerSentEventMessage(data=payload, event="order.updated", id=str(step), retry=3000)
        await asyncio.sleep(0.5)


@get("/api/orders/{order_id:str}/events", name="orders:events", sync_to_thread=False)
def stream_order_events(order_id: str) -> ServerSentEvent:
    """Stream Server-Sent Events for an order."""
    return ServerSentEvent(_order_event_generator(order_id))