Skip to content

Commit ed89a6a

Browse files
committed
fix: auto-send notifications/cancelled on anyio task cancellation
Closes #1410. Users of high-level APIs like client.call_tool() have no access to the internal equest_id, making it impossible to manually send a cancel notification when a task is aborted. The server would keep processing a request that nobody is waiting for. This fix intercepts �nyio task/scope cancellation inside send_request and automatically dispatches otifications/cancelled to the peer before re-raising. The send is shielded from the cancellation scope and guarded by a 2-second timeout to prevent deadlock if the transport buffer is full. A equest_sent flag ensures we only notify the server if the request actually hit the wire — avoiding spurious cancels for requests that were never received. This mirrors the pattern used by the TypeScript SDK with AbortSignal and applies to both client→server and server→client request directions since the change lives in BaseSession.
1 parent 3d7b311 commit ed89a6a

File tree

2 files changed

+84
-0
lines changed

2 files changed

+84
-0
lines changed

src/mcp/shared/session.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
INVALID_PARAMS,
2424
REQUEST_TIMEOUT,
2525
CancelledNotification,
26+
CancelledNotificationParams,
2627
ClientNotification,
2728
ClientRequest,
2829
ClientResult,
@@ -269,6 +270,7 @@ async def send_request(
269270
# Store the callback for this request
270271
self._progress_callbacks[request_id] = progress_callback
271272

273+
request_sent = False
272274
try:
273275
target = request_data.get("params", {}).get("name")
274276
span_name = f"MCP send {request.method} {target}" if target else f"MCP send {request.method}"
@@ -284,6 +286,7 @@ async def send_request(
284286

285287
jsonrpc_request = JSONRPCRequest(jsonrpc="2.0", id=request_id, **request_data)
286288
await self._write_stream.send(SessionMessage(message=jsonrpc_request, metadata=metadata))
289+
request_sent = True
287290

288291
# request read timeout takes precedence over session read timeout
289292
timeout = request_read_timeout_seconds or self._session_read_timeout_seconds
@@ -301,6 +304,26 @@ async def send_request(
301304
else:
302305
return result_type.model_validate(response_or_error.result, by_name=False)
303306

307+
except anyio.get_cancelled_exc_class():
308+
# Automatically notify the other side when a task/scope is cancelled,
309+
# so the peer can abort work for a request nobody is waiting for.
310+
if request_sent:
311+
with anyio.CancelScope(shield=True):
312+
try:
313+
# Add a short timeout to prevent deadlock if the transport buffer is full
314+
with anyio.move_on_after(2.0):
315+
await self.send_notification(
316+
CancelledNotification( # type: ignore[arg-type]
317+
params=CancelledNotificationParams(
318+
request_id=request_id,
319+
reason="Task cancelled",
320+
)
321+
)
322+
)
323+
except Exception:
324+
pass # Transport may already be closed
325+
raise
326+
304327
finally:
305328
self._response_streams.pop(request_id, None)
306329
self._progress_callbacks.pop(request_id, None)

tests/server/test_cancel_handling.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,3 +248,64 @@ async def run_server():
248248
# Without the fixes: RuntimeError (dict mutation) or ClosedResourceError
249249
# (respond after write-stream close) escapes run_server and this hangs.
250250
await server_run_returned.wait()
251+
252+
253+
@pytest.mark.anyio
254+
async def test_anyio_cancel_scope_sends_cancelled_notification() -> None:
255+
"""Cancelling a call_tool via anyio cancel scope should automatically
256+
send notifications/cancelled to the server, causing it to abort the handler."""
257+
258+
tool_started = anyio.Event()
259+
handler_cancelled = anyio.Event()
260+
261+
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
262+
return ListToolsResult(
263+
tools=[
264+
Tool(
265+
name="slow_tool",
266+
description="A slow tool for testing cancellation",
267+
input_schema={},
268+
)
269+
]
270+
)
271+
272+
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
273+
if params.name == "slow_tool":
274+
tool_started.set()
275+
try:
276+
await anyio.sleep_forever()
277+
except anyio.get_cancelled_exc_class():
278+
handler_cancelled.set()
279+
raise
280+
raise ValueError(f"Unknown tool: {params.name}") # pragma: no cover
281+
282+
server = Server(
283+
"test-server",
284+
on_list_tools=handle_list_tools,
285+
on_call_tool=handle_call_tool,
286+
)
287+
288+
async with Client(server) as client:
289+
# Cancel the call_tool via anyio scope cancellation.
290+
# send_request should automatically send notifications/cancelled.
291+
async with anyio.create_task_group() as tg:
292+
293+
async def do_call() -> None:
294+
with anyio.CancelScope() as scope:
295+
# Store scope so the outer task can cancel it
296+
do_call.scope = scope # type: ignore[attr-defined]
297+
await client.call_tool("slow_tool", {})
298+
299+
tg.start_soon(do_call)
300+
301+
# Wait for the server handler to start
302+
await tool_started.wait()
303+
304+
# Cancel the client-side scope — this should trigger auto-notification
305+
do_call.scope.cancel() # type: ignore[attr-defined]
306+
307+
# Give the server a moment to process the cancellation
308+
await anyio.sleep(0.1)
309+
310+
# The server handler should have been cancelled via the notification
311+
assert handler_cancelled.is_set(), "Server handler was not cancelled — notifications/cancelled was not sent"

0 commit comments

Comments
 (0)