stabilize app shutdown flow and resource cleanup

backend/src/agents/expert.py

@@ -84,3 +84,5 @@ class ExpertAgent(Agent):
                     messages.append(inference.LlamaMessage('tool', tc))
             else:
                 break
+        
+        run.close()

backend/src/app.py

@@ -4,6 +4,7 @@ import json
 import time
 import sys
 import os
+import signal
 
 from tortoise import Tortoise
 from tortoise.expressions import Q
@@ -28,6 +29,8 @@ class App(ChatMixin):
         self.remote_connections = {}
         self.streams = {}
         self.event_listeners = []
+        # active websocket tasks, for manual cancelling upon server shutdown
+        self.ws_tasks = []
 
         self.tmp = tmp.State()
 
@@ -380,12 +383,29 @@ async def main():
     )
 
     # 3. Start Server
-    # Note: start_server must be an async function or handle an event loop
-    await start_server(horde_app)
+    # manually trigger server shutdown, avoids hypercorn installing signal handlers
+    shutdown_event = asyncio.Event()
+    server_task = asyncio.create_task(start_server(horde_app, shutdown_event))
+
+    # 4. Graceful Termination
+    loop = asyncio.get_running_loop()
+    def interrupt():
+        print("Interrupt Signal")
+        shutdown_event.set()
+        # Unblock websocket handlers parked in websocket.receive()
+        for task in list(horde_app.ws_tasks):
+            task.cancel()
+
+    for sig in (signal.SIGINT, signal.SIGTERM):
+        loop.add_signal_handler(sig, interrupt)
+
+    try:
+        await server_task
+    finally:
+        print("Server stopped")
+
+        horde_app.qdrant_client.close()
+        await Tortoise.close_connections()
 
 if __name__ == "__main__":
-    try:
-        asyncio.run(main())
-    except KeyboardInterrupt:
-        # Clean shutdown for Tortoise connections
-        asyncio.run(Tortoise.close_connections())
+    asyncio.run(main())

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)