Skip to content

Commit d7e05e7

Browse files
author
Jianke LIN
committed
fix(streamable-http): bound empty SSE reconnect loops
1 parent 4b64646 commit d7e05e7

2 files changed

Lines changed: 63 additions & 3 deletions

File tree

src/mcp/client/streamable_http.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -501,12 +501,15 @@ async def _handle_reconnection(
501501
# Track for potential further reconnection
502502
reconnect_last_event_id: str = last_event_id
503503
reconnect_retry_ms = retry_interval_ms
504+
made_progress = False
504505

505506
async for sse in event_source.aiter_sse():
506507
if sse.id: # pragma: no branch
507508
reconnect_last_event_id = sse.id
508509
if sse.retry is not None:
509510
reconnect_retry_ms = sse.retry
511+
if sse.event == "message" and bool(sse.data):
512+
made_progress = True
510513

511514
is_complete = await self._handle_sse_event(
512515
sse,
@@ -518,9 +521,11 @@ async def _handle_reconnection(
518521
await event_source.response.aclose()
519522
return
520523

521-
# Stream ended again without response - reconnect again (reset attempt counter)
524+
# Stream ended again without response - reconnect again. Only reset
525+
# the retry counter when the resumed stream delivered real data.
522526
logger.info("SSE stream disconnected, reconnecting...")
523-
await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, 0)
527+
next_attempt = 0 if made_progress else attempt + 1
528+
await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, next_attempt)
524529
except Exception as e:
525530
logger.debug(f"Reconnection failed: {e}")
526531
# Try to reconnect again if we still have an event ID

tests/client/test_streamable_http.py

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,14 +92,69 @@ async def test_sse_response_disconnect_before_any_event_id_fails_request() -> No
9292

9393
async with read_stream_writer, read_stream:
9494
await transport._handle_sse_response(response, ctx)
95-
message = await read_stream.receive()
95+
with anyio.fail_after(5):
96+
message = await read_stream.receive()
9697

9798
assert isinstance(message, SessionMessage)
9899
assert isinstance(message.message, JSONRPCError)
99100
assert message.message.id == 1
100101
assert message.message.error.code == CONNECTION_CLOSED
101102

102103

104+
@pytest.mark.anyio
105+
async def test_reconnection_empty_streams_count_toward_max_attempts(monkeypatch: pytest.MonkeyPatch) -> None:
106+
class PrimingOnlyEventSource:
107+
def __init__(self) -> None:
108+
self.response = httpx.Response(200)
109+
110+
async def __aenter__(self) -> "PrimingOnlyEventSource":
111+
nonlocal reconnect_attempts
112+
reconnect_attempts += 1
113+
return self
114+
115+
async def __aexit__(self, *args: object) -> None:
116+
return None
117+
118+
async def aiter_sse(self) -> object:
119+
yield type(
120+
"SSE",
121+
(),
122+
{"event": "message", "data": "", "id": f"event-{reconnect_attempts}", "retry": 0},
123+
)()
124+
125+
def connect_sse(*args: object, **kwargs: object) -> PrimingOnlyEventSource:
126+
return PrimingOnlyEventSource()
127+
128+
reconnect_attempts = 0
129+
monkeypatch.setattr(
130+
"mcp.client.streamable_http.aconnect_sse",
131+
connect_sse,
132+
)
133+
134+
transport = StreamableHTTPTransport("http://example.com/mcp")
135+
async with httpx.AsyncClient() as client:
136+
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](1)
137+
request = JSONRPCRequest(jsonrpc="2.0", id=1, method="tools/call", params={"name": "noop", "arguments": {}})
138+
ctx = RequestContext(
139+
client=client,
140+
session_id=None,
141+
session_message=SessionMessage(request),
142+
metadata=None,
143+
read_stream_writer=read_stream_writer,
144+
)
145+
146+
async with read_stream_writer, read_stream:
147+
with anyio.fail_after(5):
148+
await transport._handle_reconnection(ctx, "event-1", retry_interval_ms=0)
149+
message = await read_stream.receive()
150+
151+
assert reconnect_attempts == 2
152+
assert isinstance(message, SessionMessage)
153+
assert isinstance(message.message, JSONRPCError)
154+
assert message.message.id == 1
155+
assert message.message.error.code == CONNECTION_CLOSED
156+
157+
103158
@pytest.mark.anyio
104159
async def test_sse_response_disconnect_ignores_closed_read_stream() -> None:
105160
transport = StreamableHTTPTransport("http://example.com/mcp")

0 commit comments

Comments
 (0)