"""Plugin-owned WebSocket and SSE streaming endpoints for queue events.
Imported only when a stream config is present and enabled, so base package
imports stay free of routing and Channels-driver imports.
The router and both relays are private: ``QueuePlugin`` is their only caller,
and applications configure streaming through
:class:`~litestar_queues.events.stream_config.EventStreamConfig` instead.
``StreamMetrics`` stays public because it describes the observability surface
an application can record against.
"""
import asyncio
import contextlib
import inspect
import time
from collections import OrderedDict
from collections.abc import Container, Sequence
from typing import TYPE_CHECKING, Any, Protocol, cast
from litestar import Request, Router, WebSocket, get, websocket
from litestar.exceptions import PermissionDeniedException, WebSocketException
from litestar.params import FromPath
from litestar.response import ServerSentEvent
from litestar_queues.events.channels import QueueChannels
from litestar_queues.events.models import QueueEvent, QueueEventScope
if TYPE_CHECKING:
from collections.abc import AsyncIterator
from litestar.types import SSEData
from litestar_queues.config import QueueConfig
from litestar_queues.events.stream_config import EventStreamConfig
from litestar_queues.events.typing import ChannelsLike, ChannelsStreamBackend, ChannelsSubscriptionBackend
from litestar_queues.namespace import QueueNamespace
from litestar_queues.observability import QueueObservabilityRuntimeProtocol
__all__ = ("StreamMetrics",)
_STREAM_DEDUP_MAX_KEYS = 1024
def _resolve_channels_backend(socket: "Any") -> "ChannelsLike | None":
if hasattr(socket, "channels_plugin"):
return cast("ChannelsLike", socket.channels_plugin)
scope = getattr(socket, "scope", None)
if isinstance(scope, dict):
scoped = scope.get("channels") or scope.get("queue_event_channels")
if scoped is not None:
return cast("ChannelsLike", scoped)
app = getattr(socket, "app", None)
state = getattr(app, "state", None)
if state is not None:
for key in ("queue_event_channels_backend", "channels", "queue_event_channels"):
with contextlib.suppress(KeyError, TypeError):
value = state[key]
if value is not None:
return cast("ChannelsLike", value)
value = getattr(state, key, None)
if value is not None:
return cast("ChannelsLike", value)
return None
@contextlib.asynccontextmanager
async def _event_stream(
backend: "ChannelsLike", channels: "Sequence[str]", *, history: "int"
) -> "AsyncIterator[AsyncIterator[bytes]]":
if hasattr(backend, "start_subscription"):
subscription_backend = cast("ChannelsSubscriptionBackend", backend)
async with subscription_backend.start_subscription(list(channels), history=history) as subscriber:
yield subscriber.iter_events()
return
if not hasattr(backend, "subscribe") or not hasattr(backend, "stream_events"):
msg = "Queue event streaming requires a ChannelsPlugin or ChannelsBackend-like object."
raise RuntimeError(msg)
stream_backend = cast("ChannelsStreamBackend", backend)
await stream_backend.subscribe(list(channels))
try:
yield _backend_events(stream_backend.stream_events(), set(channels))
finally:
await stream_backend.unsubscribe(list(channels))
async def _backend_events(events: "AsyncIterator[tuple[str, bytes]]", channels: "set[str]") -> "AsyncIterator[bytes]":
async for channel, payload in events:
if channel in channels:
yield payload
def _decode_event(raw_event: "bytes | str") -> "QueueEvent | None":
try:
return QueueEvent.from_json(raw_event)
except (KeyError, TypeError, ValueError):
return None
[docs]
class StreamMetrics(Protocol):
"""Optional metric callbacks used by the WebSocket and SSE stream relays."""
[docs]
def on_connect(self, scope: "QueueEventScope") -> None:
"""Record a stream connection."""
[docs]
def on_event(self, scope: "QueueEventScope") -> None:
"""Record an event sent to a stream client."""
[docs]
def on_heartbeat(self, scope: "QueueEventScope") -> None:
"""Record a heartbeat sent to a stream client."""
[docs]
def on_dedup_drop(self, scope: "QueueEventScope") -> None:
"""Record a deduplicated event dropped by the stream relay."""
[docs]
def on_denial(self, scope: "QueueEventScope", reason: str) -> None:
"""Record an authorization denial."""
[docs]
def on_disconnect(self, scope: "QueueEventScope", duration_seconds: float) -> None:
"""Record stream connection lifetime."""
class _RuntimeStreamMetrics:
"""Record stream metrics through the queue observability runtime."""
__slots__ = ("_runtime",)
def __init__(self, runtime: "QueueObservabilityRuntimeProtocol") -> "None":
self._runtime = runtime
def on_connect(self, scope: "QueueEventScope") -> None:
attributes = _scope_attributes(scope)
self._runtime.record_counter("litestar_queues.stream.connections", attributes=attributes)
self._runtime.record_gauge_delta("litestar_queues.stream.active", 1, attributes=attributes)
def on_event(self, scope: "QueueEventScope") -> None:
self._runtime.record_counter("litestar_queues.stream.events_sent", attributes=_scope_attributes(scope))
def on_heartbeat(self, scope: "QueueEventScope") -> None:
self._runtime.record_counter("litestar_queues.stream.heartbeats", attributes=_scope_attributes(scope))
def on_dedup_drop(self, scope: "QueueEventScope") -> None:
self._runtime.record_counter("litestar_queues.stream.dedup_drops", attributes=_scope_attributes(scope))
def on_denial(self, scope: "QueueEventScope", reason: str) -> None:
attributes = _scope_attributes(scope)
attributes["reason"] = reason
self._runtime.record_counter("litestar_queues.stream.auth_denials", attributes=attributes)
def on_disconnect(self, scope: "QueueEventScope", duration_seconds: float) -> None:
attributes = _scope_attributes(scope)
self._runtime.record_gauge_delta("litestar_queues.stream.active", -1, attributes=attributes)
self._runtime.record_duration(
"litestar_queues.stream.connection.duration", duration_seconds, attributes=attributes
)
async def _stream_queue_events_ws(
socket: Any,
channels: Sequence[str],
*,
history: int = 0,
channels_backend: "ChannelsLike | None" = None,
heartbeat_interval: float = 25.0,
stream_metrics: StreamMetrics | None = None,
scope: QueueEventScope = "task",
) -> None:
"""Stream queue events to a WebSocket with a heartbeat and serialized sends.
The caller owns route paths, guards, tenant filtering, and authorization.
Subscriber backpressure is intentionally left to the configured Channels backend.
"""
backend = channels_backend or _resolve_channels_backend(socket)
if backend is None:
msg = "A Litestar Channels backend or plugin is required to stream queue events."
raise RuntimeError(msg)
await socket.accept()
_record_metric(stream_metrics, "on_connect", scope)
started_at = time.perf_counter()
send_lock = asyncio.Lock()
stop = asyncio.Event()
try:
async with _event_stream(backend, channels, history=history) as events:
event_task = asyncio.create_task(_pump_events(socket, events, send_lock, stream_metrics, scope))
heartbeat_task = asyncio.create_task(
_pump_heartbeat(socket, send_lock, heartbeat_interval, stop, stream_metrics=stream_metrics, scope=scope)
)
await _wait_for_stream_tasks(event_task, heartbeat_task, stop)
finally:
_record_metric(stream_metrics, "on_disconnect", scope, time.perf_counter() - started_at)
async def _wait_for_stream_tasks(
event_task: "asyncio.Task[None]", heartbeat_task: "asyncio.Task[None]", stop: "asyncio.Event"
) -> None:
tasks = {event_task, heartbeat_task}
try:
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
for task in done:
task.result()
for task in pending:
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task
finally:
stop.set()
for task in tasks:
if not task.done():
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task
async def _pump_events(
socket: Any, events: "Any", send_lock: "asyncio.Lock", stream_metrics: StreamMetrics | None, scope: QueueEventScope
) -> None:
seen_dedup_keys: "OrderedDict[str, None]" = OrderedDict()
async for raw_event in events:
event = _decode_event(raw_event)
if event is None:
continue
dedup_key = event.event_key if event.event_key is not None else event.id
if dedup_key in seen_dedup_keys:
seen_dedup_keys.move_to_end(dedup_key)
_record_metric(stream_metrics, "on_dedup_drop", scope)
continue
seen_dedup_keys[dedup_key] = None
if len(seen_dedup_keys) > _STREAM_DEDUP_MAX_KEYS:
seen_dedup_keys.popitem(last=False)
if not await _send_json(socket, send_lock, event.to_dict()):
return
_record_metric(stream_metrics, "on_event", scope)
async def _pump_heartbeat(
socket: Any,
send_lock: "asyncio.Lock",
interval: float,
stop: "asyncio.Event",
*,
stream_metrics: StreamMetrics | None,
scope: QueueEventScope,
) -> None:
if interval <= 0:
await stop.wait()
return
while not stop.is_set():
try:
await asyncio.wait_for(stop.wait(), timeout=interval)
except asyncio.TimeoutError:
pass
else:
return
if not await _send_json(socket, send_lock, {"type": "ping"}):
return
_record_metric(stream_metrics, "on_heartbeat", scope)
async def _send_json(socket: Any, send_lock: "asyncio.Lock", payload: dict[str, object]) -> bool:
async with send_lock:
try:
await socket.send_json(payload)
except (OSError, RuntimeError):
return False
except Exception as exc:
if exc.__class__.__name__ == "WebSocketDisconnect":
return False
raise
return True
def _record_metric(stream_metrics: StreamMetrics | None, method_name: str, *args: object) -> None:
if stream_metrics is not None:
getattr(stream_metrics, method_name)(*args)
def _scope_attributes(scope: "QueueEventScope") -> "dict[str, str]":
return {"scope": scope}
def _resolve_observability_runtime(
connection: Any, config: "QueueConfig"
) -> "QueueObservabilityRuntimeProtocol | None":
if config.observability is None:
return None
app = getattr(connection, "app", None)
state = getattr(app, "state", None)
if state is None:
return None
key = config.observability_runtime_state_key
with contextlib.suppress(KeyError, TypeError):
runtime = state[key]
if runtime is not None:
return cast("QueueObservabilityRuntimeProtocol", runtime)
runtime = getattr(state, key, None)
if runtime is not None:
return cast("QueueObservabilityRuntimeProtocol", runtime)
return None
def _resolve_stream_metrics(connection: Any, config: "QueueConfig") -> "StreamMetrics | None":
runtime = _resolve_observability_runtime(connection, config)
if runtime is not None and runtime.enabled:
return _RuntimeStreamMetrics(runtime)
return None
def _configured_stream_channels_backend(
config: "QueueConfig", channels_backend: "ChannelsLike | None"
) -> "ChannelsLike | None":
if channels_backend is not None:
return channels_backend
if config.events is not None:
return config.events.channels
return None
def _build_stream_router(
config: "QueueConfig", stream_config: "EventStreamConfig", *, channels_backend: "ChannelsLike | None" = None
) -> "Router":
"""Build plugin-owned queue-event stream handlers for configured scopes.
Returns:
A router containing one handler per recognized configured scope for each
enabled transport (WebSocket and/or SSE).
"""
stream_config = stream_config.resolve(config.names)
authorizer = stream_config.channel_authorizer
history = stream_config.replay_limit
configured_channels_backend = _configured_stream_channels_backend(config, channels_backend)
async def _authorize(
connection: Any,
scope: "QueueEventScope",
key: str | None,
*,
websocket: bool,
stream_metrics: StreamMetrics | None,
) -> None:
if authorizer is None:
return
result = authorizer(connection, scope, key)
if inspect.isawaitable(result):
result = await result
if not result:
_record_metric(stream_metrics, "on_denial", scope, "authz")
msg = "Channel authorization denied"
if websocket:
raise WebSocketException(detail=msg, code=4003)
raise PermissionDeniedException(detail=msg)
async def _relay(socket: "WebSocket", scope: "QueueEventScope", key: str | None, channel: str) -> None:
stream_metrics = _resolve_stream_metrics(socket, config)
await _authorize(socket, scope, key, websocket=True, stream_metrics=stream_metrics)
backend = configured_channels_backend
if backend is None:
backend = _resolve_channels_backend(socket)
await _stream_queue_events_ws(
socket,
[channel],
history=history,
channels_backend=backend,
heartbeat_interval=stream_config.heartbeat_interval,
stream_metrics=stream_metrics,
scope=scope,
)
async def _sse(connection: Any, scope: "QueueEventScope", key: str | None, channel: str) -> Any:
stream_metrics = _resolve_stream_metrics(connection, config)
await _authorize(connection, scope, key, websocket=False, stream_metrics=stream_metrics)
backend = configured_channels_backend
if backend is None:
backend = _resolve_channels_backend(connection)
return _stream_queue_events_sse(
connection,
[channel],
history=history,
channels_backend=backend,
heartbeat_interval=stream_config.heartbeat_interval,
stream_metrics=stream_metrics,
scope=scope,
)
handlers: list[Any] = []
if "websocket" in stream_config.transports:
_append_task_handler(handlers, stream_config.scopes, _relay, config.names)
_append_queue_handler(handlers, stream_config.scopes, _relay, config.names)
_append_worker_handler(handlers, stream_config.scopes, _relay, config.names)
_append_global_handler(handlers, stream_config.scopes, _relay, config.names)
_append_custom_handler(handlers, stream_config.scopes, _relay, config.names)
if "sse" in stream_config.transports:
_append_sse_task_handler(handlers, stream_config.scopes, _sse, config.names)
_append_sse_queue_handler(handlers, stream_config.scopes, _sse, config.names)
_append_sse_worker_handler(handlers, stream_config.scopes, _sse, config.names)
_append_sse_global_handler(handlers, stream_config.scopes, _sse, config.names)
_append_sse_custom_handler(handlers, stream_config.scopes, _sse, config.names)
return Router(
path=stream_config.path,
route_handlers=handlers,
guards=list(stream_config.guards) if stream_config.guards else None,
opt=dict(stream_config.opt) if stream_config.opt else None,
include_in_schema=stream_config.include_in_schema,
)
def _stream_queue_events_sse(
connection: Any,
channels: Sequence[str],
*,
history: int = 0,
channels_backend: "ChannelsLike | None" = None,
heartbeat_interval: float = 25.0,
stream_metrics: StreamMetrics | None = None,
scope: QueueEventScope = "task",
) -> "ServerSentEvent":
"""Return a server-sent event stream for queue events."""
return ServerSentEvent(
_sse_events(
connection,
channels,
history=history,
channels_backend=channels_backend,
heartbeat_interval=heartbeat_interval,
stream_metrics=stream_metrics,
scope=scope,
)
)
async def _sse_events(
connection: Any,
channels: Sequence[str],
*,
history: int,
channels_backend: "ChannelsLike | None",
heartbeat_interval: float,
stream_metrics: StreamMetrics | None,
scope: QueueEventScope,
) -> "AsyncIterator[SSEData]":
backend = channels_backend or _resolve_channels_backend(connection)
if backend is None:
msg = "A Litestar Channels backend or plugin is required to stream queue events."
raise RuntimeError(msg)
_record_metric(stream_metrics, "on_connect", scope)
started_at = time.perf_counter()
stream_context = _event_stream(backend, channels, history=history)
events = await stream_context.__aenter__()
try:
async for frame in _sse_event_frames(
events, heartbeat_interval=heartbeat_interval, stream_metrics=stream_metrics, scope=scope
):
yield frame
finally:
await stream_context.__aexit__(None, None, None)
_record_metric(stream_metrics, "on_disconnect", scope, time.perf_counter() - started_at)
async def _sse_event_frames(
events: "Any", *, heartbeat_interval: float, stream_metrics: StreamMetrics | None, scope: QueueEventScope
) -> "AsyncIterator[SSEData]":
seen_dedup_keys: "OrderedDict[str, None]" = OrderedDict()
iterator = events.__aiter__()
next_event = asyncio.create_task(anext(iterator))
try:
while True:
if heartbeat_interval <= 0:
try:
raw_event = await next_event
except StopAsyncIteration:
return
else:
done, _pending = await asyncio.wait({next_event}, timeout=heartbeat_interval)
if not done:
_record_metric(stream_metrics, "on_heartbeat", scope)
yield {"comment": "ping", "data": None}
continue
try:
raw_event = next_event.result()
except StopAsyncIteration:
return
next_event = asyncio.create_task(anext(iterator))
frame = _sse_frame(raw_event, seen_dedup_keys, stream_metrics, scope)
if frame is not None:
yield frame
finally:
if not next_event.done():
next_event.cancel()
with contextlib.suppress(asyncio.CancelledError):
await next_event
def _sse_frame(
raw_event: "bytes | str",
seen_dedup_keys: "OrderedDict[str, None]",
stream_metrics: StreamMetrics | None,
scope: QueueEventScope,
) -> "dict[str, str] | None":
event = _decode_event(raw_event)
if event is None:
return None
dedup_key = event.event_key if event.event_key is not None else event.id
if dedup_key in seen_dedup_keys:
seen_dedup_keys.move_to_end(dedup_key)
_record_metric(stream_metrics, "on_dedup_drop", scope)
return None
seen_dedup_keys[dedup_key] = None
if len(seen_dedup_keys) > _STREAM_DEDUP_MAX_KEYS:
seen_dedup_keys.popitem(last=False)
_record_metric(stream_metrics, "on_event", scope)
return {"event": event.type, "data": event.to_json().decode("utf-8")}
def _append_task_handler(handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace") -> None:
if "task" not in scopes:
return
@websocket("/tasks/{task_id:str}", name=namespace.registration("event", "stream", "task"))
async def task_stream(socket: "WebSocket", task_id: FromPath[str]) -> None:
await relay(socket, "task", task_id, QueueChannels.task(task_id, namespace=namespace))
handlers.append(task_stream)
def _append_sse_task_handler(
handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace"
) -> None:
if "task" not in scopes:
return
@get(
"/sse/tasks/{task_id:str}", name=namespace.registration("event", "sse", "task"), media_type="text/event-stream"
)
async def task_sse(request: Request, task_id: FromPath[str]) -> Any:
return await relay(request, "task", task_id, QueueChannels.task(task_id, namespace=namespace))
handlers.append(task_sse)
def _append_queue_handler(handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace") -> None:
if "queue" not in scopes:
return
@websocket("/queues/{queue:str}", name=namespace.registration("event", "stream", "queue"))
async def queue_stream(socket: "WebSocket", queue: FromPath[str]) -> None:
await relay(socket, "queue", queue, QueueChannels.queue(queue, namespace=namespace))
handlers.append(queue_stream)
def _append_sse_queue_handler(
handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace"
) -> None:
if "queue" not in scopes:
return
@get(
"/sse/queues/{queue:str}", name=namespace.registration("event", "sse", "queue"), media_type="text/event-stream"
)
async def queue_sse(request: Request, queue: FromPath[str]) -> Any:
return await relay(request, "queue", queue, QueueChannels.queue(queue, namespace=namespace))
handlers.append(queue_sse)
def _append_worker_handler(
handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace"
) -> None:
if "worker" not in scopes:
return
@websocket("/workers/{worker_id:str}", name=namespace.registration("event", "stream", "worker"))
async def worker_stream(socket: "WebSocket", worker_id: FromPath[str]) -> None:
await relay(socket, "worker", worker_id, QueueChannels.worker(worker_id, namespace=namespace))
handlers.append(worker_stream)
def _append_sse_worker_handler(
handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace"
) -> None:
if "worker" not in scopes:
return
@get(
"/sse/workers/{worker_id:str}",
name=namespace.registration("event", "sse", "worker"),
media_type="text/event-stream",
)
async def worker_sse(request: Request, worker_id: FromPath[str]) -> Any:
return await relay(request, "worker", worker_id, QueueChannels.worker(worker_id, namespace=namespace))
handlers.append(worker_sse)
def _append_global_handler(
handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace"
) -> None:
if "global" not in scopes:
return
@websocket("/global", name=namespace.registration("event", "stream", "global"))
async def global_stream(socket: "WebSocket") -> None:
await relay(socket, "global", None, QueueChannels.global_channel(namespace=namespace))
handlers.append(global_stream)
def _append_sse_global_handler(
handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace"
) -> None:
if "global" not in scopes:
return
@get("/sse/global", name=namespace.registration("event", "sse", "global"), media_type="text/event-stream")
async def global_sse(request: Request) -> Any:
return await relay(request, "global", None, QueueChannels.global_channel(namespace=namespace))
handlers.append(global_sse)
def _append_custom_handler(
handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace"
) -> None:
if "custom" not in scopes:
return
@websocket("/custom/{scope_key:str}", name=namespace.registration("event", "stream", "custom"))
async def custom_stream(socket: "WebSocket", scope_key: FromPath[str]) -> None:
await relay(socket, "custom", scope_key, QueueChannels.custom(scope_key, namespace=namespace))
handlers.append(custom_stream)
def _append_sse_custom_handler(
handlers: list[Any], scopes: Container[str], relay: Any, namespace: "QueueNamespace"
) -> None:
if "custom" not in scopes:
return
@get(
"/sse/custom/{scope_key:str}",
name=namespace.registration("event", "sse", "custom"),
media_type="text/event-stream",
)
async def custom_sse(request: Request, scope_key: FromPath[str]) -> Any:
return await relay(request, "custom", scope_key, QueueChannels.custom(scope_key, namespace=namespace))
handlers.append(custom_sse)
