This commit is contained in:
moosecrap 2026-07-21 13:49:34 -07:00
parent d09378811d
commit 270ba6694a

17
main.py
View File

@ -2,6 +2,7 @@ import asyncio
import logging
import uvicorn
import signal
import os
from starlette.applications import Starlette
from starlette.routing import Route
from starlette.responses import Response, StreamingResponse
@ -33,7 +34,8 @@ async def sse_endpoint(request):
queue = mcp_logic.sessions[sid]
while mcp_logic.running:
try:
msg = await asyncio.wait_for(queue.get(), timeout=1.0)
# Use a slightly shorter timeout for faster shutdown response
msg = await asyncio.wait_for(queue.get(), timeout=0.5)
yield f"event: message\ndata: {msg}\n\n"
queue.task_done()
except asyncio.TimeoutError:
@ -91,15 +93,24 @@ if __name__ == "__main__":
host="127.0.0.1",
port=8000,
log_level="info",
timeout_graceful_shutdown=5
timeout_graceful_shutdown=2 # Reduced from 5 to 2 seconds
)
server = uvicorn.Server(config)
def signal_handler(sig, frame):
logger.info("Shutdown signal received")
mcp_logic.running = False
# Explicitly tell Uvicorn to stop the server
server.should_exit = True
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
asyncio.run(server.serve())
try:
asyncio.run(server.serve())
except KeyboardInterrupt:
pass
finally:
# Final fallback to ensure the process actually dies
# if the event loop is still hanging on a connection
os._exit(0)