@samitouri / QOSami_Agents_Zero-2 / commits / 1d81f72a

refactor: Backend core rewrite - WsHandler + WsManager + handler migration

- Add WsHandler base class, WsManager (connection tracking / event routing / buffering), WsResult - Extract network.py (is_loopback_address) and context_utils.py (use_context) to eliminate duplication - Migrate three handlers to api/ following the ws_* py naming convention - Simplify run_ui.py WebSocket init from ~170 lines to ~10 - Update import paths in api.py, plugins.py, state_monitor.py

keyboardstaff committed Mar 26, 2026 at 00:58 UTC 1d81f72a31931534d7154a2df242dbac6e439720
11 files changed +1957 -290
api/ws_dev_test.py new
+72
@@ -0,0 +1,72 @@
1 +import asyncio
2 +from typing import Any
3 +
4 +from helpers.ws import WsHandler
5 +from helpers.print_style import PrintStyle
6 +from helpers import runtime
7 +
8 +
9 +class WsDevTest(WsHandler):
10 + """Developer-only WebSocket test harness handler."""
11 +
12 + async def process(self, event: str, data: dict, sid: str) -> dict[str, Any] | None:
13 + if event == "ws_event_console_subscribe":
14 + if not runtime.is_development():
15 + return {"_error": True, "code": "NOT_AVAILABLE",
16 + "message": "Event console is available only in development mode"}
17 + registered = self.manager.register_diagnostic_watcher(self.namespace, sid)
18 + if not registered:
19 + return {"_error": True, "code": "SUBSCRIBE_FAILED",
20 + "message": "Unable to subscribe to diagnostics"}
21 + return {"status": "subscribed", "timestamp": data.get("requestedAt")}
22 +
23 + if event == "ws_event_console_unsubscribe":
24 + self.manager.unregister_diagnostic_watcher(self.namespace, sid)
25 + return {"status": "unsubscribed"}
26 +
27 + if event == "ws_tester_emit":
28 + message = data.get("message", "emit")
29 + payload = {"message": message, "echo": True, "timestamp": data.get("timestamp")}
30 + await self.broadcast("ws_tester_broadcast", payload)
31 + PrintStyle.info(f"Harness emit broadcasted message='{message}'")
32 + return None
33 +
34 + if event == "ws_tester_request":
35 + value = data.get("value")
36 + PrintStyle.debug("Harness request responded with echo %s", value)
37 + return {"echo": value, "handler": self.identifier, "status": "ok"}
38 +
39 + if event == "ws_tester_request_delayed":
40 + delay_ms = int(data.get("delay_ms", 0))
41 + await asyncio.sleep(delay_ms / 1000)
42 + PrintStyle.warning("Harness delayed request finished after %s ms", delay_ms)
43 + return {"status": "delayed", "delay_ms": delay_ms, "handler": self.identifier}
44 +
45 + if event == "ws_tester_trigger_persistence":
46 + phase = data.get("phase", "unknown")
47 + payload = {"phase": phase, "handler": self.identifier}
48 + await self.emit_to(sid, "ws_tester_persistence", payload)
49 + PrintStyle.info(f"Harness persistence event phase='{phase}' -> {sid}")
50 + return None
51 +
52 + if event == "ws_tester_broadcast_demo_trigger":
53 + payload = {"demo": True, "requested_at": data.get("requested_at")}
54 + await self.broadcast("ws_tester_broadcast_demo", payload)
55 + PrintStyle.info("Harness broadcast demo event dispatched")
56 + return None
57 +
58 + if event == "ws_tester_request_all":
59 + correlation_id = data.get("correlationId")
60 + aggregated = await self.dispatch_to_all_sids(
61 + "ws_tester_request",
62 + {"value": data.get("marker", "aggregate")},
63 + correlation_id=correlation_id,
64 + )
65 + return {"results": aggregated}
66 +
67 + # Ignore events not targeted at this handler (other activated handlers
68 + # may process them). Only warn for events that look like dev-harness
69 + # traffic so we don't spam logs with unrelated events.
70 + if event.startswith("ws_tester_"):
71 + PrintStyle.warning(f"Harness received unknown event '{event}'")
72 + return None
api/ws_hello.py new
+13
@@ -0,0 +1,13 @@
1 +from helpers.ws import WsHandler
2 +from helpers.print_style import PrintStyle
3 +
4 +
5 +class WsHello(WsHandler):
6 + """Simple echo handler used for foundational testing."""
7 +
8 + async def process(self, event: str, data: dict, sid: str) -> dict | None:
9 + if event != "hello_request":
10 + return None
11 + name = data.get("name") or "stranger"
12 + PrintStyle.info(f"hello_request from {sid} ({name})")
13 + return {"message": f"Hello, {name}!", "handler": self.identifier}
api/ws_webui.py new
+32
@@ -0,0 +1,32 @@
1 +from helpers.ws import WsHandler
2 +from helpers import extension
3 +
4 +
5 +class WsWebui(WsHandler):
6 + """State synchronisation handler — the primary WebSocket endpoint for the UI."""
7 +
8 + async def on_connect(self, sid: str) -> None:
9 + await extension.call_extensions_async(
10 + "webui_ws_connect", agent=None, instance=self, sid=sid
11 + )
12 +
13 + async def on_disconnect(self, sid: str) -> None:
14 + await extension.call_extensions_async(
15 + "webui_ws_disconnect", agent=None, instance=self, sid=sid
16 + )
17 +
18 + async def process(self, event: str, data: dict, sid: str) -> dict | None:
19 + response_data: dict = {}
20 +
21 + await extension.call_extensions_async(
22 + "webui_ws_event",
23 + agent=None,
24 + instance=self,
25 + sid=sid,
26 + event_type=event,
27 + data=data,
28 + response_data=response_data,
29 + )
30 +
31 + # Return None (fire-and-forget) when no extension populated the response.
32 + return response_data if response_data else None
helpers/api.py
+24 -58
@@ -1,11 +1,9 @@
1 from abc import abstractmethod
2 import json
3 -import socket
4 -import struct
3 import threading
4 from functools import wraps
5 from pathlib import Path
8 -from typing import Union, TypedDict, Dict, Any
6 +from typing import Union, Dict, Any
7 from flask import (
8 Request,
9 Response,
@@ -19,7 +17,6 @@ from flask import (
17 )
18 from werkzeug.wrappers.response import Response as BaseResponse
19 from agent import AgentContext
22 -from initialize import initialize_agent
20 from helpers.print_style import PrintStyle
21 from helpers.errors import format_error
22 from helpers import files, cache
@@ -30,7 +27,7 @@ CACHE_AREA = "api_handlers(api)"
27 cache.toggle_area(CACHE_AREA, False) # cache off for now
28
29 Input = dict
33 -Output = Union[Dict[str, Any], Response, TypedDict] # type: ignore
30 +Output = Union[Dict[str, Any], Response]
31
32
33 class ApiHandler:
@@ -99,59 +96,11 @@ class ApiHandler:
96
97 # get context to run agent zero in
98 def use_context(self, ctxid: str, create_if_not_exists: bool = True):
102 - with self.thread_lock:
103 - if not ctxid:
104 - first = AgentContext.first()
105 - if first:
106 - AgentContext.use(first.id)
107 - return first
108 - context = AgentContext(config=initialize_agent(), set_current=True)
109 - return context
110 - got = AgentContext.use(ctxid)
111 - if got:
112 - return got
113 - if create_if_not_exists:
114 - context = AgentContext(
115 - config=initialize_agent(), id=ctxid, set_current=True
116 - )
117 - return context
118 - else:
119 - raise Exception(f"Context {ctxid} not found")
120 -
121 -
122 -def is_loopback_address(address: str) -> bool:
123 - loopback_checker = {
124 - socket.AF_INET: lambda x: (
125 - struct.unpack("!I", socket.inet_aton(x))[0] >> (32 - 8)
126 - )
127 - == 127,
128 - socket.AF_INET6: lambda x: x == "::1",
129 - }
130 - address_type = "hostname"
131 - try:
132 - socket.inet_pton(socket.AF_INET6, address)
133 - address_type = "ipv6"
134 - except socket.error:
135 - try:
136 - socket.inet_pton(socket.AF_INET, address)
137 - address_type = "ipv4"
138 - except socket.error:
139 - address_type = "hostname"
140 -
141 - if address_type == "ipv4":
142 - return loopback_checker[socket.AF_INET](address)
143 - elif address_type == "ipv6":
144 - return loopback_checker[socket.AF_INET6](address)
145 - else:
146 - for family in (socket.AF_INET, socket.AF_INET6):
147 - try:
148 - r = socket.getaddrinfo(address, None, family, socket.SOCK_STREAM)
149 - except socket.gaierror:
150 - return False
151 - for family, _, _, _, sockaddr in r:
152 - if not loopback_checker[family](sockaddr[0]):
153 - return False
154 - return True
99 + from helpers.context_utils import use_context as _use_context
100 + return _use_context(self.thread_lock, ctxid, create_if_not_exists)
101 +
102 +
103 +from helpers.network import is_loopback_address
104
105
106 def requires_api_key(f):
@@ -301,3 +250,20 @@ def register_watchdogs():
250 patterns=["*.py"],
251 handler=on_api_change,
252 )
253 +
254 + # WS handler cache shares the same watched directories (api/, usr/api/)
255 + from helpers.ws import CACHE_AREA as WS_CACHE_AREA
256 +
257 + def on_ws_change(items: list[watchdog.WatchItem]):
258 + PrintStyle.debug("WS handler watchdog triggered:", items)
259 + cache.clear(WS_CACHE_AREA)
260 +
261 + watchdog.add_watchdog(
262 + "ws_handlers",
263 + roots=[
264 + files.get_abs_path(files.API_DIR),
265 + files.get_abs_path(files.USER_DIR, files.API_DIR),
266 + ],
267 + patterns=["ws_*.py"],
268 + handler=on_ws_change,
269 + )
helpers/context_utils.py new
+30
@@ -0,0 +1,30 @@
1 +"""Shared context helper used by both ApiHandler and WsHandler."""
2 +
3 +import threading
4 +from typing import Union
5 +
6 +ThreadLockType = Union[threading.Lock, threading.RLock]
7 +
8 +
9 +def use_context(lock: ThreadLockType, ctxid: str, create_if_not_exists: bool = True):
10 + from agent import AgentContext
11 + from initialize import initialize_agent
12 +
13 + with lock:
14 + if not ctxid:
15 + first = AgentContext.first()
16 + if first:
17 + AgentContext.use(first.id)
18 + return first
19 + context = AgentContext(config=initialize_agent(), set_current=True)
20 + return context
21 + got = AgentContext.use(ctxid)
22 + if got:
23 + return got
24 + if create_if_not_exists:
25 + context = AgentContext(
26 + config=initialize_agent(), id=ctxid, set_current=True
27 + )
28 + return context
29 + else:
30 + raise Exception(f"Context {ctxid} not found")
helpers/network.py new
+31
@@ -0,0 +1,31 @@
1 +import socket
2 +import struct
3 +
4 +
5 +def is_loopback_address(address: str) -> bool:
6 + """Check whether *address* resolves to a loopback interface."""
7 + _checkers = {
8 + socket.AF_INET: lambda x: (
9 + struct.unpack("!I", socket.inet_aton(x))[0] >> (32 - 8)
10 + ) == 127,
11 + socket.AF_INET6: lambda x: x == "::1",
12 + }
13 + try:
14 + socket.inet_pton(socket.AF_INET6, address)
15 + return _checkers[socket.AF_INET6](address)
16 + except socket.error:
17 + pass
18 + try:
19 + socket.inet_pton(socket.AF_INET, address)
20 + return _checkers[socket.AF_INET](address)
21 + except socket.error:
22 + pass
23 + for family in (socket.AF_INET, socket.AF_INET6):
24 + try:
25 + r = socket.getaddrinfo(address, None, family, socket.SOCK_STREAM)
26 + except socket.gaierror:
27 + return False
28 + for fam, _, _, _, sockaddr in r:
29 + if not _checkers[fam](sockaddr[0]):
30 + return False
31 + return True
helpers/plugins.py
+1 -1
@@ -202,7 +202,7 @@ def clear_plugin_cache(plugin_names: list[str] | None = None):
202 for area in areas:
203 cache.clear(area)
204
205 - from helpers.websocket_manager import send_data
205 + from helpers.ws_manager import send_data
206
207 DeferredTask().start_task(
208 send_data,
helpers/state_monitor.py
+4 -4
@@ -14,10 +14,10 @@ from helpers.state_snapshot import (
14 advance_state_request_after_snapshot,
15 build_snapshot_from_request,
16 )
17 -from helpers.websocket import ConnectionNotFoundError
17 +from helpers.ws import ConnectionNotFoundError
18
19 if TYPE_CHECKING: # pragma: no cover - hints only
20 - from helpers.websocket_manager import WebSocketManager
20 + from helpers.ws_manager import WsManager
21
22
23 ConnectionIdentity = tuple[str, str] # (namespace, sid)
@@ -60,12 +60,12 @@ class StateMonitor:
60 self._projections: dict[ConnectionIdentity, ConnectionProjection] = {}
61 self._debounce_handles: dict[ConnectionIdentity, asyncio.TimerHandle] = {}
62 self._push_tasks: dict[ConnectionIdentity, asyncio.Task[None]] = {}
63 - self._manager: WebSocketManager | None = None
63 + self._manager: WsManager | None = None
64 self._emit_handler_id: str | None = None
65 self._dispatcher_loop: asyncio.AbstractEventLoop | None = None
66 self._dirty_wave_seq: int = 0
67
68 - def bind_manager(self, manager: "WebSocketManager", *, handler_id: str | None = None) -> None:
68 + def bind_manager(self, manager: "WsManager", *, handler_id: str | None = None) -> None:
69 with self._lock:
70 self._manager = manager
71 if handler_id:
helpers/ws.py
+392 -55
@@ -1,22 +1,148 @@
1 import threading
2 +import uuid
3 from abc import abstractmethod
4 from dataclasses import dataclass
5 from pathlib import Path
5 -from typing import Any, Union
6 +from typing import Any, Iterable, Union, TYPE_CHECKING
7 +from urllib.parse import urlparse
8
9 import socketio
10 from flask import Flask, session, request
11
10 -from agent import AgentContext
11 -from initialize import initialize_agent
12 from helpers import files, cache
13 -from helpers.api import is_loopback_address
13 from helpers.print_style import PrintStyle
14 from helpers.errors import format_error
16 -from helpers.websocket import validate_ws_origin
15 +
16 +if TYPE_CHECKING:
17 + from helpers.ws_manager import WsManager
18 +
19 +
20 +# Utilities
21 +
22 +from helpers.network import is_loopback_address
23 +
24 +
25 +class ConnectionNotFoundError(RuntimeError):
26 + """Raised when attempting to emit to a non-existent WebSocket connection."""
27 +
28 + def __init__(self, sid: str, *, namespace: str | None = None) -> None:
29 + self.sid = sid
30 + self.namespace = namespace
31 + if namespace:
32 + super().__init__(f"Connection not found: namespace={namespace} sid={sid}")
33 + else:
34 + super().__init__(f"Connection not found: {sid}")
35 +
36 +
37 +def _default_port_for_scheme(scheme: str) -> int | None:
38 + if scheme == "http":
39 + return 80
40 + if scheme == "https":
41 + return 443
42 + return None
43 +
44 +
45 +def normalize_origin(value: Any) -> str | None:
46 + """Normalize an Origin/Referer header value to scheme://host[:port]."""
47 + if not isinstance(value, str) or not value.strip():
48 + return None
49 + parsed = urlparse(value.strip())
50 + if not parsed.scheme or not parsed.hostname:
51 + return None
52 + origin = f"{parsed.scheme}://{parsed.hostname}"
53 + if parsed.port:
54 + origin += f":{parsed.port}"
55 + return origin
56 +
57 +
58 +def _parse_host_header(value: Any) -> tuple[str | None, int | None]:
59 + if not isinstance(value, str) or not value.strip():
60 + return None, None
61 + parsed = urlparse(f"http://{value.strip()}")
62 + return parsed.hostname, parsed.port
63 +
64 +
65 +def validate_ws_origin(environ: dict[str, Any]) -> tuple[bool, str | None]:
66 + """Validate the browser Origin during the Socket.IO handshake.
67 +
68 + This is the minimum baseline recommended by RFC 6455 (Origin considerations)
69 + and OWASP (CSWSH mitigation): reject cross-origin WebSocket handshakes when
70 + the server is intended for a specific web UI origin.
71 + """
72 +
73 + raw_origin = environ.get("HTTP_ORIGIN") or environ.get("HTTP_REFERER")
74 + origin = normalize_origin(raw_origin)
75 + if origin is None:
76 + return False, "missing_origin"
77 +
78 + origin_parsed = urlparse(origin)
79 + origin_host = origin_parsed.hostname.lower() if origin_parsed.hostname else None
80 + origin_port = origin_parsed.port or _default_port_for_scheme(origin_parsed.scheme)
81 + if origin_host is None or origin_port is None:
82 + return False, "invalid_origin"
83 +
84 + raw_host = environ.get("HTTP_HOST")
85 + req_host, req_port = _parse_host_header(raw_host)
86 + if not req_host:
87 + req_host = environ.get("SERVER_NAME")
88 +
89 + if req_port is None:
90 + server_port_raw = environ.get("SERVER_PORT")
91 + try:
92 + server_port = int(server_port_raw) if server_port_raw is not None else None
93 + except (TypeError, ValueError):
94 + server_port = None
95 + if server_port is not None and server_port > 0:
96 + req_port = server_port
97 +
98 + if req_host:
99 + req_host = req_host.lower()
100 + if req_port is None:
101 + req_port = origin_port
102 +
103 + forwarded_host_raw = environ.get("HTTP_X_FORWARDED_HOST")
104 + forwarded_host = None
105 + forwarded_port = None
106 + if isinstance(forwarded_host_raw, str) and forwarded_host_raw.strip():
107 + first = forwarded_host_raw.split(",")[0].strip()
108 + forwarded_host, forwarded_port = _parse_host_header(first)
109 + if forwarded_host:
110 + forwarded_host = forwarded_host.lower()
111 +
112 + forwarded_proto_raw = environ.get("HTTP_X_FORWARDED_PROTO")
113 + forwarded_scheme = None
114 + if isinstance(forwarded_proto_raw, str) and forwarded_proto_raw.strip():
115 + forwarded_scheme = forwarded_proto_raw.split(",")[0].strip().lower()
116 + forwarded_scheme = forwarded_scheme or origin_parsed.scheme
117 + forwarded_port = (
118 + forwarded_port
119 + if forwarded_port is not None
120 + else _default_port_for_scheme(forwarded_scheme) or origin_port
121 + )
122 +
123 + candidates: list[tuple[str, int]] = []
124 + if req_host:
125 + candidates.append((req_host, int(req_port)))
126 + if forwarded_host:
127 + candidates.append((forwarded_host, int(forwarded_port)))
128 +
129 + if not candidates:
130 + return False, "missing_host"
131 +
132 + for host, port in candidates:
133 + if origin_host == host and origin_port == port:
134 + return True, None
135 +
136 + if origin_host not in {host for host, _ in candidates}:
137 + return False, "origin_host_mismatch"
138 + return False, "origin_port_mismatch"
139 +
140 +
141 +# Constants
142
143 ThreadLockType = Union[threading.Lock, threading.RLock]
144
145 +NAMESPACE = "/ws"
146 CACHE_AREA = "ws_handlers(api)(plugins)"
147 cache.toggle_area(CACHE_AREA, False) # cache off for now
148
@@ -37,9 +163,52 @@ _contexts_lock = threading.Lock()
163
164
165 class WsHandler:
40 - def __init__(self, socketio_server: socketio.AsyncServer, lock: ThreadLockType):
166 + """Base class for WebSocket handlers loaded from api/ directories.
167 +
168 + Mirrors ApiHandler conventions: declarative security flags, dynamic file-
169 + based loading, and a ``process(event, data, sid)`` entry point. Handlers
170 + are activated per-connection based on the ``auth.handlers`` list sent by the
171 + client during the Socket.IO connect handshake.
172 + """
173 +
174 + def __init__(
175 + self,
176 + socketio_server: socketio.AsyncServer,
177 + lock: ThreadLockType,
178 + *,
179 + manager: "WsManager | None" = None,
180 + namespace: str = NAMESPACE,
181 + ):
182 self.socketio = socketio_server
183 self.lock = lock
184 + self._manager = manager
185 + self._namespace = namespace
186 +
187 + # Properties
188 +
189 + @property
190 + def namespace(self) -> str:
191 + return self._namespace
192 +
193 + @property
194 + def manager(self) -> "WsManager":
195 + if self._manager is None:
196 + raise RuntimeError("WsHandler has no WsManager bound")
197 + return self._manager
198 +
199 + @property
200 + def identifier(self) -> str:
201 + return f"{self.__class__.__module__}.{self.__class__.__name__}"
202 +
203 + def bind_manager(
204 + self, manager: "WsManager", *, namespace: str | None = None
205 + ) -> None:
206 + """Late-bind (or rebind) the manager and optionally the namespace."""
207 + self._manager = manager
208 + if namespace is not None:
209 + self._namespace = namespace
210 +
211 + # Security flags (mirror ApiHandler)
212
213 @classmethod
214 def requires_loopback(cls) -> bool:
@@ -57,67 +226,164 @@ class WsHandler:
226 def requires_csrf(cls) -> bool:
227 return cls.requires_auth()
228
60 - @abstractmethod
61 - async def process(self, data: dict, sid: str) -> dict | None:
62 - pass
229 + # Lifecycle hooks
230
64 - async def on_connect(self, sid: str) -> dict | None:
65 - return None
231 + async def on_connect(self, sid: str) -> None:
232 + pass
233
234 async def on_disconnect(self, sid: str) -> None:
235 pass
236
237 + # Event processing
238 +
239 + @abstractmethod
240 + async def process(self, event: str, data: dict, sid: str) -> dict | None:
241 + """Handle an incoming event.
242 +
243 + Return a dict to include in the acknowledgement, or ``None`` for
244 + fire-and-forget semantics.
245 + """
246 +
247 + # Emit helpers (delegate to WsManager for envelope wrapping)
248 +
249 + async def emit_to(
250 + self,
251 + sid: str,
252 + event: str,
253 + data: dict,
254 + *,
255 + correlation_id: str | None = None,
256 + ) -> None:
257 + await self.manager.emit_to(
258 + self._namespace, sid, event, data,
259 + handler_id=self.identifier,
260 + correlation_id=correlation_id,
261 + )
262 +
263 + async def broadcast(
264 + self,
265 + event: str,
266 + data: dict,
267 + *,
268 + exclude_sids: str | Iterable[str] | None = None,
269 + correlation_id: str | None = None,
270 + ) -> None:
271 + await self.manager.broadcast(
272 + self._namespace, event, data,
273 + exclude_sids=exclude_sids,
274 + handler_id=self.identifier,
275 + correlation_id=correlation_id,
276 + )
277 +
278 + # Aggregation helper
279 +
280 + async def dispatch_to_all_sids(
281 + self,
282 + event: str,
283 + data: dict,
284 + *,
285 + correlation_id: str | None = None,
286 + ) -> list[dict[str, Any]]:
287 + """Dispatch *event* to every connected sid's activated handlers and
288 + aggregate the results.
289 +
290 + Returns a list of ``{sid, correlationId, results}`` dicts – one per
291 + connected sid. This mirrors the shape produced by
292 + ``WsManager.route_event_all`` so that existing frontend
293 + assertions remain valid.
294 + """
295 + cid = correlation_id or uuid.uuid4().hex
296 +
297 + with _contexts_lock:
298 + snapshot = {
299 + sid: dict(handlers)
300 + for sid, handlers in _active_handlers.items()
301 + }
302 +
303 + aggregated: list[dict[str, Any]] = []
304 + for sid, handlers in snapshot.items():
305 + ctx = _ws_contexts.get(sid)
306 + sid_results: list[dict[str, Any]] = []
307 + for _path, instance in handlers.items():
308 + if ctx is not None:
309 + error = _check_security(type(instance), ctx)
310 + if error is not None:
311 + sid_results.append({
312 + "handlerId": instance.identifier,
313 + "ok": False,
314 + "correlationId": cid,
315 + "error": error,
316 + })
317 + continue
318 + try:
319 + result = await instance.process(event, dict(data, correlationId=cid), sid)
320 + if result is not None:
321 + sid_results.append({
322 + "handlerId": instance.identifier,
323 + "ok": True,
324 + "correlationId": cid,
325 + "data": result,
326 + })
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 + })
334 + aggregated.append({
335 + "sid": sid,
336 + "correlationId": cid,
337 + "results": sid_results,
338 + })
339 + return aggregated
340 +
341 + # Context helper (shared with ApiHandler)
342 +
343 def use_context(self, ctxid: str, create_if_not_exists: bool = True):
71 - with self.lock:
72 - if not ctxid:
73 - first = AgentContext.first()
74 - if first:
75 - AgentContext.use(first.id)
76 - return first
77 - context = AgentContext(config=initialize_agent(), set_current=True)
78 - return context
79 - got = AgentContext.use(ctxid)
80 - if got:
81 - return got
82 - if create_if_not_exists:
83 - context = AgentContext(config=initialize_agent(), id=ctxid, set_current=True)
84 - return context
85 - else:
86 - raise Exception(f"Context {ctxid} not found")
344 + from helpers.context_utils import use_context as _use_context
345 + return _use_context(self.lock, ctxid, create_if_not_exists)
346 +
347
348 +# Security check (aligned with api.py decorators)
349
350 def _check_security(handler_cls: type[WsHandler], ctx: _SecurityContext) -> dict[str, Any] | None:
351 + """Return an error payload dict if the check fails, or ``None`` on success."""
352 +
353 if handler_cls.requires_loopback():
354 if not ctx.remote_addr or not is_loopback_address(ctx.remote_addr):
92 - return {"ok": False, "error": "Access denied", "code": 403}
355 + return {"code": "FORBIDDEN", "error": "Access denied"}
356
357 if handler_cls.requires_auth():
358 from helpers import login
359 user_pass_hash = login.get_credentials_hash()
360 if user_pass_hash and ctx.auth_hash != user_pass_hash:
98 - return {"ok": False, "error": "Authentication required", "code": 401}
361 + return {"code": "AUTH_REQUIRED", "error": "Authentication required"}
362
363 if handler_cls.requires_csrf():
364 if not ctx.csrf_token:
102 - return {"ok": False, "error": "CSRF token not initialised", "code": 403}
365 + return {"code": "CSRF_MISSING", "error": "CSRF token not initialised"}
366 if not ctx.client_csrf_token or ctx.client_csrf_token != ctx.csrf_token:
104 - return {"ok": False, "error": "CSRF token missing or invalid", "code": 403}
367 + return {"code": "CSRF_INVALID", "error": "CSRF token missing or invalid"}
368 if ctx.csrf_cookie != ctx.csrf_token:
106 - return {"ok": False, "error": "CSRF cookie mismatch", "code": 403}
369 + return {"code": "CSRF_COOKIE", "error": "CSRF cookie mismatch"}
370
371 if handler_cls.requires_api_key():
372 from helpers.settings import get_settings
373 valid_key = get_settings().get("mcp_server_token")
374 if not ctx.api_key or ctx.api_key != valid_key:
112 - return {"ok": False, "error": "API key required", "code": 401}
375 + return {"code": "API_KEY_REQUIRED", "error": "API key required"}
376
377 return None
378
379
380 +# Namespace registration
381 +
382 def register_ws_namespace(
383 socketio_server: socketio.AsyncServer,
384 webapp: Flask,
385 lock: ThreadLockType,
386 + manager: "WsManager | None" = None,
387 ) -> None:
388 from helpers.modules import load_classes_from_file
389 from helpers import plugins, runtime
@@ -132,6 +398,14 @@ def register_ws_namespace(
398 if classes:
399 handler_cls = classes[0]
400
401 + # Check user api/<path>.py
402 + if handler_cls is None:
403 + user_file = files.get_abs_path(files.USER_DIR, f"api/{path}.py")
404 + if files.exists(user_file):
405 + classes = load_classes_from_file(user_file, WsHandler)
406 + if classes:
407 + handler_cls = classes[0]
408 +
409 # Check plugin api/<handler>.py — path format: plugins/<plugin_name>/<handler>
410 if handler_cls is None and path.startswith("plugins/"):
411 parts = path.split("/", 2)
@@ -156,13 +430,13 @@ def register_ws_namespace(
430 cache.add(CACHE_AREA, path, handler_cls)
431 return handler_cls
432
159 - @socketio_server.on("connect", namespace="/ws") # type: ignore
433 + @socketio_server.on("connect", namespace=NAMESPACE) # type: ignore
434 async def _on_connect(sid, environ, auth):
435 with webapp.request_context(environ):
436 origin_ok, origin_reason = validate_ws_origin(environ)
437 if not origin_ok:
438 PrintStyle.warning(
165 - f"WS /ws connect rejected for {sid}: {origin_reason or 'invalid'}"
439 + f"WS connect rejected for {sid}: {origin_reason or 'invalid'}"
440 )
441 return False
442
@@ -182,11 +456,19 @@ def register_ws_namespace(
456 if isinstance(auth, dict) else None
457 ),
458 )
459 + user_id = session.get("user_id") or "single_user"
460 +
461 with _contexts_lock:
462 _ws_contexts[sid] = ctx
463
464 + # Register with WsManager first so that the dispatcher loop and
465 + # connection tracking are available before handler on_connect runs
466 + # (extensions like StateSync depend on manager._dispatcher_loop).
467 + if manager is not None:
468 + await manager.handle_connect(NAMESPACE, sid, user_id=user_id)
469 +
470 # Activate handlers declared in auth.handlers
189 - handler_paths = []
471 + handler_paths: list[str] = []
472 if isinstance(auth, dict):
473 raw = auth.get("handlers")
474 if isinstance(raw, list):
@@ -201,7 +483,10 @@ def register_ws_namespace(
483 error = _check_security(handler_cls, ctx)
484 if error is not None:
485 continue
204 - instance = handler_cls(socketio_server, lock)
486 + instance = handler_cls(
487 + socketio_server, lock,
488 + manager=manager, namespace=NAMESPACE,
489 + )
490 await instance.on_connect(sid)
491 activated[path] = instance
492 except Exception as e:
@@ -212,7 +497,7 @@ def register_ws_namespace(
497
498 return True
499
215 - @socketio_server.on("disconnect", namespace="/ws") # type: ignore
500 + @socketio_server.on("disconnect", namespace=NAMESPACE) # type: ignore
501 async def _on_disconnect(sid):
502 with _contexts_lock:
503 activated = _active_handlers.pop(sid, {})
@@ -224,31 +509,83 @@ def register_ws_namespace(
509 except Exception as e:
510 PrintStyle.error(f"WS on_disconnect error ({path}): {format_error(e)}")
511
227 - @socketio_server.on("*", namespace="/ws") # type: ignore
512 + if manager is not None:
513 + await manager.handle_disconnect(NAMESPACE, sid)
514 +
515 + @socketio_server.on("*", namespace=NAMESPACE) # type: ignore
516 async def _dispatch(event, sid, data):
229 - path = event
230 - payload = data if isinstance(data, dict) else {}
517 + incoming = data if isinstance(data, dict) else {}
518
519 try:
520 with _contexts_lock:
521 ctx = _ws_contexts.get(sid)
235 - activated = _active_handlers.get(sid, {})
236 - if ctx is None:
237 - return {"ok": False, "error": "No security context", "code": 401}
522 + activated = dict(_active_handlers.get(sid, {}))
523
239 - instance = activated.get(path)
240 - if instance is None:
241 - return {"ok": False, "error": f"WS endpoint not activated: {path}", "code": 404}
524 + correlation_id = incoming.get("correlationId") or uuid.uuid4().hex
525
243 - # Security check
244 - error = _check_security(type(instance), ctx)
245 - if error is not None:
246 - return error
526 + if ctx is None:
527 + return _error_response("AUTH_REQUIRED",
528 + "No security context", correlation_id)
529 + if not activated:
530 + return _error_response("NO_HANDLERS",
531 + "No handlers activated", correlation_id)
532 +
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
248 - # Use cached instance and process
249 - return await instance.process(payload, sid)
541 + results: list[dict[str, Any]] = []
542 + for path, instance in activated.items():
543 + error = _check_security(type(instance), ctx)
544 + if error is not None:
545 + results.append({
546 + "handlerId": instance.identifier,
547 + "ok": False,
548 + "correlationId": correlation_id,
549 + "error": error,
550 + })
551 + continue
552 + try:
553 + result = await instance.process(event, handler_payload, sid)
554 + if result is not None:
555 + results.append({
556 + "handlerId": instance.identifier,
557 + "ok": True,
558 + "correlationId": correlation_id,
559 + "data": result,
560 + })
561 + except Exception as e:
562 + error_text = format_error(e)
563 + PrintStyle.error(f"WS handler error ({path}/{event}): {error_text}")
564 + results.append({
565 + "handlerId": instance.identifier,
566 + "ok": False,
567 + "correlationId": correlation_id,
568 + "error": {"code": "HANDLER_ERROR", "error": "Internal server error"},
569 + })
570 +
571 + return {"correlationId": correlation_id, "results": results}
572
573 except Exception as e:
574 error_text = format_error(e)
253 - PrintStyle.error(f"WS handler error ({path}): {error_text}")
254 - return {"ok": False, "error": error_text, "code": 500}
\ No newline at end of file
575 + PrintStyle.error(f"WS dispatch error ({event}): {error_text}")
576 + return _error_response(
577 + "INTERNAL_ERROR", "Internal server error",
578 + incoming.get("correlationId", ""),
579 + )
580 +
581 +
582 +def _error_response(code: str, message: str,
583 + correlation_id: str) -> dict[str, Any]:
584 + return {
585 + "correlationId": correlation_id,
586 + "results": [{
587 + "handlerId": "ws.dispatch",
588 + "ok": False,
589 + "error": {"code": code, "error": message},
590 + }],
591 + }
\ No newline at end of file
helpers/ws_manager.py new
+1351
@@ -0,0 +1,1351 @@
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 ConnectionNotFoundError, WsHandler
19 +
20 +
21 +def _ws_debug_enabled() -> bool:
22 + """Check A0_WS_DEBUG env var — no heavyweight imports needed."""
23 + value = os.getenv("A0_WS_DEBUG", "").strip().lower()
24 + return value in {"1", "true", "yes", "on"}
25 +
26 +
27 +# Event validation
28 +
29 +_EVENT_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$")
30 +_RESERVED_EVENT_NAMES: set[str] = {
31 + "connect",
32 + "disconnect",
33 + "error",
34 + "ping",
35 + "pong",
36 + "connect_error",
37 + "reconnect",
38 + "reconnect_attempt",
39 + "reconnect_error",
40 + "reconnect_failed",
41 +}
42 +
43 +
44 +# WsResult – standardized handler return value
45 +
46 +class WsResult:
47 + """Helper wrapper for standardized handler results.
48 +
49 + Instances are converted to the canonical ``RequestResultItem`` shape by
50 + :class:`WsManager`. Helper constructors enforce payload validation
51 + so handlers no longer need to hand-craft dictionaries.
52 + """
53 +
54 + __slots__ = ("_ok", "_data", "_error", "_correlation_id", "_duration_ms")
55 +
56 + def __init__(
57 + self,
58 + ok: bool,
59 + data: dict[str, Any] | None = None,
60 + error: dict[str, Any] | None = None,
61 + correlation_id: str | None = None,
62 + duration_ms: float | None = None,
63 + ) -> None:
64 + if ok and error:
65 + raise ValueError("Cannot be both ok and have an error")
66 + if not ok and not error:
67 + raise ValueError("Must either be ok or have an error")
68 + if data is not None and not isinstance(data, dict):
69 + raise TypeError("Data payload must be a dictionary or None")
70 + if error is not None and not isinstance(error, dict):
71 + raise TypeError("Error payload must be a dictionary or None")
72 + if correlation_id is not None and not isinstance(correlation_id, str):
73 + raise TypeError("Correlation ID must be a string or None")
74 + if duration_ms is not None and not isinstance(duration_ms, (int, float)):
75 + raise TypeError("Duration must be a number or None")
76 +
77 + self._ok = bool(ok)
78 + self._data = dict(data) if data is not None else None
79 + self._error = dict(error) if error is not None else None
80 + self._correlation_id = correlation_id
81 + self._duration_ms = float(duration_ms) if duration_ms is not None else None
82 +
83 + @classmethod
84 + def ok(
85 + cls,
86 + data: dict[str, Any] | None = None,
87 + *,
88 + correlation_id: str | None = None,
89 + duration_ms: float | None = None,
90 + ) -> "WsResult":
91 + if data is not None and not isinstance(data, dict):
92 + raise TypeError("WsResult.ok data must be a dict or None")
93 + payload = dict(data) if data is not None else None
94 + return cls(
95 + ok=True,
96 + data=payload,
97 + correlation_id=correlation_id,
98 + duration_ms=duration_ms,
99 + )
100 +
101 + @classmethod
102 + def error(
103 + cls,
104 + *,
105 + code: str,
106 + message: str,
107 + details: Any | None = None,
108 + correlation_id: str | None = None,
109 + duration_ms: float | None = None,
110 + ) -> "WsResult":
111 + if not isinstance(code, str) or not code.strip():
112 + raise ValueError("Error code must be a non-empty string")
113 + if not isinstance(message, str) or not message.strip():
114 + raise ValueError("Error message must be a non-empty string")
115 +
116 + error_payload: dict[str, Any] = {"code": code, "error": message}
117 + if details is not None:
118 + error_payload["details"] = details
119 + return cls(
120 + ok=False,
121 + error=error_payload,
122 + correlation_id=correlation_id,
123 + duration_ms=duration_ms,
124 + )
125 +
126 + def as_result(
127 + self,
128 + *,
129 + handler_id: str,
130 + fallback_correlation_id: str | None,
131 + duration_ms: float | None = None,
132 + ) -> dict[str, Any]:
133 + result: dict[str, Any] = {
134 + "handlerId": handler_id,
135 + "ok": self._ok,
136 + }
137 +
138 + effective_duration = (
139 + self._duration_ms if self._duration_ms is not None else duration_ms
140 + )
141 + if effective_duration is not None:
142 + result["durationMs"] = round(effective_duration, 4)
143 +
144 + correlation = (
145 + self._correlation_id
146 + if self._correlation_id is not None
147 + else fallback_correlation_id
148 + )
149 + if correlation is not None:
150 + result["correlationId"] = correlation
151 +
152 + if self._ok:
153 + result["data"] = dict(self._data) if self._data is not None else {}
154 + else:
155 + result["error"] = dict(self._error) if self._error is not None else {
156 + "code": "INTERNAL_ERROR",
157 + "error": "Internal server error",
158 + }
159 + return result
160 +
161 +
162 +def validate_event_type(event_type: str) -> str:
163 + """Validate an event name: must be lowercase_snake_case and not reserved."""
164 + if not isinstance(event_type, str):
165 + raise TypeError("Event type must be a string")
166 + if not _EVENT_NAME_PATTERN.fullmatch(event_type):
167 + raise ValueError(
168 + f"Invalid event type '{event_type}' – must match lowercase_snake_case"
169 + )
170 + if event_type in _RESERVED_EVENT_NAMES:
171 + raise ValueError(
172 + f"Event type '{event_type}' is reserved by Socket.IO and cannot be used"
173 + )
174 + return event_type
175 +
176 +
177 +BUFFER_MAX_SIZE = 100
178 +BUFFER_TTL = timedelta(hours=1)
179 +_shared_ws_manager: WsManager | None = None
180 +
181 +
182 +async def send_data(
183 + event_type: str,
184 + data: dict[str, Any],
185 + endpoint_name: str = "/ws",
186 + connection_id: str | None = None,
187 +) -> None:
188 + manager = get_shared_ws_manager()
189 + await manager.send_data(endpoint_name, event_type, data, connection_id)
190 +
191 +
192 +def _utcnow() -> datetime:
193 + return datetime.now(timezone.utc)
194 +
195 +
196 +def set_shared_ws_manager(manager: "WsManager") -> None:
197 + global _shared_ws_manager
198 + _shared_ws_manager = manager
199 +
200 +
201 +def get_shared_ws_manager() -> "WsManager":
202 + manager = _shared_ws_manager
203 + if manager is None:
204 + raise RuntimeError("Shared WsManager has not been initialized")
205 + return manager
206 +
207 +
208 +@dataclass
209 +class BufferedEvent:
210 + event_type: str
211 + data: dict[str, Any]
212 + handler_id: str | None = None
213 + correlation_id: str | None = None
214 + timestamp: datetime = field(default_factory=_utcnow)
215 +
216 +
217 +@dataclass
218 +class ConnectionInfo:
219 + namespace: str
220 + sid: str
221 + connected_at: datetime = field(default_factory=_utcnow)
222 + last_activity: datetime = field(default_factory=_utcnow)
223 +
224 +
225 +ConnectionIdentity = tuple[str, str] # (namespace, sid)
226 +
227 +
228 +@dataclass
229 +class _HandlerExecution:
230 + handler: WsHandler
231 + value: Any
232 + duration_ms: float | None
233 +
234 +
235 +DIAGNOSTIC_EVENT = "ws_dev_console_event"
236 +LIFECYCLE_CONNECT_EVENT = "ws_lifecycle_connect"
237 +LIFECYCLE_DISCONNECT_EVENT = "ws_lifecycle_disconnect"
238 +
239 +
240 +class WsManager:
241 + def __init__(self, socketio: socketio.AsyncServer, lock) -> None:
242 + self.socketio = socketio
243 + self.lock = lock
244 + self.handlers: defaultdict[str, List[WsHandler]] = defaultdict(list)
245 + self.connections: Dict[ConnectionIdentity, ConnectionInfo] = {}
246 + self.buffers: defaultdict[ConnectionIdentity, Deque[BufferedEvent]] = (
247 + defaultdict(deque)
248 + )
249 + self._known_sids: Set[ConnectionIdentity] = set()
250 + self._identifier: str = f"{self.__class__.__module__}.{self.__class__.__name__}"
251 + # Session tracking (single-user default)
252 + self.user_to_sids: defaultdict[str, Set[ConnectionIdentity]] = defaultdict(set)
253 + self.sid_to_user: Dict[ConnectionIdentity, str | None] = {}
254 + self._ALL_USERS_BUCKET = "allUsers"
255 + self._server_restart_enabled: bool = False
256 + self._diagnostic_watchers: Set[ConnectionIdentity] = set()
257 + self._diagnostics_enabled: bool = runtime.is_development()
258 + self._dispatcher_loop: asyncio.AbstractEventLoop | None = None
259 + self._handler_worker: DeferredTask | None = None
260 +
261 + # Internal: development-only debug logging to avoid noise in production
262 + def _debug(self, message: str) -> None:
263 + if _ws_debug_enabled():
264 + PrintStyle.debug(message)
265 +
266 + def _ensure_dispatcher_loop(self) -> None:
267 + if self._dispatcher_loop is None:
268 + try:
269 + self._dispatcher_loop = asyncio.get_running_loop()
270 + except RuntimeError:
271 + return
272 +
273 + def _get_handler_worker(self) -> DeferredTask:
274 + if self._handler_worker is None:
275 + self._handler_worker = DeferredTask(thread_name="WsHandlers")
276 + return self._handler_worker
277 +
278 + async def _run_on_dispatcher_loop(self, coro: Any) -> Any:
279 + self._ensure_dispatcher_loop()
280 + dispatcher_loop = self._dispatcher_loop
281 + if dispatcher_loop is None:
282 + return await coro
283 + if dispatcher_loop.is_closed():
284 + try:
285 + coro.close()
286 + except Exception: # pragma: no cover - best-effort cleanup
287 + pass
288 + raise RuntimeError("Dispatcher event loop is closed")
289 +
290 + try:
291 + running_loop = asyncio.get_running_loop()
292 + except RuntimeError:
293 + running_loop = None
294 +
295 + if running_loop is dispatcher_loop:
296 + return await coro
297 +
298 + future = asyncio.run_coroutine_threadsafe(coro, dispatcher_loop)
299 + return await asyncio.wrap_future(future)
300 +
301 + def _diagnostics_active(self) -> bool:
302 + if not self._diagnostics_enabled:
303 + return False
304 + with self.lock:
305 + return bool(self._diagnostic_watchers)
306 +
307 + def _copy_diagnostic_watchers(self) -> list[ConnectionIdentity]:
308 + with self.lock:
309 + return list(self._diagnostic_watchers)
310 +
311 + def register_diagnostic_watcher(self, namespace: str, sid: str) -> bool:
312 + if not self._diagnostics_enabled:
313 + return False
314 + identity: ConnectionIdentity = (namespace, sid)
315 + with self.lock:
316 + if identity not in self.connections:
317 + return False
318 + self._diagnostic_watchers.add(identity)
319 + return True
320 +
321 + def unregister_diagnostic_watcher(self, namespace: str, sid: str) -> None:
322 + identity: ConnectionIdentity = (namespace, sid)
323 + with self.lock:
324 + self._diagnostic_watchers.discard(identity)
325 +
326 + def _timestamp(self) -> str:
327 + return _utcnow().isoformat(timespec="milliseconds").replace("+00:00", "Z")
328 +
329 + def _summarize_payload(self, payload: dict[str, Any] | None) -> dict[str, Any]:
330 + if not isinstance(payload, dict):
331 + return {}
332 + summary: dict[str, Any] = {}
333 + for key in list(payload.keys())[:5]:
334 + value = payload[key]
335 + if isinstance(value, (str, int, float, bool)) or value is None:
336 + preview = value
337 + elif isinstance(value, dict):
338 + preview = f"dict({len(value)})"
339 + elif isinstance(value, list):
340 + preview = f"list({len(value)})"
341 + else:
342 + preview = value.__class__.__name__
343 + summary[key] = preview
344 + summary["__sizeBytes__"] = len(str(payload).encode("utf-8"))
345 + return summary
346 +
347 + def _summarize_results(self, results: List[dict[str, Any]]) -> dict[str, Any]:
348 + summary = {"ok": 0, "error": 0, "handlers": []}
349 + for result in results:
350 + handler_id = result.get("handlerId")
351 + ok = bool(result.get("ok"))
352 + if ok:
353 + summary["ok"] += 1
354 + else:
355 + summary["error"] += 1
356 + summary["handlers"].append(
357 + {
358 + "handlerId": handler_id,
359 + "ok": ok,
360 + "errorCode": (result.get("error") or {}).get("code"),
361 + "durationMs": result.get("durationMs"),
362 + }
363 + )
364 + summary["handlerCount"] = len(summary["handlers"])
365 + return summary
366 +
367 + async def _publish_diagnostic_event(
368 + self, payload: dict[str, Any] | Callable[[], dict[str, Any]]
369 + ) -> None:
370 + if not self._diagnostics_enabled:
371 + return
372 + watchers = self._copy_diagnostic_watchers()
373 + if not watchers:
374 + return
375 + effective_payload = payload() if callable(payload) else payload
376 + if (
377 + isinstance(effective_payload, dict)
378 + and "sourceNamespace" not in effective_payload
379 + ):
380 + origin = effective_payload.get("namespace")
381 + if isinstance(origin, str) and origin.strip():
382 + effective_payload = {
383 + **effective_payload,
384 + "sourceNamespace": origin.strip(),
385 + }
386 +
387 + async def _emit_to_watcher(identity: ConnectionIdentity) -> None:
388 + namespace, sid = identity
389 + try:
390 + await self.emit_to(
391 + namespace,
392 + sid,
393 + DIAGNOSTIC_EVENT,
394 + effective_payload,
395 + handler_id=self._identifier,
396 + diagnostic=True,
397 + )
398 + except ConnectionNotFoundError:
399 + self.unregister_diagnostic_watcher(namespace, sid)
400 +
401 + await asyncio.gather(*(_emit_to_watcher(identity) for identity in watchers))
402 +
403 + def _schedule_lifecycle_broadcast(
404 + self, namespace: str, event_type: str, payload: dict[str, Any]
405 + ) -> None:
406 + async def _broadcast() -> None:
407 + try:
408 + await self.broadcast(
409 + namespace,
410 + event_type,
411 + payload,
412 + diagnostic=True,
413 + )
414 + except Exception as exc: # pragma: no cover - diagnostic
415 + self._debug(f"Failed to broadcast lifecycle event {event_type}: {exc}")
416 +
417 + asyncio.create_task(_broadcast())
418 +
419 + def _normalize_handler_filter(self, value: Any, field_name: str) -> Set[str] | None:
420 + if value is None:
421 + return None
422 + if isinstance(value, str):
423 + return {value}
424 + try:
425 + iterator = iter(value)
426 + except TypeError as exc: # pragma: no cover - defensive
427 + raise ValueError(
428 + f"{field_name} must be an array of handler identifiers"
429 + ) from exc
430 +
431 + normalized: Set[str] = set()
432 + for item in iterator:
433 + if not isinstance(item, str):
434 + raise ValueError(
435 + f"{field_name} values must be handler identifier strings"
436 + )
437 + normalized.add(item)
438 + return normalized
439 +
440 + def _normalize_sid_filter(self, value: str | Iterable[str] | None) -> Set[str]:
441 + if value is None:
442 + return set()
443 + if isinstance(value, str):
444 + return {value}
445 + normalized: Set[str] = set()
446 + for item in value:
447 + normalized.add(str(item))
448 + return normalized
449 +
450 + def _select_handlers(
451 + self,
452 + namespace: str,
453 + *,
454 + include: Set[str] | None,
455 + exclude: Set[str] | None,
456 + ) -> tuple[list[WsHandler], Set[str]]:
457 + registered = self.handlers.get(namespace, [])
458 + available_ids = {handler.identifier for handler in registered}
459 +
460 + if include is not None:
461 + unknown = include - available_ids
462 + if unknown:
463 + raise ValueError(
464 + f"Unknown handler(s) in includeHandlers for namespace '{namespace}': "
465 + f"{', '.join(sorted(unknown))}"
466 + )
467 + if exclude is not None:
468 + unknown = exclude - available_ids
469 + if unknown:
470 + raise ValueError(
471 + f"Unknown handler(s) in excludeHandlers for namespace '{namespace}': "
472 + f"{', '.join(sorted(unknown))}"
473 + )
474 +
475 + selected: list[WsHandler] = []
476 + for handler in registered:
477 + ident = handler.identifier
478 + if include is not None and ident not in include:
479 + continue
480 + if exclude is not None and ident in exclude:
481 + continue
482 + selected.append(handler)
483 +
484 + return selected, available_ids
485 +
486 + def _resolve_correlation_id(self, payload: dict[str, Any]) -> str:
487 + value = payload.get("correlationId")
488 + if isinstance(value, str) and value.strip():
489 + correlation_id = value.strip()
490 + else:
491 + correlation_id = uuid.uuid4().hex
492 + payload["correlationId"] = correlation_id
493 + return correlation_id
494 +
495 + def register_handlers(
496 + self, handlers_by_namespace: dict[str, Iterable[WsHandler]]
497 + ) -> None:
498 + for namespace, handlers in handlers_by_namespace.items():
499 + for handler in handlers:
500 + handler.bind_manager(self, namespace=namespace)
501 + if _ws_debug_enabled():
502 + PrintStyle.info(
503 + "Registered WebSocket handler %s namespace=%s"
504 + % (handler.identifier, namespace)
505 + )
506 + existing = self.handlers.get(namespace, [])
507 + if handler in existing:
508 + PrintStyle.warning(
509 + f"Duplicate handler registration for namespace '{namespace}'"
510 + )
511 + self.handlers[namespace].append(handler)
512 + self._debug(
513 + f"Registered handler {handler.identifier} namespace={namespace}"
514 + )
515 +
516 + def iter_event_types(self, namespace: str) -> Iterable[str]:
517 + return []
518 +
519 + def iter_namespaces(self) -> list[str]:
520 + return list(self.handlers.keys())
521 +
522 + async def _invoke_handler(
523 + self,
524 + handler: WsHandler,
525 + event_type: str,
526 + payload: dict[str, Any],
527 + sid: str,
528 + ) -> _HandlerExecution:
529 + instrument = self._diagnostics_active()
530 + start = time.perf_counter() if instrument else None
531 + try:
532 + value = await self._get_handler_worker().execute_inside(
533 + handler.process, event_type, payload, sid
534 + )
535 + except Exception as exc: # pragma: no cover - handled by caller
536 + duration_ms = (
537 + (time.perf_counter() - start) * 1000 if start is not None else None
538 + )
539 + return _HandlerExecution(handler, exc, duration_ms)
540 + duration_ms = (
541 + (time.perf_counter() - start) * 1000 if start is not None else None
542 + )
543 + return _HandlerExecution(handler, value, duration_ms)
544 +
545 + async def handle_connect(
546 + self, namespace: str, sid: str, user_id: str | None = None
547 + ) -> None:
548 + self._ensure_dispatcher_loop()
549 + user_bucket = user_id or "single_user"
550 + identity: ConnectionIdentity = (namespace, sid)
551 + with self.lock:
552 + self.connections[identity] = ConnectionInfo(namespace=namespace, sid=sid)
553 + self._known_sids.add(identity)
554 + self.sid_to_user[identity] = user_bucket
555 + self.user_to_sids[self._ALL_USERS_BUCKET].add(identity)
556 + self.user_to_sids[user_bucket].add(identity)
557 + connection_count = sum(
558 + 1 for conn_identity in self.connections if conn_identity[0] == namespace
559 + )
560 + if _ws_debug_enabled():
561 + PrintStyle.info(f"WebSocket connected: namespace={namespace} sid={sid}")
562 + await self._run_lifecycle(namespace, lambda h: h.on_connect(sid))
563 + await self._flush_buffer(identity)
564 + if self._server_restart_enabled:
565 + await self.emit_to(
566 + namespace,
567 + sid,
568 + "server_restart",
569 + {
570 + "emittedAt": _utcnow()
571 + .isoformat(timespec="milliseconds")
572 + .replace("+00:00", "Z"),
573 + "runtimeId": runtime.get_runtime_id(),
574 + },
575 + handler_id=self._identifier,
576 + )
577 + if _ws_debug_enabled():
578 + PrintStyle.info(
579 + f"server_restart broadcast emitted to namespace={namespace} sid={sid}"
580 + )
581 + lifecycle_payload = {
582 + "namespace": namespace,
583 + "sid": sid,
584 + "connectionCount": connection_count,
585 + "timestamp": self._timestamp(),
586 + }
587 + await self._publish_diagnostic_event(
588 + {
589 + "kind": "lifecycle",
590 + "event": "connect",
591 + **lifecycle_payload,
592 + }
593 + )
594 + self._schedule_lifecycle_broadcast(
595 + namespace, LIFECYCLE_CONNECT_EVENT, lifecycle_payload
596 + )
597 +
598 + async def handle_disconnect(self, namespace: str, sid: str) -> None:
599 + self._ensure_dispatcher_loop()
600 + identity: ConnectionIdentity = (namespace, sid)
601 + with self.lock:
602 + self.connections.pop(identity, None)
603 + # Keep identity in _known_sids so emit_to buffers instead of raising
604 + # session tracking cleanup
605 + user_bucket = self.sid_to_user.pop(identity, None)
606 + if self._ALL_USERS_BUCKET in self.user_to_sids:
607 + self.user_to_sids[self._ALL_USERS_BUCKET].discard(identity)
608 + if not self.user_to_sids[self._ALL_USERS_BUCKET]:
609 + self.user_to_sids.pop(self._ALL_USERS_BUCKET, None)
610 + if user_bucket and user_bucket in self.user_to_sids:
611 + self.user_to_sids[user_bucket].discard(identity)
612 + if not self.user_to_sids[user_bucket]:
613 + self.user_to_sids.pop(user_bucket, None)
614 + connection_count = sum(
615 + 1 for conn_identity in self.connections if conn_identity[0] == namespace
616 + )
617 + self.unregister_diagnostic_watcher(namespace, sid)
618 + PrintStyle.info(f"WebSocket disconnected: namespace={namespace} sid={sid}")
619 + await self._run_lifecycle(namespace, lambda h: h.on_disconnect(sid))
620 + lifecycle_payload = {
621 + "namespace": namespace,
622 + "sid": sid,
623 + "connectionCount": connection_count,
624 + "timestamp": self._timestamp(),
625 + }
626 + await self._publish_diagnostic_event(
627 + {
628 + "kind": "lifecycle",
629 + "event": "disconnect",
630 + **lifecycle_payload,
631 + }
632 + )
633 + self._schedule_lifecycle_broadcast(
634 + namespace, LIFECYCLE_DISCONNECT_EVENT, lifecycle_payload
635 + )
636 +
637 + async def route_event(
638 + self,
639 + namespace: str,
640 + event_type: str,
641 + data: dict[str, Any],
642 + sid: str,
643 + ack: Optional[Callable[[Any], None]] = None,
644 + *,
645 + include_handlers: Set[str] | None = None,
646 + exclude_handlers: Set[str] | None = None,
647 + allow_exclude: bool = False,
648 + handler_id: str | None = None,
649 + ) -> dict[str, Any]:
650 + self._ensure_dispatcher_loop()
651 + incoming = dict(data or {})
652 + correlation_id = self._resolve_correlation_id(incoming)
653 + self._debug(
654 + f"Routing event namespace={namespace} '{event_type}' sid={sid} correlation={correlation_id}"
655 + )
656 +
657 + include_meta_raw = incoming.pop("includeHandlers", None)
658 + exclude_meta_raw = incoming.pop("excludeHandlers", None)
659 +
660 + if "data" in incoming and isinstance(incoming.get("data"), dict):
661 + handler_payload = dict(incoming.get("data") or {})
662 + if "excludeSids" in incoming:
663 + handler_payload["excludeSids"] = incoming.get("excludeSids")
664 + else:
665 + handler_payload = dict(incoming)
666 +
667 + handler_payload["correlationId"] = correlation_id
668 +
669 + try:
670 + include_meta = self._normalize_handler_filter(
671 + include_meta_raw, "includeHandlers"
672 + )
673 + except ValueError as exc:
674 + error = self._build_error_result(
675 + handler_id=handler_id or self._identifier,
676 + code="INVALID_FILTER",
677 + message=str(exc),
678 + correlation_id=correlation_id,
679 + )
680 + if ack:
681 + ack({"correlationId": correlation_id, "results": [error]})
682 + return {"correlationId": correlation_id, "results": [error]}
683 +
684 + try:
685 + exclude_meta = self._normalize_handler_filter(
686 + exclude_meta_raw, "excludeHandlers"
687 + )
688 + except ValueError as exc:
689 + error = self._build_error_result(
690 + handler_id=handler_id or self._identifier,
691 + code="INVALID_FILTER",
692 + message=str(exc),
693 + correlation_id=correlation_id,
694 + )
695 + payload_error = {"correlationId": correlation_id, "results": [error]}
696 + if ack:
697 + ack(payload_error)
698 + return payload_error
699 +
700 + if exclude_meta_raw is not None and not allow_exclude:
701 + error = self._build_error_result(
702 + handler_id=handler_id or self._identifier,
703 + code="INVALID_FILTER",
704 + message="excludeHandlers is not supported for this operation",
705 + correlation_id=correlation_id,
706 + )
707 + if ack:
708 + ack({"correlationId": correlation_id, "results": [error]})
709 + return {"correlationId": correlation_id, "results": [error]}
710 +
711 + if include_handlers is not None and include_meta is not None:
712 + if include_handlers != include_meta:
713 + error = self._build_error_result(
714 + handler_id=handler_id or self._identifier,
715 + code="INVALID_FILTER",
716 + message="Conflicting includeHandlers filters supplied",
717 + correlation_id=correlation_id,
718 + )
719 + if ack:
720 + ack({"correlationId": correlation_id, "results": [error]})
721 + return {"correlationId": correlation_id, "results": [error]}
722 +
723 + if allow_exclude and exclude_handlers is not None and exclude_meta is not None:
724 + if exclude_handlers != exclude_meta:
725 + error = self._build_error_result(
726 + handler_id=handler_id or self._identifier,
727 + code="INVALID_FILTER",
728 + message="Conflicting excludeHandlers filters supplied",
729 + correlation_id=correlation_id,
730 + )
731 + if ack:
732 + ack({"correlationId": correlation_id, "results": [error]})
733 + return {"correlationId": correlation_id, "results": [error]}
734 +
735 + include = include_handlers or include_meta
736 + exclude = exclude_handlers or (exclude_meta if allow_exclude else None)
737 +
738 + try:
739 + validate_event_type(event_type)
740 + except (TypeError, ValueError) as exc:
741 + error = self._build_error_result(
742 + handler_id=handler_id or self._identifier,
743 + code="INVALID_EVENT",
744 + message=str(exc),
745 + correlation_id=correlation_id,
746 + )
747 + if ack:
748 + ack({"correlationId": correlation_id, "results": [error]})
749 + return {"correlationId": correlation_id, "results": [error]}
750 +
751 + registered = self.handlers.get(namespace, [])
752 + if not registered:
753 + PrintStyle.warning(f"No handlers registered for namespace '{namespace}'")
754 + error = self._build_error_result(
755 + handler_id=handler_id or self._identifier,
756 + code="NO_HANDLERS",
757 + message=f"No handler for namespace '{namespace}'",
758 + correlation_id=correlation_id,
759 + )
760 + if ack:
761 + ack({"correlationId": correlation_id, "results": [error]})
762 + return {"correlationId": correlation_id, "results": [error]}
763 +
764 + try:
765 + selected_handlers, _ = self._select_handlers(
766 + namespace, include=include, exclude=exclude
767 + )
768 + except ValueError as exc:
769 + error = self._build_error_result(
770 + handler_id=handler_id or self._identifier,
771 + code="INVALID_FILTER",
772 + message=str(exc),
773 + correlation_id=correlation_id,
774 + )
775 + if ack:
776 + ack({"correlationId": correlation_id, "results": [error]})
777 + return {"correlationId": correlation_id, "results": [error]}
778 +
779 + if not selected_handlers:
780 + error = self._build_error_result(
781 + handler_id=handler_id or self._identifier,
782 + code="NO_HANDLERS",
783 + message=f"No handler for '{event_type}' after applying filters",
784 + correlation_id=correlation_id,
785 + )
786 + if ack:
787 + ack({"correlationId": correlation_id, "results": [error]})
788 + return {"correlationId": correlation_id, "results": [error]}
789 +
790 + with self.lock:
791 + info = self.connections.get((namespace, sid))
792 + if info:
793 + info.last_activity = _utcnow()
794 +
795 + executions = await asyncio.gather(
796 + *[
797 + self._invoke_handler(handler, event_type, dict(handler_payload), sid)
798 + for handler in selected_handlers
799 + ]
800 + )
801 +
802 + results: List[dict[str, Any]] = []
803 + for execution in executions:
804 + handler = execution.handler
805 + value = execution.value
806 + duration_ms = execution.duration_ms
807 +
808 + if isinstance(value, Exception): # pragma: no cover - defensive logging
809 + PrintStyle.error(
810 + f"Error in handler {handler.identifier} for '{event_type}' (correlation {correlation_id}): {value}"
811 + )
812 + results.append(
813 + self._build_error_result(
814 + handler_id=handler.identifier,
815 + code="HANDLER_ERROR",
816 + message="Internal server error",
817 + details=str(value),
818 + correlation_id=correlation_id,
819 + duration_ms=duration_ms,
820 + )
821 + )
822 + continue
823 +
824 + if isinstance(value, WsResult):
825 + results.append(
826 + value.as_result(
827 + handler_id=handler.identifier,
828 + fallback_correlation_id=correlation_id,
829 + duration_ms=duration_ms,
830 + )
831 + )
832 + continue
833 +
834 + if value is None:
835 + helper_result = WsResult(ok=True)
836 + elif isinstance(value, dict):
837 + helper_result = WsResult(ok=True, data=value)
838 + else:
839 + helper_result = WsResult(ok=True, data={"result": value})
840 +
841 + results.append(
842 + helper_result.as_result(
843 + handler_id=handler.identifier,
844 + fallback_correlation_id=correlation_id,
845 + duration_ms=duration_ms,
846 + )
847 + )
848 +
849 + await self._publish_diagnostic_event(
850 + lambda: {
851 + "kind": "inbound",
852 + "sourceNamespace": namespace,
853 + "namespace": namespace,
854 + "eventType": event_type,
855 + "sid": sid,
856 + "correlationId": correlation_id,
857 + "timestamp": self._timestamp(),
858 + "handlerCount": len(selected_handlers),
859 + "durationMs": sum((exec.duration_ms or 0.0) for exec in executions),
860 + "resultSummary": self._summarize_results(results),
861 + "payloadSummary": self._summarize_payload(handler_payload),
862 + }
863 + )
864 +
865 + response_payload = {"correlationId": correlation_id, "results": results}
866 + if ack:
867 + ack(response_payload)
868 + self._debug(
869 + f"Completed event namespace={namespace} '{event_type}' sid={sid} correlation={correlation_id}"
870 + )
871 + return response_payload
872 +
873 + async def request_for_sid(
874 + self,
875 + *,
876 + namespace: str,
877 + sid: str,
878 + event_type: str,
879 + data: dict[str, Any],
880 + timeout_ms: int = 0,
881 + handler_id: str | None = None,
882 + include_handlers: Set[str] | None = None,
883 + ) -> dict[str, Any]:
884 + payload = dict(data or {})
885 + correlation_id = self._resolve_correlation_id(payload)
886 +
887 + with self.lock:
888 + connected = (namespace, sid) in self.connections
889 + if not connected:
890 + return {
891 + "correlationId": correlation_id,
892 + "results": [
893 + self._build_error_result(
894 + handler_id=handler_id or self._identifier,
895 + code="CONNECTION_NOT_FOUND",
896 + message=f"Connection '{sid}' not found in namespace '{namespace}'",
897 + correlation_id=correlation_id,
898 + )
899 + ],
900 + }
901 +
902 + async def _invoke() -> dict[str, Any]:
903 + return await self.route_event(
904 + namespace,
905 + event_type,
906 + payload,
907 + sid,
908 + include_handlers=include_handlers,
909 + handler_id=handler_id,
910 + )
911 +
912 + if timeout_ms and timeout_ms > 0:
913 + try:
914 + return await asyncio.wait_for(_invoke(), timeout=timeout_ms / 1000)
915 + except asyncio.TimeoutError:
916 + PrintStyle.warning(
917 + f"request timeout for sid {sid} event '{event_type}'"
918 + )
919 + return {
920 + "correlationId": correlation_id,
921 + "results": [
922 + self._build_error_result(
923 + handler_id=handler_id or self._identifier,
924 + code="TIMEOUT",
925 + message="Request timeout",
926 + correlation_id=correlation_id,
927 + )
928 + ],
929 + }
930 + return await _invoke()
931 +
932 + async def route_event_all(
933 + self,
934 + namespace: str,
935 + event_type: str,
936 + data: dict[str, Any],
937 + *,
938 + timeout_ms: int = 0,
939 + exclude_handlers: Set[str] | None = None,
940 + handler_id: str | None = None,
941 + ) -> list[dict[str, Any]]:
942 + """Fan-out a request to all active connections and aggregate responses."""
943 +
944 + base_payload = dict(data or {})
945 + exclude_meta_raw = base_payload.pop("excludeHandlers", None)
946 + exclude_combined: Set[str] | None = exclude_handlers
947 + correlation_id = self._resolve_correlation_id(base_payload)
948 +
949 + if exclude_meta_raw is not None:
950 + try:
951 + exclude_meta = self._normalize_handler_filter(
952 + exclude_meta_raw, "excludeHandlers"
953 + )
954 + except ValueError as exc:
955 + error = self._build_error_result(
956 + handler_id=handler_id or self._identifier,
957 + code="INVALID_FILTER",
958 + message=str(exc),
959 + correlation_id=correlation_id,
960 + )
961 + return [
962 + {
963 + "sid": "__invalid__",
964 + "correlationId": correlation_id,
965 + "results": [error],
966 + }
967 + ]
968 +
969 + if exclude_combined is None:
970 + exclude_combined = exclude_meta
971 + elif exclude_meta is not None and exclude_combined != exclude_meta:
972 + error = self._build_error_result(
973 + handler_id=handler_id or self._identifier,
974 + code="INVALID_FILTER",
975 + message="Conflicting excludeHandlers filters supplied",
976 + correlation_id=correlation_id,
977 + )
978 + return [
979 + {
980 + "sid": "__invalid__",
981 + "correlationId": correlation_id,
982 + "results": [error],
983 + }
984 + ]
985 +
986 + self._debug(
987 + f"Starting requestAll namespace={namespace} for '{event_type}' correlation={correlation_id}"
988 + )
989 +
990 + with self.lock:
991 + active_sids = [
992 + conn_identity[1]
993 + for conn_identity in self.connections.keys()
994 + if conn_identity[0] == namespace
995 + ]
996 + if not active_sids:
997 + self._debug(
998 + f"No active connections for requestAll namespace={namespace} '{event_type}' correlation={correlation_id}"
999 + )
1000 + return []
1001 +
1002 + timeout_seconds = timeout_ms / 1000 if timeout_ms and timeout_ms > 0 else None
1003 +
1004 + async def _invoke_for_sid(target_sid: str) -> dict[str, Any]:
1005 + async def _dispatch() -> dict[str, Any]:
1006 + return await self.route_event(
1007 + namespace,
1008 + event_type,
1009 + base_payload,
1010 + target_sid,
1011 + allow_exclude=True,
1012 + exclude_handlers=exclude_combined,
1013 + handler_id=handler_id,
1014 + )
1015 +
1016 + if timeout_seconds is None:
1017 + return await _dispatch()
1018 +
1019 + try:
1020 + task = asyncio.create_task(_dispatch())
1021 + return await asyncio.wait_for(
1022 + asyncio.shield(task), timeout=timeout_seconds
1023 + )
1024 + except asyncio.TimeoutError:
1025 + PrintStyle.warning(
1026 + f"requestAll timeout for sid {target_sid} correlation={correlation_id}"
1027 + )
1028 + # Ensure any late exceptions are observed so asyncio does not log
1029 + # "Task exception was never retrieved".
1030 + try:
1031 + task.add_done_callback(lambda t: t.exception()) # type: ignore[arg-type]
1032 + except Exception: # pragma: no cover - defensive
1033 + pass
1034 + return {
1035 + "correlationId": correlation_id,
1036 + "results": [
1037 + self._build_error_result(
1038 + handler_id=handler_id or self._identifier,
1039 + code="TIMEOUT",
1040 + message="Request timeout",
1041 + correlation_id=correlation_id,
1042 + )
1043 + ],
1044 + }
1045 +
1046 + tasks = {sid: asyncio.create_task(_invoke_for_sid(sid)) for sid in active_sids}
1047 +
1048 + aggregated: list[dict[str, Any]] = []
1049 + for sid, task in tasks.items():
1050 + result = await task
1051 + if isinstance(result, dict):
1052 + aggregated.append(
1053 + {
1054 + "sid": sid,
1055 + "correlationId": result.get("correlationId", correlation_id),
1056 + "results": result.get("results", []),
1057 + }
1058 + )
1059 + else:
1060 + aggregated.append(
1061 + {
1062 + "sid": sid,
1063 + "correlationId": correlation_id,
1064 + "results": result,
1065 + }
1066 + )
1067 +
1068 + self._debug(
1069 + f"Completed requestAll namespace={namespace} for '{event_type}' correlation={correlation_id}"
1070 + )
1071 + return aggregated
1072 +
1073 + def _wrap_envelope(
1074 + self,
1075 + handler_id: str | None,
1076 + data: dict[str, Any],
1077 + *,
1078 + correlation_id: str | None = None,
1079 + ) -> dict[str, Any]:
1080 + hid = handler_id or self._identifier
1081 + ts = _utcnow().isoformat(timespec="milliseconds").replace("+00:00", "Z")
1082 + event_id = str(uuid.uuid4())
1083 + correlation = correlation_id or str(uuid.uuid4())
1084 + return {
1085 + "handlerId": hid,
1086 + "eventId": event_id,
1087 + "correlationId": correlation,
1088 + "ts": ts,
1089 + "data": data or {},
1090 + }
1091 +
1092 + async def emit_to(
1093 + self,
1094 + namespace: str,
1095 + sid: str,
1096 + event_type: str,
1097 + data: dict[str, Any],
1098 + *,
1099 + handler_id: str | None = None,
1100 + correlation_id: str | None = None,
1101 + diagnostic: bool = False,
1102 + ) -> None:
1103 + envelope = self._wrap_envelope(
1104 + handler_id,
1105 + data,
1106 + correlation_id=correlation_id,
1107 + )
1108 + delivered = False
1109 + buffered = False
1110 + identity: ConnectionIdentity = (namespace, sid)
1111 +
1112 + with self.lock:
1113 + connected = identity in self.connections
1114 + known = identity in self._known_sids or identity in self.buffers
1115 +
1116 + if connected:
1117 + self._debug(
1118 + "Emit to namespace=%s sid=%s event=%s eventId=%s correlationId=%s handlerId=%s"
1119 + % (
1120 + namespace,
1121 + sid,
1122 + event_type,
1123 + envelope.get("eventId"),
1124 + envelope.get("correlationId"),
1125 + envelope.get("handlerId"),
1126 + )
1127 + )
1128 + await self._run_on_dispatcher_loop(
1129 + self.socketio.emit(event_type, envelope, to=sid, namespace=namespace)
1130 + )
1131 + delivered = True
1132 + else:
1133 + if not known:
1134 + raise ConnectionNotFoundError(sid, namespace=namespace)
1135 + with self.lock:
1136 + self._buffer_event(
1137 + identity,
1138 + event_type,
1139 + data,
1140 + handler_id,
1141 + envelope["correlationId"],
1142 + )
1143 + buffered = True
1144 +
1145 + if not diagnostic:
1146 + await self._publish_diagnostic_event(
1147 + lambda: {
1148 + "kind": "outbound",
1149 + "direction": "emit_to",
1150 + "eventType": event_type,
1151 + "namespace": namespace,
1152 + "sid": sid,
1153 + "correlationId": envelope["correlationId"],
1154 + "handlerId": envelope["handlerId"],
1155 + "timestamp": self._timestamp(),
1156 + "delivered": delivered,
1157 + "buffered": buffered,
1158 + "payloadSummary": self._summarize_payload(data),
1159 + }
1160 + )
1161 +
1162 + async def send_data(
1163 + self,
1164 + endpoint_name: str,
1165 + event_type: str,
1166 + data: dict[str, Any],
1167 + connection_id: str | None = None,
1168 + ) -> None:
1169 + if connection_id is not None:
1170 + await self.emit_to(endpoint_name, connection_id, event_type, data)
1171 + return
1172 + await self.broadcast(endpoint_name, event_type, data)
1173 +
1174 + async def broadcast(
1175 + self,
1176 + namespace: str,
1177 + event_type: str,
1178 + data: dict[str, Any],
1179 + *,
1180 + exclude_sids: str | Iterable[str] | None = None,
1181 + handler_id: str | None = None,
1182 + correlation_id: str | None = None,
1183 + diagnostic: bool = False,
1184 + ) -> None:
1185 + excluded = self._normalize_sid_filter(exclude_sids)
1186 +
1187 + targets: list[str] = []
1188 + with self.lock:
1189 + current_identities = list(self.connections.keys())
1190 + for conn_identity in current_identities:
1191 + if conn_identity[0] != namespace:
1192 + continue
1193 + sid = conn_identity[1]
1194 + if sid in excluded:
1195 + continue
1196 + targets.append(sid)
1197 +
1198 + if targets:
1199 + envelope = self._wrap_envelope(
1200 + handler_id,
1201 + data,
1202 + correlation_id=correlation_id,
1203 + )
1204 + coros = [
1205 + self._run_on_dispatcher_loop(
1206 + self.socketio.emit(event_type, envelope, to=sid, namespace=namespace)
1207 + )
1208 + for sid in targets
1209 + ]
1210 + await asyncio.gather(*coros)
1211 +
1212 + if not diagnostic:
1213 + await self._publish_diagnostic_event(
1214 + lambda: {
1215 + "kind": "outbound",
1216 + "direction": "broadcast",
1217 + "eventType": event_type,
1218 + "namespace": namespace,
1219 + "targets": targets[:10],
1220 + "targetCount": len(targets),
1221 + "correlationId": correlation_id,
1222 + "handlerId": handler_id or self._identifier,
1223 + "timestamp": self._timestamp(),
1224 + "payloadSummary": self._summarize_payload(data),
1225 + }
1226 + )
1227 +
1228 + async def _run_lifecycle(
1229 + self, namespace: str, fn: Callable[[WsHandler], Any]
1230 + ) -> None:
1231 + seen: Set[WsHandler] = set()
1232 + coros: list[Any] = []
1233 + for handler in self.handlers.get(namespace, []):
1234 + if handler in seen:
1235 + continue
1236 + seen.add(handler)
1237 + coros.append(self._get_handler_worker().execute_inside(fn, handler))
1238 + if coros:
1239 + await asyncio.gather(*coros, return_exceptions=True)
1240 +
1241 + def _buffer_event(
1242 + self,
1243 + identity: ConnectionIdentity,
1244 + event_type: str,
1245 + data: dict[str, Any],
1246 + handler_id: str | None,
1247 + correlation_id: str | None,
1248 + ) -> None:
1249 + namespace, sid = identity
1250 + buffer = self.buffers[identity]
1251 + buffer.append(
1252 + BufferedEvent(
1253 + event_type=event_type,
1254 + data=data,
1255 + handler_id=handler_id,
1256 + correlation_id=correlation_id,
1257 + )
1258 + )
1259 + while len(buffer) > BUFFER_MAX_SIZE:
1260 + dropped = buffer.popleft()
1261 + PrintStyle.warning(
1262 + f"Dropping buffered event '{dropped.event_type}' for namespace={namespace} sid={sid} (overflow)"
1263 + )
1264 + self._debug(
1265 + f"Buffered event namespace={namespace} '{event_type}' sid={sid} (queue length={len(buffer)})"
1266 + )
1267 +
1268 + async def _flush_buffer(self, identity: ConnectionIdentity) -> None:
1269 + self._ensure_dispatcher_loop()
1270 + buffer = self.buffers.get(identity)
1271 + if not buffer:
1272 + return
1273 + namespace, sid = identity
1274 + now = _utcnow()
1275 + delivered = 0
1276 + while buffer:
1277 + event = buffer.popleft()
1278 + if now - event.timestamp > BUFFER_TTL:
1279 + self._debug(
1280 + f"Discarding expired buffered event '{event.event_type}' for sid {sid}"
1281 + )
1282 + continue
1283 + envelope = self._wrap_envelope(
1284 + event.handler_id,
1285 + event.data,
1286 + correlation_id=event.correlation_id,
1287 + )
1288 + self._debug(
1289 + "Flush to sid=%s event=%s eventId=%s correlationId=%s handlerId=%s"
1290 + % (
1291 + sid,
1292 + event.event_type,
1293 + envelope.get("eventId"),
1294 + envelope.get("correlationId"),
1295 + envelope.get("handlerId"),
1296 + )
1297 + )
1298 + await self._run_on_dispatcher_loop(
1299 + self.socketio.emit(
1300 + event.event_type, envelope, to=sid, namespace=namespace
1301 + )
1302 + )
1303 + delivered += 1
1304 + if identity in self.buffers:
1305 + self.buffers.pop(identity, None)
1306 + if delivered:
1307 + PrintStyle.info(
1308 + f"Flushed {delivered} buffered event(s) to namespace={namespace} sid={sid}"
1309 + )
1310 +
1311 + def _build_error_result(
1312 + self,
1313 + *,
1314 + handler_id: str | None = None,
1315 + code: str,
1316 + message: str,
1317 + details: str | None = None,
1318 + correlation_id: str | None = None,
1319 + duration_ms: float | None = None,
1320 + ) -> dict[str, Any]:
1321 + error_payload = {"code": code, "error": message}
1322 + if details:
1323 + error_payload["details"] = details
1324 + result: dict[str, Any] = {
1325 + "handlerId": handler_id or self._identifier,
1326 + "ok": False,
1327 + "error": error_payload,
1328 + }
1329 + if correlation_id is not None:
1330 + result["correlationId"] = correlation_id
1331 + if duration_ms is not None:
1332 + result["durationMs"] = round(duration_ms, 4)
1333 + return result
1334 +
1335 + # Session tracking helpers (single-user defaults)
1336 + def get_sids_for_user(self, user: str | None = None) -> list[ConnectionIdentity]:
1337 + """Return connection identities for a user; single-user default returns all."""
1338 + with self.lock:
1339 + bucket = self._ALL_USERS_BUCKET if user is None else user
1340 + return list(self.user_to_sids.get(bucket, set()))
1341 +
1342 + def get_user_for_sid(self, namespace: str, sid: str) -> str | None:
1343 + """Return user identifier for a connection or None."""
1344 + identity: ConnectionIdentity = (namespace, sid)
1345 + with self.lock:
1346 + return self.sid_to_user.get(identity)
1347 +
1348 + def set_server_restart_broadcast(self, enabled: bool) -> None:
1349 + """Enable or disable automatic server restart broadcasts."""
1350 +
1351 + self._server_restart_enabled = bool(enabled)
run_ui.py
+7 -172
@@ -14,9 +14,8 @@ import initialize
14 from helpers import files, git, mcp_server, fasta2a_server, settings as settings_helper, extension
15 from helpers.files import get_abs_path
16 from helpers import runtime, dotenv, process
17 -from helpers.websocket import WebSocketHandler, validate_ws_origin
17 from helpers.api import register_api_route, requires_auth, csrf_protect
19 -from helpers.ws import register_ws_namespace
18 +from helpers.ws import register_ws_namespace, validate_ws_origin
19 from helpers.print_style import PrintStyle
20 from helpers import login
21 import socketio # type: ignore[import-untyped]
@@ -24,8 +23,7 @@ from socketio import ASGIApp, packet
23 from starlette.applications import Starlette
24 from starlette.routing import Mount
25 from uvicorn.middleware.wsgi import WSGIMiddleware
27 -from helpers.websocket_manager import WebSocketManager, set_shared_websocket_manager
28 -from helpers.websocket_namespace_discovery import discover_websocket_namespaces
26 +from helpers.ws_manager import WsManager, set_shared_ws_manager
27 from flask import send_file
28
29 # disable logging
@@ -73,11 +71,11 @@ socketio_server = socketio.AsyncServer(
71 max_http_buffer_size=50 * 1024 * 1024,
72 )
73
76 -websocket_manager = WebSocketManager(socketio_server, lock)
77 -set_shared_websocket_manager(websocket_manager)
74 +ws_manager = WsManager(socketio_server, lock)
75 +set_shared_ws_manager(ws_manager)
76 _settings = settings_helper.get_settings()
77 settings_helper.set_runtime_settings_snapshot(_settings)
80 -websocket_manager.set_server_restart_broadcast(
78 +ws_manager.set_server_restart_broadcast(
79 _settings.get("websocket_server_restart_enabled", True)
80 )
81
@@ -191,160 +189,6 @@ async def _serve_plugin_asset(plugin_name, asset_path):
189 return Response("Error serving asset", 500)
190
191
194 -def _build_websocket_handlers_by_namespace(
195 - socketio_server: socketio.AsyncServer,
196 - lock: threading.RLock,
197 -) -> dict[str, list[WebSocketHandler]]:
198 - discoveries = discover_websocket_namespaces(
199 - handlers_folder="python/websocket_handlers",
200 - include_root_default=True,
201 - )
202 -
203 - handlers_by_namespace: dict[str, list[WebSocketHandler]] = {}
204 - for discovery in discoveries:
205 - namespace = discovery.namespace
206 - for handler_cls in discovery.handler_classes:
207 - handler = handler_cls.get_instance(socketio_server, lock)
208 - handlers_by_namespace.setdefault(namespace, []).append(handler)
209 -
210 - return handlers_by_namespace
211 -
212 -
213 -def configure_websocket_namespaces(
214 - *,
215 - webapp: Flask,
216 - socketio_server: socketio.AsyncServer,
217 - websocket_manager: WebSocketManager,
218 - handlers_by_namespace: dict[str, list[WebSocketHandler]],
219 -) -> set[str]:
220 - namespace_map: dict[str, list[WebSocketHandler]] = {
221 - namespace: list(handlers) for namespace, handlers in handlers_by_namespace.items()
222 - }
223 -
224 - # Always include the reserved root namespace. It is unhandled for application events by
225 - # default, but request-style calls must resolve deterministically with NO_HANDLERS.
226 - namespace_map.setdefault("/", [])
227 -
228 - websocket_manager.register_handlers(namespace_map)
229 -
230 - allowed_namespaces = set(namespace_map.keys())
231 - original_handle_connect = socketio_server._handle_connect # type: ignore[attr-defined]
232 -
233 - async def _handle_connect_with_namespace_gatekeeper(eio_sid, namespace, data):
234 - requested = namespace or "/"
235 - if requested not in allowed_namespaces:
236 - await socketio_server._send_packet(
237 - eio_sid,
238 - socketio_server.packet_class(
239 - packet.CONNECT_ERROR,
240 - data={
241 - "message": "UNKNOWN_NAMESPACE",
242 - "data": {"code": "UNKNOWN_NAMESPACE", "namespace": requested},
243 - },
244 - namespace=requested,
245 - ),
246 - )
247 - return
248 - await original_handle_connect(eio_sid, namespace, data)
249 -
250 - socketio_server._handle_connect = _handle_connect_with_namespace_gatekeeper # type: ignore[assignment]
251 -
252 - def _register_namespace_handlers(
253 - namespace: str, namespace_handlers: list[WebSocketHandler]
254 - ) -> None:
255 - # A namespace is the WebSocket equivalent of an API endpoint.
256 - # Security requirements must be consistent within the namespace (no any()-based union).
257 - auth_required = False
258 - csrf_required = False
259 - if namespace_handlers:
260 - auth_required = bool(namespace_handlers[0].requires_auth())
261 - csrf_required = bool(namespace_handlers[0].requires_csrf())
262 - for handler in namespace_handlers[1:]:
263 - if (
264 - bool(handler.requires_auth()) != auth_required
265 - or bool(handler.requires_csrf()) != csrf_required
266 - ):
267 - raise ValueError(
268 - f"WebSocket namespace {namespace!r} has mixed auth/csrf requirements across handlers"
269 - )
270 -
271 - @socketio_server.on("connect", namespace=namespace)
272 - async def _connect( # type: ignore[override]
273 - sid,
274 - environ,
275 - _auth,
276 - _namespace: str = namespace,
277 - _auth_required: bool = auth_required,
278 - _csrf_required: bool = csrf_required,
279 - ):
280 - with webapp.request_context(environ):
281 - origin_ok, origin_reason = validate_ws_origin(environ)
282 - if not origin_ok:
283 - PrintStyle.warning(
284 - f"WebSocket origin validation failed for {_namespace} {sid}: {origin_reason or 'invalid'}"
285 - )
286 - return False
287 -
288 - if _auth_required:
289 - credentials_hash = login.get_credentials_hash()
290 - if credentials_hash:
291 - if session.get("authentication") != credentials_hash:
292 - PrintStyle.warning(
293 - f"WebSocket authentication failed for {_namespace} {sid}: session not valid"
294 - )
295 - return False
296 - else:
297 - PrintStyle.debug(
298 - "WebSocket authentication required but credentials not configured; proceeding"
299 - )
300 -
301 - if _csrf_required:
302 - expected_token = session.get("csrf_token")
303 - if not isinstance(expected_token, str) or not expected_token:
304 - PrintStyle.warning(
305 - f"WebSocket CSRF validation failed for {_namespace} {sid}: csrf_token not initialized"
306 - )
307 - return False
308 -
309 - auth_token = None
310 - if isinstance(_auth, dict):
311 - auth_token = _auth.get("csrf_token") or _auth.get("csrfToken")
312 - if not isinstance(auth_token, str) or not auth_token:
313 - PrintStyle.warning(
314 - f"WebSocket CSRF validation failed for {_namespace} {sid}: missing csrf_token in auth"
315 - )
316 - return False
317 - if auth_token != expected_token:
318 - PrintStyle.warning(
319 - f"WebSocket CSRF validation failed for {_namespace} {sid}: csrf_token mismatch"
320 - )
321 - return False
322 -
323 - cookie_name = f"csrf_token_{runtime.get_runtime_id()}"
324 - cookie_token = request.cookies.get(cookie_name)
325 - if cookie_token != expected_token:
326 - PrintStyle.warning(
327 - f"WebSocket CSRF validation failed for {_namespace} {sid}: csrf cookie mismatch"
328 - )
329 - return False
330 -
331 - user_id = session.get("user_id") or "single_user"
332 - await websocket_manager.handle_connect(_namespace, sid, user_id=user_id)
333 - return True
334 -
335 - @socketio_server.on("disconnect", namespace=namespace)
336 - async def _disconnect(sid, _namespace: str = namespace): # type: ignore[override]
337 - await websocket_manager.handle_disconnect(_namespace, sid)
338 -
339 - @socketio_server.on("*", namespace=namespace)
340 - async def _catch_all(event, sid, data, _namespace: str = namespace):
341 - payload = data or {}
342 - return await websocket_manager.route_event(_namespace, event, payload, sid)
343 -
344 - for namespace, namespace_handlers in namespace_map.items():
345 - _register_namespace_handlers(namespace, namespace_handlers)
346 -
347 - return allowed_namespaces
192
193
194 def run():
@@ -373,16 +217,7 @@ def run():
217
218 register_api_route(webapp, lock)
219
376 - handlers_by_namespace = _build_websocket_handlers_by_namespace(socketio_server, lock)
377 - allowed_namespaces = configure_websocket_namespaces(
378 - webapp=webapp,
379 - socketio_server=socketio_server,
380 - websocket_manager=websocket_manager,
381 - handlers_by_namespace=handlers_by_namespace,
382 - )
383 -
384 - register_ws_namespace(socketio_server, webapp, lock)
385 - allowed_namespaces.add("/ws")
220 + register_ws_namespace(socketio_server, webapp, lock, manager=ws_manager)
221
222 init_a0()
223
@@ -443,7 +278,7 @@ def run():
278
279
280 def wait_for_health(host: str, port: int):
446 - url = f"http://{host}:{port}/health"
281 + url = f"http://{host}:{port}/api/health"
282 while True:
283 try:
284 with urllib.request.urlopen(url, timeout=2) as resp: