stabilize app shutdown flow and resource cleanup

backend/src/server.py

@@ -25,7 +25,7 @@ def encode_json(value):
             return repr(value)
     return json.dumps(value, default=encode_unknown)
 
-async def start_server(app):
+async def start_server(app, shutdown_event):
     server = Quart(__name__, static_folder=frontend_path)
     server.secret_key = 'horde-secret'
     auth_manager = QuartAuth(server, mode="bearer")
@@ -105,6 +105,7 @@ async def start_server(app):
 
         # Authenticate before dedicating any resources
         auth_req = json.loads(await websocket.receive())
+        loop = asyncio.get_running_loop()
 
         if 'token' in auth_req:
             try:
@@ -132,7 +133,7 @@ async def start_server(app):
         await websocket.send(json.dumps({'message': 'authentication successful'}))
 
         send_queue = asyncio.Queue()
-        async def sender_task():
+        async def sender_impl():
             nonlocal send_queue
             while True:
                 msg = await send_queue.get()
@@ -140,12 +141,17 @@ async def start_server(app):
         async def send(msg):
             await send_queue.put(msg)
 
+        asyncio.create_task(sender_impl())
+        ws_task = asyncio.current_task()
+        app.ws_tasks.append(ws_task)
+
         try:
+            # TODO: remotes aren't separate connections
+            # instead, any connection can create a remotes and listen for requests sent to them
             if 'remotename' in auth_req:
                 app.add_remote_connection(user_id, authreq['remotename'], authreq['tools'], send_queue)
             else:
                 app.add_user_connection(user_id, send_queue)
-            asyncio.create_task(sender_task())
             
             # Listen for events and broadcast them
             async def on_event(name, data):
@@ -328,6 +334,8 @@ async def start_server(app):
                     await send({'error': 'unknown action: ' + action, 'nonce': nonce})
         finally:
             print("User logged out:", user_id)
+            if ws_task in app.ws_tasks:
+                app.ws_tasks.remove(ws_task)
             app.remove_user_connection(user_id, send_queue)
 
     # https://github.com/hrgarber/openai-compatible-api/blob/main/app/main.py
@@ -584,4 +592,4 @@ async def start_server(app):
         }
 
 
-    await server.run_task(debug=True)
+    await server.run_task(debug=True, shutdown_trigger=shutdown_event.wait)