feat(event): add typed channel event system with static type checking

Author: huynguyengl99

chanx/generic/websocket.py

@@ -5,12 +5,13 @@
 offering a robust framework for building real-time applications with Django
 Channels and Django REST Framework. The AsyncJsonWebsocketConsumer serves as the
 foundation for WebSocket connections with integrated authentication, permissions,
-structured message handling, and group messaging capabilities.
+structured message handling, group messaging capabilities, and typed channel events.
 
 Key features:
 - DRF-style authentication and permission checking
 - Structured message handling with Pydantic validation
 - Automatic group management for pub/sub messaging
+- Typed channel event system with discriminated unions
 - Comprehensive error handling and reporting
 - Configurable logging and message completion signals
 - Support for object-level permissions and retrieval
@@ -25,13 +26,14 @@
 import sys
 import uuid
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Sequence
+from collections.abc import Awaitable, Callable, Iterable, Sequence
 from types import ModuleType
 from typing import (
     Annotated,
     Any,
     Generic,
     Literal,
+    TypedDict,
     cast,
     get_args,
     get_origin,
@@ -40,6 +42,7 @@
 from channels.generic.websocket import (
     AsyncJsonWebsocketConsumer as BaseAsyncJsonWebsocketConsumer,
 )
+from channels.layers import get_channel_layer
 from django.contrib.auth.models import AnonymousUser, User
 from django.db.models import Model
 from rest_framework.authentication import BaseAuthentication
@@ -50,13 +53,15 @@
 )
 
 import structlog
+from asgiref.sync import async_to_sync
 from asgiref.typing import WebSocketConnectEvent, WebSocketDisconnectEvent
 from pydantic import Field, TypeAdapter, ValidationError
 from typing_extensions import TypeVar, get_original_bases
 
 from chanx.constants import MISSING_PYHUMPS_ERROR
 from chanx.generic.authenticator import ChanxWebsocketAuthenticator, QuerysetLike
 from chanx.messages.base import (
+    BaseChannelEvent,
     BaseGroupMessage,
     BaseMessage,
 )
@@ -83,10 +88,15 @@
     "OG", bound=BaseGroupMessage | None, default=None
 )  # Outgoing group messages
 M = TypeVar("M", bound=Model | None, default=None)  # Object model
+Event = TypeVar("Event", bound=BaseChannelEvent | None, default=None)  # Channel Events
+
+
+class EventPayload(TypedDict):
+    event_data: BaseChannelEvent
 
 
 class AsyncJsonWebsocketConsumer(
-    Generic[IC, OG, M], BaseAsyncJsonWebsocketConsumer, ABC
+    Generic[IC, Event, OG, M], BaseAsyncJsonWebsocketConsumer, ABC
 ):
     """
     Base class for asynchronous JSON WebSocket consumers with authentication and permissions.
@@ -99,6 +109,9 @@ class AsyncJsonWebsocketConsumer(
     `OUTGOING_GROUP_MESSAGE_SCHEMA` to enable proper validation and handling
     of group message broadcasts.
 
+    For typed channel events, subclasses can define a union type of channel events
+    and use the generic type parameter Event to enable type-safe channel event handling.
+
     Attributes:
         authentication_classes: DRF authentication classes for connection verification
         permission_classes: DRF permission classes for connection authorization
@@ -135,6 +148,7 @@ class AsyncJsonWebsocketConsumer(
 
     # Message schemas
     _INCOMING_MESSAGE_SCHEMA: IC
+    _EVENT_SCHEMA: Event
     _OUTGOING_GROUP_MESSAGE_SCHEMA: OG
 
     # Object instance
@@ -164,14 +178,21 @@ def __init_subclass__(cls, *args: Any, **kwargs: Any):
                         if i < len(type_var_vals) and var is not None:
                             type_var_vals[i] = var
 
-                    incoming_message_schema, outgoing_group_message_schema, _model = (
-                        type_var_vals
-                    )
+                    (
+                        incoming_message_schema,
+                        event_schema,
+                        outgoing_group_message_schema,
+                        _model,
+                    ) = type_var_vals
                 else:
-                    incoming_message_schema, outgoing_group_message_schema, _model = (
-                        get_args(base)
-                    )
+                    (
+                        incoming_message_schema,
+                        event_schema,
+                        outgoing_group_message_schema,
+                        _model,
+                    ) = get_args(base)
                 cls._INCOMING_MESSAGE_SCHEMA = incoming_message_schema
+                cls._EVENT_SCHEMA = event_schema
                 cls._OUTGOING_GROUP_MESSAGE_SCHEMA = outgoing_group_message_schema
                 break
 
@@ -237,6 +258,13 @@ def _setup_message_adapters(self) -> None:
             ]
         )
 
+        self.event_adapter: TypeAdapter[Event] = TypeAdapter(
+            Annotated[
+                self._EVENT_SCHEMA,
+                Field(discriminator="handler"),
+            ]
+        )
+
         self.outgoing_group_message_adapter: TypeAdapter[OG] = TypeAdapter(
             Annotated[
                 self._OUTGOING_GROUP_MESSAGE_SCHEMA,
@@ -582,6 +610,123 @@ async def send_group_member(self, event: GroupMemberEvent) -> None:
         if self.send_completion:
             await self.send_message(GroupCompleteMessage())
 
+    # Channel event system methods
+    @classmethod
+    async def asend_channel_event(
+        cls,
+        group_name: str,
+        event: Event,
+    ) -> None:
+        """
+        Send a typed channel event to one or more channel groups.
+
+        This is a class method that provides a type-safe way to send events through
+        the channel layer to consumers. It can be called from tasks, views, or other
+        places where you don't have a consumer instance.
+
+        Args:
+            event: The typed event to send (constrained by the consumer's Event type)
+            group_name: Group name to send to (required)
+
+        Example:
+            ```python
+            # From a Django task or view:
+            await ChatDetailConsumer.send_channel_event(
+                MemberAddedEvent(
+                    type="member_added",
+                    payload={"member_id": 123, "email": "user@example.com"}
+                ),
+                groups=["chat_room_1"],
+                from_user_pk=request.user.pk
+            )
+            ```
+        """
+        channel_layer = get_channel_layer()
+        assert channel_layer is not None
+
+        assert event is not None
+        await channel_layer.group_send(
+            group_name,
+            {
+                "type": "handle_channel_event",
+                "event_data": event.model_dump(),
+            },
+        )
+
+    @classmethod
+    def send_channel_event(
+        cls,
+        group_name: str,
+        event: Event,
+    ) -> None:
+        """
+        Synchronous version of send_channel_event for use in Django tasks/views.
+
+        This method provides the same functionality as send_channel_event but
+        can be called from synchronous code like Django tasks, views, or signals.
+
+        Args:
+            group_name: Group name to send to (required)
+            event: The typed event to send (constrained by the consumer's Event type)
+
+        Example:
+            ```python
+            # From a Django task:
+            ChatDetailConsumer.send_channel_event_sync(
+                "chat_room_1",
+                MemberAddedEvent(
+                    type="member_added",
+                    payload={"member_id": 123, "email": "user@example.com"}
+                ),
+            )
+            ```
+        """
+        async_to_sync(cls.asend_channel_event)(group_name, event)
+
+    async def handle_channel_event(self, event_payload: EventPayload) -> None:
+        """
+        Internal dispatcher for channel events with completion signal.
+
+        This method is called by the channel layer when an event is sent to a group
+        this consumer belongs to. It extracts the event data, checks exclusion rules,
+        finds the appropriate handler method, and calls it with proper error handling.
+
+        Args:
+            event_payload: The message from the channel layer containing event data
+                         and metadata about the sender
+        """
+        event_data_dict = event_payload.get("event_data", {})
+        event_data = self.event_adapter.validate_python(event_data_dict)
+
+        assert event_data is not None
+
+        try:
+            handler_name = event_data.handler
+            if not handler_name:
+                await logger.awarning("Received channel event without handler field")
+                return
+
+            # Find and call the handler method
+            handler_method: Callable[[Event], Awaitable[None]] | None = getattr(
+                self, handler_name, None
+            )
+            if not callable(handler_method):
+                await logger.ainfo(
+                    f"Handler method '{handler_name}' is not available for sending event"
+                )
+                return
+
+                # Handler is async, await it
+            await handler_method(event_data)
+
+        except Exception as e:
+            await logger.aexception(f"Failed to process channel event: {str(e)}")
+            # Don't re-raise to avoid breaking the channel layer
+        finally:
+            # Send completion signal if configured
+            if self.send_completion:
+                await self.send_message(CompleteMessage())
+
     # Helper methods
 
     async def _handle_receive_json_and_signal_complete(

chanx/messages/base.py

@@ -21,11 +21,7 @@
 """
 
 import abc
-from typing import (
-    Any,
-    Literal,
-    get_origin,
-)
+from typing import Any, Literal, get_origin
 
 from pydantic import BaseModel, ConfigDict
 from typing_extensions import Unpack
@@ -93,3 +89,8 @@ class BaseGroupMessage(BaseMessage, abc.ABC):
 
     is_mine: bool = False
     is_current: bool = False
+
+
+class BaseChannelEvent(BaseModel, abc.ABC):
+    handler: Any
+    payload: Any

docs/_static/interrogate_badge.svg

@@ -1,7 +1,5 @@
 from typing import Any
 
-from rest_framework.permissions import IsAuthenticated
-
 from chanx.generic.websocket import AsyncJsonWebsocketConsumer
 from chanx.messages.incoming import PingMessage
 from chanx.messages.outgoing import PongMessage
@@ -17,7 +15,8 @@
 class AssistantConsumer(AsyncJsonWebsocketConsumer[AssistantIncomingMessage]):
     """Websocket to chat with server, like chat with chatbot system"""
 
-    permission_classes = [IsAuthenticated]
+    permission_classes = []
+    authentication_classes = []
 
     async def receive_message(
         self, message: AssistantIncomingMessage, **kwargs: Any

sandbox/assistants/consumers.py

@@ -3,6 +3,7 @@
 from unittest.mock import patch
 from uuid import uuid4
 
+from django.conf import settings
 from rest_framework import status
 
 import pytest
@@ -288,3 +289,26 @@ async def test_ignore_actions(self) -> None:
 
             # Should not log received message for silent action
             assert "ping" not in str(mock_logger.call_args_list)
+
+    async def test_unauthorized_connection(self) -> None:
+        # Test anonymous connection
+        anonymous_communicator = self.create_communicator(
+            headers=[
+                (b"origin", settings.SERVER_URL.encode()),
+            ]
+        )
+        await anonymous_communicator.connect()
+        await anonymous_communicator.assert_authenticated_status_ok()
+
+        # Test chat functionality
+        message_content = "My message content"
+        await anonymous_communicator.send_message(
+            NewMessage(payload=MessagePayload(content=message_content))
+        )
+
+        all_messages = await anonymous_communicator.receive_all_json()
+        assert all_messages == [
+            ReplyMessage(
+                payload=MessagePayload(content=f"Reply: {message_content}")
+            ).model_dump()
+        ]

sandbox/assistants/tests/test_consumers.py

@@ -17,7 +17,9 @@
 
 
 class ChatDetailConsumer(
-    AsyncJsonWebsocketConsumer[ChatIncomingMessage, OutgoingGroupMessage, GroupChat]
+    AsyncJsonWebsocketConsumer[
+        ChatIncomingMessage, None, OutgoingGroupMessage, GroupChat
+    ]
 ):
     permission_classes = [IsGroupChatMember]
     queryset = GroupChat.objects.get_queryset()

sandbox/chat/consumers/chat_detail.py

@@ -21,6 +21,7 @@
 from django.urls import include, path
 from django.views.generic import RedirectView
 
+from discussion.views import NotifyView
 from drf_spectacular.views import (
     SpectacularAPIView,
     SpectacularRedocView,
@@ -31,6 +32,7 @@
     path("accounts/", include("accounts.urls")),
     path("chat/", include("chat.urls")),
     path("playground/", include("chanx.playground.urls")),
+    path("notify/", NotifyView.as_view(), name="notify_people"),
     path("schema/", SpectacularAPIView.as_view(), name="schema"),
     path(
         "schema/swg/",