Skip to content

Commit b93cdb8

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

File tree

3 files changed

+68
-1
lines changed

3 files changed

+68
-1
lines changed

src/mcp/shared/session.py

+8
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
ClientNotification,
1919
ClientRequest,
2020
ClientResult,
21+
CONNECTION_CLOSED,
2122
ErrorData,
2223
JSONRPCError,
2324
JSONRPCMessage,
@@ -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

+56-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,55 @@ 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+
"""Test that pending requests are cancelled when the connection is closed remotely."""
135+
136+
ev_closed = anyio.Event()
137+
ev_response = anyio.Event()
138+
139+
async with create_client_server_memory_streams() as (
140+
client_streams,
141+
server_streams,
142+
):
143+
client_read, client_write = client_streams
144+
server_read, server_write = server_streams
145+
146+
async def make_request(client_session):
147+
"""Send a request in a separate task"""
148+
nonlocal ev_response
149+
try:
150+
# any request will do
151+
await client_session.initialize()
152+
pytest.fail("Request should have errored")
153+
except McpError as e:
154+
# Expected - request errored
155+
assert "Connection closed" in str(e)
156+
ev_response.set()
157+
158+
async def mock_server():
159+
"""Wait for a request, then close the connection"""
160+
nonlocal ev_closed
161+
# Wait for a request
162+
await server_read.receive()
163+
# Close the connection, as if the server exited
164+
server_write.close()
165+
server_read.close()
166+
ev_closed.set()
167+
168+
async with (
169+
anyio.create_task_group() as tg,
170+
ClientSession(
171+
read_stream=client_read,
172+
write_stream=client_write,
173+
) as client_session,
174+
):
175+
tg.start_soon(make_request, client_session)
176+
tg.start_soon(mock_server)
177+
178+
with anyio.fail_after(1):
179+
await ev_closed.wait()
180+
with anyio.fail_after(1):
181+
await ev_response.wait()

0 commit comments

Comments
 (0)