From 2b44dcebf411a8b6d779846cb57cad793eacc314 Mon Sep 17 00:00:00 2001 From: Shivering Ocelot Date: Tue, 30 Jun 2026 10:59:46 -0500 Subject: [PATCH] stabilize app shutdown flow and resource cleanup --- backend/src/agents/expert.py | 2 ++ backend/src/app.py | 34 +++++++++++++++++++++++++++------- backend/src/server.py | 16 ++++++++++++---- 3 files changed, 41 insertions(+), 11 deletions(-) diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index a3f5b72..76ed296 100644 --- a/backend/src/agents/expert.py +++ b/backend/src/agents/expert.py @@ -84,3 +84,5 @@ class ExpertAgent(Agent): messages.append(inference.LlamaMessage('tool', tc)) else: break + + run.close() diff --git a/backend/src/app.py b/backend/src/app.py index e335ce8..ffeb0ff 100644 --- a/backend/src/app.py +++ b/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()) \ No newline at end of file + asyncio.run(main()) \ No newline at end of file diff --git a/backend/src/server.py b/backend/src/server.py index 769239c..7005e8b 100644 --- a/backend/src/server.py +++ b/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)