fix: stop reconnecting after terminal SSE drain errors

Author: he-yufeng

src/mcp/client/streamable_http.py

@@ -342,9 +342,9 @@ async def _handle_sse_response(
         assert isinstance(ctx.session_message.message, JSONRPCRequest)
         original_request_id = ctx.session_message.message.id
 
+        response_complete = False
         try:
             event_source = EventSource(response)
-            response_complete = False
             async for sse in event_source.aiter_sse():  # pragma: no branch
                 # Track last event ID for potential reconnection
                 if sse.id:
@@ -371,7 +371,9 @@ async def _handle_sse_response(
             if response_complete:
                 return  # Normal completion, no reconnect needed
         except Exception:
-            logger.debug("SSE stream ended", exc_info=True)  # pragma: no cover
+            logger.debug("SSE stream ended", exc_info=True)
+            if response_complete:
+                return
 
         # Stream ended without response - reconnect if we received an event with ID
         if last_event_id is not None:  # pragma: no branch
@@ -403,6 +405,7 @@ async def _handle_reconnection(
         if isinstance(ctx.session_message.message, JSONRPCRequest):  # pragma: no branch
             original_request_id = ctx.session_message.message.id
 
+        response_complete = False
         try:
             async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source:
                 event_source.response.raise_for_status()
@@ -411,7 +414,6 @@ async def _handle_reconnection(
                 # Track for potential further reconnection
                 reconnect_last_event_id: str = last_event_id
                 reconnect_retry_ms = retry_interval_ms
-                response_complete = False
 
                 async for sse in event_source.aiter_sse():
                     if sse.id:  # pragma: no branch
@@ -438,6 +440,8 @@ async def _handle_reconnection(
                 await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, 0)
         except Exception as e:
             logger.debug(f"Reconnection failed: {e}")
+            if response_complete:
+                return
             # Try to reconnect again if we still have an event ID
             await self._handle_reconnection(ctx, last_event_id, retry_interval_ms, attempt + 1)
 

tests/shared/test_streamable_http.py

@@ -49,11 +49,6 @@
 from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared._context import RequestContext
 from mcp.shared._context_streams import ContextSendStream, create_context_streams
-from mcp.shared._httpx_utils import (
-    MCP_DEFAULT_SSE_READ_TIMEOUT,
-    MCP_DEFAULT_TIMEOUT,
-    create_mcp_http_client,
-)
 from mcp.shared.message import ClientMessageMetadata, ServerMessageMetadata, SessionMessage
 from mcp.shared.session import RequestResponder
 from mcp.types import (
@@ -1663,6 +1658,76 @@ async def fail_reconnect(*args: Any, **kwargs: Any) -> None:  # pragma: no cover
             await read_stream.aclose()
 
 
+@pytest.mark.anyio
+async def test_sse_response_does_not_reconnect_after_terminal_then_drain_error(monkeypatch: pytest.MonkeyPatch):
+    transport = StreamableHTTPTransport(url="http://localhost:8000/mcp")
+    response = _FakeStreamResponse()
+
+    class FakeEventSource:
+        def __init__(self, response: _FakeStreamResponse) -> None:
+            self.response = response
+
+        async def aiter_sse(self):
+            yield _response_sse(1)
+            raise RuntimeError("drain failed after terminal response")
+
+    async def fail_reconnect(*args: Any, **kwargs: Any) -> None:  # pragma: no cover
+        raise AssertionError("completed responses should not reconnect after drain errors")
+
+    monkeypatch.setattr(streamable_http_module, "EventSource", FakeEventSource)
+    monkeypatch.setattr(transport, "_handle_reconnection", fail_reconnect)
+
+    write_stream, read_stream = create_context_streams[SessionMessage | Exception](2)
+    async with httpx.AsyncClient() as client:
+        try:
+            ctx = _make_streamable_http_request_context(1, client, write_stream)
+            await transport._handle_sse_response(response, ctx)
+
+            message = await read_stream.receive()
+            assert isinstance(message, SessionMessage)
+            assert isinstance(message.message, types.JSONRPCResponse)
+            assert message.message.id == 1
+        finally:
+            await write_stream.aclose()
+            await read_stream.aclose()
+
+
+@pytest.mark.anyio
+async def test_sse_response_reconnects_after_pre_terminal_drain_error(monkeypatch: pytest.MonkeyPatch):
+    transport = StreamableHTTPTransport(url="http://localhost:8000/mcp")
+    response = _FakeStreamResponse()
+    reconnects: list[tuple[str, int | None]] = []
+
+    class FakeEventSource:
+        def __init__(self, response: _FakeStreamResponse) -> None:
+            self.response = response
+
+        async def aiter_sse(self):
+            yield ServerSentEvent(event="message", data="", id="resume-from-here")
+            raise RuntimeError("stream failed before terminal response")
+
+    async def record_reconnect(
+        ctx: StreamableHTTPClientRequestContext,
+        last_event_id: str,
+        retry_interval_ms: int | None = None,
+    ) -> None:
+        reconnects.append((last_event_id, retry_interval_ms))
+
+    monkeypatch.setattr(streamable_http_module, "EventSource", FakeEventSource)
+    monkeypatch.setattr(transport, "_handle_reconnection", record_reconnect)
+
+    write_stream, read_stream = create_context_streams[SessionMessage | Exception](2)
+    async with httpx.AsyncClient() as client:
+        try:
+            ctx = _make_streamable_http_request_context(1, client, write_stream)
+            await transport._handle_sse_response(response, ctx)
+
+            assert reconnects == [("resume-from-here", None)]
+        finally:
+            await write_stream.aclose()
+            await read_stream.aclose()
+
+
 @pytest.mark.anyio
 async def test_reconnection_drains_after_terminal_response(monkeypatch: pytest.MonkeyPatch):
     """Resumed GET responses use EOF draining instead of response.aclose()."""
@@ -1699,6 +1764,44 @@ async def fake_aconnect_sse(*args: Any, **kwargs: Any):
             await read_stream.aclose()
 
 
+@pytest.mark.anyio
+async def test_reconnection_does_not_retry_after_terminal_then_drain_error(monkeypatch: pytest.MonkeyPatch):
+    transport = StreamableHTTPTransport(url="http://localhost:8000/mcp")
+    response = _FakeStreamResponse()
+    attempts = 0
+
+    class FakeReconnectionEventSource:
+        def __init__(self, response: _FakeStreamResponse) -> None:
+            self.response = response
+
+        async def aiter_sse(self):
+            yield _response_sse("abc")
+            raise RuntimeError("drain failed after terminal response")
+
+    @asynccontextmanager
+    async def fake_aconnect_sse(*args: Any, **kwargs: Any):
+        nonlocal attempts
+        attempts += 1
+        yield FakeReconnectionEventSource(response)
+
+    monkeypatch.setattr(streamable_http_module, "aconnect_sse", fake_aconnect_sse)
+
+    write_stream, read_stream = create_context_streams[SessionMessage | Exception](2)
+    async with httpx.AsyncClient() as client:
+        try:
+            ctx = _make_streamable_http_request_context("abc", client, write_stream)
+            await transport._handle_reconnection(ctx, "previous-event", retry_interval_ms=0)
+
+            assert attempts == 1
+            message = await read_stream.receive()
+            assert isinstance(message, SessionMessage)
+            assert isinstance(message.message, types.JSONRPCResponse)
+            assert message.message.id == "abc"
+        finally:
+            await write_stream.aclose()
+            await read_stream.aclose()
+
+
 @pytest.mark.anyio
 async def test_reconnection_retries_after_failed_resume(monkeypatch: pytest.MonkeyPatch):
     """A failed resume attempt falls back to the next reconnection attempt."""