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

api/ws_dev_test.py

@@ -0,0 +1,72 @@
+import asyncio
+from typing import Any
+
+from helpers.ws import WsHandler
+from helpers.print_style import PrintStyle
+from helpers import runtime
+
+
+class WsDevTest(WsHandler):
+    """Developer-only WebSocket test harness handler."""
+
+    async def process(self, event: str, data: dict, sid: str) -> dict[str, Any] | None:
+        if event == "ws_event_console_subscribe":
+            if not runtime.is_development():
+                return {"_error": True, "code": "NOT_AVAILABLE",
+                        "message": "Event console is available only in development mode"}
+            registered = self.manager.register_diagnostic_watcher(self.namespace, sid)
+            if not registered:
+                return {"_error": True, "code": "SUBSCRIBE_FAILED",
+                        "message": "Unable to subscribe to diagnostics"}
+            return {"status": "subscribed", "timestamp": data.get("requestedAt")}
+
+        if event == "ws_event_console_unsubscribe":
+            self.manager.unregister_diagnostic_watcher(self.namespace, sid)
+            return {"status": "unsubscribed"}
+
+        if event == "ws_tester_emit":
+            message = data.get("message", "emit")
+            payload = {"message": message, "echo": True, "timestamp": data.get("timestamp")}
+            await self.broadcast("ws_tester_broadcast", payload)
+            PrintStyle.info(f"Harness emit broadcasted message='{message}'")
+            return None
+
+        if event == "ws_tester_request":
+            value = data.get("value")
+            PrintStyle.debug("Harness request responded with echo %s", value)
+            return {"echo": value, "handler": self.identifier, "status": "ok"}
+
+        if event == "ws_tester_request_delayed":
+            delay_ms = int(data.get("delay_ms", 0))
+            await asyncio.sleep(delay_ms / 1000)
+            PrintStyle.warning("Harness delayed request finished after %s ms", delay_ms)
+            return {"status": "delayed", "delay_ms": delay_ms, "handler": self.identifier}
+
+        if event == "ws_tester_trigger_persistence":
+            phase = data.get("phase", "unknown")
+            payload = {"phase": phase, "handler": self.identifier}
+            await self.emit_to(sid, "ws_tester_persistence", payload)
+            PrintStyle.info(f"Harness persistence event phase='{phase}' -> {sid}")
+            return None
+
+        if event == "ws_tester_broadcast_demo_trigger":
+            payload = {"demo": True, "requested_at": data.get("requested_at")}
+            await self.broadcast("ws_tester_broadcast_demo", payload)
+            PrintStyle.info("Harness broadcast demo event dispatched")
+            return None
+
+        if event == "ws_tester_request_all":
+            correlation_id = data.get("correlationId")
+            aggregated = await self.dispatch_to_all_sids(
+                "ws_tester_request",
+                {"value": data.get("marker", "aggregate")},
+                correlation_id=correlation_id,
+            )
+            return {"results": aggregated}
+
+        # Ignore events not targeted at this handler (other activated handlers
+        # may process them).  Only warn for events that look like dev-harness
+        # traffic so we don't spam logs with unrelated events.
+        if event.startswith("ws_tester_"):
+            PrintStyle.warning(f"Harness received unknown event '{event}'")
+        return None

api/ws_hello.py

@@ -0,0 +1,13 @@
+from helpers.ws import WsHandler
+from helpers.print_style import PrintStyle
+
+
+class WsHello(WsHandler):
+    """Simple echo handler used for foundational testing."""
+
+    async def process(self, event: str, data: dict, sid: str) -> dict | None:
+        if event != "hello_request":
+            return None
+        name = data.get("name") or "stranger"
+        PrintStyle.info(f"hello_request from {sid} ({name})")
+        return {"message": f"Hello, {name}!", "handler": self.identifier}

api/ws_webui.py

@@ -0,0 +1,32 @@
+from helpers.ws import WsHandler
+from helpers import extension
+
+
+class WsWebui(WsHandler):
+    """State synchronisation handler — the primary WebSocket endpoint for the UI."""
+
+    async def on_connect(self, sid: str) -> None:
+        await extension.call_extensions_async(
+            "webui_ws_connect", agent=None, instance=self, sid=sid
+        )
+
+    async def on_disconnect(self, sid: str) -> None:
+        await extension.call_extensions_async(
+            "webui_ws_disconnect", agent=None, instance=self, sid=sid
+        )
+
+    async def process(self, event: str, data: dict, sid: str) -> dict | None:
+        response_data: dict = {}
+
+        await extension.call_extensions_async(
+            "webui_ws_event",
+            agent=None,
+            instance=self,
+            sid=sid,
+            event_type=event,
+            data=data,
+            response_data=response_data,
+        )
+
+        # Return None (fire-and-forget) when no extension populated the response.
+        return response_data if response_data else None

helpers/api.py

@@ -1,11 +1,9 @@
 from abc import abstractmethod
 import json
-import socket
-import struct
 import threading
 from functools import wraps
 from pathlib import Path
-from typing import Union, TypedDict, Dict, Any
+from typing import Union, Dict, Any
 from flask import (
     Request,
     Response,
@@ -19,7 +17,6 @@ from flask import (
 )
 from werkzeug.wrappers.response import Response as BaseResponse
 from agent import AgentContext
-from initialize import initialize_agent
 from helpers.print_style import PrintStyle
 from helpers.errors import format_error
 from helpers import files, cache
@@ -30,7 +27,7 @@ CACHE_AREA = "api_handlers(api)"
 cache.toggle_area(CACHE_AREA, False)  # cache off for now
 
 Input = dict
-Output = Union[Dict[str, Any], Response, TypedDict]  # type: ignore
+Output = Union[Dict[str, Any], Response]
 
 
 class ApiHandler:
@@ -99,59 +96,11 @@ class ApiHandler:
 
     # get context to run agent zero in
     def use_context(self, ctxid: str, create_if_not_exists: bool = True):
-        with self.thread_lock:
+        from helpers.context_utils import use_context as _use_context
-            if not ctxid:
+        return _use_context(self.thread_lock, ctxid, create_if_not_exists)
-                first = AgentContext.first()
-                if first:
-                    AgentContext.use(first.id)
-                    return first
-                context = AgentContext(config=initialize_agent(), set_current=True)
-                return context
-            got = AgentContext.use(ctxid)
-            if got:
-                return got
-            if create_if_not_exists:
-                context = AgentContext(
-                    config=initialize_agent(), id=ctxid, set_current=True
-                )
-                return context
-            else:
-                raise Exception(f"Context {ctxid} not found")
 
 
-def is_loopback_address(address: str) -> bool:
+from helpers.network import is_loopback_address
-    loopback_checker = {
-        socket.AF_INET: lambda x: (
-            struct.unpack("!I", socket.inet_aton(x))[0] >> (32 - 8)
-        )
-        == 127,
-        socket.AF_INET6: lambda x: x == "::1",
-    }
-    address_type = "hostname"
-    try:
-        socket.inet_pton(socket.AF_INET6, address)
-        address_type = "ipv6"
-    except socket.error:
-        try:
-            socket.inet_pton(socket.AF_INET, address)
-            address_type = "ipv4"
-        except socket.error:
-            address_type = "hostname"
-
-    if address_type == "ipv4":
-        return loopback_checker[socket.AF_INET](address)
-    elif address_type == "ipv6":
-        return loopback_checker[socket.AF_INET6](address)
-    else:
-        for family in (socket.AF_INET, socket.AF_INET6):
-            try:
-                r = socket.getaddrinfo(address, None, family, socket.SOCK_STREAM)
-            except socket.gaierror:
-                return False
-            for family, _, _, _, sockaddr in r:
-                if not loopback_checker[family](sockaddr[0]):
-                    return False
-        return True
 
 
 def requires_api_key(f):
@@ -301,3 +250,20 @@ def register_watchdogs():
         patterns=["*.py"],
         handler=on_api_change,
     )
+
+    # WS handler cache shares the same watched directories (api/, usr/api/)
+    from helpers.ws import CACHE_AREA as WS_CACHE_AREA
+
+    def on_ws_change(items: list[watchdog.WatchItem]):
+        PrintStyle.debug("WS handler watchdog triggered:", items)
+        cache.clear(WS_CACHE_AREA)
+
+    watchdog.add_watchdog(
+        "ws_handlers",
+        roots=[
+            files.get_abs_path(files.API_DIR),
+            files.get_abs_path(files.USER_DIR, files.API_DIR),
+        ],
+        patterns=["ws_*.py"],
+        handler=on_ws_change,
+    )

helpers/context_utils.py

@@ -0,0 +1,30 @@
+"""Shared context helper used by both ApiHandler and WsHandler."""
+
+import threading
+from typing import Union
+
+ThreadLockType = Union[threading.Lock, threading.RLock]
+
+
+def use_context(lock: ThreadLockType, ctxid: str, create_if_not_exists: bool = True):
+    from agent import AgentContext
+    from initialize import initialize_agent
+
+    with lock:
+        if not ctxid:
+            first = AgentContext.first()
+            if first:
+                AgentContext.use(first.id)
+                return first
+            context = AgentContext(config=initialize_agent(), set_current=True)
+            return context
+        got = AgentContext.use(ctxid)
+        if got:
+            return got
+        if create_if_not_exists:
+            context = AgentContext(
+                config=initialize_agent(), id=ctxid, set_current=True
+            )
+            return context
+        else:
+            raise Exception(f"Context {ctxid} not found")

helpers/network.py

@@ -0,0 +1,31 @@
+import socket
+import struct
+
+
+def is_loopback_address(address: str) -> bool:
+    """Check whether *address* resolves to a loopback interface."""
+    _checkers = {
+        socket.AF_INET: lambda x: (
+            struct.unpack("!I", socket.inet_aton(x))[0] >> (32 - 8)
+        ) == 127,
+        socket.AF_INET6: lambda x: x == "::1",
+    }
+    try:
+        socket.inet_pton(socket.AF_INET6, address)
+        return _checkers[socket.AF_INET6](address)
+    except socket.error:
+        pass
+    try:
+        socket.inet_pton(socket.AF_INET, address)
+        return _checkers[socket.AF_INET](address)
+    except socket.error:
+        pass
+    for family in (socket.AF_INET, socket.AF_INET6):
+        try:
+            r = socket.getaddrinfo(address, None, family, socket.SOCK_STREAM)
+        except socket.gaierror:
+            return False
+        for fam, _, _, _, sockaddr in r:
+            if not _checkers[fam](sockaddr[0]):
+                return False
+    return True

helpers/plugins.py

@@ -202,7 +202,7 @@ def clear_plugin_cache(plugin_names: list[str] | None = None):
     for area in areas:
         cache.clear(area)
 
-    from helpers.websocket_manager import send_data
+    from helpers.ws_manager import send_data
 
     DeferredTask().start_task(
         send_data,

helpers/state_monitor.py

@@ -14,10 +14,10 @@ from helpers.state_snapshot import (
     advance_state_request_after_snapshot,
     build_snapshot_from_request,
 )
-from helpers.websocket import ConnectionNotFoundError
+from helpers.ws import ConnectionNotFoundError
 
 if TYPE_CHECKING:  # pragma: no cover - hints only
-    from helpers.websocket_manager import WebSocketManager
+    from helpers.ws_manager import WsManager
 
 
 ConnectionIdentity = tuple[str, str]  # (namespace, sid)
@@ -60,12 +60,12 @@ class StateMonitor:
         self._projections: dict[ConnectionIdentity, ConnectionProjection] = {}
         self._debounce_handles: dict[ConnectionIdentity, asyncio.TimerHandle] = {}
         self._push_tasks: dict[ConnectionIdentity, asyncio.Task[None]] = {}
-        self._manager: WebSocketManager | None = None
+        self._manager: WsManager | None = None
         self._emit_handler_id: str | None = None
         self._dispatcher_loop: asyncio.AbstractEventLoop | None = None
         self._dirty_wave_seq: int = 0
 
-    def bind_manager(self, manager: "WebSocketManager", *, handler_id: str | None = None) -> None:
+    def bind_manager(self, manager: "WsManager", *, handler_id: str | None = None) -> None:
         with self._lock:
             self._manager = manager
             if handler_id:

helpers/ws.py

@@ -1,22 +1,148 @@
 import threading
+import uuid
 from abc import abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Union
+from typing import Any, Iterable, Union, TYPE_CHECKING
+from urllib.parse import urlparse
 
 import socketio
 from flask import Flask, session, request
 
-from agent import AgentContext
-from initialize import initialize_agent
 from helpers import files, cache
-from helpers.api import is_loopback_address
 from helpers.print_style import PrintStyle
 from helpers.errors import format_error
-from helpers.websocket import validate_ws_origin
+
+if TYPE_CHECKING:
+    from helpers.ws_manager import WsManager
+
+
+# Utilities
+
+from helpers.network import is_loopback_address
+
+
+class ConnectionNotFoundError(RuntimeError):
+    """Raised when attempting to emit to a non-existent WebSocket connection."""
+
+    def __init__(self, sid: str, *, namespace: str | None = None) -> None:
+        self.sid = sid
+        self.namespace = namespace
+        if namespace:
+            super().__init__(f"Connection not found: namespace={namespace} sid={sid}")
+        else:
+            super().__init__(f"Connection not found: {sid}")
+
+
+def _default_port_for_scheme(scheme: str) -> int | None:
+    if scheme == "http":
+        return 80
+    if scheme == "https":
+        return 443
+    return None
+
+
+def normalize_origin(value: Any) -> str | None:
+    """Normalize an Origin/Referer header value to scheme://host[:port]."""
+    if not isinstance(value, str) or not value.strip():
+        return None
+    parsed = urlparse(value.strip())
+    if not parsed.scheme or not parsed.hostname:
+        return None
+    origin = f"{parsed.scheme}://{parsed.hostname}"
+    if parsed.port:
+        origin += f":{parsed.port}"
+    return origin
+
+
+def _parse_host_header(value: Any) -> tuple[str | None, int | None]:
+    if not isinstance(value, str) or not value.strip():
+        return None, None
+    parsed = urlparse(f"http://{value.strip()}")
+    return parsed.hostname, parsed.port
+
+
+def validate_ws_origin(environ: dict[str, Any]) -> tuple[bool, str | None]:
+    """Validate the browser Origin during the Socket.IO handshake.
+
+    This is the minimum baseline recommended by RFC 6455 (Origin considerations)
+    and OWASP (CSWSH mitigation): reject cross-origin WebSocket handshakes when
+    the server is intended for a specific web UI origin.
+    """
+
+    raw_origin = environ.get("HTTP_ORIGIN") or environ.get("HTTP_REFERER")
+    origin = normalize_origin(raw_origin)
+    if origin is None:
+        return False, "missing_origin"
+
+    origin_parsed = urlparse(origin)
+    origin_host = origin_parsed.hostname.lower() if origin_parsed.hostname else None
+    origin_port = origin_parsed.port or _default_port_for_scheme(origin_parsed.scheme)
+    if origin_host is None or origin_port is None:
+        return False, "invalid_origin"
+
+    raw_host = environ.get("HTTP_HOST")
+    req_host, req_port = _parse_host_header(raw_host)
+    if not req_host:
+        req_host = environ.get("SERVER_NAME")
+
+    if req_port is None:
+        server_port_raw = environ.get("SERVER_PORT")
+        try:
+            server_port = int(server_port_raw) if server_port_raw is not None else None
+        except (TypeError, ValueError):
+            server_port = None
+        if server_port is not None and server_port > 0:
+            req_port = server_port
+
+    if req_host:
+        req_host = req_host.lower()
+    if req_port is None:
+        req_port = origin_port
+
+    forwarded_host_raw = environ.get("HTTP_X_FORWARDED_HOST")
+    forwarded_host = None
+    forwarded_port = None
+    if isinstance(forwarded_host_raw, str) and forwarded_host_raw.strip():
+        first = forwarded_host_raw.split(",")[0].strip()
+        forwarded_host, forwarded_port = _parse_host_header(first)
+        if forwarded_host:
+            forwarded_host = forwarded_host.lower()
+
+    forwarded_proto_raw = environ.get("HTTP_X_FORWARDED_PROTO")
+    forwarded_scheme = None
+    if isinstance(forwarded_proto_raw, str) and forwarded_proto_raw.strip():
+        forwarded_scheme = forwarded_proto_raw.split(",")[0].strip().lower()
+    forwarded_scheme = forwarded_scheme or origin_parsed.scheme
+    forwarded_port = (
+        forwarded_port
+        if forwarded_port is not None
+        else _default_port_for_scheme(forwarded_scheme) or origin_port
+    )
+
+    candidates: list[tuple[str, int]] = []
+    if req_host:
+        candidates.append((req_host, int(req_port)))
+    if forwarded_host:
+        candidates.append((forwarded_host, int(forwarded_port)))
+
+    if not candidates:
+        return False, "missing_host"
+
+    for host, port in candidates:
+        if origin_host == host and origin_port == port:
+            return True, None
+
+    if origin_host not in {host for host, _ in candidates}:
+        return False, "origin_host_mismatch"
+    return False, "origin_port_mismatch"
+
+
+# Constants
 
 ThreadLockType = Union[threading.Lock, threading.RLock]
 
+NAMESPACE = "/ws"
 CACHE_AREA = "ws_handlers(api)(plugins)"
 cache.toggle_area(CACHE_AREA, False)  # cache off for now
 
@@ -37,9 +163,52 @@ _contexts_lock = threading.Lock()
 
 
 class WsHandler:
-    def __init__(self, socketio_server: socketio.AsyncServer, lock: ThreadLockType):
+    """Base class for WebSocket handlers loaded from api/ directories.
+
+    Mirrors ApiHandler conventions: declarative security flags, dynamic file-
+    based loading, and a ``process(event, data, sid)`` entry point.  Handlers
+    are activated per-connection based on the ``auth.handlers`` list sent by the
+    client during the Socket.IO connect handshake.
+    """
+
+    def __init__(
+        self,
+        socketio_server: socketio.AsyncServer,
+        lock: ThreadLockType,
+        *,
+        manager: "WsManager | None" = None,
+        namespace: str = NAMESPACE,
+    ):
         self.socketio = socketio_server
         self.lock = lock
+        self._manager = manager
+        self._namespace = namespace
+
+    # Properties
+
+    @property
+    def namespace(self) -> str:
+        return self._namespace
+
+    @property
+    def manager(self) -> "WsManager":
+        if self._manager is None:
+            raise RuntimeError("WsHandler has no WsManager bound")
+        return self._manager
+
+    @property
+    def identifier(self) -> str:
+        return f"{self.__class__.__module__}.{self.__class__.__name__}"
+
+    def bind_manager(
+        self, manager: "WsManager", *, namespace: str | None = None
+    ) -> None:
+        """Late-bind (or rebind) the manager and optionally the namespace."""
+        self._manager = manager
+        if namespace is not None:
+            self._namespace = namespace
+
+    # Security flags (mirror ApiHandler)
 
     @classmethod
     def requires_loopback(cls) -> bool:
@@ -57,67 +226,164 @@ class WsHandler:
     def requires_csrf(cls) -> bool:
         return cls.requires_auth()
 
-    @abstractmethod
+    # Lifecycle hooks
-    async def process(self, data: dict, sid: str) -> dict | None:
-        pass
 
-    async def on_connect(self, sid: str) -> dict | None:
+    async def on_connect(self, sid: str) -> None:
-        return None
+        pass
 
     async def on_disconnect(self, sid: str) -> None:
         pass
 
-    def use_context(self, ctxid: str, create_if_not_exists: bool = True):
+    # Event processing
-        with self.lock:
-            if not ctxid:
-                first = AgentContext.first()
-                if first:
-                    AgentContext.use(first.id)
-                    return first
-                context = AgentContext(config=initialize_agent(), set_current=True)
-                return context
-            got = AgentContext.use(ctxid)
-            if got:
-                return got
-            if create_if_not_exists:
-                context = AgentContext(config=initialize_agent(), id=ctxid, set_current=True)
-                return context
-            else:
-                raise Exception(f"Context {ctxid} not found")
 
+    @abstractmethod
+    async def process(self, event: str, data: dict, sid: str) -> dict | None:
+        """Handle an incoming event.
+
+        Return a dict to include in the acknowledgement, or ``None`` for
+        fire-and-forget semantics.
+        """
+
+    # Emit helpers (delegate to WsManager for envelope wrapping)
+
+    async def emit_to(
+        self,
+        sid: str,
+        event: str,
+        data: dict,
+        *,
+        correlation_id: str | None = None,
+    ) -> None:
+        await self.manager.emit_to(
+            self._namespace, sid, event, data,
+            handler_id=self.identifier,
+            correlation_id=correlation_id,
+        )
+
+    async def broadcast(
+        self,
+        event: str,
+        data: dict,
+        *,
+        exclude_sids: str | Iterable[str] | None = None,
+        correlation_id: str | None = None,
+    ) -> None:
+        await self.manager.broadcast(
+            self._namespace, event, data,
+            exclude_sids=exclude_sids,
+            handler_id=self.identifier,
+            correlation_id=correlation_id,
+        )
+
+    # Aggregation helper
+
+    async def dispatch_to_all_sids(
+        self,
+        event: str,
+        data: dict,
+        *,
+        correlation_id: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Dispatch *event* to every connected sid's activated handlers and
+        aggregate the results.
+
+        Returns a list of ``{sid, correlationId, results}`` dicts – one per
+        connected sid.  This mirrors the shape produced by
+        ``WsManager.route_event_all`` so that existing frontend
+        assertions remain valid.
+        """
+        cid = correlation_id or uuid.uuid4().hex
+
+        with _contexts_lock:
+            snapshot = {
+                sid: dict(handlers)
+                for sid, handlers in _active_handlers.items()
+            }
+
+        aggregated: list[dict[str, Any]] = []
+        for sid, handlers in snapshot.items():
+            ctx = _ws_contexts.get(sid)
+            sid_results: list[dict[str, Any]] = []
+            for _path, instance in handlers.items():
+                if ctx is not None:
+                    error = _check_security(type(instance), ctx)
+                    if error is not None:
+                        sid_results.append({
+                            "handlerId": instance.identifier,
+                            "ok": False,
+                            "correlationId": cid,
+                            "error": error,
+                        })
+                        continue
+                try:
+                    result = await instance.process(event, dict(data, correlationId=cid), sid)
+                    if result is not None:
+                        sid_results.append({
+                            "handlerId": instance.identifier,
+                            "ok": True,
+                            "correlationId": cid,
+                            "data": result,
+                        })
+                except Exception as e:
+                    sid_results.append({
+                        "handlerId": instance.identifier,
+                        "ok": False,
+                        "correlationId": cid,
+                        "error": {"code": "HANDLER_ERROR", "error": str(e)},
+                    })
+            aggregated.append({
+                "sid": sid,
+                "correlationId": cid,
+                "results": sid_results,
+            })
+        return aggregated
+
+    # Context helper (shared with ApiHandler)
+
+    def use_context(self, ctxid: str, create_if_not_exists: bool = True):
+        from helpers.context_utils import use_context as _use_context
+        return _use_context(self.lock, ctxid, create_if_not_exists)
+
+
+# Security check (aligned with api.py decorators)
 
 def _check_security(handler_cls: type[WsHandler], ctx: _SecurityContext) -> dict[str, Any] | None:
+    """Return an error payload dict if the check fails, or ``None`` on success."""
+
     if handler_cls.requires_loopback():
         if not ctx.remote_addr or not is_loopback_address(ctx.remote_addr):
-            return {"ok": False, "error": "Access denied", "code": 403}
+            return {"code": "FORBIDDEN", "error": "Access denied"}
 
     if handler_cls.requires_auth():
         from helpers import login
         user_pass_hash = login.get_credentials_hash()
         if user_pass_hash and ctx.auth_hash != user_pass_hash:
-            return {"ok": False, "error": "Authentication required", "code": 401}
+            return {"code": "AUTH_REQUIRED", "error": "Authentication required"}
 
     if handler_cls.requires_csrf():
         if not ctx.csrf_token:
-            return {"ok": False, "error": "CSRF token not initialised", "code": 403}
+            return {"code": "CSRF_MISSING", "error": "CSRF token not initialised"}
         if not ctx.client_csrf_token or ctx.client_csrf_token != ctx.csrf_token:
-            return {"ok": False, "error": "CSRF token missing or invalid", "code": 403}
+            return {"code": "CSRF_INVALID", "error": "CSRF token missing or invalid"}
         if ctx.csrf_cookie != ctx.csrf_token:
-            return {"ok": False, "error": "CSRF cookie mismatch", "code": 403}
+            return {"code": "CSRF_COOKIE", "error": "CSRF cookie mismatch"}
 
     if handler_cls.requires_api_key():
         from helpers.settings import get_settings
         valid_key = get_settings().get("mcp_server_token")
         if not ctx.api_key or ctx.api_key != valid_key:
-            return {"ok": False, "error": "API key required", "code": 401}
+            return {"code": "API_KEY_REQUIRED", "error": "API key required"}
 
     return None
 
 
+# Namespace registration
+
 def register_ws_namespace(
     socketio_server: socketio.AsyncServer,
     webapp: Flask,
     lock: ThreadLockType,
+    manager: "WsManager | None" = None,
 ) -> None:
     from helpers.modules import load_classes_from_file
     from helpers import plugins, runtime
@@ -132,6 +398,14 @@ def register_ws_namespace(
             if classes:
                 handler_cls = classes[0]
 
+        # Check user api/<path>.py
+        if handler_cls is None:
+            user_file = files.get_abs_path(files.USER_DIR, f"api/{path}.py")
+            if files.exists(user_file):
+                classes = load_classes_from_file(user_file, WsHandler)
+                if classes:
+                    handler_cls = classes[0]
+
         # Check plugin api/<handler>.py — path format: plugins/<plugin_name>/<handler>
         if handler_cls is None and path.startswith("plugins/"):
             parts = path.split("/", 2)
@@ -156,13 +430,13 @@ def register_ws_namespace(
             cache.add(CACHE_AREA, path, handler_cls)
         return handler_cls
 
-    @socketio_server.on("connect", namespace="/ws") # type: ignore
+    @socketio_server.on("connect", namespace=NAMESPACE)  # type: ignore
     async def _on_connect(sid, environ, auth):
         with webapp.request_context(environ):
             origin_ok, origin_reason = validate_ws_origin(environ)
             if not origin_ok:
                 PrintStyle.warning(
-                    f"WS /ws connect rejected for {sid}: {origin_reason or 'invalid'}"
+                    f"WS connect rejected for {sid}: {origin_reason or 'invalid'}"
                 )
                 return False
 
@@ -182,11 +456,19 @@ def register_ws_namespace(
                     if isinstance(auth, dict) else None
                 ),
             )
+            user_id = session.get("user_id") or "single_user"
+
             with _contexts_lock:
                 _ws_contexts[sid] = ctx
 
+        # Register with WsManager first so that the dispatcher loop and
+        # connection tracking are available before handler on_connect runs
+        # (extensions like StateSync depend on manager._dispatcher_loop).
+        if manager is not None:
+            await manager.handle_connect(NAMESPACE, sid, user_id=user_id)
+
         # Activate handlers declared in auth.handlers
-        handler_paths = []
+        handler_paths: list[str] = []
         if isinstance(auth, dict):
             raw = auth.get("handlers")
             if isinstance(raw, list):
@@ -201,7 +483,10 @@ def register_ws_namespace(
                 error = _check_security(handler_cls, ctx)
                 if error is not None:
                     continue
-                instance = handler_cls(socketio_server, lock)
+                instance = handler_cls(
+                    socketio_server, lock,
+                    manager=manager, namespace=NAMESPACE,
+                )
                 await instance.on_connect(sid)
                 activated[path] = instance
             except Exception as e:
@@ -212,7 +497,7 @@ def register_ws_namespace(
 
         return True
 
-    @socketio_server.on("disconnect", namespace="/ws") # type: ignore
+    @socketio_server.on("disconnect", namespace=NAMESPACE)  # type: ignore
     async def _on_disconnect(sid):
         with _contexts_lock:
             activated = _active_handlers.pop(sid, {})
@@ -224,31 +509,83 @@ def register_ws_namespace(
             except Exception as e:
                 PrintStyle.error(f"WS on_disconnect error ({path}): {format_error(e)}")
 
-    @socketio_server.on("*", namespace="/ws") # type: ignore
+        if manager is not None:
+            await manager.handle_disconnect(NAMESPACE, sid)
+
+    @socketio_server.on("*", namespace=NAMESPACE)  # type: ignore
     async def _dispatch(event, sid, data):
-        path = event
+        incoming = data if isinstance(data, dict) else {}
-        payload = data if isinstance(data, dict) else {}
 
         try:
             with _contexts_lock:
                 ctx = _ws_contexts.get(sid)
-                activated = _active_handlers.get(sid, {})
+                activated = dict(_active_handlers.get(sid, {}))
+
+            correlation_id = incoming.get("correlationId") or uuid.uuid4().hex
+
             if ctx is None:
-                return {"ok": False, "error": "No security context", "code": 401}
+                return _error_response("AUTH_REQUIRED",
+                                       "No security context", correlation_id)
+            if not activated:
+                return _error_response("NO_HANDLERS",
+                                       "No handlers activated", correlation_id)
 
-            instance = activated.get(path)
+            # Unwrap nested payload (mirrors WsManager.route_event):
-            if instance is None:
+            # frontend sends {ts, data: {actual fields...}, correlationId}
-                return {"ok": False, "error": f"WS endpoint not activated: {path}", "code": 404}
+            if "data" in incoming and isinstance(incoming.get("data"), dict):
+                handler_payload = dict(incoming["data"])
+            else:
+                handler_payload = dict(incoming)
+            handler_payload["correlationId"] = correlation_id
 
-            # Security check
+            results: list[dict[str, Any]] = []
-            error = _check_security(type(instance), ctx)
+            for path, instance in activated.items():
-            if error is not None:
+                error = _check_security(type(instance), ctx)
-                return error
+                if error is not None:
+                    results.append({
+                        "handlerId": instance.identifier,
+                        "ok": False,
+                        "correlationId": correlation_id,
+                        "error": error,
+                    })
+                    continue
+                try:
+                    result = await instance.process(event, handler_payload, sid)
+                    if result is not None:
+                        results.append({
+                            "handlerId": instance.identifier,
+                            "ok": True,
+                            "correlationId": correlation_id,
+                            "data": result,
+                        })
+                except Exception as e:
+                    error_text = format_error(e)
+                    PrintStyle.error(f"WS handler error ({path}/{event}): {error_text}")
+                    results.append({
+                        "handlerId": instance.identifier,
+                        "ok": False,
+                        "correlationId": correlation_id,
+                        "error": {"code": "HANDLER_ERROR", "error": "Internal server error"},
+                    })
 
-            # Use cached instance and process
+            return {"correlationId": correlation_id, "results": results}
-            return await instance.process(payload, sid)
 
         except Exception as e:
             error_text = format_error(e)
-            PrintStyle.error(f"WS handler error ({path}): {error_text}")
+            PrintStyle.error(f"WS dispatch error ({event}): {error_text}")
-            return {"ok": False, "error": error_text, "code": 500}
+            return _error_response(
+                "INTERNAL_ERROR", "Internal server error",
+                incoming.get("correlationId", ""),
+            )
+
+
+def _error_response(code: str, message: str,
+                    correlation_id: str) -> dict[str, Any]:
+    return {
+        "correlationId": correlation_id,
+        "results": [{
+            "handlerId": "ws.dispatch",
+            "ok": False,
+            "error": {"code": code, "error": message},
+        }],
+    }

run_ui.py

@@ -14,9 +14,8 @@ import initialize
 from helpers import files, git, mcp_server, fasta2a_server, settings as settings_helper, extension
 from helpers.files import get_abs_path
 from helpers import runtime, dotenv, process
-from helpers.websocket import WebSocketHandler, validate_ws_origin
 from helpers.api import register_api_route, requires_auth, csrf_protect
-from helpers.ws import register_ws_namespace
+from helpers.ws import register_ws_namespace, validate_ws_origin
 from helpers.print_style import PrintStyle
 from helpers import login
 import socketio  # type: ignore[import-untyped]
@@ -24,8 +23,7 @@ from socketio import ASGIApp, packet
 from starlette.applications import Starlette
 from starlette.routing import Mount
 from uvicorn.middleware.wsgi import WSGIMiddleware
-from helpers.websocket_manager import WebSocketManager, set_shared_websocket_manager
+from helpers.ws_manager import WsManager, set_shared_ws_manager
-from helpers.websocket_namespace_discovery import discover_websocket_namespaces
 from flask import send_file
 
 # disable logging
@@ -73,11 +71,11 @@ socketio_server = socketio.AsyncServer(
     max_http_buffer_size=50 * 1024 * 1024,
 )
 
-websocket_manager = WebSocketManager(socketio_server, lock)
+ws_manager = WsManager(socketio_server, lock)
-set_shared_websocket_manager(websocket_manager)
+set_shared_ws_manager(ws_manager)
 _settings = settings_helper.get_settings()
 settings_helper.set_runtime_settings_snapshot(_settings)
-websocket_manager.set_server_restart_broadcast(
+ws_manager.set_server_restart_broadcast(
     _settings.get("websocket_server_restart_enabled", True)
 )
 
@@ -191,160 +189,6 @@ async def _serve_plugin_asset(plugin_name, asset_path):
         return Response("Error serving asset", 500)
 
 
-def _build_websocket_handlers_by_namespace(
-    socketio_server: socketio.AsyncServer,
-    lock: threading.RLock,
-) -> dict[str, list[WebSocketHandler]]:
-    discoveries = discover_websocket_namespaces(
-        handlers_folder="python/websocket_handlers",
-        include_root_default=True,
-    )
-
-    handlers_by_namespace: dict[str, list[WebSocketHandler]] = {}
-    for discovery in discoveries:
-        namespace = discovery.namespace
-        for handler_cls in discovery.handler_classes:
-            handler = handler_cls.get_instance(socketio_server, lock)
-            handlers_by_namespace.setdefault(namespace, []).append(handler)
-
-    return handlers_by_namespace
-
-
-def configure_websocket_namespaces(
-    *,
-    webapp: Flask,
-    socketio_server: socketio.AsyncServer,
-    websocket_manager: WebSocketManager,
-    handlers_by_namespace: dict[str, list[WebSocketHandler]],
-) -> set[str]:
-    namespace_map: dict[str, list[WebSocketHandler]] = {
-        namespace: list(handlers) for namespace, handlers in handlers_by_namespace.items()
-    }
-
-    # Always include the reserved root namespace. It is unhandled for application events by
-    # default, but request-style calls must resolve deterministically with NO_HANDLERS.
-    namespace_map.setdefault("/", [])
-
-    websocket_manager.register_handlers(namespace_map)
-
-    allowed_namespaces = set(namespace_map.keys())
-    original_handle_connect = socketio_server._handle_connect  # type: ignore[attr-defined]
-
-    async def _handle_connect_with_namespace_gatekeeper(eio_sid, namespace, data):
-        requested = namespace or "/"
-        if requested not in allowed_namespaces:
-            await socketio_server._send_packet(
-                eio_sid,
-                socketio_server.packet_class(
-                    packet.CONNECT_ERROR,
-                    data={
-                        "message": "UNKNOWN_NAMESPACE",
-                        "data": {"code": "UNKNOWN_NAMESPACE", "namespace": requested},
-                    },
-                    namespace=requested,
-                ),
-            )
-            return
-        await original_handle_connect(eio_sid, namespace, data)
-
-    socketio_server._handle_connect = _handle_connect_with_namespace_gatekeeper  # type: ignore[assignment]
-
-    def _register_namespace_handlers(
-        namespace: str, namespace_handlers: list[WebSocketHandler]
-    ) -> None:
-        # A namespace is the WebSocket equivalent of an API endpoint.
-        # Security requirements must be consistent within the namespace (no any()-based union).
-        auth_required = False
-        csrf_required = False
-        if namespace_handlers:
-            auth_required = bool(namespace_handlers[0].requires_auth())
-            csrf_required = bool(namespace_handlers[0].requires_csrf())
-            for handler in namespace_handlers[1:]:
-                if (
-                    bool(handler.requires_auth()) != auth_required
-                    or bool(handler.requires_csrf()) != csrf_required
-                ):
-                    raise ValueError(
-                        f"WebSocket namespace {namespace!r} has mixed auth/csrf requirements across handlers"
-                    )
-
-        @socketio_server.on("connect", namespace=namespace)
-        async def _connect(  # type: ignore[override]
-            sid,
-            environ,
-            _auth,
-            _namespace: str = namespace,
-            _auth_required: bool = auth_required,
-            _csrf_required: bool = csrf_required,
-        ):
-            with webapp.request_context(environ):
-                origin_ok, origin_reason = validate_ws_origin(environ)
-                if not origin_ok:
-                    PrintStyle.warning(
-                        f"WebSocket origin validation failed for {_namespace} {sid}: {origin_reason or 'invalid'}"
-                    )
-                    return False
-
-                if _auth_required:
-                    credentials_hash = login.get_credentials_hash()
-                    if credentials_hash:
-                        if session.get("authentication") != credentials_hash:
-                            PrintStyle.warning(
-                                f"WebSocket authentication failed for {_namespace} {sid}: session not valid"
-                            )
-                            return False
-                    else:
-                        PrintStyle.debug(
-                            "WebSocket authentication required but credentials not configured; proceeding"
-                        )
-
-                if _csrf_required:
-                    expected_token = session.get("csrf_token")
-                    if not isinstance(expected_token, str) or not expected_token:
-                        PrintStyle.warning(
-                            f"WebSocket CSRF validation failed for {_namespace} {sid}: csrf_token not initialized"
-                        )
-                        return False
-
-                    auth_token = None
-                    if isinstance(_auth, dict):
-                        auth_token = _auth.get("csrf_token") or _auth.get("csrfToken")
-                    if not isinstance(auth_token, str) or not auth_token:
-                        PrintStyle.warning(
-                            f"WebSocket CSRF validation failed for {_namespace} {sid}: missing csrf_token in auth"
-                        )
-                        return False
-                    if auth_token != expected_token:
-                        PrintStyle.warning(
-                            f"WebSocket CSRF validation failed for {_namespace} {sid}: csrf_token mismatch"
-                        )
-                        return False
-
-                    cookie_name = f"csrf_token_{runtime.get_runtime_id()}"
-                    cookie_token = request.cookies.get(cookie_name)
-                    if cookie_token != expected_token:
-                        PrintStyle.warning(
-                            f"WebSocket CSRF validation failed for {_namespace} {sid}: csrf cookie mismatch"
-                        )
-                        return False
-
-                user_id = session.get("user_id") or "single_user"
-                await websocket_manager.handle_connect(_namespace, sid, user_id=user_id)
-                return True
-
-        @socketio_server.on("disconnect", namespace=namespace)
-        async def _disconnect(sid, _namespace: str = namespace):  # type: ignore[override]
-            await websocket_manager.handle_disconnect(_namespace, sid)
-
-        @socketio_server.on("*", namespace=namespace)
-        async def _catch_all(event, sid, data, _namespace: str = namespace):
-            payload = data or {}
-            return await websocket_manager.route_event(_namespace, event, payload, sid)
-
-    for namespace, namespace_handlers in namespace_map.items():
-        _register_namespace_handlers(namespace, namespace_handlers)
-
-    return allowed_namespaces
 
 
 def run():
@@ -373,16 +217,7 @@ def run():
 
     register_api_route(webapp, lock)
 
-    handlers_by_namespace = _build_websocket_handlers_by_namespace(socketio_server, lock)
+    register_ws_namespace(socketio_server, webapp, lock, manager=ws_manager)
-    allowed_namespaces = configure_websocket_namespaces(
-        webapp=webapp,
-        socketio_server=socketio_server,
-        websocket_manager=websocket_manager,
-        handlers_by_namespace=handlers_by_namespace,
-    )
-
-    register_ws_namespace(socketio_server, webapp, lock)
-    allowed_namespaces.add("/ws")
 
     init_a0()
 
@@ -443,7 +278,7 @@ def run():
 
 
 def wait_for_health(host: str, port: int):
-    url = f"http://{host}:{port}/health"
+    url = f"http://{host}:{port}/api/health"
     while True:
         try:
             with urllib.request.urlopen(url, timeout=2) as resp: