| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio, os |
| 4 | import re |
| 5 | import time |
| 6 | import threading |
| 7 | from collections import defaultdict, deque |
| 8 | from dataclasses import dataclass, field |
| 9 | from datetime import datetime, timedelta, timezone |
| 10 | from typing import Any, Callable, Deque, Dict, Iterable, List, Optional, Set |
| 11 | |
| 12 | import socketio |
| 13 | import uuid |
| 14 | |
| 15 | from helpers.defer import DeferredTask |
| 16 | from helpers.print_style import PrintStyle |
| 17 | from helpers import runtime |
| 18 | from helpers.ws import ConnectionIdentity, ConnectionNotFoundError, WsHandler, _ws_debug_enabled, ws_debug |
| 19 | |
| 20 | |
| 21 | # Event validation |
| 22 | |
| 23 | _EVENT_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$") |
| 24 | _RESERVED_EVENT_NAMES: set[str] = { |
| 25 | "connect", |
| 26 | "disconnect", |
| 27 | "error", |
| 28 | "ping", |
| 29 | "pong", |
| 30 | "connect_error", |
| 31 | "reconnect", |
| 32 | "reconnect_attempt", |
| 33 | "reconnect_error", |
| 34 | "reconnect_failed", |
| 35 | } |
| 36 | |
| 37 | |
| 38 | # WsResult – standardized handler return value |
| 39 | |
| 40 | class WsResult: |
| 41 | """Helper wrapper for standardized handler results. |
| 42 | |
| 43 | Instances are converted to the canonical ``RequestResultItem`` shape by |
| 44 | :class:`WsManager`. Helper constructors enforce payload validation |
| 45 | so handlers no longer need to hand-craft dictionaries. |
| 46 | """ |
| 47 | |
| 48 | __slots__ = ("_ok", "_data", "_error", "_correlation_id", "_duration_ms") |
| 49 | |
| 50 | def __init__( |
| 51 | self, |
| 52 | ok: bool, |
| 53 | data: dict[str, Any] | None = None, |
| 54 | error: dict[str, Any] | None = None, |
| 55 | correlation_id: str | None = None, |
| 56 | duration_ms: float | None = None, |
| 57 | ) -> None: |
| 58 | if ok and error: |
| 59 | raise ValueError("Cannot be both ok and have an error") |
| 60 | if not ok and not error: |
| 61 | raise ValueError("Must either be ok or have an error") |
| 62 | if data is not None and not isinstance(data, dict): |
| 63 | raise TypeError("Data payload must be a dictionary or None") |
| 64 | if error is not None and not isinstance(error, dict): |
| 65 | raise TypeError("Error payload must be a dictionary or None") |
| 66 | if correlation_id is not None and not isinstance(correlation_id, str): |
| 67 | raise TypeError("Correlation ID must be a string or None") |
| 68 | if duration_ms is not None and not isinstance(duration_ms, (int, float)): |
| 69 | raise TypeError("Duration must be a number or None") |
| 70 | |
| 71 | self._ok = bool(ok) |
| 72 | self._data = dict(data) if data is not None else None |
| 73 | self._error = dict(error) if error is not None else None |
| 74 | self._correlation_id = correlation_id |
| 75 | self._duration_ms = float(duration_ms) if duration_ms is not None else None |
| 76 | |
| 77 | @classmethod |
| 78 | def ok( |
| 79 | cls, |
| 80 | data: dict[str, Any] | None = None, |
| 81 | *, |
| 82 | correlation_id: str | None = None, |
| 83 | duration_ms: float | None = None, |
| 84 | ) -> "WsResult": |
| 85 | if data is not None and not isinstance(data, dict): |
| 86 | raise TypeError("WsResult.ok data must be a dict or None") |
| 87 | payload = dict(data) if data is not None else None |
| 88 | return cls( |
| 89 | ok=True, |
| 90 | data=payload, |
| 91 | correlation_id=correlation_id, |
| 92 | duration_ms=duration_ms, |
| 93 | ) |
| 94 | |
| 95 | @classmethod |
| 96 | def error( |
| 97 | cls, |
| 98 | *, |
| 99 | code: str, |
| 100 | message: str, |
| 101 | details: Any | None = None, |
| 102 | correlation_id: str | None = None, |
| 103 | duration_ms: float | None = None, |
| 104 | ) -> "WsResult": |
| 105 | if not isinstance(code, str) or not code.strip(): |
| 106 | raise ValueError("Error code must be a non-empty string") |
| 107 | if not isinstance(message, str) or not message.strip(): |
| 108 | raise ValueError("Error message must be a non-empty string") |
| 109 | |
| 110 | error_payload: dict[str, Any] = {"code": code, "error": message} |
| 111 | if details is not None: |
| 112 | error_payload["details"] = details |
| 113 | return cls( |
| 114 | ok=False, |
| 115 | error=error_payload, |
| 116 | correlation_id=correlation_id, |
| 117 | duration_ms=duration_ms, |
| 118 | ) |
| 119 | |
| 120 | def as_result( |
| 121 | self, |
| 122 | *, |
| 123 | handler_id: str, |
| 124 | fallback_correlation_id: str | None, |
| 125 | duration_ms: float | None = None, |
| 126 | ) -> dict[str, Any]: |
| 127 | result: dict[str, Any] = { |
| 128 | "handlerId": handler_id, |
| 129 | "ok": self._ok, |
| 130 | } |
| 131 | |
| 132 | effective_duration = ( |
| 133 | self._duration_ms if self._duration_ms is not None else duration_ms |
| 134 | ) |
| 135 | if effective_duration is not None: |
| 136 | result["durationMs"] = round(effective_duration, 4) |
| 137 | |
| 138 | correlation = ( |
| 139 | self._correlation_id |
| 140 | if self._correlation_id is not None |
| 141 | else fallback_correlation_id |
| 142 | ) |
| 143 | if correlation is not None: |
| 144 | result["correlationId"] = correlation |
| 145 | |
| 146 | if self._ok: |
| 147 | result["data"] = dict(self._data) if self._data is not None else {} |
| 148 | else: |
| 149 | result["error"] = dict(self._error) if self._error is not None else { |
| 150 | "code": "INTERNAL_ERROR", |
| 151 | "error": "Internal server error", |
| 152 | } |
| 153 | return result |
| 154 | |
| 155 | |
| 156 | def validate_event_type(event_type: str) -> str: |
| 157 | """Validate an event name: must be lowercase_snake_case and not reserved.""" |
| 158 | if not isinstance(event_type, str): |
| 159 | raise TypeError("Event type must be a string") |
| 160 | if not _EVENT_NAME_PATTERN.fullmatch(event_type): |
| 161 | raise ValueError( |
| 162 | f"Invalid event type '{event_type}' – must match lowercase_snake_case" |
| 163 | ) |
| 164 | if event_type in _RESERVED_EVENT_NAMES: |
| 165 | raise ValueError( |
| 166 | f"Event type '{event_type}' is reserved by Socket.IO and cannot be used" |
| 167 | ) |
| 168 | return event_type |
| 169 | |
| 170 | |
| 171 | BUFFER_MAX_SIZE = 100 |
| 172 | BUFFER_TTL = timedelta(hours=1) |
| 173 | _shared_ws_manager: WsManager | None = None |
| 174 | |
| 175 | |
| 176 | async def send_data( |
| 177 | event_type: str, |
| 178 | data: dict[str, Any], |
| 179 | *, |
| 180 | endpoint_name: str = "/ws", |
| 181 | connection_id: str | None = None, |
| 182 | ) -> None: |
| 183 | """Convenience wrapper around :pymeth:`WsManager.send_data`. |
| 184 | |
| 185 | All optional parameters are keyword-only to match the instance method's |
| 186 | ``(endpoint_name, event_type, data, connection_id)`` order and avoid |
| 187 | positional confusion between the two signatures. |
| 188 | """ |
| 189 | manager = get_shared_ws_manager() |
| 190 | await manager.send_data(endpoint_name, event_type, data, connection_id) |
| 191 | |
| 192 | |
| 193 | def _utcnow() -> datetime: |
| 194 | return datetime.now(timezone.utc) |
| 195 | |
| 196 | |
| 197 | def set_shared_ws_manager(manager: "WsManager") -> None: |
| 198 | global _shared_ws_manager |
| 199 | _shared_ws_manager = manager |
| 200 | |
| 201 | |
| 202 | def get_shared_ws_manager() -> "WsManager": |
| 203 | manager = _shared_ws_manager |
| 204 | if manager is None: |
| 205 | raise RuntimeError("Shared WsManager has not been initialized") |
| 206 | return manager |
| 207 | |
| 208 | |
| 209 | @dataclass |
| 210 | class BufferedEvent: |
| 211 | event_type: str |
| 212 | data: dict[str, Any] |
| 213 | handler_id: str | None = None |
| 214 | correlation_id: str | None = None |
| 215 | timestamp: datetime = field(default_factory=_utcnow) |
| 216 | |
| 217 | |
| 218 | @dataclass |
| 219 | class ConnectionInfo: |
| 220 | namespace: str |
| 221 | sid: str |
| 222 | connected_at: datetime = field(default_factory=_utcnow) |
| 223 | last_activity: datetime = field(default_factory=_utcnow) |
| 224 | |
| 225 | |
| 226 | @dataclass |
| 227 | class _HandlerExecution: |
| 228 | handler: WsHandler |
| 229 | value: Any |
| 230 | duration_ms: float | None |
| 231 | |
| 232 | |
| 233 | DIAGNOSTIC_EVENT = "ws_dev_console_event" |
| 234 | LIFECYCLE_CONNECT_EVENT = "ws_lifecycle_connect" |
| 235 | LIFECYCLE_DISCONNECT_EVENT = "ws_lifecycle_disconnect" |
| 236 | STATE_PUSH_EVENT = "state_push" |
| 237 | SERVER_RESTART_EVENT = "server_restart" |
| 238 | |
| 239 | # Error codes returned by _build_error_result |
| 240 | ERR_NO_HANDLERS = "NO_HANDLERS" |
| 241 | ERR_HANDLER_ERROR = "HANDLER_ERROR" |
| 242 | ERR_INVALID_FILTER = "INVALID_FILTER" |
| 243 | ERR_INVALID_EVENT = "INVALID_EVENT" |
| 244 | ERR_CONNECTION_NOT_FOUND = "CONNECTION_NOT_FOUND" |
| 245 | ERR_TIMEOUT = "TIMEOUT" |
| 246 | |
| 247 | |
| 248 | class WsManager: |
| 249 | def __init__(self, socketio: socketio.AsyncServer, lock) -> None: |
| 250 | self.socketio = socketio |
| 251 | self.lock = lock |
| 252 | self.handlers: defaultdict[str, List[WsHandler]] = defaultdict(list) |
| 253 | self.connections: Dict[ConnectionIdentity, ConnectionInfo] = {} |
| 254 | self.buffers: defaultdict[ConnectionIdentity, Deque[BufferedEvent]] = ( |
| 255 | defaultdict(deque) |
| 256 | ) |
| 257 | self._known_sids: Set[ConnectionIdentity] = set() |
| 258 | self._disconnect_times: Dict[ConnectionIdentity, datetime] = {} |
| 259 | self._identifier: str = f"{self.__class__.__module__}.{self.__class__.__name__}" |
| 260 | # Session tracking (single-user default) |
| 261 | self.user_to_sids: defaultdict[str, Set[ConnectionIdentity]] = defaultdict(set) |
| 262 | self.sid_to_user: Dict[ConnectionIdentity, str | None] = {} |
| 263 | self._ALL_USERS_BUCKET = "allUsers" |
| 264 | self._server_restart_enabled: bool = False |
| 265 | self._diagnostic_watchers: Set[ConnectionIdentity] = set() |
| 266 | self._diagnostics_enabled: bool = runtime.is_development() |
| 267 | self._dispatcher_loop: asyncio.AbstractEventLoop | None = None |
| 268 | self._handler_worker: DeferredTask | None = None |
| 269 | self._lifecycle_tasks: Set[asyncio.Task] = set() |
| 270 | |
| 271 | # Internal: development-only debug logging to avoid noise in production |
| 272 | def _debug(self, message: str) -> None: |
| 273 | ws_debug(message) |
| 274 | |
| 275 | def _ensure_dispatcher_loop(self) -> None: |
| 276 | if self._dispatcher_loop is None: |
| 277 | try: |
| 278 | self._dispatcher_loop = asyncio.get_running_loop() |
| 279 | except RuntimeError: |
| 280 | return |
| 281 | |
| 282 | def _get_handler_worker(self) -> DeferredTask: |
| 283 | if self._handler_worker is None: |
| 284 | self._handler_worker = DeferredTask(thread_name="WsHandlers") |
| 285 | return self._handler_worker |
| 286 | |
| 287 | async def _run_on_dispatcher_loop(self, coro: Any) -> Any: |
| 288 | self._ensure_dispatcher_loop() |
| 289 | dispatcher_loop = self._dispatcher_loop |
| 290 | if dispatcher_loop is None: |
| 291 | return await coro |
| 292 | if dispatcher_loop.is_closed(): |
| 293 | try: |
| 294 | coro.close() |
| 295 | except Exception: # pragma: no cover - best-effort cleanup |
| 296 | pass |
| 297 | raise RuntimeError("Dispatcher event loop is closed") |
| 298 | |
| 299 | try: |
| 300 | running_loop = asyncio.get_running_loop() |
| 301 | except RuntimeError: |
| 302 | running_loop = None |
| 303 | |
| 304 | if running_loop is dispatcher_loop: |
| 305 | return await coro |
| 306 | |
| 307 | future = asyncio.run_coroutine_threadsafe(coro, dispatcher_loop) |
| 308 | return await asyncio.wrap_future(future) |
| 309 | |
| 310 | def _diagnostics_active(self) -> bool: |
| 311 | if not self._diagnostics_enabled: |
| 312 | return False |
| 313 | with self.lock: |
| 314 | return bool(self._diagnostic_watchers) |
| 315 | |
| 316 | def _copy_diagnostic_watchers(self) -> list[ConnectionIdentity]: |
| 317 | with self.lock: |
| 318 | return list(self._diagnostic_watchers) |
| 319 | |
| 320 | def register_diagnostic_watcher(self, namespace: str, sid: str) -> bool: |
| 321 | if not self._diagnostics_enabled: |
| 322 | return False |
| 323 | identity: ConnectionIdentity = (namespace, sid) |
| 324 | with self.lock: |
| 325 | if identity not in self.connections: |
| 326 | return False |
| 327 | self._diagnostic_watchers.add(identity) |
| 328 | return True |
| 329 | |
| 330 | def unregister_diagnostic_watcher(self, namespace: str, sid: str) -> None: |
| 331 | identity: ConnectionIdentity = (namespace, sid) |
| 332 | with self.lock: |
| 333 | self._diagnostic_watchers.discard(identity) |
| 334 | |
| 335 | def _timestamp(self) -> str: |
| 336 | return _utcnow().isoformat(timespec="milliseconds").replace("+00:00", "Z") |
| 337 | |
| 338 | def _summarize_payload(self, payload: dict[str, Any] | None) -> dict[str, Any]: |
| 339 | if not isinstance(payload, dict): |
| 340 | return {} |
| 341 | summary: dict[str, Any] = {} |
| 342 | for key in list(payload.keys())[:5]: |
| 343 | value = payload[key] |
| 344 | if isinstance(value, (str, int, float, bool)) or value is None: |
| 345 | preview = value |
| 346 | elif isinstance(value, dict): |
| 347 | preview = f"dict({len(value)})" |
| 348 | elif isinstance(value, list): |
| 349 | preview = f"list({len(value)})" |
| 350 | else: |
| 351 | preview = value.__class__.__name__ |
| 352 | summary[key] = preview |
| 353 | summary["__sizeBytes__"] = len(str(payload).encode("utf-8")) |
| 354 | return summary |
| 355 | |
| 356 | def _summarize_results(self, results: List[dict[str, Any]]) -> dict[str, Any]: |
| 357 | summary = {"ok": 0, "error": 0, "handlers": []} |
| 358 | for result in results: |
| 359 | handler_id = result.get("handlerId") |
| 360 | ok = bool(result.get("ok")) |
| 361 | if ok: |
| 362 | summary["ok"] += 1 |
| 363 | else: |
| 364 | summary["error"] += 1 |
| 365 | summary["handlers"].append( |
| 366 | { |
| 367 | "handlerId": handler_id, |
| 368 | "ok": ok, |
| 369 | "errorCode": (result.get("error") or {}).get("code"), |
| 370 | "durationMs": result.get("durationMs"), |
| 371 | } |
| 372 | ) |
| 373 | summary["handlerCount"] = len(summary["handlers"]) |
| 374 | return summary |
| 375 | |
| 376 | async def _publish_diagnostic_event( |
| 377 | self, payload: dict[str, Any] | Callable[[], dict[str, Any]] |
| 378 | ) -> None: |
| 379 | if not self._diagnostics_enabled: |
| 380 | return |
| 381 | watchers = self._copy_diagnostic_watchers() |
| 382 | if not watchers: |
| 383 | return |
| 384 | effective_payload = payload() if callable(payload) else payload |
| 385 | if ( |
| 386 | isinstance(effective_payload, dict) |
| 387 | and "sourceNamespace" not in effective_payload |
| 388 | ): |
| 389 | origin = effective_payload.get("namespace") |
| 390 | if isinstance(origin, str) and origin.strip(): |
| 391 | effective_payload = { |
| 392 | **effective_payload, |
| 393 | "sourceNamespace": origin.strip(), |
| 394 | } |
| 395 | |
| 396 | async def _emit_to_watcher(identity: ConnectionIdentity) -> None: |
| 397 | namespace, sid = identity |
| 398 | try: |
| 399 | await self.emit_to( |
| 400 | namespace, |
| 401 | sid, |
| 402 | DIAGNOSTIC_EVENT, |
| 403 | effective_payload, |
| 404 | handler_id=self._identifier, |
| 405 | diagnostic=True, |
| 406 | ) |
| 407 | except ConnectionNotFoundError: |
| 408 | self.unregister_diagnostic_watcher(namespace, sid) |
| 409 | |
| 410 | await asyncio.gather(*(_emit_to_watcher(identity) for identity in watchers)) |
| 411 | |
| 412 | def _schedule_lifecycle_broadcast( |
| 413 | self, namespace: str, event_type: str, payload: dict[str, Any] |
| 414 | ) -> None: |
| 415 | async def _broadcast() -> None: |
| 416 | try: |
| 417 | await self.broadcast( |
| 418 | namespace, |
| 419 | event_type, |
| 420 | payload, |
| 421 | diagnostic=True, |
| 422 | ) |
| 423 | except Exception as exc: # pragma: no cover - diagnostic |
| 424 | self._debug(f"Failed to broadcast lifecycle event {event_type}: {exc}") |
| 425 | |
| 426 | task = asyncio.create_task(_broadcast()) |
| 427 | self._lifecycle_tasks.add(task) |
| 428 | task.add_done_callback(self._lifecycle_tasks.discard) |
| 429 | |
| 430 | def _sweep_stale_sids(self) -> None: |
| 431 | """Remove _known_sids entries whose disconnect exceeds BUFFER_TTL.""" |
| 432 | now = _utcnow() |
| 433 | with self.lock: |
| 434 | stale = [ |
| 435 | identity |
| 436 | for identity, dt in self._disconnect_times.items() |
| 437 | if identity not in self.connections and (now - dt) > BUFFER_TTL |
| 438 | ] |
| 439 | for identity in stale: |
| 440 | self._known_sids.discard(identity) |
| 441 | self._disconnect_times.pop(identity, None) |
| 442 | self.buffers.pop(identity, None) |
| 443 | |
| 444 | def _normalize_handler_filter(self, value: Any, field_name: str) -> Set[str] | None: |
| 445 | if value is None: |
| 446 | return None |
| 447 | if isinstance(value, str): |
| 448 | return {value} |
| 449 | try: |
| 450 | iterator = iter(value) |
| 451 | except TypeError as exc: # pragma: no cover - defensive |
| 452 | raise ValueError( |
| 453 | f"{field_name} must be an array of handler identifiers" |
| 454 | ) from exc |
| 455 | |
| 456 | normalized: Set[str] = set() |
| 457 | for item in iterator: |
| 458 | if not isinstance(item, str): |
| 459 | raise ValueError( |
| 460 | f"{field_name} values must be handler identifier strings" |
| 461 | ) |
| 462 | normalized.add(item) |
| 463 | return normalized |
| 464 | |
| 465 | def _normalize_sid_filter(self, value: str | Iterable[str] | None) -> Set[str]: |
| 466 | if value is None: |
| 467 | return set() |
| 468 | if isinstance(value, str): |
| 469 | return {value} |
| 470 | normalized: Set[str] = set() |
| 471 | for item in value: |
| 472 | normalized.add(str(item)) |
| 473 | return normalized |
| 474 | |
| 475 | def _select_handlers( |
| 476 | self, |
| 477 | namespace: str, |
| 478 | *, |
| 479 | include: Set[str] | None, |
| 480 | exclude: Set[str] | None, |
| 481 | ) -> tuple[list[WsHandler], Set[str]]: |
| 482 | registered = self.handlers.get(namespace, []) |
| 483 | available_ids = {handler.identifier for handler in registered} |
| 484 | |
| 485 | if include is not None: |
| 486 | unknown = include - available_ids |
| 487 | if unknown: |
| 488 | raise ValueError( |
| 489 | f"Unknown handler(s) in includeHandlers for namespace '{namespace}': " |
| 490 | f"{', '.join(sorted(unknown))}" |
| 491 | ) |
| 492 | if exclude is not None: |
| 493 | unknown = exclude - available_ids |
| 494 | if unknown: |
| 495 | raise ValueError( |
| 496 | f"Unknown handler(s) in excludeHandlers for namespace '{namespace}': " |
| 497 | f"{', '.join(sorted(unknown))}" |
| 498 | ) |
| 499 | |
| 500 | selected: list[WsHandler] = [] |
| 501 | for handler in registered: |
| 502 | ident = handler.identifier |
| 503 | if include is not None and ident not in include: |
| 504 | continue |
| 505 | if exclude is not None and ident in exclude: |
| 506 | continue |
| 507 | selected.append(handler) |
| 508 | |
| 509 | return selected, available_ids |
| 510 | |
| 511 | def _resolve_correlation_id(self, payload: dict[str, Any]) -> str: |
| 512 | value = payload.get("correlationId") |
| 513 | if isinstance(value, str) and value.strip(): |
| 514 | correlation_id = value.strip() |
| 515 | else: |
| 516 | correlation_id = uuid.uuid4().hex |
| 517 | payload["correlationId"] = correlation_id |
| 518 | return correlation_id |
| 519 | |
| 520 | def register_handlers( |
| 521 | self, handlers_by_namespace: dict[str, Iterable[WsHandler]] |
| 522 | ) -> None: |
| 523 | for namespace, handlers in handlers_by_namespace.items(): |
| 524 | for handler in handlers: |
| 525 | handler.bind_manager(self, namespace=namespace) |
| 526 | if _ws_debug_enabled(): |
| 527 | PrintStyle.info( |
| 528 | "Registered WebSocket handler %s namespace=%s" |
| 529 | % (handler.identifier, namespace) |
| 530 | ) |
| 531 | existing = self.handlers.get(namespace, []) |
| 532 | if handler in existing: |
| 533 | PrintStyle.warning( |
| 534 | f"Duplicate handler registration for namespace '{namespace}'" |
| 535 | ) |
| 536 | self.handlers[namespace].append(handler) |
| 537 | self._debug( |
| 538 | f"Registered handler {handler.identifier} namespace={namespace}" |
| 539 | ) |
| 540 | |
| 541 | def iter_namespaces(self) -> list[str]: |
| 542 | return list(self.handlers.keys()) |
| 543 | |
| 544 | async def process_client_event( |
| 545 | self, |
| 546 | namespace: str, |
| 547 | event_type: str, |
| 548 | data: dict[str, Any], |
| 549 | sid: str, |
| 550 | *, |
| 551 | handlers: list[WsHandler], |
| 552 | ) -> dict[str, Any]: |
| 553 | """Process a client-originated event through provided handler instances. |
| 554 | |
| 555 | Unlike ``route_event`` which selects from globally registered handlers, |
| 556 | this accepts pre-selected instances (e.g. per-connection activated |
| 557 | handlers that have already passed security checks). |
| 558 | """ |
| 559 | self._ensure_dispatcher_loop() |
| 560 | incoming = dict(data or {}) |
| 561 | correlation_id = self._resolve_correlation_id(incoming) |
| 562 | |
| 563 | if "data" in incoming and isinstance(incoming.get("data"), dict): |
| 564 | handler_payload = dict(incoming["data"]) |
| 565 | if "excludeSids" in incoming: |
| 566 | handler_payload["excludeSids"] = incoming["excludeSids"] |
| 567 | else: |
| 568 | handler_payload = dict(incoming) |
| 569 | handler_payload["correlationId"] = correlation_id |
| 570 | |
| 571 | if not handlers: |
| 572 | return self._ack_error( |
| 573 | handler_id=self._identifier, |
| 574 | code=ERR_NO_HANDLERS, |
| 575 | message="No handlers available after security filtering", |
| 576 | correlation_id=correlation_id, |
| 577 | ) |
| 578 | |
| 579 | with self.lock: |
| 580 | info = self.connections.get((namespace, sid)) |
| 581 | if info: |
| 582 | info.last_activity = _utcnow() |
| 583 | |
| 584 | executions = await asyncio.gather( |
| 585 | *[ |
| 586 | self._invoke_handler(handler, event_type, dict(handler_payload), sid) |
| 587 | for handler in handlers |
| 588 | ] |
| 589 | ) |
| 590 | |
| 591 | results = self._collect_results( |
| 592 | executions, event_type, correlation_id, skip_none=True, |
| 593 | ) |
| 594 | |
| 595 | await self._publish_diagnostic_event( |
| 596 | lambda: { |
| 597 | "kind": "inbound", |
| 598 | "sourceNamespace": namespace, |
| 599 | "namespace": namespace, |
| 600 | "eventType": event_type, |
| 601 | "sid": sid, |
| 602 | "correlationId": correlation_id, |
| 603 | "timestamp": self._timestamp(), |
| 604 | "handlerCount": len(handlers), |
| 605 | "durationMs": sum( |
| 606 | (exec.duration_ms or 0.0) for exec in executions |
| 607 | ), |
| 608 | "resultSummary": self._summarize_results(results), |
| 609 | "payloadSummary": self._summarize_payload(handler_payload), |
| 610 | } |
| 611 | ) |
| 612 | |
| 613 | response = {"correlationId": correlation_id, "results": results} |
| 614 | self._debug( |
| 615 | f"Completed client event namespace={namespace} '{event_type}' " |
| 616 | f"sid={sid} correlation={correlation_id}" |
| 617 | ) |
| 618 | return response |
| 619 | |
| 620 | def _collect_results( |
| 621 | self, |
| 622 | executions: list[_HandlerExecution], |
| 623 | event_type: str, |
| 624 | correlation_id: str, |
| 625 | *, |
| 626 | skip_none: bool = False, |
| 627 | ) -> List[dict[str, Any]]: |
| 628 | """Build a result list from handler executions. |
| 629 | |
| 630 | Args: |
| 631 | skip_none: When ``True``, handlers that return ``None`` are omitted |
| 632 | from the results (fire-and-forget / opt-out semantics used by |
| 633 | ``process_client_event``). When ``False``, ``None`` is converted |
| 634 | to ``WsResult(ok=True)`` (default ``route_event`` behaviour). |
| 635 | """ |
| 636 | results: List[dict[str, Any]] = [] |
| 637 | for execution in executions: |
| 638 | handler = execution.handler |
| 639 | value = execution.value |
| 640 | duration_ms = execution.duration_ms |
| 641 | |
| 642 | if isinstance(value, Exception): |
| 643 | PrintStyle.error( |
| 644 | f"Error in handler {handler.identifier} for '{event_type}' " |
| 645 | f"(correlation {correlation_id}): {value}" |
| 646 | ) |
| 647 | results.append( |
| 648 | self._build_error_result( |
| 649 | handler_id=handler.identifier, |
| 650 | code=ERR_HANDLER_ERROR, |
| 651 | message="Internal server error", |
| 652 | details=str(value), |
| 653 | correlation_id=correlation_id, |
| 654 | duration_ms=duration_ms, |
| 655 | ) |
| 656 | ) |
| 657 | continue |
| 658 | |
| 659 | if isinstance(value, WsResult): |
| 660 | results.append( |
| 661 | value.as_result( |
| 662 | handler_id=handler.identifier, |
| 663 | fallback_correlation_id=correlation_id, |
| 664 | duration_ms=duration_ms, |
| 665 | ) |
| 666 | ) |
| 667 | continue |
| 668 | |
| 669 | if value is None: |
| 670 | if skip_none: |
| 671 | continue |
| 672 | helper_result = WsResult(ok=True) |
| 673 | elif isinstance(value, dict): |
| 674 | helper_result = WsResult(ok=True, data=value) |
| 675 | else: |
| 676 | helper_result = WsResult(ok=True, data={"result": value}) |
| 677 | |
| 678 | results.append( |
| 679 | helper_result.as_result( |
| 680 | handler_id=handler.identifier, |
| 681 | fallback_correlation_id=correlation_id, |
| 682 | duration_ms=duration_ms, |
| 683 | ) |
| 684 | ) |
| 685 | return results |
| 686 | |
| 687 | async def _invoke_handler( |
| 688 | self, |
| 689 | handler: WsHandler, |
| 690 | event_type: str, |
| 691 | payload: dict[str, Any], |
| 692 | sid: str, |
| 693 | ) -> _HandlerExecution: |
| 694 | instrument = self._diagnostics_active() |
| 695 | start = time.perf_counter() if instrument else None |
| 696 | try: |
| 697 | value = await self._get_handler_worker().execute_inside( |
| 698 | handler.process, event_type, payload, sid |
| 699 | ) |
| 700 | except Exception as exc: # pragma: no cover - handled by caller |
| 701 | duration_ms = ( |
| 702 | (time.perf_counter() - start) * 1000 if start is not None else None |
| 703 | ) |
| 704 | return _HandlerExecution(handler, exc, duration_ms) |
| 705 | duration_ms = ( |
| 706 | (time.perf_counter() - start) * 1000 if start is not None else None |
| 707 | ) |
| 708 | return _HandlerExecution(handler, value, duration_ms) |
| 709 | |
| 710 | async def handle_connect( |
| 711 | self, namespace: str, sid: str, user_id: str | None = None |
| 712 | ) -> None: |
| 713 | self._ensure_dispatcher_loop() |
| 714 | user_bucket = user_id or "single_user" |
| 715 | identity: ConnectionIdentity = (namespace, sid) |
| 716 | with self.lock: |
| 717 | self.connections[identity] = ConnectionInfo(namespace=namespace, sid=sid) |
| 718 | self._known_sids.add(identity) |
| 719 | self._disconnect_times.pop(identity, None) |
| 720 | self.sid_to_user[identity] = user_bucket |
| 721 | self.user_to_sids[self._ALL_USERS_BUCKET].add(identity) |
| 722 | self.user_to_sids[user_bucket].add(identity) |
| 723 | connection_count = sum( |
| 724 | 1 for conn_identity in self.connections if conn_identity[0] == namespace |
| 725 | ) |
| 726 | if _ws_debug_enabled(): |
| 727 | PrintStyle.info(f"WebSocket connected: namespace={namespace} sid={sid}") |
| 728 | await self._run_lifecycle(namespace, lambda h: h.on_connect(sid)) |
| 729 | await self._flush_buffer(identity) |
| 730 | if self._server_restart_enabled: |
| 731 | await self.emit_to( |
| 732 | namespace, |
| 733 | sid, |
| 734 | SERVER_RESTART_EVENT, |
| 735 | { |
| 736 | "emittedAt": self._timestamp(), |
| 737 | "runtimeId": runtime.get_runtime_id(), |
| 738 | }, |
| 739 | handler_id=self._identifier, |
| 740 | ) |
| 741 | if _ws_debug_enabled(): |
| 742 | PrintStyle.info( |
| 743 | f"server_restart broadcast emitted to namespace={namespace} sid={sid}" |
| 744 | ) |
| 745 | lifecycle_payload = { |
| 746 | "namespace": namespace, |
| 747 | "sid": sid, |
| 748 | "connectionCount": connection_count, |
| 749 | "timestamp": self._timestamp(), |
| 750 | } |
| 751 | await self._publish_diagnostic_event( |
| 752 | { |
| 753 | "kind": "lifecycle", |
| 754 | "event": "connect", |
| 755 | **lifecycle_payload, |
| 756 | } |
| 757 | ) |
| 758 | self._schedule_lifecycle_broadcast( |
| 759 | namespace, LIFECYCLE_CONNECT_EVENT, lifecycle_payload |
| 760 | ) |
| 761 | |
| 762 | async def handle_disconnect(self, namespace: str, sid: str) -> None: |
| 763 | self._ensure_dispatcher_loop() |
| 764 | identity: ConnectionIdentity = (namespace, sid) |
| 765 | with self.lock: |
| 766 | self.connections.pop(identity, None) |
| 767 | # Keep identity in _known_sids so emit_to buffers instead of raising; |
| 768 | # record disconnect time for TTL-based cleanup |
| 769 | self._disconnect_times[identity] = _utcnow() |
| 770 | # session tracking cleanup |
| 771 | user_bucket = self.sid_to_user.pop(identity, None) |
| 772 | if self._ALL_USERS_BUCKET in self.user_to_sids: |
| 773 | self.user_to_sids[self._ALL_USERS_BUCKET].discard(identity) |
| 774 | if not self.user_to_sids[self._ALL_USERS_BUCKET]: |
| 775 | self.user_to_sids.pop(self._ALL_USERS_BUCKET, None) |
| 776 | if user_bucket and user_bucket in self.user_to_sids: |
| 777 | self.user_to_sids[user_bucket].discard(identity) |
| 778 | if not self.user_to_sids[user_bucket]: |
| 779 | self.user_to_sids.pop(user_bucket, None) |
| 780 | connection_count = sum( |
| 781 | 1 for conn_identity in self.connections if conn_identity[0] == namespace |
| 782 | ) |
| 783 | self.unregister_diagnostic_watcher(namespace, sid) |
| 784 | PrintStyle.info(f"WebSocket disconnected: namespace={namespace} sid={sid}") |
| 785 | await self._run_lifecycle(namespace, lambda h: h.on_disconnect(sid)) |
| 786 | lifecycle_payload = { |
| 787 | "namespace": namespace, |
| 788 | "sid": sid, |
| 789 | "connectionCount": connection_count, |
| 790 | "timestamp": self._timestamp(), |
| 791 | } |
| 792 | await self._publish_diagnostic_event( |
| 793 | { |
| 794 | "kind": "lifecycle", |
| 795 | "event": "disconnect", |
| 796 | **lifecycle_payload, |
| 797 | } |
| 798 | ) |
| 799 | self._schedule_lifecycle_broadcast( |
| 800 | namespace, LIFECYCLE_DISCONNECT_EVENT, lifecycle_payload |
| 801 | ) |
| 802 | self._sweep_stale_sids() |
| 803 | |
| 804 | async def route_event( |
| 805 | self, |
| 806 | namespace: str, |
| 807 | event_type: str, |
| 808 | data: dict[str, Any], |
| 809 | sid: str, |
| 810 | ack: Optional[Callable[[Any], None]] = None, |
| 811 | *, |
| 812 | include_handlers: Set[str] | None = None, |
| 813 | exclude_handlers: Set[str] | None = None, |
| 814 | allow_exclude: bool = False, |
| 815 | handler_id: str | None = None, |
| 816 | ) -> dict[str, Any]: |
| 817 | self._ensure_dispatcher_loop() |
| 818 | incoming = dict(data or {}) |
| 819 | correlation_id = self._resolve_correlation_id(incoming) |
| 820 | self._debug( |
| 821 | f"Routing event namespace={namespace} '{event_type}' sid={sid} correlation={correlation_id}" |
| 822 | ) |
| 823 | |
| 824 | include_meta_raw = incoming.pop("includeHandlers", None) |
| 825 | exclude_meta_raw = incoming.pop("excludeHandlers", None) |
| 826 | |
| 827 | if "data" in incoming and isinstance(incoming.get("data"), dict): |
| 828 | handler_payload = dict(incoming.get("data") or {}) |
| 829 | if "excludeSids" in incoming: |
| 830 | handler_payload["excludeSids"] = incoming.get("excludeSids") |
| 831 | else: |
| 832 | handler_payload = dict(incoming) |
| 833 | |
| 834 | handler_payload["correlationId"] = correlation_id |
| 835 | |
| 836 | try: |
| 837 | include_meta = self._normalize_handler_filter( |
| 838 | include_meta_raw, "includeHandlers" |
| 839 | ) |
| 840 | except ValueError as exc: |
| 841 | return self._ack_error( |
| 842 | handler_id=handler_id or self._identifier, |
| 843 | code=ERR_INVALID_FILTER, |
| 844 | message=str(exc), |
| 845 | correlation_id=correlation_id, |
| 846 | ack=ack, |
| 847 | ) |
| 848 | |
| 849 | try: |
| 850 | exclude_meta = self._normalize_handler_filter( |
| 851 | exclude_meta_raw, "excludeHandlers" |
| 852 | ) |
| 853 | except ValueError as exc: |
| 854 | return self._ack_error( |
| 855 | handler_id=handler_id or self._identifier, |
| 856 | code=ERR_INVALID_FILTER, |
| 857 | message=str(exc), |
| 858 | correlation_id=correlation_id, |
| 859 | ack=ack, |
| 860 | ) |
| 861 | |
| 862 | if exclude_meta_raw is not None and not allow_exclude: |
| 863 | return self._ack_error( |
| 864 | handler_id=handler_id or self._identifier, |
| 865 | code=ERR_INVALID_FILTER, |
| 866 | message="excludeHandlers is not supported for this operation", |
| 867 | correlation_id=correlation_id, |
| 868 | ack=ack, |
| 869 | ) |
| 870 | |
| 871 | if include_handlers is not None and include_meta is not None: |
| 872 | if include_handlers != include_meta: |
| 873 | return self._ack_error( |
| 874 | handler_id=handler_id or self._identifier, |
| 875 | code=ERR_INVALID_FILTER, |
| 876 | message="Conflicting includeHandlers filters supplied", |
| 877 | correlation_id=correlation_id, |
| 878 | ack=ack, |
| 879 | ) |
| 880 | |
| 881 | if allow_exclude and exclude_handlers is not None and exclude_meta is not None: |
| 882 | if exclude_handlers != exclude_meta: |
| 883 | return self._ack_error( |
| 884 | handler_id=handler_id or self._identifier, |
| 885 | code=ERR_INVALID_FILTER, |
| 886 | message="Conflicting excludeHandlers filters supplied", |
| 887 | correlation_id=correlation_id, |
| 888 | ack=ack, |
| 889 | ) |
| 890 | |
| 891 | include = include_handlers or include_meta |
| 892 | exclude = exclude_handlers or (exclude_meta if allow_exclude else None) |
| 893 | |
| 894 | try: |
| 895 | validate_event_type(event_type) |
| 896 | except (TypeError, ValueError) as exc: |
| 897 | return self._ack_error( |
| 898 | handler_id=handler_id or self._identifier, |
| 899 | code=ERR_INVALID_EVENT, |
| 900 | message=str(exc), |
| 901 | correlation_id=correlation_id, |
| 902 | ack=ack, |
| 903 | ) |
| 904 | |
| 905 | registered = self.handlers.get(namespace, []) |
| 906 | if not registered: |
| 907 | PrintStyle.warning(f"No handlers registered for namespace '{namespace}'") |
| 908 | return self._ack_error( |
| 909 | handler_id=handler_id or self._identifier, |
| 910 | code=ERR_NO_HANDLERS, |
| 911 | message=f"No handler for namespace '{namespace}'", |
| 912 | correlation_id=correlation_id, |
| 913 | ack=ack, |
| 914 | ) |
| 915 | |
| 916 | try: |
| 917 | selected_handlers, _ = self._select_handlers( |
| 918 | namespace, include=include, exclude=exclude |
| 919 | ) |
| 920 | except ValueError as exc: |
| 921 | return self._ack_error( |
| 922 | handler_id=handler_id or self._identifier, |
| 923 | code=ERR_INVALID_FILTER, |
| 924 | message=str(exc), |
| 925 | correlation_id=correlation_id, |
| 926 | ack=ack, |
| 927 | ) |
| 928 | |
| 929 | if not selected_handlers: |
| 930 | return self._ack_error( |
| 931 | handler_id=handler_id or self._identifier, |
| 932 | code=ERR_NO_HANDLERS, |
| 933 | message=f"No handler for '{event_type}' after applying filters", |
| 934 | correlation_id=correlation_id, |
| 935 | ack=ack, |
| 936 | ) |
| 937 | |
| 938 | with self.lock: |
| 939 | info = self.connections.get((namespace, sid)) |
| 940 | if info: |
| 941 | info.last_activity = _utcnow() |
| 942 | |
| 943 | executions = await asyncio.gather( |
| 944 | *[ |
| 945 | self._invoke_handler(handler, event_type, dict(handler_payload), sid) |
| 946 | for handler in selected_handlers |
| 947 | ] |
| 948 | ) |
| 949 | |
| 950 | # NOTE: skip_none=False here — route_event converts None to ok=True, |
| 951 | # unlike process_client_event which skips None (fire-and-forget). |
| 952 | # This is intentional: route_event is server-initiated and callers |
| 953 | # expect a result entry for every handler. |
| 954 | results = self._collect_results( |
| 955 | executions, event_type, correlation_id, skip_none=False, |
| 956 | ) |
| 957 | |
| 958 | await self._publish_diagnostic_event( |
| 959 | lambda: { |
| 960 | "kind": "inbound", |
| 961 | "sourceNamespace": namespace, |
| 962 | "namespace": namespace, |
| 963 | "eventType": event_type, |
| 964 | "sid": sid, |
| 965 | "correlationId": correlation_id, |
| 966 | "timestamp": self._timestamp(), |
| 967 | "handlerCount": len(selected_handlers), |
| 968 | "durationMs": sum((exec.duration_ms or 0.0) for exec in executions), |
| 969 | "resultSummary": self._summarize_results(results), |
| 970 | "payloadSummary": self._summarize_payload(handler_payload), |
| 971 | } |
| 972 | ) |
| 973 | |
| 974 | response_payload = {"correlationId": correlation_id, "results": results} |
| 975 | if ack: |
| 976 | ack(response_payload) |
| 977 | self._debug( |
| 978 | f"Completed event namespace={namespace} '{event_type}' sid={sid} correlation={correlation_id}" |
| 979 | ) |
| 980 | return response_payload |
| 981 | |
| 982 | async def request_for_sid( |
| 983 | self, |
| 984 | *, |
| 985 | namespace: str, |
| 986 | sid: str, |
| 987 | event_type: str, |
| 988 | data: dict[str, Any], |
| 989 | timeout_ms: int = 0, |
| 990 | handler_id: str | None = None, |
| 991 | include_handlers: Set[str] | None = None, |
| 992 | ) -> dict[str, Any]: |
| 993 | payload = dict(data or {}) |
| 994 | correlation_id = self._resolve_correlation_id(payload) |
| 995 | |
| 996 | with self.lock: |
| 997 | connected = (namespace, sid) in self.connections |
| 998 | if not connected: |
| 999 | return { |
| 1000 | "correlationId": correlation_id, |
| 1001 | "results": [ |
| 1002 | self._build_error_result( |
| 1003 | handler_id=handler_id or self._identifier, |
| 1004 | code=ERR_CONNECTION_NOT_FOUND, |
| 1005 | message=f"Connection '{sid}' not found in namespace '{namespace}'", |
| 1006 | correlation_id=correlation_id, |
| 1007 | ) |
| 1008 | ], |
| 1009 | } |
| 1010 | |
| 1011 | async def _invoke() -> dict[str, Any]: |
| 1012 | return await self.route_event( |
| 1013 | namespace, |
| 1014 | event_type, |
| 1015 | payload, |
| 1016 | sid, |
| 1017 | include_handlers=include_handlers, |
| 1018 | handler_id=handler_id, |
| 1019 | ) |
| 1020 | |
| 1021 | if timeout_ms and timeout_ms > 0: |
| 1022 | try: |
| 1023 | return await asyncio.wait_for(_invoke(), timeout=timeout_ms / 1000) |
| 1024 | except asyncio.TimeoutError: |
| 1025 | PrintStyle.warning( |
| 1026 | f"request timeout for sid {sid} event '{event_type}'" |
| 1027 | ) |
| 1028 | return { |
| 1029 | "correlationId": correlation_id, |
| 1030 | "results": [ |
| 1031 | self._build_error_result( |
| 1032 | handler_id=handler_id or self._identifier, |
| 1033 | code=ERR_TIMEOUT, |
| 1034 | message="Request timeout", |
| 1035 | correlation_id=correlation_id, |
| 1036 | ) |
| 1037 | ], |
| 1038 | } |
| 1039 | return await _invoke() |
| 1040 | |
| 1041 | async def route_event_all( |
| 1042 | self, |
| 1043 | namespace: str, |
| 1044 | event_type: str, |
| 1045 | data: dict[str, Any], |
| 1046 | *, |
| 1047 | timeout_ms: int = 0, |
| 1048 | exclude_handlers: Set[str] | None = None, |
| 1049 | handler_id: str | None = None, |
| 1050 | ) -> list[dict[str, Any]]: |
| 1051 | """Fan-out a request to all active connections and aggregate responses.""" |
| 1052 | |
| 1053 | base_payload = dict(data or {}) |
| 1054 | exclude_meta_raw = base_payload.pop("excludeHandlers", None) |
| 1055 | exclude_combined: Set[str] | None = exclude_handlers |
| 1056 | correlation_id = self._resolve_correlation_id(base_payload) |
| 1057 | |
| 1058 | if exclude_meta_raw is not None: |
| 1059 | try: |
| 1060 | exclude_meta = self._normalize_handler_filter( |
| 1061 | exclude_meta_raw, "excludeHandlers" |
| 1062 | ) |
| 1063 | except ValueError as exc: |
| 1064 | error = self._build_error_result( |
| 1065 | handler_id=handler_id or self._identifier, |
| 1066 | code=ERR_INVALID_FILTER, |
| 1067 | message=str(exc), |
| 1068 | correlation_id=correlation_id, |
| 1069 | ) |
| 1070 | return [ |
| 1071 | { |
| 1072 | "sid": "__invalid__", |
| 1073 | "correlationId": correlation_id, |
| 1074 | "results": [error], |
| 1075 | } |
| 1076 | ] |
| 1077 | |
| 1078 | if exclude_combined is None: |
| 1079 | exclude_combined = exclude_meta |
| 1080 | elif exclude_meta is not None and exclude_combined != exclude_meta: |
| 1081 | error = self._build_error_result( |
| 1082 | handler_id=handler_id or self._identifier, |
| 1083 | code=ERR_INVALID_FILTER, |
| 1084 | message="Conflicting excludeHandlers filters supplied", |
| 1085 | correlation_id=correlation_id, |
| 1086 | ) |
| 1087 | return [ |
| 1088 | { |
| 1089 | "sid": "__invalid__", |
| 1090 | "correlationId": correlation_id, |
| 1091 | "results": [error], |
| 1092 | } |
| 1093 | ] |
| 1094 | |
| 1095 | self._debug( |
| 1096 | f"Starting requestAll namespace={namespace} for '{event_type}' correlation={correlation_id}" |
| 1097 | ) |
| 1098 | |
| 1099 | with self.lock: |
| 1100 | active_sids = [ |
| 1101 | conn_identity[1] |
| 1102 | for conn_identity in self.connections.keys() |
| 1103 | if conn_identity[0] == namespace |
| 1104 | ] |
| 1105 | if not active_sids: |
| 1106 | self._debug( |
| 1107 | f"No active connections for requestAll namespace={namespace} '{event_type}' correlation={correlation_id}" |
| 1108 | ) |
| 1109 | return [] |
| 1110 | |
| 1111 | timeout_seconds = timeout_ms / 1000 if timeout_ms and timeout_ms > 0 else None |
| 1112 | |
| 1113 | async def _invoke_for_sid(target_sid: str) -> dict[str, Any]: |
| 1114 | async def _dispatch() -> dict[str, Any]: |
| 1115 | return await self.route_event( |
| 1116 | namespace, |
| 1117 | event_type, |
| 1118 | base_payload, |
| 1119 | target_sid, |
| 1120 | allow_exclude=True, |
| 1121 | exclude_handlers=exclude_combined, |
| 1122 | handler_id=handler_id, |
| 1123 | ) |
| 1124 | |
| 1125 | if timeout_seconds is None: |
| 1126 | return await _dispatch() |
| 1127 | |
| 1128 | try: |
| 1129 | task = asyncio.create_task(_dispatch()) |
| 1130 | return await asyncio.wait_for( |
| 1131 | asyncio.shield(task), timeout=timeout_seconds |
| 1132 | ) |
| 1133 | except asyncio.TimeoutError: |
| 1134 | PrintStyle.warning( |
| 1135 | f"requestAll timeout for sid {target_sid} correlation={correlation_id}" |
| 1136 | ) |
| 1137 | # Ensure any late exceptions are observed so asyncio does not log |
| 1138 | # "Task exception was never retrieved". |
| 1139 | try: |
| 1140 | task.add_done_callback(lambda t: t.exception()) # type: ignore[arg-type] |
| 1141 | except Exception: # pragma: no cover - defensive |
| 1142 | pass |
| 1143 | return { |
| 1144 | "correlationId": correlation_id, |
| 1145 | "results": [ |
| 1146 | self._build_error_result( |
| 1147 | handler_id=handler_id or self._identifier, |
| 1148 | code=ERR_TIMEOUT, |
| 1149 | message="Request timeout", |
| 1150 | correlation_id=correlation_id, |
| 1151 | ) |
| 1152 | ], |
| 1153 | } |
| 1154 | |
| 1155 | tasks = {sid: asyncio.create_task(_invoke_for_sid(sid)) for sid in active_sids} |
| 1156 | |
| 1157 | aggregated: list[dict[str, Any]] = [] |
| 1158 | for sid, task in tasks.items(): |
| 1159 | result = await task |
| 1160 | if isinstance(result, dict): |
| 1161 | aggregated.append( |
| 1162 | { |
| 1163 | "sid": sid, |
| 1164 | "correlationId": result.get("correlationId", correlation_id), |
| 1165 | "results": result.get("results", []), |
| 1166 | } |
| 1167 | ) |
| 1168 | else: |
| 1169 | aggregated.append( |
| 1170 | { |
| 1171 | "sid": sid, |
| 1172 | "correlationId": correlation_id, |
| 1173 | "results": result, |
| 1174 | } |
| 1175 | ) |
| 1176 | |
| 1177 | self._debug( |
| 1178 | f"Completed requestAll namespace={namespace} for '{event_type}' correlation={correlation_id}" |
| 1179 | ) |
| 1180 | return aggregated |
| 1181 | |
| 1182 | def _wrap_envelope( |
| 1183 | self, |
| 1184 | handler_id: str | None, |
| 1185 | data: dict[str, Any], |
| 1186 | *, |
| 1187 | correlation_id: str | None = None, |
| 1188 | ) -> dict[str, Any]: |
| 1189 | hid = handler_id or self._identifier |
| 1190 | ts = self._timestamp() |
| 1191 | event_id = str(uuid.uuid4()) |
| 1192 | correlation = correlation_id or str(uuid.uuid4()) |
| 1193 | return { |
| 1194 | "handlerId": hid, |
| 1195 | "eventId": event_id, |
| 1196 | "correlationId": correlation, |
| 1197 | "ts": ts, |
| 1198 | "data": data or {}, |
| 1199 | } |
| 1200 | |
| 1201 | async def emit_to( |
| 1202 | self, |
| 1203 | namespace: str, |
| 1204 | sid: str, |
| 1205 | event_type: str, |
| 1206 | data: dict[str, Any], |
| 1207 | *, |
| 1208 | handler_id: str | None = None, |
| 1209 | correlation_id: str | None = None, |
| 1210 | diagnostic: bool = False, |
| 1211 | ) -> None: |
| 1212 | envelope = self._wrap_envelope( |
| 1213 | handler_id, |
| 1214 | data, |
| 1215 | correlation_id=correlation_id, |
| 1216 | ) |
| 1217 | delivered = False |
| 1218 | buffered = False |
| 1219 | identity: ConnectionIdentity = (namespace, sid) |
| 1220 | |
| 1221 | with self.lock: |
| 1222 | connected = identity in self.connections |
| 1223 | known = identity in self._known_sids or identity in self.buffers |
| 1224 | # Evict if disconnect has exceeded BUFFER_TTL |
| 1225 | if not connected and known: |
| 1226 | dt = self._disconnect_times.get(identity) |
| 1227 | if dt is not None and (_utcnow() - dt) > BUFFER_TTL: |
| 1228 | self._known_sids.discard(identity) |
| 1229 | self._disconnect_times.pop(identity, None) |
| 1230 | self.buffers.pop(identity, None) |
| 1231 | known = False |
| 1232 | |
| 1233 | if connected: |
| 1234 | self._debug( |
| 1235 | "Emit to namespace=%s sid=%s event=%s eventId=%s correlationId=%s handlerId=%s" |
| 1236 | % ( |
| 1237 | namespace, |
| 1238 | sid, |
| 1239 | event_type, |
| 1240 | envelope.get("eventId"), |
| 1241 | envelope.get("correlationId"), |
| 1242 | envelope.get("handlerId"), |
| 1243 | ) |
| 1244 | ) |
| 1245 | await self._run_on_dispatcher_loop( |
| 1246 | self.socketio.emit(event_type, envelope, to=sid, namespace=namespace) |
| 1247 | ) |
| 1248 | delivered = True |
| 1249 | else: |
| 1250 | if not known: |
| 1251 | raise ConnectionNotFoundError(sid, namespace=namespace) |
| 1252 | with self.lock: |
| 1253 | self._buffer_event( |
| 1254 | identity, |
| 1255 | event_type, |
| 1256 | data, |
| 1257 | handler_id, |
| 1258 | envelope["correlationId"], |
| 1259 | ) |
| 1260 | buffered = True |
| 1261 | |
| 1262 | if not diagnostic: |
| 1263 | await self._publish_diagnostic_event( |
| 1264 | lambda: { |
| 1265 | "kind": "outbound", |
| 1266 | "direction": "emit_to", |
| 1267 | "eventType": event_type, |
| 1268 | "namespace": namespace, |
| 1269 | "sid": sid, |
| 1270 | "correlationId": envelope["correlationId"], |
| 1271 | "handlerId": envelope["handlerId"], |
| 1272 | "timestamp": self._timestamp(), |
| 1273 | "delivered": delivered, |
| 1274 | "buffered": buffered, |
| 1275 | "payloadSummary": self._summarize_payload(data), |
| 1276 | } |
| 1277 | ) |
| 1278 | |
| 1279 | async def send_data( |
| 1280 | self, |
| 1281 | endpoint_name: str, |
| 1282 | event_type: str, |
| 1283 | data: dict[str, Any], |
| 1284 | connection_id: str | None = None, |
| 1285 | ) -> None: |
| 1286 | if connection_id is not None: |
| 1287 | await self.emit_to(endpoint_name, connection_id, event_type, data) |
| 1288 | return |
| 1289 | await self.broadcast(endpoint_name, event_type, data) |
| 1290 | |
| 1291 | async def broadcast( |
| 1292 | self, |
| 1293 | namespace: str, |
| 1294 | event_type: str, |
| 1295 | data: dict[str, Any], |
| 1296 | *, |
| 1297 | exclude_sids: str | Iterable[str] | None = None, |
| 1298 | handler_id: str | None = None, |
| 1299 | correlation_id: str | None = None, |
| 1300 | diagnostic: bool = False, |
| 1301 | ) -> None: |
| 1302 | excluded = self._normalize_sid_filter(exclude_sids) |
| 1303 | |
| 1304 | targets: list[str] = [] |
| 1305 | with self.lock: |
| 1306 | current_identities = list(self.connections.keys()) |
| 1307 | for conn_identity in current_identities: |
| 1308 | if conn_identity[0] != namespace: |
| 1309 | continue |
| 1310 | sid = conn_identity[1] |
| 1311 | if sid in excluded: |
| 1312 | continue |
| 1313 | targets.append(sid) |
| 1314 | |
| 1315 | if targets: |
| 1316 | envelope = self._wrap_envelope( |
| 1317 | handler_id, |
| 1318 | data, |
| 1319 | correlation_id=correlation_id, |
| 1320 | ) |
| 1321 | coros = [ |
| 1322 | self._run_on_dispatcher_loop( |
| 1323 | self.socketio.emit(event_type, envelope, to=sid, namespace=namespace) |
| 1324 | ) |
| 1325 | for sid in targets |
| 1326 | ] |
| 1327 | await asyncio.gather(*coros) |
| 1328 | |
| 1329 | if not diagnostic: |
| 1330 | await self._publish_diagnostic_event( |
| 1331 | lambda: { |
| 1332 | "kind": "outbound", |
| 1333 | "direction": "broadcast", |
| 1334 | "eventType": event_type, |
| 1335 | "namespace": namespace, |
| 1336 | "targets": targets[:10], |
| 1337 | "targetCount": len(targets), |
| 1338 | "correlationId": correlation_id, |
| 1339 | "handlerId": handler_id or self._identifier, |
| 1340 | "timestamp": self._timestamp(), |
| 1341 | "payloadSummary": self._summarize_payload(data), |
| 1342 | } |
| 1343 | ) |
| 1344 | |
| 1345 | async def _run_lifecycle( |
| 1346 | self, namespace: str, fn: Callable[[WsHandler], Any] |
| 1347 | ) -> None: |
| 1348 | seen: Set[WsHandler] = set() |
| 1349 | coros: list[Any] = [] |
| 1350 | for handler in self.handlers.get(namespace, []): |
| 1351 | if handler in seen: |
| 1352 | continue |
| 1353 | seen.add(handler) |
| 1354 | coros.append(self._get_handler_worker().execute_inside(fn, handler)) |
| 1355 | if coros: |
| 1356 | await asyncio.gather(*coros, return_exceptions=True) |
| 1357 | |
| 1358 | def _buffer_event( |
| 1359 | self, |
| 1360 | identity: ConnectionIdentity, |
| 1361 | event_type: str, |
| 1362 | data: dict[str, Any], |
| 1363 | handler_id: str | None, |
| 1364 | correlation_id: str | None, |
| 1365 | ) -> None: |
| 1366 | namespace, sid = identity |
| 1367 | buffer = self.buffers[identity] |
| 1368 | buffer.append( |
| 1369 | BufferedEvent( |
| 1370 | event_type=event_type, |
| 1371 | data=data, |
| 1372 | handler_id=handler_id, |
| 1373 | correlation_id=correlation_id, |
| 1374 | ) |
| 1375 | ) |
| 1376 | while len(buffer) > BUFFER_MAX_SIZE: |
| 1377 | dropped = buffer.popleft() |
| 1378 | PrintStyle.warning( |
| 1379 | f"Dropping buffered event '{dropped.event_type}' for namespace={namespace} sid={sid} (overflow)" |
| 1380 | ) |
| 1381 | self._debug( |
| 1382 | f"Buffered event namespace={namespace} '{event_type}' sid={sid} (queue length={len(buffer)})" |
| 1383 | ) |
| 1384 | |
| 1385 | async def _flush_buffer(self, identity: ConnectionIdentity) -> None: |
| 1386 | self._ensure_dispatcher_loop() |
| 1387 | buffer = self.buffers.get(identity) |
| 1388 | if not buffer: |
| 1389 | return |
| 1390 | namespace, sid = identity |
| 1391 | now = _utcnow() |
| 1392 | delivered = 0 |
| 1393 | while buffer: |
| 1394 | event = buffer.popleft() |
| 1395 | if now - event.timestamp > BUFFER_TTL: |
| 1396 | self._debug( |
| 1397 | f"Discarding expired buffered event '{event.event_type}' for sid {sid}" |
| 1398 | ) |
| 1399 | continue |
| 1400 | envelope = self._wrap_envelope( |
| 1401 | event.handler_id, |
| 1402 | event.data, |
| 1403 | correlation_id=event.correlation_id, |
| 1404 | ) |
| 1405 | self._debug( |
| 1406 | "Flush to sid=%s event=%s eventId=%s correlationId=%s handlerId=%s" |
| 1407 | % ( |
| 1408 | sid, |
| 1409 | event.event_type, |
| 1410 | envelope.get("eventId"), |
| 1411 | envelope.get("correlationId"), |
| 1412 | envelope.get("handlerId"), |
| 1413 | ) |
| 1414 | ) |
| 1415 | await self._run_on_dispatcher_loop( |
| 1416 | self.socketio.emit( |
| 1417 | event.event_type, envelope, to=sid, namespace=namespace |
| 1418 | ) |
| 1419 | ) |
| 1420 | delivered += 1 |
| 1421 | if identity in self.buffers: |
| 1422 | self.buffers.pop(identity, None) |
| 1423 | if delivered: |
| 1424 | PrintStyle.info( |
| 1425 | f"Flushed {delivered} buffered event(s) to namespace={namespace} sid={sid}" |
| 1426 | ) |
| 1427 | |
| 1428 | def _build_error_result( |
| 1429 | self, |
| 1430 | *, |
| 1431 | handler_id: str | None = None, |
| 1432 | code: str, |
| 1433 | message: str, |
| 1434 | details: str | None = None, |
| 1435 | correlation_id: str | None = None, |
| 1436 | duration_ms: float | None = None, |
| 1437 | ) -> dict[str, Any]: |
| 1438 | error_payload = {"code": code, "error": message} |
| 1439 | if details: |
| 1440 | error_payload["details"] = details |
| 1441 | result: dict[str, Any] = { |
| 1442 | "handlerId": handler_id or self._identifier, |
| 1443 | "ok": False, |
| 1444 | "error": error_payload, |
| 1445 | } |
| 1446 | if correlation_id is not None: |
| 1447 | result["correlationId"] = correlation_id |
| 1448 | if duration_ms is not None: |
| 1449 | result["durationMs"] = round(duration_ms, 4) |
| 1450 | return result |
| 1451 | |
| 1452 | def _ack_error( |
| 1453 | self, |
| 1454 | *, |
| 1455 | handler_id: str | None = None, |
| 1456 | code: str, |
| 1457 | message: str, |
| 1458 | correlation_id: str | None = None, |
| 1459 | ack: Callable | None = None, |
| 1460 | ) -> dict[str, Any]: |
| 1461 | """Build an error response, optionally invoke the ack callback, and return.""" |
| 1462 | error = self._build_error_result( |
| 1463 | handler_id=handler_id, |
| 1464 | code=code, |
| 1465 | message=message, |
| 1466 | correlation_id=correlation_id, |
| 1467 | ) |
| 1468 | response = {"correlationId": correlation_id, "results": [error]} |
| 1469 | if ack: |
| 1470 | ack(response) |
| 1471 | return response |
| 1472 | |
| 1473 | # Session tracking helpers (single-user defaults) |
| 1474 | def get_sids_for_user(self, user: str | None = None) -> list[ConnectionIdentity]: |
| 1475 | """Return connection identities for a user; single-user default returns all.""" |
| 1476 | with self.lock: |
| 1477 | bucket = self._ALL_USERS_BUCKET if user is None else user |
| 1478 | return list(self.user_to_sids.get(bucket, set())) |
| 1479 | |
| 1480 | def get_user_for_sid(self, namespace: str, sid: str) -> str | None: |
| 1481 | """Return user identifier for a connection or None.""" |
| 1482 | identity: ConnectionIdentity = (namespace, sid) |
| 1483 | with self.lock: |
| 1484 | return self.sid_to_user.get(identity) |
| 1485 | |
| 1486 | def set_server_restart_broadcast(self, enabled: bool) -> None: |
| 1487 | """Enable or disable automatic server restart broadcasts.""" |
| 1488 | |
| 1489 | self._server_restart_enabled = bool(enabled) |