Coverage for src/lexigram/web/websocket/handler.py: 33%
126 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 04:37 +0800
1"""WebSocket Handler Base Class.
3Guidance — Handler vs Gateway
4------------------------------
6:class:`AbstractWebSocketHandler` is for **application-level WebSocket endpoints** that
7are owned and processed by the current Lexigram application:
9- Chat rooms, live dashboards, notifications, collaborative editing
10- When the connection logic (auth, rooms, broadcast) lives *in this app*
11- Each handler class maps to one URL pattern
13**When to use a WebSocket Gateway instead:**
15A gateway (reverse-proxy pattern) is appropriate when the Lexigram application
16acts as a *protocol translator or hub* rather than the final consumer:
18- Proxying WebSocket connections to a backend service
19- Multiplexing many logical channels over one underlying connection
20- Bridging WebSocket↔AMQP/Redis Pub-Sub
22In those cases, do not use :class:`AbstractWebSocketHandler` — use a raw ASGI handler
23or the messaging layer (:mod:`lexigram.messaging`) with a pub/sub backend.
25"""
27from __future__ import annotations
29from abc import ABC, abstractmethod
30from typing import Any, ClassVar
32from starlette.websockets import WebSocketDisconnect
34from lexigram import serialization as json
35from lexigram.logging import get_logger
36from lexigram.web.transport.websockets import WebSocket
37from lexigram.web.websocket.connection_id import ConnectionIDManager
38from lexigram.web.websocket.rate_limiter import WebSocketRateLimiter
40logger = get_logger(__name__)
43class AbstractWebSocketHandler(ABC):
44 """Base class for WebSocket handlers.
46 Subclass this to create WebSocket endpoints with connection lifecycle
47 management, message handling, and optional room/broadcast support.
49 Attributes:
50 ping_interval: Seconds between ping frames (default: 30)
51 ping_timeout: Seconds to wait for pong response (default: 10)
52 max_connections: Global cap on concurrent connections for this handler.
53 New connections are rejected when the cap is reached (default: None).
54 max_connections_per_user: Limit concurrent connections per user identity.
55 Requires ``get_user_id()`` to return a non-None value (default: None).
56 max_messages_per_second: Per-connection message rate limit using a
57 token-bucket algorithm (default: None — unlimited).
59 Example:
60 ```python
61 @websocket_handler("/ws/chat/{room_id}")
62 class ChatHandler(AbstractWebSocketHandler):
63 rooms: dict[str, set[WebSocket]] = {}
65 async def on_connect(self, websocket):
66 room_id = websocket.path_params["room_id"]
67 if room_id not in self.rooms:
68 self.rooms[room_id] = set()
69 self.rooms[room_id].add(websocket)
70 await websocket.accept()
72 async def on_message(self, websocket, message):
73 room_id = websocket.path_params["room_id"]
74 for ws in self.rooms.get(room_id, []):
75 await ws.send_json(message)
77 async def on_disconnect(self, websocket):
78 room_id = websocket.path_params["room_id"]
79 self.rooms.get(room_id, set()).discard(websocket)
80 ```
81 """
83 # Configuration
84 ping_interval: int = 30
85 ping_timeout: int = 10
86 max_connections: int | None = None
87 max_connections_per_user: int | None = None
88 max_messages_per_second: float | None = None
90 # Route metadata (set by decorator)
91 _path: str | None = None
92 _guards: ClassVar[list[Any]] = []
94 # Connection tracking
95 _connections: set[WebSocket]
97 def __init__(
98 self,
99 connection_id_manager: ConnectionIDManager | None = None,
100 ) -> None:
101 self._connections = set()
102 # Per-user connection count tracking: user_id → set of connection UUIDs
103 self._user_connections: dict[str, set[str]] = {}
104 self._connection_id_manager = (
105 connection_id_manager
106 if connection_id_manager is not None
107 else ConnectionIDManager()
108 )
109 self._rate_limiter: WebSocketRateLimiter | None = (
110 WebSocketRateLimiter(self.max_messages_per_second)
111 if self.max_messages_per_second is not None
112 else None
113 )
115 @abstractmethod
116 async def on_connect(self, websocket: WebSocket) -> None:
117 """Called when a client connects.
119 Override to handle connection setup. You MUST call
120 `await websocket.accept()` to accept the connection.
122 Args:
123 websocket: The WebSocket connection
124 """
126 @abstractmethod
127 async def on_message(self, websocket: WebSocket, message: dict[str, Any]) -> None:
128 """Called when a message is received.
130 Override to handle incoming messages.
132 Args:
133 websocket: The WebSocket connection
134 message: The parsed JSON message
135 """
136 # pragma: no cover
138 @abstractmethod
139 async def on_disconnect(self, websocket: WebSocket) -> None:
140 """Called when a client disconnects.
142 Override to handle cleanup.
144 Args:
145 websocket: The WebSocket connection
146 """
147 # pragma: no cover
149 async def on_error(self, websocket: WebSocket, error: Exception) -> None:
150 """Called when an error occurs.
152 Override to handle errors. Default implementation logs and closes.
154 Args:
155 websocket: The WebSocket connection
156 error: The exception that occurred
157 """
158 try:
159 await websocket.send_json(
160 {
161 "type": "error",
162 "error": str(error),
163 },
164 )
165 except (RuntimeError, ConnectionError, OSError, AttributeError):
166 # Log and swallow to avoid crashing the WebSocket handler for operational errors
167 logger.exception("Failed to send error message to websocket")
169 async def send_json(self, websocket: WebSocket, data: dict[str, Any]) -> None:
170 """Send JSON data to a websocket.
172 Args:
173 websocket: The WebSocket connection
174 data: Data to send as JSON
175 """
176 await websocket.send_json(data)
178 async def send_text(self, websocket: WebSocket, text: str) -> None:
179 """Send text data to a websocket.
181 Args:
182 websocket: The WebSocket connection
183 text: Text to send
184 """
185 await websocket.send_text(text)
187 async def send_error(self, websocket: WebSocket, error: str) -> None:
188 """Send an error message to a websocket.
190 Args:
191 websocket: The WebSocket connection
192 error: Error message
193 """
194 await websocket.send_json({"type": "error", "error": error})
196 async def broadcast(
197 self,
198 message: dict[str, Any],
199 exclude: WebSocket | None = None,
200 ) -> None:
201 """Broadcast message to all connected clients.
203 Args:
204 message: Message to broadcast
205 exclude: Optional WebSocket to exclude from broadcast
206 """
207 for ws in self._connections:
208 if ws != exclude:
209 try:
210 await ws.send_json(message)
211 except (RuntimeError, ConnectionError, OSError):
212 logger.exception(
213 "Error broadcasting JSON message to websocket client",
214 )
216 async def broadcast_text(
217 self,
218 text: str,
219 exclude: WebSocket | None = None,
220 ) -> None:
221 """Broadcast text to all connected clients.
223 Args:
224 text: Text to broadcast
225 exclude: Optional WebSocket to exclude from broadcast
226 """
227 for ws in self._connections:
228 if ws != exclude:
229 try:
230 await ws.send_text(text)
231 except (RuntimeError, ConnectionError, OSError):
232 logger.exception(
233 "Error broadcasting text message to websocket client",
234 )
236 @property
237 def connection_count(self) -> int:
238 """Get current number of connections."""
239 return len(self._connections)
241 @property
242 def count(self) -> int:
243 """Alias for :attr:`connection_count` (satisfies :class:`ConnectionManagerProtocol`)."""
244 return len(self._connections)
246 async def add(self, connection: WebSocket) -> None:
247 """Register a connection.
249 Called automatically during :meth:`handle`. Can also be called
250 directly when managing connections outside the default loop.
251 """
252 self._connections.add(connection)
254 async def remove(self, connection: WebSocket) -> None:
255 """Unregister a connection.
257 Called automatically during :meth:`handle`. Can also be called
258 directly when managing connections outside the default loop.
259 """
260 self._connections.discard(connection)
262 def get_user_id(self, websocket: WebSocket) -> str | None:
263 """Return a stable identifier for the connecting user.
265 Override this method to enable ``max_connections_per_user``
266 enforcement. The default implementation checks for a ``user``
267 attribute on ``websocket.state`` (commonly set by authentication
268 middleware) and falls back to ``None`` (enforcement disabled).
270 Args:
271 websocket: The incoming WebSocket connection (before accept).
273 Returns:
274 A hashable string user identifier, or ``None`` to skip
275 per-user limit enforcement for this connection.
276 """
277 state = getattr(websocket, "state", None)
278 user = getattr(state, "user", None) if state is not None else None
279 if user is None:
280 return None
281 return str(getattr(user, "id", None) or getattr(user, "sub", None) or str(user))
283 async def handle(self, websocket: WebSocket) -> None:
284 """Handle a WebSocket connection.
286 This is called by the router to process the WebSocket.
287 It manages the connection lifecycle and message loop.
289 Args:
290 websocket: The WebSocket connection
291 """
292 # --- Connection-level rate limiting (before accept) ---
294 # Global connection cap
295 if (
296 self.max_connections is not None
297 and len(self._connections) >= self.max_connections
298 ):
299 logger.warning(
300 "websocket.connection_limit_reached",
301 max_connections=self.max_connections,
302 current=len(self._connections),
303 )
304 await websocket.close(code=1008, reason="connection_limit_exceeded") # type: ignore[call-arg]
305 return
307 # Per-user connection cap
308 user_id: str | None = None
309 if self.max_connections_per_user is not None:
310 user_id = self.get_user_id(websocket)
311 if user_id is not None:
312 user_conn_ids = self._user_connections.get(user_id, set())
313 if len(user_conn_ids) >= self.max_connections_per_user:
314 logger.warning(
315 "websocket.user_connection_limit_reached",
316 user_id=user_id,
317 max_connections_per_user=self.max_connections_per_user,
318 )
319 await websocket.close(
320 code=1008, reason="user_connection_limit_exceeded"
321 ) # type: ignore[call-arg]
322 return
324 try:
325 # Register connection before on_connect so broadcast includes it
326 await self.add(websocket)
327 if user_id is not None:
328 connection_uuid = self._connection_id_manager.get_or_create(websocket) # type: ignore[arg-type]
329 self._user_connections.setdefault(user_id, set()).add(connection_uuid)
331 # Call on_connect - user must call accept()
332 await self.on_connect(websocket)
334 # Message loop
335 while True:
336 try:
337 # Receive message
338 data = await websocket.receive_json()
340 # Enforce per-connection rate limit if configured
341 if self._rate_limiter is not None:
342 connection_uuid = self._connection_id_manager.get_or_create(
343 websocket # type: ignore[arg-type]
344 )
345 if not await self._rate_limiter.check(connection_uuid):
346 await self.send_error(websocket, "rate_limit_exceeded")
347 continue
349 await self.on_message(websocket, data)
350 except json.JSONDecodeError as e:
351 await self.on_error(websocket, e)
352 except WebSocketDisconnect:
353 # Client disconnected, stop the loop gracefully
354 break
355 except (RuntimeError, ConnectionError, OSError):
356 # Connection closed or other operational error; log and stop the loop
357 logger.exception(
358 "Error while receiving messages from websocket; closing connection",
359 )
360 break
362 except Exception as e: # noqa: BLE001 — websocket handler must catch any error to invoke on_error and clean up the connection
363 logger.exception("Unhandled websocket handler error")
364 await self.on_error(websocket, e)
365 finally:
366 await self.remove(websocket)
367 # Clean up per-user tracking
368 if user_id is not None and user_id in self._user_connections:
369 connection_uuid = self._connection_id_manager.get(websocket) # type: ignore[arg-type,assignment]
370 if connection_uuid is not None:
371 self._user_connections[user_id].discard(connection_uuid)
372 if not self._user_connections[user_id]:
373 del self._user_connections[user_id]
374 if self._rate_limiter is not None:
375 connection_uuid = self._connection_id_manager.get(websocket) # type: ignore[arg-type,assignment]
376 if connection_uuid is not None:
377 await self._rate_limiter.remove(connection_uuid)
378 try:
379 await self.on_disconnect(websocket)
380 except (RuntimeError, ConnectionError, OSError, AttributeError):
381 logger.exception("Error during websocket disconnect handler")
384__all__ = ["AbstractWebSocketHandler"]