@samitouri / QOSami_Agents_Zero / commits / b351de45

fix: resolve option whitelist, memory leak, task tracking, and dispatch unification

- Fix Memory Leaks: Resolved SID retention in _known_sids after disconnection and cleaned up unreferenced broadcast tasks in _schedule_lifecycle_broadcast. - Unify Dispatching Paths: Unified client and server event dispatching through the process_client_event() method to ensure diagnostic consistency. - Optimization & Cleanup: Expanded the _OPTION_KEYS whitelist, removed dead code (iter_event_types), and deleted unused websocket exports. - Robustness: Added handling for None responses in process_client_event to prevent cluttering responses with empty results. - Testing: Added test cases to verify SID TTL expiration and stale SID cleanup on disconnect.

keyboardstaff committed Mar 27, 2026 at 01:21 UTC b351de456e20e0c9b2de553a2d44c60208347598
4 files changed +273 -32
helpers/ws.py
+66 -24
@@ -300,37 +300,56 @@ class WsHandler:
300 for sid, handlers in _active_handlers.items()
301 }
302
303 + mgr = self._manager
304 aggregated: list[dict[str, Any]] = []
305 for sid, handlers in snapshot.items():
306 ctx = _ws_contexts.get(sid)
306 - sid_results: list[dict[str, Any]] = []
307 + security_errors: list[dict[str, Any]] = []
308 + passing: list[WsHandler] = []
309 for _path, instance in handlers.items():
310 if ctx is not None:
311 error = _check_security(type(instance), ctx)
312 if error is not None:
311 - sid_results.append({
313 + security_errors.append({
314 "handlerId": instance.identifier,
315 "ok": False,
316 "correlationId": cid,
317 "error": error,
318 })
319 continue
318 - try:
319 - result = await instance.process(event, dict(data, correlationId=cid), sid)
320 - if result is not None:
320 + passing.append(instance)
321 +
322 + if mgr is not None and passing:
323 + result = await mgr.process_client_event(
324 + self._namespace, event,
325 + dict(data, correlationId=cid), sid,
326 + handlers=passing,
327 + )
328 + sid_results = security_errors + result.get("results", [])
329 + else:
330 + # Fallback: inline processing
331 + sid_results = list(security_errors)
332 + for _path, instance in handlers.items():
333 + if instance not in passing:
334 + continue
335 + try:
336 + result = await instance.process(
337 + event, dict(data, correlationId=cid), sid,
338 + )
339 + if result is not None:
340 + sid_results.append({
341 + "handlerId": instance.identifier,
342 + "ok": True,
343 + "correlationId": cid,
344 + "data": result,
345 + })
346 + except Exception as e:
347 sid_results.append({
348 "handlerId": instance.identifier,
323 - "ok": True,
349 + "ok": False,
350 "correlationId": cid,
325 - "data": result,
351 + "error": {"code": "HANDLER_ERROR", "error": str(e)},
352 })
327 - except Exception as e:
328 - sid_results.append({
329 - "handlerId": instance.identifier,
330 - "ok": False,
331 - "correlationId": cid,
332 - "error": {"code": "HANDLER_ERROR", "error": str(e)},
333 - })
353 aggregated.append({
354 "sid": sid,
355 "correlationId": cid,
@@ -530,24 +549,47 @@ def register_ws_namespace(
549 return _error_response("NO_HANDLERS",
550 "No handlers activated", correlation_id)
551
533 - # Unwrap nested payload (mirrors WsManager.route_event):
534 - # frontend sends {ts, data: {actual fields...}, correlationId}
535 - if "data" in incoming and isinstance(incoming.get("data"), dict):
536 - handler_payload = dict(incoming["data"])
537 - else:
538 - handler_payload = dict(incoming)
539 - handler_payload["correlationId"] = correlation_id
540 -
541 - results: list[dict[str, Any]] = []
552 + # Pre-filter handlers through security checks
553 + passing_handlers: list[WsHandler] = []
554 + security_errors: list[dict[str, Any]] = []
555 for path, instance in activated.items():
556 error = _check_security(type(instance), ctx)
557 if error is not None:
545 - results.append({
558 + security_errors.append({
559 "handlerId": instance.identifier,
560 "ok": False,
561 "correlationId": correlation_id,
562 "error": error,
563 })
564 + else:
565 + passing_handlers.append(instance)
566 +
567 + # Delegate to WsManager for unified processing pipeline
568 + # (worker thread isolation, diagnostic events, WsResult support)
569 + if manager is not None and passing_handlers:
570 + result = await manager.process_client_event(
571 + NAMESPACE, event, incoming, sid,
572 + handlers=passing_handlers,
573 + )
574 + if security_errors:
575 + result["results"] = security_errors + result.get("results", [])
576 + return result
577 +
578 + # All handlers failed security or no manager — return collected errors
579 + if not passing_handlers:
580 + return {"correlationId": correlation_id, "results": security_errors}
581 +
582 + # Fallback: inline processing (no manager — should not happen in practice)
583 + handler_payload: dict[str, Any]
584 + if "data" in incoming and isinstance(incoming.get("data"), dict):
585 + handler_payload = dict(incoming["data"])
586 + else:
587 + handler_payload = dict(incoming)
588 + handler_payload["correlationId"] = correlation_id
589 +
590 + results: list[dict[str, Any]] = list(security_errors)
591 + for path, instance in activated.items():
592 + if instance not in passing_handlers:
593 continue
594 try:
595 result = await instance.process(event, handler_payload, sid)
helpers/ws_manager.py
+156 -5
@@ -247,6 +247,7 @@ class WsManager:
247 defaultdict(deque)
248 )
249 self._known_sids: Set[ConnectionIdentity] = set()
250 + self._disconnect_times: Dict[ConnectionIdentity, datetime] = {}
251 self._identifier: str = f"{self.__class__.__module__}.{self.__class__.__name__}"
252 # Session tracking (single-user default)
253 self.user_to_sids: defaultdict[str, Set[ConnectionIdentity]] = defaultdict(set)
@@ -257,6 +258,7 @@ class WsManager:
258 self._diagnostics_enabled: bool = runtime.is_development()
259 self._dispatcher_loop: asyncio.AbstractEventLoop | None = None
260 self._handler_worker: DeferredTask | None = None
261 + self._lifecycle_tasks: Set[asyncio.Task] = set()
262
263 # Internal: development-only debug logging to avoid noise in production
264 def _debug(self, message: str) -> None:
@@ -414,7 +416,23 @@ class WsManager:
416 except Exception as exc: # pragma: no cover - diagnostic
417 self._debug(f"Failed to broadcast lifecycle event {event_type}: {exc}")
418
417 - asyncio.create_task(_broadcast())
419 + task = asyncio.create_task(_broadcast())
420 + self._lifecycle_tasks.add(task)
421 + task.add_done_callback(self._lifecycle_tasks.discard)
422 +
423 + def _sweep_stale_sids(self) -> None:
424 + """Remove _known_sids entries whose disconnect exceeds BUFFER_TTL."""
425 + now = _utcnow()
426 + with self.lock:
427 + stale = [
428 + identity
429 + for identity, dt in self._disconnect_times.items()
430 + if identity not in self.connections and (now - dt) > BUFFER_TTL
431 + ]
432 + for identity in stale:
433 + self._known_sids.discard(identity)
434 + self._disconnect_times.pop(identity, None)
435 + self.buffers.pop(identity, None)
436
437 def _normalize_handler_filter(self, value: Any, field_name: str) -> Set[str] | None:
438 if value is None:
@@ -513,12 +531,133 @@ class WsManager:
531 f"Registered handler {handler.identifier} namespace={namespace}"
532 )
533
516 - def iter_event_types(self, namespace: str) -> Iterable[str]:
517 - return []
518 -
534 def iter_namespaces(self) -> list[str]:
535 return list(self.handlers.keys())
536
537 + async def process_client_event(
538 + self,
539 + namespace: str,
540 + event_type: str,
541 + data: dict[str, Any],
542 + sid: str,
543 + *,
544 + handlers: list[WsHandler],
545 + ) -> dict[str, Any]:
546 + """Process a client-originated event through provided handler instances.
547 +
548 + Unlike ``route_event`` which selects from globally registered handlers,
549 + this accepts pre-selected instances (e.g. per-connection activated
550 + handlers that have already passed security checks).
551 + """
552 + self._ensure_dispatcher_loop()
553 + incoming = dict(data or {})
554 + correlation_id = self._resolve_correlation_id(incoming)
555 +
556 + if "data" in incoming and isinstance(incoming.get("data"), dict):
557 + handler_payload = dict(incoming["data"])
558 + if "excludeSids" in incoming:
559 + handler_payload["excludeSids"] = incoming["excludeSids"]
560 + else:
561 + handler_payload = dict(incoming)
562 + handler_payload["correlationId"] = correlation_id
563 +
564 + if not handlers:
565 + error = self._build_error_result(
566 + handler_id=self._identifier,
567 + code="NO_HANDLERS",
568 + message="No handlers available after security filtering",
569 + correlation_id=correlation_id,
570 + )
571 + return {"correlationId": correlation_id, "results": [error]}
572 +
573 + with self.lock:
574 + info = self.connections.get((namespace, sid))
575 + if info:
576 + info.last_activity = _utcnow()
577 +
578 + executions = await asyncio.gather(
579 + *[
580 + self._invoke_handler(handler, event_type, dict(handler_payload), sid)
581 + for handler in handlers
582 + ]
583 + )
584 +
585 + results: List[dict[str, Any]] = []
586 + for execution in executions:
587 + handler = execution.handler
588 + value = execution.value
589 + duration_ms = execution.duration_ms
590 +
591 + if isinstance(value, Exception):
592 + PrintStyle.error(
593 + f"Error in handler {handler.identifier} for '{event_type}' "
594 + f"(correlation {correlation_id}): {value}"
595 + )
596 + results.append(
597 + self._build_error_result(
598 + handler_id=handler.identifier,
599 + code="HANDLER_ERROR",
600 + message="Internal server error",
601 + details=str(value),
602 + correlation_id=correlation_id,
603 + duration_ms=duration_ms,
604 + )
605 + )
606 + continue
607 +
608 + if isinstance(value, WsResult):
609 + results.append(
610 + value.as_result(
611 + handler_id=handler.identifier,
612 + fallback_correlation_id=correlation_id,
613 + duration_ms=duration_ms,
614 + )
615 + )
616 + continue
617 +
618 + # Skip handlers that return None — they opted out of contributing
619 + # a result (fire-and-forget semantics, matching legacy _dispatch).
620 + if value is None:
621 + continue
622 +
623 + if isinstance(value, dict):
624 + helper_result = WsResult(ok=True, data=value)
625 + else:
626 + helper_result = WsResult(ok=True, data={"result": value})
627 +
628 + results.append(
629 + helper_result.as_result(
630 + handler_id=handler.identifier,
631 + fallback_correlation_id=correlation_id,
632 + duration_ms=duration_ms,
633 + )
634 + )
635 +
636 + await self._publish_diagnostic_event(
637 + lambda: {
638 + "kind": "inbound",
639 + "sourceNamespace": namespace,
640 + "namespace": namespace,
641 + "eventType": event_type,
642 + "sid": sid,
643 + "correlationId": correlation_id,
644 + "timestamp": self._timestamp(),
645 + "handlerCount": len(handlers),
646 + "durationMs": sum(
647 + (exec.duration_ms or 0.0) for exec in executions
648 + ),
649 + "resultSummary": self._summarize_results(results),
650 + "payloadSummary": self._summarize_payload(handler_payload),
651 + }
652 + )
653 +
654 + response = {"correlationId": correlation_id, "results": results}
655 + self._debug(
656 + f"Completed client event namespace={namespace} '{event_type}' "
657 + f"sid={sid} correlation={correlation_id}"
658 + )
659 + return response
660 +
661 async def _invoke_handler(
662 self,
663 handler: WsHandler,
@@ -551,6 +690,7 @@ class WsManager:
690 with self.lock:
691 self.connections[identity] = ConnectionInfo(namespace=namespace, sid=sid)
692 self._known_sids.add(identity)
693 + self._disconnect_times.pop(identity, None)
694 self.sid_to_user[identity] = user_bucket
695 self.user_to_sids[self._ALL_USERS_BUCKET].add(identity)
696 self.user_to_sids[user_bucket].add(identity)
@@ -600,7 +740,9 @@ class WsManager:
740 identity: ConnectionIdentity = (namespace, sid)
741 with self.lock:
742 self.connections.pop(identity, None)
603 - # Keep identity in _known_sids so emit_to buffers instead of raising
743 + # Keep identity in _known_sids so emit_to buffers instead of raising;
744 + # record disconnect time for TTL-based cleanup
745 + self._disconnect_times[identity] = _utcnow()
746 # session tracking cleanup
747 user_bucket = self.sid_to_user.pop(identity, None)
748 if self._ALL_USERS_BUCKET in self.user_to_sids:
@@ -633,6 +775,7 @@ class WsManager:
775 self._schedule_lifecycle_broadcast(
776 namespace, LIFECYCLE_DISCONNECT_EVENT, lifecycle_payload
777 )
778 + self._sweep_stale_sids()
779
780 async def route_event(
781 self,
@@ -1112,6 +1255,14 @@ class WsManager:
1255 with self.lock:
1256 connected = identity in self.connections
1257 known = identity in self._known_sids or identity in self.buffers
1258 + # Evict if disconnect has exceeded BUFFER_TTL
1259 + if not connected and known:
1260 + dt = self._disconnect_times.get(identity)
1261 + if dt is not None and (_utcnow() - dt) > BUFFER_TTL:
1262 + self._known_sids.discard(identity)
1263 + self._disconnect_times.pop(identity, None)
1264 + self.buffers.pop(identity, None)
1265 + known = False
1266
1267 if connected:
1268 self._debug(
tests/test_ws_manager.py
+50
@@ -396,6 +396,56 @@ async def test_flush_buffer_delivers_and_logs(monkeypatch):
396 assert (NAMESPACE, "sid-1") not in manager.buffers
397
398
399 +@pytest.mark.asyncio
400 +async def test_known_sid_expires_after_buffer_ttl(monkeypatch):
401 + """After BUFFER_TTL, a disconnected sid is swept from _known_sids and emit_to raises."""
402 + socketio = FakeSocketIOServer()
403 + manager = WsManager(socketio, threading.RLock())
404 +
405 + await manager.handle_connect(NAMESPACE, "sid-stale")
406 + await manager.handle_disconnect(NAMESPACE, "sid-stale")
407 +
408 + # Immediately after disconnect, buffering still works
409 + await manager.emit_to(NAMESPACE, "sid-stale", "event", {"x": 1})
410 + assert (NAMESPACE, "sid-stale") in manager.buffers
411 +
412 + from datetime import timedelta, timezone, datetime
413 +
414 + future = datetime.now(timezone.utc) + BUFFER_TTL + timedelta(seconds=10)
415 + monkeypatch.setattr("helpers.ws_manager._utcnow", lambda: future)
416 +
417 + # After TTL, emit_to should raise because the sid is no longer known
418 + with pytest.raises(ConnectionNotFoundError):
419 + await manager.emit_to(NAMESPACE, "sid-stale", "event", {"x": 2})
420 +
421 + # _known_sids and buffers should be cleaned
422 + assert (NAMESPACE, "sid-stale") not in manager._known_sids
423 + assert (NAMESPACE, "sid-stale") not in manager.buffers
424 + assert (NAMESPACE, "sid-stale") not in manager._disconnect_times
425 +
426 +
427 +@pytest.mark.asyncio
428 +async def test_sweep_cleans_stale_sids_on_disconnect(monkeypatch):
429 + """_sweep_stale_sids runs during handle_disconnect and cleans expired entries."""
430 + socketio = FakeSocketIOServer()
431 + manager = WsManager(socketio, threading.RLock())
432 +
433 + await manager.handle_connect(NAMESPACE, "old-sid")
434 + await manager.handle_disconnect(NAMESPACE, "old-sid")
435 +
436 + from datetime import timedelta, timezone, datetime
437 +
438 + future = datetime.now(timezone.utc) + BUFFER_TTL + timedelta(seconds=10)
439 + monkeypatch.setattr("helpers.ws_manager._utcnow", lambda: future)
440 +
441 + # A new connect/disconnect triggers sweep which cleans old-sid
442 + await manager.handle_connect(NAMESPACE, "new-sid")
443 + await manager.handle_disconnect(NAMESPACE, "new-sid")
444 +
445 + assert (NAMESPACE, "old-sid") not in manager._known_sids
446 + assert (NAMESPACE, "old-sid") not in manager._disconnect_times
447 +
448 +
449 @pytest.mark.asyncio
450 async def test_broadcast_excludes_multiple_sids():
451 socketio = FakeSocketIOServer()
webui/js/websocket.js
+1 -3
@@ -5,7 +5,7 @@ const MAX_PAYLOAD_BYTES = 50 * 1024 * 1024; // 50MB hard cap per contract
5 const DEFAULT_TIMEOUT_MS = 0;
6
7 const _UUID_HEX = [..."0123456789abcdef"];
8 -const _OPTION_KEYS = new Set(["correlationId"]);
8 +const _OPTION_KEYS = new Set(["correlationId", "includeHandlers", "excludeHandlers", "excludeSids"]);
9
10 /**
11 * @param {unknown} value
@@ -744,5 +744,3 @@ export function getNamespacedClient(namespace) {
744 _namespacedClients.set(key, client);
745 return client;
746 }
747 -
748 -export const websocket = getNamespacedClient("/");