121 lines
3.4 KiB
Python
121 lines
3.4 KiB
Python
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
|
|
from starlette.middleware.cors import CORSMiddleware
|
|
|
|
from mcp_logic import MCPServer
|
|
from tools import TOOL_REGISTRY
|
|
import config
|
|
|
|
|
|
# --- Configuration & Logging ---
|
|
logging.basicConfig(
|
|
level=getattr(logging, config.LOG_LEVEL),
|
|
filename=config.LOG_FILE,
|
|
|
|
filemode="a",
|
|
format="%(asctime)s - %(levelname)s - %(message)s"
|
|
)
|
|
logger = logging.getLogger("MooseCP")
|
|
|
|
mcp_logic = MCPServer()
|
|
|
|
# --- Transport Layer (Starlette/SSE) ---
|
|
|
|
async def sse_endpoint(request):
|
|
sid = mcp_logic.create_session()
|
|
logger.debug(f"Session started: {sid}")
|
|
|
|
async def event_generator():
|
|
try:
|
|
yield f"event: endpoint\ndata: /messages?session_id={sid}\n\n"
|
|
queue = mcp_logic.sessions[sid]
|
|
while mcp_logic.running:
|
|
try:
|
|
# 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:
|
|
continue
|
|
except asyncio.CancelledError:
|
|
logger.debug(f"Session cancelled: {sid}")
|
|
finally:
|
|
mcp_logic.sessions.pop(sid, None)
|
|
logger.debug(f"Session closed: {sid}")
|
|
|
|
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
|
|
|
async def messages_endpoint(request):
|
|
sid = request.query_params.get("session_id")
|
|
if not sid or sid not in mcp_logic.sessions:
|
|
return Response("Session not found", status_code=404)
|
|
|
|
try:
|
|
body = await request.json()
|
|
except Exception as e:
|
|
return Response(f"Invalid JSON: {e}", status_code=400)
|
|
|
|
method = body.get("method")
|
|
params = body.get("params", {})
|
|
req_id = body.get("id")
|
|
|
|
logger.debug(f"Request: {method} (ID: {req_id})")
|
|
|
|
result = await mcp_logic.handle_request(method, params, req_id, TOOL_REGISTRY)
|
|
|
|
if isinstance(result, Response):
|
|
return result
|
|
|
|
await mcp_logic.send_to_session(sid, {"jsonrpc": "2.0", "id": req_id, "result": result})
|
|
return Response(status_code=202)
|
|
|
|
app = Starlette(
|
|
routes=[
|
|
Route("/sse", endpoint=sse_endpoint),
|
|
Route("/messages", endpoint=messages_endpoint, methods=["POST"]),
|
|
]
|
|
)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
config = uvicorn.Config(
|
|
app=app,
|
|
host=config.HOST,
|
|
port=config.PORT,
|
|
log_level="info",
|
|
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)
|
|
|
|
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)
|