Skip to content

Commit dcfc757

Browse files
author
Jianke LIN
committed
fix(streamable-http): close SSE responses on errors
1 parent a4f4ccd commit dcfc757

2 files changed

Lines changed: 136 additions & 14 deletions

File tree

src/mcp/client/streamable_http.py

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -253,20 +253,25 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
253253
if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch
254254
original_request_id = ctx.session_message.message.id
255255

256-
async with ctx.client.sse(self.url, headers=headers) as event_source:
257-
event_source.response.raise_for_status()
258-
logger.debug("Resumption GET SSE connection established")
256+
event_source: EventSource | None = None
257+
try:
258+
async with ctx.client.sse(self.url, headers=headers) as es:
259+
event_source = es
260+
event_source.response.raise_for_status()
261+
logger.debug("Resumption GET SSE connection established")
259262

260-
async for sse in event_source: # pragma: no branch
261-
is_complete = await self._handle_sse_event(
262-
sse,
263-
ctx.read_stream_writer,
264-
original_request_id,
265-
ctx.metadata.on_resumption_token_update if ctx.metadata else None,
266-
)
267-
if is_complete:
268-
await event_source.response.aclose()
269-
break
263+
async for sse in event_source: # pragma: no branch
264+
is_complete = await self._handle_sse_event(
265+
sse,
266+
ctx.read_stream_writer,
267+
original_request_id,
268+
ctx.metadata.on_resumption_token_update if ctx.metadata else None,
269+
)
270+
if is_complete:
271+
break
272+
finally:
273+
if event_source is not None:
274+
await event_source.response.aclose()
270275

271276
def _consume_modern_cancellation(self, session_message: SessionMessage) -> bool:
272277
"""Translate an outbound `notifications/cancelled` at 2026; True means "do not POST".
@@ -442,10 +447,11 @@ async def _handle_sse_response(
442447
# If the SSE event indicates completion, like returning response/error
443448
# break the loop
444449
if is_complete:
445-
await response.aclose()
446450
return # Normal completion, no reconnect needed
447451
except Exception:
448452
logger.debug("SSE stream ended", exc_info=True) # pragma: lax no cover
453+
finally:
454+
await response.aclose()
449455

450456
# Stream ended without response - reconnect if we received an event with ID
451457
if last_event_id is not None:
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
import contextlib
2+
3+
import anyio
4+
import httpx
5+
import pytest
6+
from httpx_sse import ServerSentEvent
7+
8+
from mcp.client.streamable_http import RequestContext, StreamableHTTPTransport
9+
from mcp.shared.message import ClientMessageMetadata, SessionMessage
10+
from mcp.types import JSONRPCRequest
11+
12+
13+
class _RaiseEventSource:
14+
def __init__(self, response: httpx.Response) -> None:
15+
self.response = response
16+
17+
async def aiter_sse(self):
18+
yield ServerSentEvent(event="message", data="", id=None, retry=None)
19+
raise RuntimeError("boom")
20+
21+
22+
@pytest.mark.anyio
23+
async def test_handle_sse_response_closes_response_on_exception(monkeypatch: pytest.MonkeyPatch) -> None:
24+
closed = False
25+
26+
async def spy_aclose() -> None:
27+
nonlocal closed
28+
closed = True
29+
30+
response = httpx.Response(200, headers={"content-type": "text/event-stream"})
31+
response.aclose = spy_aclose # type: ignore[method-assign]
32+
33+
monkeypatch.setattr("mcp.client.streamable_http.EventSource", _RaiseEventSource)
34+
35+
send_stream, receive_stream = anyio.create_memory_object_stream[SessionMessage | Exception](1)
36+
async with send_stream, receive_stream:
37+
transport = StreamableHTTPTransport("http://example.invalid/mcp")
38+
async with httpx.AsyncClient(transport=httpx.MockTransport(lambda _: httpx.Response(200))) as client:
39+
ctx = RequestContext(
40+
client=client,
41+
session_id=None,
42+
session_message=SessionMessage(JSONRPCRequest(method="initialize", params={}, jsonrpc="2.0", id=1)),
43+
metadata=ClientMessageMetadata(),
44+
read_stream_writer=send_stream,
45+
)
46+
await transport._handle_sse_response(response, ctx)
47+
48+
assert closed
49+
50+
51+
@pytest.mark.anyio
52+
async def test_handle_resumption_request_closes_response_when_aconnect_sse_raises(
53+
monkeypatch: pytest.MonkeyPatch,
54+
) -> None:
55+
@contextlib.asynccontextmanager
56+
async def fake_aconnect_sse(*_args, **_kwargs):
57+
raise RuntimeError("connect failed")
58+
yield
59+
60+
monkeypatch.setattr("mcp.client.streamable_http.aconnect_sse", fake_aconnect_sse)
61+
62+
send_stream, receive_stream = anyio.create_memory_object_stream[SessionMessage | Exception](1)
63+
async with send_stream, receive_stream:
64+
transport = StreamableHTTPTransport("http://example.invalid/mcp")
65+
metadata = ClientMessageMetadata(resumption_token="1")
66+
async with httpx.AsyncClient(transport=httpx.MockTransport(lambda _: httpx.Response(200))) as client:
67+
ctx = RequestContext(
68+
client=client,
69+
session_id=None,
70+
session_message=SessionMessage(JSONRPCRequest(method="initialize", params={}, jsonrpc="2.0", id=1)),
71+
metadata=metadata,
72+
read_stream_writer=send_stream,
73+
)
74+
75+
with pytest.raises(RuntimeError, match="connect failed"):
76+
await transport._handle_resumption_request(ctx)
77+
78+
79+
@pytest.mark.anyio
80+
async def test_handle_resumption_request_closes_response_on_exception(monkeypatch: pytest.MonkeyPatch) -> None:
81+
closed = False
82+
83+
async def spy_aclose() -> None:
84+
nonlocal closed
85+
closed = True
86+
87+
response = httpx.Response(
88+
200,
89+
headers={"content-type": "text/event-stream"},
90+
request=httpx.Request("GET", "http://example.invalid/mcp"),
91+
)
92+
response.aclose = spy_aclose # type: ignore[method-assign]
93+
94+
@contextlib.asynccontextmanager
95+
async def fake_aconnect_sse(*_args, **_kwargs):
96+
yield _RaiseEventSource(response)
97+
98+
monkeypatch.setattr("mcp.client.streamable_http.aconnect_sse", fake_aconnect_sse)
99+
100+
send_stream, receive_stream = anyio.create_memory_object_stream[SessionMessage | Exception](1)
101+
async with send_stream, receive_stream:
102+
transport = StreamableHTTPTransport("http://example.invalid/mcp")
103+
metadata = ClientMessageMetadata(resumption_token="1")
104+
async with httpx.AsyncClient(transport=httpx.MockTransport(lambda _: httpx.Response(200))) as client:
105+
ctx = RequestContext(
106+
client=client,
107+
session_id=None,
108+
session_message=SessionMessage(JSONRPCRequest(method="initialize", params={}, jsonrpc="2.0", id=1)),
109+
metadata=metadata,
110+
read_stream_writer=send_stream,
111+
)
112+
113+
with pytest.raises(RuntimeError, match="boom"):
114+
await transport._handle_resumption_request(ctx)
115+
116+
assert closed

0 commit comments

Comments
 (0)