refactor: ws handlers on connect/disconnect + api/ unification
keyboardstaff committed
Mar 16, 2026 at 09:25 UTC
1979212ddb554b19891a1c39f4fb7e863e4e1bd9
1 file changed
+60
-17
helpers/ws.py
+60
-17
@@ -17,7 +17,7 @@ from helpers.websocket import validate_ws_origin
17
18
ThreadLockType = Union[threading.Lock, threading.RLock]
19
20
-CACHE_AREA = "ws_handlers(ws)(plugins)"
20
+CACHE_AREA = "ws_handlers(api)(plugins)"
21
cache.toggle_area(CACHE_AREA, False) # cache off for now
22
23
@@ -32,6 +32,7 @@ class _SecurityContext:
32
33
34
_ws_contexts: dict[str, _SecurityContext] = {}
35
+_active_handlers: dict[str, list[tuple[str, type["WsHandler"]]]] = {}
36
_contexts_lock = threading.Lock()
37
38
@@ -60,6 +61,12 @@ class WsHandler:
61
async def process(self, data: dict, sid: str) -> dict | None:
62
pass
63
64
+ async def on_connect(self, sid: str) -> dict | None:
65
+ return None
66
+
67
+ async def on_disconnect(self, sid: str) -> None:
68
+ pass
69
+
70
def use_context(self, ctxid: str, create_if_not_exists: bool = True):
71
with self.lock:
72
if not ctxid:
@@ -118,21 +125,21 @@ def register_ws_namespace(
125
def _resolve_handler(path: str) -> type[WsHandler] | None:
126
handler_cls: type[WsHandler] | None = None
127
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):
128
+ # Check built-in api/<path>.py
129
+ builtin_file = files.get_abs_path(f"api/{path}.py")
130
+ if files.is_in_dir(builtin_file, files.get_abs_path("api")) and files.exists(builtin_file):
131
classes = load_classes_from_file(builtin_file, WsHandler)
132
if classes:
133
handler_cls = classes[0]
134
128
- # Plugin ws/<handler>.py — path format: plugins/<plugin_name>/<handler>
135
+ # Check plugin api/<handler>.py — path format: plugins/<plugin_name>/<handler>
136
if handler_cls is None and path.startswith("plugins/"):
137
parts = path.split("/", 2)
138
if len(parts) == 3:
139
_, plugin_name, handler_name = parts
140
plugin_dir = plugins.find_plugin_dir(plugin_name)
141
if plugin_dir:
135
- plugin_file = Path(plugin_dir) / "ws" / f"{handler_name}.py"
142
+ plugin_file = Path(plugin_dir) / "api" / f"{handler_name}.py"
143
if plugin_file.is_file():
144
classes = load_classes_from_file(str(plugin_file), WsHandler)
145
if classes:
@@ -140,6 +147,15 @@ def register_ws_namespace(
147
148
return handler_cls
149
150
+ def _resolve_cached(path: str) -> type[WsHandler] | None:
151
+ cached = cache.get(CACHE_AREA, path)
152
+ if cached is not None:
153
+ return cached
154
+ handler_cls = _resolve_handler(path)
155
+ if handler_cls is not None:
156
+ cache.add(CACHE_AREA, path, handler_cls)
157
+ return handler_cls
158
+
159
@socketio_server.on("connect", namespace="/ws")
160
async def _on_connect(sid, environ, auth):
161
with webapp.request_context(environ):
@@ -169,13 +185,46 @@ def register_ws_namespace(
185
with _contexts_lock:
186
_ws_contexts[sid] = ctx
187
172
- return True
188
+ # Activate handlers declared in auth.handlers
189
+ handler_paths = []
190
+ if isinstance(auth, dict):
191
+ raw = auth.get("handlers")
192
+ if isinstance(raw, list):
193
+ handler_paths = [p for p in raw if isinstance(p, str)]
194
+
195
+ activated: list[tuple[str, type[WsHandler]]] = []
196
+ for path in handler_paths:
197
+ try:
198
+ handler_cls = _resolve_cached(path)
199
+ if handler_cls is None:
200
+ continue
201
+ error = _check_security(handler_cls, ctx)
202
+ if error is not None:
203
+ continue
204
+ instance = handler_cls(socketio_server, lock)
205
+ await instance.on_connect(sid)
206
+ activated.append((path, handler_cls))
207
+ except Exception as e:
208
+ PrintStyle.error(f"WS on_connect error ({path}): {format_error(e)}")
209
+
210
+ with _contexts_lock:
211
+ _active_handlers[sid] = activated
212
+
213
+ return True
214
215
@socketio_server.on("disconnect", namespace="/ws")
216
async def _on_disconnect(sid):
217
with _contexts_lock:
218
+ activated = _active_handlers.pop(sid, [])
219
_ws_contexts.pop(sid, None)
220
221
+ for path, handler_cls in activated:
222
+ try:
223
+ instance = handler_cls(socketio_server, lock)
224
+ await instance.on_disconnect(sid)
225
+ except Exception as e:
226
+ PrintStyle.error(f"WS on_disconnect error ({path}): {format_error(e)}")
227
+
228
@socketio_server.on("*", namespace="/ws")
229
async def _dispatch(event, sid, data):
230
path = event
@@ -187,15 +236,9 @@ def register_ws_namespace(
236
if ctx is None:
237
return {"ok": False, "error": "No security context", "code": 401}
238
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)
239
+ handler_cls = _resolve_cached(path)
240
+ if handler_cls is None:
241
+ return {"ok": False, "error": f"WS endpoint not found: {path}", "code": 404}
242
243
# Security check
244
error = _check_security(handler_cls, ctx)
@@ -209,4 +252,4 @@ def register_ws_namespace(
252
except Exception as e:
253
error_text = format_error(e)
254
PrintStyle.error(f"WS handler error ({path}): {error_text}")
212
- return {"ok": False, "error": error_text, "code": 500}
255
+ return {"ok": False, "error": error_text, "code": 500}
\ No newline at end of file