Skip to content

Commit 37addeb

Browse files
Fix: Enhance ASGI Compliance and Robustness in router.py SSE Handling (#164)
* Fix: Ensure ASGI compliance in transport.py for POST message handling - Modified mcpm.router.transport.py: - Refactored handle_post_message to correctly send a 202 Accepted response as a raw ASGI application before proceeding with internal message forwarding. It now implicitly returns None as expected. (Note: Complementary ASGI fixes for mcpm.router.router.py's handle_sse were already present in the base branch 076a6bf) * Fix: Enhance ASGI compliance in router.py SSE handling --------- Co-authored-by: Jonathan Wang <[email protected]>
1 parent a278722 commit 37addeb

File tree

1 file changed

+39
-11
lines changed

1 file changed

+39
-11
lines changed

src/mcpm/router/router.py

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from collections import defaultdict
88
from contextlib import asynccontextmanager
99
from typing import Literal, Optional, Sequence, TextIO
10+
import asyncio
1011

1112
import uvicorn
1213
from deprecated import deprecated
@@ -46,6 +47,21 @@
4647
logger = logging.getLogger(__name__)
4748

4849

50+
class NoOpsResponse(Response):
51+
def __init__(self):
52+
super().__init__(content=b"", status_code=204)
53+
54+
async def __call__(self, scope, receive, send):
55+
await send(
56+
{
57+
"type": "http.response.start",
58+
"status": self.status_code,
59+
"headers": self.render_headers(),
60+
}
61+
)
62+
await send({"type": "http.response.body", "body": b"", "more_body": False})
63+
64+
4965
class MCPRouter:
5066
"""
5167
A router that aggregates multiple MCP servers (SSE/STDIO) and
@@ -615,17 +631,29 @@ async def get_sse_server_app(
615631
sse = RouterSseTransport("/messages/", api_key=api_key)
616632

617633
async def handle_sse(request: Request) -> Response:
618-
async with sse.connect_sse(
619-
request.scope,
620-
request.receive,
621-
request._send, # noqa: SLF001
622-
) as (read_stream, write_stream):
623-
await self.aggregated_server.run(
624-
read_stream,
625-
write_stream,
626-
self.aggregated_server.initialization_options,
627-
)
628-
return Response()
634+
try:
635+
async with sse.connect_sse(
636+
request.scope,
637+
request.receive,
638+
request._send, # noqa: SLF001
639+
) as (read_stream, write_stream):
640+
await self.aggregated_server.run(
641+
read_stream,
642+
write_stream,
643+
self.aggregated_server.initialization_options,
644+
)
645+
# Keep alive while client connected.
646+
# EventSourceResponse (inside connect_sse) manages the stream,
647+
# but this loop ensures this handler itself stays alive until disconnect.
648+
while not await request.is_disconnected():
649+
await asyncio.sleep(0.1)
650+
651+
except asyncio.CancelledError:
652+
raise
653+
except Exception as e:
654+
logger.error(f"Unexpected error in handle_sse (router.py): {e}", exc_info=True)
655+
finally:
656+
return NoOpsResponse()
629657

630658
lifespan_handler: t.Optional[Lifespan[Starlette]] = None
631659
if include_lifespan:

0 commit comments

Comments
 (0)