refactor: decouple transport and session

calmini · calmini · commit b5cfb1976d24 · 2025-04-29T20:27:04.000+08:00
diff --git a/src/mcpm/core/router/app.py b/src/mcpm/core/router/app.py
@@ -0,0 +1,63 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.routing import Mount, Route
+from starlette.types import Receive, Scope, Send
+
+from mcpm.monitor.event import monitor
+from mcpm.router.router import MCPRouter
+
+from .middleware import SessionMiddleware
+from .session import SessionManager
+from .transport import SseTransport
+
+logger = logging.getLogger("mcpm.router")
+
+session_manager = SessionManager()
+transport = SseTransport(endpoint="/messages/", session_manager=session_manager)
+
+router = MCPRouter(reload_server=False)
+
+class NoOpsResponse(Response):
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        # To comply with Starlette's ASGI application design, this method must return a response.
+        # Since no further client interaction is needed after server shutdown, we provide a no-operation response
+        # that allows the application to exit gracefully when cancelled by Uvicorn.
+        # No content is sent back to the client as EventSourceResponse has already returned a 200 status code.
+        pass
+
+
+async def handle_sse(request: Request):
+    try:
+        async with transport.connect_sse(request.scope, request.receive, request._send) as (read, write):
+            await router.aggregated_server.run(read, write, router.aggregated_server.initialization_options) # type: ignore
+    except asyncio.CancelledError:
+        return NoOpsResponse()
+
+
+@asynccontextmanager
+async def lifespan(app):
+    logger.info("Starting MCPRouter...")
+    await router.initialize_router()
+    await monitor.initialize_storage()
+
+    yield
+
+    logger.info("Shutting down MCPRouter...")
+    await router.shutdown()
+    await monitor.close()
+
+app = Starlette(
+    debug=True,
+    routes=[
+        Route("/sse", endpoint=handle_sse),
+        Mount("/messages/", app=transport.handle_post_message)
+    ],
+    middleware=[Middleware(SessionMiddleware, session_manager=session_manager)],
+    lifespan=lifespan
+)
diff --git a/src/mcpm/core/router/middleware.py b/src/mcpm/core/router/middleware.py
@@ -0,0 +1,76 @@
+import logging
+from typing import Any
+from uuid import UUID
+
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from .session import SessionManager
+
+logger = logging.getLogger(__name__)
+
+META_DATA_KEYS = ["profile", "client"]
+
+def get_meta(request: Request) -> dict[str, Any]:
+    meta: dict[str, Any] = {}
+    for key in META_DATA_KEYS:
+        value = request.query_params.get(key)
+        if not value:
+            value = request.headers.get(key)
+
+        if value:
+            meta[key] = value
+
+    return meta
+
+class SessionMiddleware:
+
+    def __init__(self, app: ASGIApp, session_manager: SessionManager) -> None:
+        self.app = app
+        self.session_manager = session_manager
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        # we related metadata with session through this middleware, so that in the transport layer we only need to handle
+        # session_id and dispatch message to the correct memory stream
+
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        request = Request(scope)
+
+        if scope["path"] == "/sse":
+            # retrieve metadata from query params or header
+            session = await self.session_manager.create_session(meta=get_meta(request))
+            logger.debug(f"Created new session with ID: {session['id']}")
+
+            scope["session_id"] = session["id"].hex
+
+        if scope["path"] == "/messages/":
+            session_id_param = request.query_params.get("session_id")
+            if not session_id_param:
+                logger.debug("Missing session_id")
+                response = Response("session_id is required", status_code=400)
+                await response(scope, receive, send)
+                return
+
+            # validate session_id
+            try:
+                session_id = UUID(hex=session_id_param)
+            except ValueError:
+                logger.warning(f"Received invalid session ID: {session_id_param}")
+                response = Response("invalid session ID", status_code=400)
+                await response(scope, receive, send)
+                return
+
+            # if session_id is not in session manager, return 404
+            if not self.session_manager.exist(session_id):
+                logger.debug(f"session {session_id} not found")
+                response = Response("session not found", status_code=404)
+                await response(scope, receive, send)
+                return
+
+            scope["session_id"] = session_id.hex
+
+        await self.app(scope, receive, send)
diff --git a/src/mcpm/core/router/session.py b/src/mcpm/core/router/session.py
@@ -0,0 +1,95 @@
+from typing import Any, Protocol, TypedDict
+from uuid import UUID, uuid4
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp import types
+
+
+class Session(TypedDict):
+    id: UUID
+    # some read,write streams related with session
+    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+    # any meta data is saved here
+    meta: dict[str, Any]
+
+
+class SessionStore(Protocol):
+
+    def exist(self, session_id: UUID) -> bool:
+        ...
+
+    async def put(self, session: Session) -> None:
+        ...
+
+    async def get(self, session_id: UUID) -> Session:
+        ...
+
+    async def remove(self, session_id: UUID):
+        ...
+
+    async def cleanup(self):
+        ...
+
+
+class LocalSessionStore:
+
+    def __init__(self):
+        self._store: dict[UUID, Session] = {}
+
+    def exist(self, session_id: UUID) -> bool:
+        return session_id in self._store
+
+    async def put(self, session: Session) -> None:
+        self._store[session["id"]] = session
+
+    async def get(self, session_id: UUID) -> Session:
+        return self._store[session_id]
+
+    async def remove(self, session_id: UUID):
+        session = self._store.pop(session_id, None)
+        if session:
+            await session["read_stream_writer"].aclose()
+            await session["write_stream"].aclose()
+
+    async def cleanup(self):
+        keys = list(self._store.keys())
+        for session_id in keys:
+            await self.remove(session_id)
+
+
+class SessionManager:
+
+    def __init__(self):
+        self.session_store: SessionStore = LocalSessionStore()
+
+    def exist(self, session_id: UUID) -> bool:
+        return self.session_store.exist(session_id)
+
+    async def create_session(self, meta: dict[str, Any] = {}) -> Session:
+        read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+        write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+        session_id = uuid4()
+        session = Session(
+            id=session_id,
+            read_stream=read_stream,
+            read_stream_writer=read_stream_writer,
+            write_stream=write_stream,
+            write_stream_reader=write_stream_reader,
+            meta=meta
+        )
+        await self.session_store.put(session)
+        return session
+
+    async def get_session(self, session_id: UUID) -> Session:
+        return await self.session_store.get(session_id)
+
+    async def close_session(self, session_id: UUID):
+        await self.session_store.remove(session_id)
+
+    async def cleanup_resources(self):
+        await self.session_store.cleanup()
diff --git a/src/mcpm/core/router/transport.py b/src/mcpm/core/router/transport.py
@@ -0,0 +1,126 @@
+import asyncio
+import json
+import logging
+from contextlib import asynccontextmanager
+from typing import Any
+from urllib.parse import quote
+from uuid import UUID
+
+import anyio
+from mcp import types
+from pydantic import ValidationError
+from sse_starlette import EventSourceResponse
+from starlette.background import BackgroundTask
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.types import Receive, Scope, Send
+
+from .session import Session, SessionManager
+
+logger = logging.getLogger(__name__)
+
+def patch_meta_data(body: bytes, **kwargs) -> bytes:
+    data = json.loads(body.decode("utf-8"))
+    if "params" not in data:
+        data["params"] = {}
+
+    for key, value in kwargs.items():
+        data["params"].setdefault("_meta", {})[key] = value
+    return json.dumps(data).encode("utf-8")
+
+class SseTransport:
+
+    def __init__(self, endpoint: str, session_manager: SessionManager) -> None:
+        self.session_manager = session_manager
+        self._endpoint = endpoint
+
+    @asynccontextmanager
+    async def connect_sse(self, scope: Scope, receive: Receive, send: Send):
+        session_id_hex = scope["session_id"]
+        session_id: UUID = UUID(hex=session_id_hex)
+        session = await self.session_manager.get_session(session_id)
+
+        session_uri = f"{quote(self._endpoint)}?session_id={session_id.hex}"
+
+        sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[dict[str, Any]](0)
+
+        async def sse_writer():
+            logger.debug("Starting SSE writer")
+            async with sse_stream_writer, session["write_stream_reader"]:
+                await sse_stream_writer.send({"event": "endpoint", "data": session_uri})
+                logger.debug(f"Sent endpoint event: {session_uri}")
+
+                async for message in session["write_stream_reader"]:
+                    logger.debug(f"Sending message via SSE: {message}")
+                    await sse_stream_writer.send(
+                        {
+                            "event": "message",
+                            "data": message.model_dump_json(by_alias=True, exclude_none=True),
+                        }
+                    )
+
+        async with anyio.create_task_group() as tg:
+            async def on_client_disconnect():
+                await self.session_manager.close_session(session_id)
+
+            try:
+                response = EventSourceResponse(
+                    content=sse_stream_reader,
+                    data_sender_callable=sse_writer,
+                    background=BackgroundTask(on_client_disconnect),
+                )
+                logger.debug("Starting SSE response task")
+                tg.start_soon(response, scope, receive, send)
+
+                logger.debug("Yielding read and write streams")
+                # Due to limitations with interrupting the MCP server run operation,
+                # this will always block here regardless of client disconnection status
+                yield (session["read_stream"], session["write_stream"])
+            except asyncio.CancelledError as exc:
+                logger.warning(f"SSE connection for session {session_id} was cancelled")
+                tg.cancel_scope.cancel()
+                # raise the exception again so that to interrupt mcp server run operation
+                raise exc
+            finally:
+                # for server shutdown
+                await self.session_manager.cleanup_resources()
+
+    async def handle_post_message(self, scope: Scope, receive: Receive, send: Send):
+
+        session_id = scope["session_id"]
+        session: Session = await self.session_manager.get_session(UUID(hex=session_id))
+
+        request = Request(scope, receive)
+        body = await request.body()
+        # patch meta data
+        body = patch_meta_data(body, **session["meta"])
+
+        # send message to writer
+        writer = session["read_stream_writer"]
+        try:
+            message = types.JSONRPCMessage.model_validate_json(body)
+            logger.debug(f"Validated client message: {message}")
+        except ValidationError as err:
+            logger.error(f"Failed to parse message: {err}")
+            response = Response("Could not parse message", status_code=400)
+            await response(scope, receive, send)
+            try:
+                await writer.send(err)
+            except (BrokenPipeError, ConnectionError, OSError) as pipe_err:
+                logger.warning(f"Failed to send error due to pipe issue: {pipe_err}")
+            return
+
+        logger.debug(f"Sending message to writer: {message}")
+        response = Response("Accepted", status_code=202)
+        await response(scope, receive, send)
+
+        # add error handling, catch possible pipe errors
+        try:
+            await writer.send(message)
+        except (BrokenPipeError, ConnectionError, OSError) as e:
+            # if it's EPIPE error or other connection error, log it but don't throw an exception
+            if isinstance(e, OSError) and e.errno == 32:  # EPIPE
+                logger.warning(f"EPIPE error when sending message to session {session_id}, connection may be closing")
+            else:
+                logger.warning(f"Connection error when sending message to session {session_id}: {e}")
+                await self.session_manager.close_session(session_id)
diff --git a/src/mcpm/utils/errlog_manager.py b/src/mcpm/utils/errlog_manager.py
@@ -30,5 +30,6 @@ def close_errlog_file(self, server_id: str) -> None:
             del self._log_files[server_id]
 
     def close_all(self) -> None:
-        for server_id in self._log_files:
+        keys = list(self._log_files.keys())
+        for server_id in keys:
             self.close_errlog_file(server_id)
