|
7 | 7 | from collections import defaultdict |
8 | 8 | from contextlib import asynccontextmanager |
9 | 9 | from typing import Literal, Optional, Sequence, TextIO |
| 10 | +import asyncio |
10 | 11 |
|
11 | 12 | import uvicorn |
12 | 13 | from deprecated import deprecated |
|
46 | 47 | logger = logging.getLogger(__name__) |
47 | 48 |
|
48 | 49 |
|
| 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 | + |
49 | 65 | class MCPRouter: |
50 | 66 | """ |
51 | 67 | A router that aggregates multiple MCP servers (SSE/STDIO) and |
@@ -615,17 +631,29 @@ async def get_sse_server_app( |
615 | 631 | sse = RouterSseTransport("/messages/", api_key=api_key) |
616 | 632 |
|
617 | 633 | 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() |
629 | 657 |
|
630 | 658 | lifespan_handler: t.Optional[Lifespan[Starlette]] = None |
631 | 659 | if include_lifespan: |
|
0 commit comments