fix: Handle SSE Disconnects Properly When use starlette middleware

[wenxuwan]

Motivation and Context

When using the Starlette middleware capabilities, if the SSE client disconnects, the current server will return starlette.responses.Response, and the Response will send a message of type "http.response.start", which will cause the following error：

File "/opt/anaconda3/lib/python3.12/contextlib.py", line 158, in __exit__
    self.gen.throw(value)
  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/_utils.py", line 83, in collapse_excgroups
    raise exc
  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/middleware/base.py", line 185, in __call__
    await response(scope, wrapped_receive, send)
  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/middleware/base.py", line 223, in __call__
    async for chunk in self.body_iterator:
  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/middleware/base.py", line 169, in body_stream
    assert message["type"] == "http.response.body", f"Unexpected message: {message}"
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: Unexpected message: {'type': 'http.response.start', 'status': 200, 'headers': [(b'content-length', b'0')]}

There are several issues that should be related to this PR, such as:

#883
jlowin/fastmcp#858

How Has This Been Tested?

Tested manually and ensured that the route was closed when clients close the connection.

Breaking Changes

No

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation update

Checklist

• I have read the MCP Documentation
• My code follows the repository's style guidelines
• New and existing tests pass locally
• I have added appropriate error handling
• I have added or updated documentation as needed

Additional context

[ochafik]

Thanks @wenxuwan !

Tested it w/ a variant of @llamantino's code snippet (just updated import to from mcp.server.fastmcp import FastMCP) and it seems to fix the issue 🎉

Two questions:

• Could you try and add some test case?
• Would/should this also apply to sse streams from streamable http?

[wenxuwan]

Thanks @wenxuwan !

Tested it w/ a variant of @llamantino's code snippet (just updated import to from mcp.server.fastmcp import FastMCP) and it seems to fix the issue 🎉

Two questions:

• Could you try and add some test case?
• Would/should this also apply to sse streams from streamable http?

@ochafik Thank you for your comments! , For your questions:

• Could you try and add some test case?
  Currently, the errors on the server side are caught by Uvicorn, making it challenging for us to add corresponding test cases.

• Would/should this also apply to sse streams from streamable http?
  The handling logic for SSE streams in streamable HTTP is different from that of the older SSE streams. In the case of streamable HTTP, the server closes the SSE connection after responding to the client’s DELETE request, so the issue we’re discussing does not apply in this context.

[Kludex]

I don't think this is the solution. Why are we sending 2 http.response.start events? The solution should be to avoid sending that, not doing a workaround with the response class.

[wenxuwan]

I don't think this is the solution. Why are we sending 2 http.response.start events? The solution should be to avoid sending that, not doing a workaround with the response class.

The operation for sending the response was introduced in #612

[Kludex]

I don't think this is the solution. Why are we sending 2 http.response.start events? The solution should be to avoid sending that, not doing a workaround with the response class.

The operation for sending the response was introduced in #612

Ah... I see the issue.

The problem is that we are mixing Starlette way with pure ASGI way. 😞

---

This PR is not the solution, it just hides the conceptual issue from above. I guess the solution is to decide what to do... We probably want to leverage Starlette instead of use the ASGI Interface directly...

[wenxuwan]

I don't think this is the solution. Why are we sending 2 http.response.start events? The solution should be to avoid sending that, not doing a workaround with the response class.

The operation for sending the response was introduced in #612

Ah... I see the issue.

The problem is that we are mixing Starlette way with pure ASGI way. 😞

This PR is not the solution, it just hides the conceptual issue from above. I guess the solution is to decide what to do... We probably want to leverage Starlette instead of use the ASGI Interface directly...

I don't think this is the solution. Why are we sending 2 http.response.start events? The solution should be to avoid sending that, not doing a workaround with the response class.

The operation for sending the response was introduced in #612

Ah... I see the issue.

The problem is that we are mixing Starlette way with pure ASGI way. 😞

This PR is not the solution, it just hides the conceptual issue from above. I guess the solution is to decide what to do... We probably want to leverage Starlette instead of use the ASGI Interface directly...

@Kludex Thank you for your suggestion. As you mentioned, the current legacy implementation of MCP's SSE mixes ASGI and Starlette, which poses challenges for the seamless integration of the handle_sse implementation with Starlette. I see one potential solution to address this issue.Maintain our existing implementation as is, but wrap the current handle_sse implementation in a Response, allowing handle_sse to directly return a McpResponse, For example:

	class McpResponse(Response):
	    """A response that does not send any HTTP response back to the client."""
	
	    def __init__(self, sse: SseServerTransport, mcp_server: MCPServer) -> None:
	        self.sse = sse
	        self._mcp_server = mcp_server
	        super().__init__()
	
	    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
	        async with self.sse.connect_sse(
	                scope,
	                receive,
	                send,
	        ) as streams:
	            await self._mcp_server.run(
	                streams[0],
	                streams[1],
	                self._mcp_server.create_initialization_options(),
	            )
	        return
        
	
	async def handle_sse(scope: Scope, receive: Receive, send: Send):
            # Add client ID from auth context into request context if available
            return SilentResponse(sse, self._mcp_server)	

In this way, our implementation will comply with the requirements of the Starlette framework, even though we internally use ASGI directly. At the same time, it does not require significant changes to the existing architecture.What do you think?

[Kludex]

Can't we just remove those Response for now?

[wenxuwan]

Can't we just remove those Response for now?

No，if we directly remove return Response, the Starlette framework will throw an exception.

	  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/routing.py", line 290, in handle
	    await self.app(scope, receive, send)
	  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/routing.py", line 78, in app
	    await wrap_app_handling_exceptions(app, request)(scope, receive, send)
	  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
	    raise exc
	  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/_exception_handler.py", line 42, in wrapped_app
	    await app(scope, receive, sender)
	  File "/Users/wangwenxue/code/gitlab/sofapyapp/myenv/lib/python3.12/site-packages/starlette/routing.py", line 76, in app
	    await response(scope, receive, send)
	          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Because the Starlette framework expects the handler to return a Response, as shown below:

    async def app(scope: Scope, receive: Receive, send: Send) -> None:
        request = Request(scope, receive, send)

        async def app(scope: Scope, receive: Receive, send: Send) -> None:
            response = await f(request)
            await response(scope, receive, send)

        await wrap_app_handling_exceptions(app, request)(scope, receive, send)

    return app

Therefore, the handle_sse function needs to return a Response.

[Kludex]

Ah, I see the issue.

We can't really mix things here. We need to remove the way we handle with pure ASGI. Which means refactoring quite a lot here, but it needs to be done.