Skip to content

Commit d7e7b46

Browse files
committed
send errors to pending requests if server closes
1 parent ae77772 commit d7e7b46

File tree

3 files changed

+70
-1
lines changed

3 files changed

+70
-1
lines changed

src/mcp/shared/session.py

+8
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
from mcp.shared.exceptions import McpError
1616
from mcp.types import (
17+
CONNECTION_CLOSED,
1718
CancelledNotification,
1819
ClientNotification,
1920
ClientRequest,
@@ -374,6 +375,13 @@ async def _receive_loop(self) -> None:
374375
)
375376
)
376377

378+
# after the read stream is closed, we need to send errors
379+
# to any pending requests
380+
for id, stream in self._response_streams.items():
381+
error = ErrorData(code=CONNECTION_CLOSED, message="Connection closed")
382+
await stream.send(JSONRPCError(jsonrpc="2.0", id=id, error=error))
383+
await stream.aclose()
384+
377385
async def _received_request(
378386
self, responder: RequestResponder[ReceiveRequestT, SendResultT]
379387
) -> None:

src/mcp/types.py

+4
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,10 @@ class JSONRPCResponse(BaseModel):
137137
model_config = ConfigDict(extra="allow")
138138

139139

140+
# SDK error codes
141+
CONNECTION_CLOSED = -32000
142+
# REQUEST_TIMEOUT = -32001 # the typescript sdk uses this
143+
140144
# Standard JSON-RPC error codes
141145
PARSE_ERROR = -32700
142146
INVALID_REQUEST = -32600

tests/shared/test_session.py

+58-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@
77
from mcp.client.session import ClientSession
88
from mcp.server.lowlevel.server import Server
99
from mcp.shared.exceptions import McpError
10-
from mcp.shared.memory import create_connected_server_and_client_session
10+
from mcp.shared.memory import (
11+
create_client_server_memory_streams,
12+
create_connected_server_and_client_session,
13+
)
1114
from mcp.types import (
1215
CancelledNotification,
1316
CancelledNotificationParams,
@@ -124,3 +127,57 @@ async def make_request(client_session):
124127
# Give cancellation time to process
125128
with anyio.fail_after(1):
126129
await ev_cancelled.wait()
130+
131+
132+
@pytest.mark.anyio
133+
async def test_connection_closed():
134+
"""
135+
Test that pending requests are cancelled when the connection is closed remotely.
136+
"""
137+
138+
ev_closed = anyio.Event()
139+
ev_response = anyio.Event()
140+
141+
async with create_client_server_memory_streams() as (
142+
client_streams,
143+
server_streams,
144+
):
145+
client_read, client_write = client_streams
146+
server_read, server_write = server_streams
147+
148+
async def make_request(client_session):
149+
"""Send a request in a separate task"""
150+
nonlocal ev_response
151+
try:
152+
# any request will do
153+
await client_session.initialize()
154+
pytest.fail("Request should have errored")
155+
except McpError as e:
156+
# Expected - request errored
157+
assert "Connection closed" in str(e)
158+
ev_response.set()
159+
160+
async def mock_server():
161+
"""Wait for a request, then close the connection"""
162+
nonlocal ev_closed
163+
# Wait for a request
164+
await server_read.receive()
165+
# Close the connection, as if the server exited
166+
server_write.close()
167+
server_read.close()
168+
ev_closed.set()
169+
170+
async with (
171+
anyio.create_task_group() as tg,
172+
ClientSession(
173+
read_stream=client_read,
174+
write_stream=client_write,
175+
) as client_session,
176+
):
177+
tg.start_soon(make_request, client_session)
178+
tg.start_soon(mock_server)
179+
180+
with anyio.fail_after(1):
181+
await ev_closed.wait()
182+
with anyio.fail_after(1):
183+
await ev_response.wait()

0 commit comments

Comments
 (0)