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

15
main.py
View File

@ -2,6 +2,7 @@ import asyncio
import logging import logging
import uvicorn import uvicorn
import signal import signal
import os
from starlette.applications import Starlette from starlette.applications import Starlette
from starlette.routing import Route from starlette.routing import Route
from starlette.responses import Response, StreamingResponse from starlette.responses import Response, StreamingResponse
@ -33,7 +34,8 @@ async def sse_endpoint(request):
queue = mcp_logic.sessions[sid] queue = mcp_logic.sessions[sid]
while mcp_logic.running: while mcp_logic.running:
try: 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" yield f"event: message\ndata: {msg}\n\n"
queue.task_done() queue.task_done()
except asyncio.TimeoutError: except asyncio.TimeoutError:
@ -91,15 +93,24 @@ if __name__ == "__main__":
host="127.0.0.1", host="127.0.0.1",
port=8000, port=8000,
log_level="info", log_level="info",
timeout_graceful_shutdown=5 timeout_graceful_shutdown=2 # Reduced from 5 to 2 seconds
) )
server = uvicorn.Server(config) server = uvicorn.Server(config)
def signal_handler(sig, frame): def signal_handler(sig, frame):
logger.info("Shutdown signal received") logger.info("Shutdown signal received")
mcp_logic.running = False mcp_logic.running = False
# Explicitly tell Uvicorn to stop the server
server.should_exit = True
signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGTERM, signal_handler)
try:
asyncio.run(server.serve()) 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)