Skip to content

Handle SSE Disconnects Properly #612

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
May 2, 2025
Merged

Handle SSE Disconnects Properly #612

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b3d0657
Fix
akash329d May 2, 2025
4672eea
merge main
akash329d May 2, 2025
74dec3a
trailing slashes for mount
akash329d May 2, 2025
79f79b3
remove exit stack cleanup now that we have context mgr
akash329d May 2, 2025
41ba669
fix
akash329d May 2, 2025
952b434
misc fixes
akash329d May 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 4 additions & 6 deletions examples/servers/simple-prompt/mcp_simple_prompt/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,22 +90,20 @@ async def get_prompt(
if transport == "sse":
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.routing import Mount, Route
from starlette.routing import Mount

sse = SseServerTransport("/messages/")

async def handle_sse(request):
async with sse.connect_sse(
request.scope, request.receive, request._send
) as streams:
async def handle_sse(scope, receive, send):
async with sse.connect_sse(scope, receive, send) as streams:
await app.run(
streams[0], streams[1], app.create_initialization_options()
)

starlette_app = Starlette(
debug=True,
routes=[
Route("/sse", endpoint=handle_sse),
Mount("/sse", app=handle_sse),
Mount("/messages/", app=sse.handle_post_message),
],
)
Expand Down
10 changes: 4 additions & 6 deletions examples/servers/simple-resource/mcp_simple_resource/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,22 +46,20 @@ async def read_resource(uri: FileUrl) -> str | bytes:
if transport == "sse":
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.routing import Mount, Route
from starlette.routing import Mount

sse = SseServerTransport("/messages/")

async def handle_sse(request):
async with sse.connect_sse(
request.scope, request.receive, request._send
) as streams:
async def handle_sse(scope, receive, send):
async with sse.connect_sse(scope, receive, send) as streams:
await app.run(
streams[0], streams[1], app.create_initialization_options()
)

starlette_app = Starlette(
debug=True,
routes=[
Route("/sse", endpoint=handle_sse),
Mount("/sse", app=handle_sse),
Mount("/messages/", app=sse.handle_post_message),
],
)
Expand Down
10 changes: 4 additions & 6 deletions examples/servers/simple-tool/mcp_simple_tool/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,22 +60,20 @@ async def list_tools() -> list[types.Tool]:
if transport == "sse":
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.routing import Mount, Route
from starlette.routing import Mount

sse = SseServerTransport("/messages/")

async def handle_sse(request):
async with sse.connect_sse(
request.scope, request.receive, request._send
) as streams:
async def handle_sse(scope, receive, send):
async with sse.connect_sse(scope, receive, send) as streams:
await app.run(
streams[0], streams[1], app.create_initialization_options()
)

starlette_app = Starlette(
debug=True,
routes=[
Route("/sse", endpoint=handle_sse),
Mount("/sse", app=handle_sse),
Mount("/messages/", app=sse.handle_post_message),
],
)
Expand Down
2 changes: 1 addition & 1 deletion src/mcp/server/fastmcp/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ class Settings(BaseSettings, Generic[LifespanResultT]):
# HTTP settings
host: str = "0.0.0.0"
port: int = 8000
sse_path: str = "/sse"
sse_path: str = "/sse/"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why the trailing /?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Apparently Mount makes it trailing by default.

Starlette adds a trailing slash to the mount path by default and redirects requests without a trailing slash to the path with a trailing slash.

Our test client doesn't follow redirects so didn't like it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussed offline; we don't want to break current users so switched back to Route but return an empty Response to get around this.

message_path: str = "/messages/"

# resource settings
Expand Down
4 changes: 4 additions & 0 deletions src/mcp/server/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,10 @@ def check_client_capability(self, capability: types.ClientCapabilities) -> bool:

return True

async def _receive_loop(self) -> None:
async with self._incoming_message_stream_writer:
await super()._receive_loop()

async def _received_request(
self, responder: RequestResponder[types.ClientRequest, types.ServerResult]
):
Expand Down
31 changes: 22 additions & 9 deletions src/mcp/server/sse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,13 @@

# Create Starlette routes for SSE and message handling
routes = [
Route("/sse", endpoint=handle_sse),
Mount("/sse", app=handle_sse),
Mount("/messages/", app=sse.handle_post_message),
]

# Define handler functions
async def handle_sse(request):
async with sse.connect_sse(
request.scope, request.receive, request._send
) as streams:
async def handle_sse(scope, receive, send):
async with sse.connect_sse(scope, receive, send) as streams:
await app.run(
streams[0], streams[1], app.create_initialization_options()
)
Expand All @@ -28,6 +26,11 @@ async def handle_sse(request):
uvicorn.run(starlette_app, host="0.0.0.0", port=port)
```

Note: If you get a "TypeError: 'NoneType' object is not callable" error,
you need to change from Route("/sse", endpoint=handle_sse) to
Mount("/sse", app=handle_sse) and update the handle_sse signature to
accept (scope, receive, send) instead of (request). See examples for details.

See SseServerTransport class documentation for more details.
"""

Expand Down Expand Up @@ -121,11 +124,21 @@ async def sse_writer():
)

async with anyio.create_task_group() as tg:
response = EventSourceResponse(
content=sse_stream_reader, data_sender_callable=sse_writer
)
async def response_wrapper(scope: Scope, receive: Receive, send: Send):
"""
The EventSourceResponse returning signals a client close / disconnect.
In this case we close our side of the streams to signal the client that
the connection has been closed.
"""
await EventSourceResponse(
content=sse_stream_reader, data_sender_callable=sse_writer
)(scope, receive, send)
await read_stream_writer.aclose()
await write_stream_reader.aclose()
logging.debug(f"Client session disconnected {session_id}")

logger.debug("Starting SSE response task")
tg.start_soon(response, scope, receive, send)
tg.start_soon(response_wrapper, scope, receive, send)

logger.debug("Yielding read and write streams")
yield (read_stream, write_stream)
Expand Down
17 changes: 7 additions & 10 deletions tests/shared/test_sse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,7 @@
import uvicorn
from pydantic import AnyUrl
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.routing import Mount, Route
from starlette.routing import Mount

from mcp.client.session import ClientSession
from mcp.client.sse import sse_client
Expand Down Expand Up @@ -83,17 +82,15 @@ def make_server_app() -> Starlette:
sse = SseServerTransport("/messages/")
server = ServerTest()

async def handle_sse(request: Request) -> None:
async with sse.connect_sse(
request.scope, request.receive, request._send
) as streams:
async def handle_sse(scope, receive, send) -> None:
async with sse.connect_sse(scope, receive, send) as streams:
await server.run(
streams[0], streams[1], server.create_initialization_options()
)

app = Starlette(
routes=[
Route("/sse", endpoint=handle_sse),
Mount("/sse", app=handle_sse),
Mount("/messages/", app=sse.handle_post_message),
]
)
Expand Down Expand Up @@ -164,7 +161,7 @@ async def test_raw_sse_connection(http_client: httpx.AsyncClient) -> None:
async with anyio.create_task_group():

async def connection_test() -> None:
async with http_client.stream("GET", "/sse") as response:
async with http_client.stream("GET", "/sse/") as response:
assert response.status_code == 200
assert (
response.headers["content-type"]
Expand All @@ -188,7 +185,7 @@ async def connection_test() -> None:

@pytest.mark.anyio
async def test_sse_client_basic_connection(server: None, server_url: str) -> None:
async with sse_client(server_url + "/sse") as streams:
async with sse_client(server_url + "/sse/") as streams:
async with ClientSession(*streams) as session:
# Test initialization
result = await session.initialize()
Expand All @@ -204,7 +201,7 @@ async def test_sse_client_basic_connection(server: None, server_url: str) -> Non
async def initialized_sse_client_session(
server, server_url: str
) -> AsyncGenerator[ClientSession, None]:
async with sse_client(server_url + "/sse", sse_read_timeout=0.5) as streams:
async with sse_client(server_url + "/sse/", sse_read_timeout=0.5) as streams:
async with ClientSession(*streams) as session:
await session.initialize()
yield session
Expand Down
Loading