feat(backend): refactor websocket handler

backend/src/server/websocket.py

@@ -68,14 +68,54 @@ def register(server, app, auth_manager):
             stop_listening = app.on_event(on_event)
 
             while True:
-                req = json.loads(await websocket.receive())
-                if not 'action' in req:
+                asyncio.create_task(handle_request(
+                    app, user_id, session_connections, send, await websocket.receive()
+                ))
+
+        finally:
+            print("User logged out:", user_id)
+            sender_task.cancel()
+            if ws_task in app.ws_tasks:
+                app.ws_tasks.remove(ws_task)
+            # Unregister the event listener: the send closure would
+            # otherwise keep queueing events for a dead socket forever
+            if stop_listening:
+                stop_listening()
+            # All remotes served by this session are disconnected with it.
+            # disconnectRemote is idempotent, so remotes that were already
+            # disconnected manually (possibly from another session) are safe.
+            for remotename in list(session_connections):
+                app.disconnectRemote(user_id, remotename)
+            app.remove_user_connection(user_id, send_queue)
+
+# Error handling
+async def handle_request(app, user_id, session_connections, send, req):
+    try:
+        req = json.loads(req)
+        await process_request(app, user_id, session_connections, send, req)
+
+    except Exception as err:
+        action = req.get('action') if isinstance(req, dict) else None
+        nonce  = req.get('nonce') if isinstance(req, dict) else None
+
+        # Maintainers
+        print(f"[ws] req error @{user_id} {action!r}: {err!r}")
+        print(traceback.format_exc())
+        # Client — nonce makes this routable to the caller's promise/callback
+        try:
+            await send({'error': str(err), 'nonce': nonce})
+        except Exception:
+            print(traceback.format_exc())
+
+# Actual processing
+async def process_request(app, user_id, session_connections, send, req):
+    if 'action' not in req:
         await send({'error': 'no action'})
-                    continue
+        return
     action = req['action']
     body = req.get('body', {})
     nonce = req.get('nonce', None)
-                print(f"[ws] action @{user_id} {action}")
+    print(f"[ws] req @{user_id} {action}")
 
     async def agent_id_from_body (allow_none=False):
         if 'agentId' in body:
@@ -88,7 +128,7 @@ def register(server, app, auth_manager):
             raise Exception("field 'agentId' or 'agentname' required")
 
     if action == "ping":
-                    await send({'message': "pong"})
+        await send({'message': "pong", 'nonce': nonce})
 
     elif action == "create-agent":
         name = body['name']
@@ -300,7 +340,7 @@ def register(server, app, auth_manager):
     elif action == "prompt-agent":
         if 'instruction' not in body and 'input' not in body:
             await send({'error': "instruction or input required", 'nonce': nonce})
-                        continue
+            return
 
         # Multiple-result action, same shape as read-stream: every
         # generated message is its own packet carrying the request
@@ -337,18 +377,3 @@ def register(server, app, auth_manager):
 
     else:
         await send({'error': 'unknown action: ' + action, 'nonce': nonce})
-        finally:
-            print("User logged out:", user_id)
-            sender_task.cancel()
-            if ws_task in app.ws_tasks:
-                app.ws_tasks.remove(ws_task)
-            # Unregister the event listener: the send closure would
-            # otherwise keep queueing events for a dead socket forever
-            if stop_listening:
-                stop_listening()
-            # All remotes served by this session are disconnected with it.
-            # disconnectRemote is idempotent, so remotes that were already
-            # disconnected manually (possibly from another session) are safe.
-            for remotename in list(session_connections):
-                app.disconnectRemote(user_id, remotename)
-            app.remove_user_connection(user_id, send_queue)