refactor: websocket dynamic endpoints
keyboardstaff committed
Mar 16, 2026 at 07:37 UTC
28002047e65959c3142d8b69837810e5ca522b5d
2 files changed
+217
-1
helpers/ws.py
new
+212
@@ -0,0 +1,212 @@
1
+import threading
2
+from abc import abstractmethod
3
+from dataclasses import dataclass
4
+from pathlib import Path
5
+from typing import Any, Union
6
+
7
+import socketio
8
+from flask import Flask, session, request
9
+
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
14
+from helpers.print_style import PrintStyle
15
+from helpers.errors import format_error
16
+from helpers.websocket import validate_ws_origin
17
+
18
+ThreadLockType = Union[threading.Lock, threading.RLock]
19
+
20
+CACHE_AREA = "ws_handlers(ws)(plugins)"
21
+cache.toggle_area(CACHE_AREA, False) # cache off for now
22
+
23
+
24
+@dataclass
25
+class _SecurityContext:
26
+ auth_hash: str | None
27
+ csrf_token: str | None
28
+ client_csrf_token: str | None
29
+ csrf_cookie: str | None
30
+ remote_addr: str | None
31
+ api_key: str | None
32
+
33
+
34
+_ws_contexts: dict[str, _SecurityContext] = {}
35
+_contexts_lock = threading.Lock()
36
+
37
+
38
+class WsHandler:
39
+ def __init__(self, socketio_server: socketio.AsyncServer, lock: ThreadLockType):
40
+ self.socketio = socketio_server
41
+ self.lock = lock
42
+
43
+ @classmethod
44
+ def requires_loopback(cls) -> bool:
45
+ return False
46
+
47
+ @classmethod
48
+ def requires_api_key(cls) -> bool:
49
+ return False
50
+
51
+ @classmethod
52
+ def requires_auth(cls) -> bool:
53
+ return True
54
+
55
+ @classmethod
56
+ def requires_csrf(cls) -> bool:
57
+ return cls.requires_auth()
58
+
59
+ @abstractmethod
60
+ async def process(self, data: dict, sid: str) -> dict | None:
61
+ pass
62
+
63
+ def use_context(self, ctxid: str, create_if_not_exists: bool = True):
64
+ with self.lock:
65
+ if not ctxid:
66
+ first = AgentContext.first()
67
+ if first:
68
+ AgentContext.use(first.id)
69
+ return first
70
+ context = AgentContext(config=initialize_agent(), set_current=True)
71
+ return context
72
+ got = AgentContext.use(ctxid)
73
+ if got:
74
+ return got
75
+ if create_if_not_exists:
76
+ context = AgentContext(config=initialize_agent(), id=ctxid, set_current=True)
77
+ return context
78
+ else:
79
+ raise Exception(f"Context {ctxid} not found")
80
+
81
+
82
+def _check_security(handler_cls: type[WsHandler], ctx: _SecurityContext) -> dict[str, Any] | None:
83
+ if handler_cls.requires_loopback():
84
+ if not ctx.remote_addr or not is_loopback_address(ctx.remote_addr):
85
+ return {"ok": False, "error": "Access denied", "code": 403}
86
+
87
+ if handler_cls.requires_auth():
88
+ from helpers import login
89
+ user_pass_hash = login.get_credentials_hash()
90
+ if user_pass_hash and ctx.auth_hash != user_pass_hash:
91
+ return {"ok": False, "error": "Authentication required", "code": 401}
92
+
93
+ if handler_cls.requires_csrf():
94
+ if not ctx.csrf_token:
95
+ return {"ok": False, "error": "CSRF token not initialised", "code": 403}
96
+ if not ctx.client_csrf_token or ctx.client_csrf_token != ctx.csrf_token:
97
+ return {"ok": False, "error": "CSRF token missing or invalid", "code": 403}
98
+ if ctx.csrf_cookie != ctx.csrf_token:
99
+ return {"ok": False, "error": "CSRF cookie mismatch", "code": 403}
100
+
101
+ if handler_cls.requires_api_key():
102
+ from helpers.settings import get_settings
103
+ valid_key = get_settings().get("mcp_server_token")
104
+ if not ctx.api_key or ctx.api_key != valid_key:
105
+ return {"ok": False, "error": "API key required", "code": 401}
106
+
107
+ return None
108
+
109
+
110
+def register_ws_namespace(
111
+ socketio_server: socketio.AsyncServer,
112
+ webapp: Flask,
113
+ lock: ThreadLockType,
114
+) -> None:
115
+ from helpers.extract_tools import load_classes_from_file
116
+ from helpers import plugins, runtime
117
+
118
+ def _resolve_handler(path: str) -> type[WsHandler] | None:
119
+ handler_cls: type[WsHandler] | None = None
120
+
121
+ # Built-in ws/<path>.py
122
+ builtin_file = files.get_abs_path(f"ws/{path}.py")
123
+ if files.is_in_dir(builtin_file, files.get_abs_path("ws")) and files.exists(builtin_file):
124
+ classes = load_classes_from_file(builtin_file, WsHandler)
125
+ if classes:
126
+ handler_cls = classes[0]
127
+
128
+ # Plugin ws/<handler>.py — path format: plugins/<plugin_name>/<handler>
129
+ if handler_cls is None and path.startswith("plugins/"):
130
+ parts = path.split("/", 2)
131
+ if len(parts) == 3:
132
+ _, plugin_name, handler_name = parts
133
+ plugin_dir = plugins.find_plugin_dir(plugin_name)
134
+ if plugin_dir:
135
+ plugin_file = Path(plugin_dir) / "ws" / f"{handler_name}.py"
136
+ if plugin_file.is_file():
137
+ classes = load_classes_from_file(str(plugin_file), WsHandler)
138
+ if classes:
139
+ handler_cls = classes[0]
140
+
141
+ return handler_cls
142
+
143
+ @socketio_server.on("connect", namespace="/ws")
144
+ async def _on_connect(sid, environ, auth):
145
+ with webapp.request_context(environ):
146
+ origin_ok, origin_reason = validate_ws_origin(environ)
147
+ if not origin_ok:
148
+ PrintStyle.warning(
149
+ f"WS /ws connect rejected for {sid}: {origin_reason or 'invalid'}"
150
+ )
151
+ return False
152
+
153
+ ctx = _SecurityContext(
154
+ auth_hash=session.get("authentication"),
155
+ csrf_token=session.get("csrf_token"),
156
+ client_csrf_token=(
157
+ (auth.get("csrf_token") or auth.get("csrfToken"))
158
+ if isinstance(auth, dict) else None
159
+ ),
160
+ csrf_cookie=request.cookies.get(
161
+ f"csrf_token_{runtime.get_runtime_id()}"
162
+ ),
163
+ remote_addr=str(request.remote_addr) if request.remote_addr else None,
164
+ api_key=(
165
+ (auth.get("api_key") or auth.get("apiKey"))
166
+ if isinstance(auth, dict) else None
167
+ ),
168
+ )
169
+ with _contexts_lock:
170
+ _ws_contexts[sid] = ctx
171
+
172
+ return True
173
+
174
+ @socketio_server.on("disconnect", namespace="/ws")
175
+ async def _on_disconnect(sid):
176
+ with _contexts_lock:
177
+ _ws_contexts.pop(sid, None)
178
+
179
+ @socketio_server.on("*", namespace="/ws")
180
+ async def _dispatch(event, sid, data):
181
+ path = event
182
+ payload = data if isinstance(data, dict) else {}
183
+
184
+ try:
185
+ with _contexts_lock:
186
+ ctx = _ws_contexts.get(sid)
187
+ if ctx is None:
188
+ return {"ok": False, "error": "No security context", "code": 401}
189
+
190
+ # Resolve handler (cache-aware)
191
+ cached = cache.get(CACHE_AREA, path)
192
+ if cached is not None:
193
+ handler_cls = cached
194
+ else:
195
+ handler_cls = _resolve_handler(path)
196
+ if handler_cls is None:
197
+ return {"ok": False, "error": f"WS endpoint not found: {path}", "code": 404}
198
+ cache.add(CACHE_AREA, path, handler_cls)
199
+
200
+ # Security check
201
+ error = _check_security(handler_cls, ctx)
202
+ if error is not None:
203
+ return error
204
+
205
+ # Instantiate and process
206
+ instance = handler_cls(socketio_server, lock)
207
+ return await instance.process(payload, sid)
208
+
209
+ except Exception as e:
210
+ error_text = format_error(e)
211
+ PrintStyle.error(f"WS handler error ({path}): {error_text}")
212
+ return {"ok": False, "error": error_text, "code": 500}
run_ui.py
+5
-1
@@ -16,6 +16,7 @@ from helpers.files import get_abs_path
16
from helpers import runtime, dotenv, process
17
from helpers.websocket import WebSocketHandler, validate_ws_origin
18
from helpers.api import register_api_route, requires_auth, csrf_protect
19
+from helpers.ws import register_ws_namespace
20
from helpers.print_style import PrintStyle
21
from helpers import login
22
import socketio # type: ignore[import-untyped]
@@ -387,13 +388,16 @@ def run():
388
register_api_route(webapp, lock)
389
390
handlers_by_namespace = _build_websocket_handlers_by_namespace(socketio_server, lock)
390
- configure_websocket_namespaces(
391
+ allowed_namespaces = configure_websocket_namespaces(
392
webapp=webapp,
393
socketio_server=socketio_server,
394
websocket_manager=websocket_manager,
395
handlers_by_namespace=handlers_by_namespace,
396
)
397
398
+ register_ws_namespace(socketio_server, webapp, lock)
399
+ allowed_namespaces.add("/ws")
400
+
401
init_a0()
402
403
wsgi_app = WSGIMiddleware(webapp)