MooseCP/main.py
2026-07-21 03:37:45 -07:00

106 lines
3.0 KiB
Python

import asyncio
import logging
import uvicorn
import signal
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
# --- Configuration & Logging ---
logging.basicConfig(
level=logging.DEBUG,
filename="debug.log",
filemode="a",
format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("MattCP")
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:
msg = await asyncio.wait_for(queue.get(), timeout=1.0)
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="127.0.0.1",
port=8000,
log_level="info",
timeout_graceful_shutdown=5
)
server = uvicorn.Server(config)
def signal_handler(sig, frame):
logger.info("Shutdown signal received")
mcp_logic.running = False
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
asyncio.run(server.serve())