-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Sending cancellation notification to server based on client anyio.CancelScope status #628
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
davemssavage
wants to merge
33
commits into
modelcontextprotocol:main
Choose a base branch
from
davemssavage:cancellation
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+169
−46
Open
Changes from all commits
Commits
Show all changes
33 commits
Select commit
Hold shift + click to select a range
f9598ad
attempt to support cancelation
davemssavage dfb3686
fix sending cancel
davemssavage abda067
add decorator for handling cancelation
davemssavage 24553c6
fix capitalisation on error message
davemssavage fe49931
try to check notification event received
davemssavage efd0ffd
fix test
davemssavage 1364b7a
fix debug message
davemssavage 17ae44c
prevent cancel on intialisation as per https://modelcontextprotocol.i…
davemssavage 06f4b3c
remove unused imports
davemssavage 07e1a52
fixed long line
davemssavage 92f806b
fixed long line
davemssavage e46c693
updates from ruff format
davemssavage 2a24e0c
removed dev print statements committed by mistake
davemssavage 87722f8
add constant for request cancelled and use it in case of cancelled task
davemssavage f0782d2
fix long line
davemssavage a0164f9
fix import order
davemssavage bf220d5
add assert that cancel scope triggered on server, add comments for re…
davemssavage 45ac52a
trivial update to comment capitalisation
davemssavage 8b7f1cd
simplify test code
davemssavage ac4b822
add tests to assert cancellable=False behaves as expected
davemssavage d86b4a5
tidy up ruff check
davemssavage 6f4ae44
ruff format update
davemssavage 11d2e52
fixed test description
davemssavage 235df35
use defined types in decorator
davemssavage f96aaa5
Add docstring info for cancellable and add TODO note for initialised …
davemssavage bd73448
set read_timeout_seconds to avoid test blocking for ever in case some…
davemssavage 3817fe2
remove unnecessary shielded cancel scope
davemssavage 2fd27a2
whitespace
davemssavage 2e86d32
clarified doc string comment
davemssavage 1039b99
fixed doc string and added comment on internal notification method
davemssavage b926970
Merge branch 'main' into cancellation
davemssavage cf1bb2b
merge from upstream main
davemssavage 39f7739
ignore type warning for internal use of send_notification
davemssavage File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,4 +1,5 @@ | ||
from collections.abc import AsyncGenerator | ||
from datetime import timedelta | ||
|
||
import anyio | ||
import pytest | ||
|
@@ -9,10 +10,6 @@ | |
from mcp.shared.exceptions import McpError | ||
from mcp.shared.memory import create_connected_server_and_client_session | ||
from mcp.types import ( | ||
CancelledNotification, | ||
CancelledNotificationParams, | ||
ClientNotification, | ||
ClientRequest, | ||
EmptyResult, | ||
) | ||
|
||
|
@@ -46,11 +43,11 @@ async def test_in_flight_requests_cleared_after_completion( | |
@pytest.mark.anyio | ||
async def test_request_cancellation(): | ||
"""Test that requests can be cancelled while in-flight.""" | ||
# The tool is already registered in the fixture | ||
|
||
ev_tool_called = anyio.Event() | ||
ev_tool_cancelled = anyio.Event() | ||
ev_cancelled = anyio.Event() | ||
request_id = None | ||
ev_cancel_notified = anyio.Event() | ||
|
||
# Start the request in a separate task so we can cancel it | ||
def make_server() -> Server: | ||
|
@@ -59,14 +56,24 @@ def make_server() -> Server: | |
# Register the tool handler | ||
@server.call_tool() | ||
async def handle_call_tool(name: str, arguments: dict | None) -> list: | ||
nonlocal request_id, ev_tool_called | ||
nonlocal ev_tool_called, ev_tool_cancelled | ||
if name == "slow_tool": | ||
request_id = server.request_context.request_id | ||
ev_tool_called.set() | ||
await anyio.sleep(10) # Long enough to ensure we can cancel | ||
return [] | ||
with anyio.CancelScope(): | ||
try: | ||
await anyio.sleep(10) # Long enough to ensure we can cancel | ||
return [] | ||
except anyio.get_cancelled_exc_class() as err: | ||
ev_tool_cancelled.set() | ||
raise err | ||
|
||
raise ValueError(f"Unknown tool: {name}") | ||
|
||
@server.cancel_notification() | ||
async def handle_cancel(requestId: str | int, reason: str | None): | ||
nonlocal ev_cancel_notified | ||
ev_cancel_notified.set() | ||
|
||
# Register the tool so it shows up in list_tools | ||
@server.list_tools() | ||
async def handle_list_tools() -> list[types.Tool]: | ||
|
@@ -80,20 +87,10 @@ async def handle_list_tools() -> list[types.Tool]: | |
|
||
return server | ||
|
||
async def make_request(client_session): | ||
async def make_request(client_session: ClientSession): | ||
nonlocal ev_cancelled | ||
try: | ||
await client_session.send_request( | ||
ClientRequest( | ||
types.CallToolRequest( | ||
method="tools/call", | ||
params=types.CallToolRequestParams( | ||
name="slow_tool", arguments={} | ||
), | ||
) | ||
), | ||
types.CallToolResult, | ||
) | ||
await client_session.call_tool("slow_tool") | ||
pytest.fail("Request should have been cancelled") | ||
except McpError as e: | ||
# Expected - request was cancelled | ||
|
@@ -110,17 +107,87 @@ async def make_request(client_session): | |
with anyio.fail_after(1): # Timeout after 1 second | ||
await ev_tool_called.wait() | ||
|
||
# Send cancellation notification | ||
assert request_id is not None | ||
await client_session.send_notification( | ||
ClientNotification( | ||
CancelledNotification( | ||
method="notifications/cancelled", | ||
params=CancelledNotificationParams(requestId=request_id), | ||
) | ||
) | ||
) | ||
# Cancel the task via task group | ||
tg.cancel_scope.cancel() | ||
|
||
# Give cancellation time to process | ||
with anyio.fail_after(1): | ||
await ev_cancelled.wait() | ||
|
||
# Check server cancel notification received | ||
with anyio.fail_after(1): | ||
await ev_cancel_notified.wait() | ||
|
||
# Give cancellation time to process on server | ||
with anyio.fail_after(1): | ||
await ev_tool_cancelled.wait() | ||
|
||
|
||
@pytest.mark.anyio | ||
async def test_request_cancellation_uncancellable(): | ||
"""Test that asserts a call with cancellable=False is not cancelled on | ||
server when cancel scope on client is set.""" | ||
|
||
ev_tool_called = anyio.Event() | ||
ev_tool_commplete = anyio.Event() | ||
ev_cancelled = anyio.Event() | ||
|
||
# Start the request in a separate task so we can cancel it | ||
def make_server() -> Server: | ||
server = Server(name="TestSessionServer") | ||
|
||
# Register the tool handler | ||
@server.call_tool() | ||
async def handle_call_tool(name: str, arguments: dict | None) -> list: | ||
nonlocal ev_tool_called, ev_tool_commplete | ||
if name == "slow_tool": | ||
ev_tool_called.set() | ||
with anyio.CancelScope(): | ||
with anyio.fail_after(10): # Long enough to ensure we can cancel | ||
await ev_cancelled.wait() | ||
ev_tool_commplete.set() | ||
return [] | ||
|
||
raise ValueError(f"Unknown tool: {name}") | ||
|
||
# Register the tool so it shows up in list_tools | ||
@server.list_tools() | ||
async def handle_list_tools() -> list[types.Tool]: | ||
return [ | ||
types.Tool( | ||
name="slow_tool", | ||
description="A slow tool that takes 10 seconds to complete", | ||
inputSchema={}, | ||
) | ||
] | ||
|
||
return server | ||
|
||
async def make_request(client_session: ClientSession): | ||
nonlocal ev_cancelled | ||
try: | ||
await client_session.call_tool( | ||
"slow_tool", | ||
cancellable=False, | ||
read_timeout_seconds=timedelta(seconds=10), | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. A test that validates a short timeout with an uncancelable call would be sensible |
||
) | ||
except McpError: | ||
pytest.fail("Request should not have been cancelled") | ||
|
||
async with create_connected_server_and_client_session( | ||
make_server() | ||
) as client_session: | ||
async with anyio.create_task_group() as tg: | ||
tg.start_soon(make_request, client_session) | ||
|
||
# Wait for the request to be in-flight | ||
with anyio.fail_after(1): # Timeout after 1 second | ||
await ev_tool_called.wait() | ||
|
||
# Cancel the task via task group | ||
tg.cancel_scope.cancel() | ||
ev_cancelled.set() | ||
|
||
# Check server completed regardless | ||
with anyio.fail_after(1): | ||
await ev_tool_commplete.wait() |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Alternative to this might be to separate client cancelation and server cancelation, e.g. client can be cancelled and server cancellation event is only set if a flag such as 'propagate_client_cancelation_to_server' (shorter names available) is set on the request.