Fix #883: AssertionError with Starlette middleware and Fix #156: Jupyter logging support

Author: Ashutosh0x

src/mcp/client/stdio.py

@@ -21,6 +21,7 @@
     terminate_windows_process_tree,
 )
 from mcp.shared.message import SessionMessage
+from mcp.shared.jupyter import is_jupyter
 
 logger = logging.getLogger(__name__)
 
@@ -118,12 +119,18 @@ async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stder
     try:
         command = _get_executable_command(server.command)
 
+        # In Jupyter, we pipe stderr to read it in the main process
+        # because sys.stderr redirection might not be reliable for subprocesses.
+        actual_errlog = errlog
+        if is_jupyter():
+            actual_errlog = anyio.lowlevel.PIPE
+
         # Open process with stderr piped for capture
         process = await _create_platform_compatible_process(
             command=command,
             args=server.args,
             env=({**get_default_environment(), **server.env} if server.env is not None else get_default_environment()),
-            errlog=errlog,
+            errlog=actual_errlog,
             cwd=server.cwd,
         )
     except OSError:
@@ -177,9 +184,29 @@ async def stdin_writer():
         except anyio.ClosedResourceError:  # pragma: no cover
             await anyio.lowlevel.checkpoint()
 
+    async def stderr_reader():
+        if not process.stderr:
+            return
+
+        try:
+            async for chunk in TextReceiveStream(
+                process.stderr,
+                encoding=server.encoding,
+                errors=server.encoding_error_handler,
+            ):
+                if is_jupyter():
+                    # In Jupyter, we use print() for better output handling
+                    # Red text for stderr
+                    print(f"\033[91m{chunk}\033[0m", end="", flush=True)
+                else:
+                    print(chunk, file=sys.stderr, end="", flush=True)
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
     async with anyio.create_task_group() as tg, process:
         tg.start_soon(stdout_reader)
         tg.start_soon(stdin_writer)
+        tg.start_soon(stderr_reader)
         try:
             yield read_stream, write_stream
         finally:

src/mcp/server/mcpserver/server.py

@@ -805,7 +805,6 @@ async def handle_sse(scope: Scope, receive: Receive, send: Send):  # pragma: no
                 await self._lowlevel_server.run(
                     streams[0], streams[1], self._lowlevel_server.create_initialization_options()
                 )
-            return Response()
 
         # Create routes
         routes: list[Route | Mount] = []
@@ -869,15 +868,18 @@ async def handle_sse(scope: Scope, receive: Receive, send: Send):  # pragma: no
             )
         else:
             # Auth is disabled, no need for RequireAuthMiddleware
-            # Since handle_sse is an ASGI app, we need to create a compatible endpoint
-            async def sse_endpoint(request: Request) -> Response:  # pragma: no cover
-                # Convert the Starlette request to ASGI parameters
-                return await handle_sse(request.scope, request.receive, request._send)  # type: ignore[reportPrivateUsage]
+
+
+            # Use an ASGI-compatible wrapper to avoid Starlette's high-level route wrapping
+            # which expects a Response object and causes double-sending.
+            class HandleSseAsgi:
+                async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+                    await handle_sse(scope, receive, send)
 
             routes.append(
                 Route(
                     sse_path,
-                    endpoint=sse_endpoint,
+                    endpoint=HandleSseAsgi(),
                     methods=["GET"],
                 )
             )

src/mcp/shared/jupyter.py

@@ -0,0 +1,15 @@
+
+import sys
+
+def is_jupyter() -> bool:
+    """Check if we are running in a Jupyter notebook environment."""
+    try:
+        shell = get_ipython().__class__.__name__ # type: ignore
+        if shell == 'ZMQInteractiveShell':
+            return True   # Jupyter notebook or qtconsole
+        elif shell == 'TerminalInteractiveShell':
+            return False  # Terminal running IPython
+        else:
+            return False  # Other type (?)
+    except NameError:
+        return False      # Probably standard Python interpreter

tests/issues/test_883_middleware.py

@@ -0,0 +1,50 @@
+
+import pytest
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.testclient import TestClient
+from mcp.server.mcpserver import MCPServer
+from mcp.server.transport_security import TransportSecuritySettings
+
+class MockMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request, call_next):
+        return await call_next(request)
+
+def test_883_middleware_sse_no_assertion_error():
+    """Test that using MCP SSE with Starlette middleware doesn't cause double-response error."""
+    mcp_server = MCPServer("test-server")
+    transport_security = TransportSecuritySettings(enable_dns_rebinding_protection=False)
+    sse_app = mcp_server.sse_app(transport_security=transport_security)
+    
+    app = Starlette(middleware=[Middleware(MockMiddleware)])
+    app.mount("/", sse_app)
+    
+    client = TestClient(app)
+    
+    # We use a context manager to ensure the stream is closed quickly
+    with client.stream("GET", "/sse") as response:
+        assert response.status_code == 200
+        # Just check headers are there
+        assert "text/event-stream" in response.headers["content-type"]
+
+def test_883_middleware_post_accepted():
+    """Test that POST messages work with middleware."""
+    mcp_server = MCPServer("test-server")
+    transport_security = TransportSecuritySettings(enable_dns_rebinding_protection=False)
+    sse_app = mcp_server.sse_app(transport_security=transport_security)
+    
+    app = Starlette(middleware=[Middleware(MockMiddleware)])
+    app.mount("/", sse_app)
+    
+    client = TestClient(app)
+    
+    # POST to /messages/ (with invalid session, but should not AssertionError)
+    response = client.post("/messages/?session_id=00000000000000000000000000000000", json={
+        "jsonrpc": "2.0",
+        "method": "notifications/initialized",
+        "params": {}
+    })
+    
+    # 404 is expected here as we didn't establish a real session
+    assert response.status_code == 404