refactor: decouple transport and session management

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
+)

src/mcpm/core/router/extra.py

@@ -0,0 +1,49 @@
+
+from typing import Any, Protocol
+
+from mcp.types import ServerResult
+from starlette.requests import Request
+
+from mcpm.core.router.session import Session
+
+
+class MetaRequestProcessor(Protocol):
+
+    def process(self, request: Request, session: Session):
+        ...
+
+class MetaResponseProcessor(Protocol):
+
+    def process(self, response: ServerResult, request_context: dict[str, Any], response_context: dict[str, Any]) -> ServerResult:
+        ...
+
+class ProfileMetaRequestProcessor:
+
+    def process(self, request: Request, session: Session):
+        profile = request.query_params.get("profile")
+        if not profile:
+            # fallback to headers
+            profile = request.headers.get("profile")
+
+        if profile:
+            session["meta"]["profile"] = profile
+
+class ClientMetaRequestProcessor:
+
+    def process(self, request: Request, session: Session):
+        client = request.query_params.get("client")
+        if client:
+            session["meta"]["client_id"] = client
+
+
+class MCPResponseProcessor:
+
+    def process(self, response: ServerResult, request_context: dict[str, Any], response_context: dict[str, Any]) -> ServerResult:
+        if not response.root.meta:
+            response.root.meta = {}
+
+        response.root.meta.update({
+            "request_context": request_context,
+            "response_context": response_context,
+        })
+        return response

src/mcpm/core/router/middleware.py

@@ -0,0 +1,77 @@
+import logging
+from uuid import UUID
+
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from .extra import ClientMetaRequestProcessor, MetaRequestProcessor, ProfileMetaRequestProcessor
+from .session import SessionManager
+
+logger = logging.getLogger(__name__)
+
+class SessionMiddleware:
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        session_manager: SessionManager,
+        meta_request_processors: list[MetaRequestProcessor] = [
+            ProfileMetaRequestProcessor(),
+            ClientMetaRequestProcessor(),
+        ]
+    ) -> None:
+        self.app = app
+        self.session_manager = session_manager
+        # patch meta data from request to session
+        self.meta_request_processors = meta_request_processors
+
+    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()
+            # session.meta will identically copied to JSONRPCMessage
+            if self.meta_request_processors:
+                for processor in self.meta_request_processors:
+                    processor.process(request, session)
+
+            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)

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()