fix: wrap MCP call_tool coroutine in a Task so CancellationToken can cancel it - #8266
Open
Pujitha Paladugu (pujitha24) wants to merge 1 commit into
Open
Pujitha Paladugu (pujitha24) wants to merge 1 commit into
Pujitha Paladugu (pujitha24) wants to merge 1 commit into
Conversation
…cancel it
McpWorkbench.call_tool linked a CancellationToken to the object returned
by McpSessionActor.call("call_tool", ...), but McpSessionActor._run_actor
handed back a bare coroutine (session.call_tool(...) was never wrapped in
a Task/Future). CancellationToken.link_future calls future.cancel(),
which a coroutine object does not implement, so cancelling the token
mid-call raised AttributeError instead of cancelling, and the in-flight
tool call ran to completion anyway. A token cancelled before the call hit
the same AttributeError inside call_tool's try block, surfacing as a
confusing is_error=True result containing the AttributeError string.
Wrap the coroutine in asyncio.ensure_future(...) in _run_actor's
call_tool branch, matching the working pattern already used by the
sibling McpToolAdapter._run in _base.py. Also check
cancellation_token.is_cancelled() up front in McpWorkbench.call_tool,
mirroring _base.py, so an already-cancelled token raises
asyncio.CancelledError immediately instead of dispatching a doomed call.
CancelledError is a BaseException and is not caught by call_tool's
existing except Exception clause, so it propagates to the caller,
matching the already-shipped behavior of McpToolAdapter._run and of
StaticWorkbench.call_tool in autogen_core.
The other _run_actor command branches (list_tools, list_prompts,
list_resources, list_resource_templates, read_resource, get_prompt) are
left unchanged: none of them accept a cancellation_token in
McpWorkbench and none call link_future, so they cannot trigger this bug.
Validation: added test_mcp_cancellation.py with two regression tests,
confirmed to fail with the exact reported AttributeError before this fix
(reverted via git stash) and pass after. Ran `poe format`, `poe lint`,
`poe mypy`, and `poe pyright` for autogen-ext (all clean), and
`pytest packages/autogen-ext/tests/tools/test_mcp_cancellation.py
packages/autogen-ext/tests/tools/test_mcp_actor.py
packages/autogen-ext/tests/tools/test_mcp_workbench_warnings_and_errors.py
packages/autogen-ext/tests/tools/test_mcp_tools.py` (83 passed, 1 skipped,
pre-existing unrelated skip) plus the full tests/tools mcp-filtered suite
(147 passed, 1 skipped).
Report: microsoft#8265
Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Assisted-by: claude-sonnet-5 (via Claude Code)
This branch has not been deployed
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Why are these changes needed?
McpWorkbench.call_toollinks aCancellationTokento the object returned byMcpSessionActor.call("call_tool", ...). That object was a bare coroutine (McpSessionActor._run_actordidresult = session.call_tool(...)without ever wrapping it in aTask/Future), butCancellationToken.link_futurecallsfuture.cancel(), which a coroutine object does not implement. As a result:AttributeError: 'coroutine' object has no attribute 'cancel'out ofCancellationToken.cancel(), and the in-flight MCP tool call was never actually cancelled — it ran to completion anyway.call_toolraised the sameAttributeErrorinsidecall_tool'stryblock, which was swallowed by the genericexcept Exceptionhandler and surfaced to the caller as a confusingis_error=Trueresult containing theAttributeErrorstring instead of a clean cancellation signal.The sibling implementation
McpToolAdapter._run(same package,_base.py) already does this correctly by wrapping the coroutine withasyncio.ensure_future(...)before linking the cancellation token, and by checkingcancellation_token.is_cancelled()up front.Approach
_actor.py: inMcpSessionActor._run_actor'scall_toolbranch, wrapsession.call_tool(...)inasyncio.ensure_future(...)before handing it back, so the object linked to theCancellationTokenis an actualTaskwith a working.cancel()._workbench.py: inMcpWorkbench.call_tool, checkcancellation_token.is_cancelled()before dispatching to the actor and raiseasyncio.CancelledErrorimmediately, mirroringMcpToolAdapter._run's existing pattern, instead of starting a doomed tool call.asyncio.CancelledErroris aBaseException(Python 3.8+), so it is not caught by the existingexcept Exceptionclause and propagates to the caller, matching the already-shipped behavior ofMcpToolAdapter._run(see existing teststest_mcp_tool_adapter_run_cancelled_before_call/_during_call) and ofStaticWorkbench.call_tool(autogen_core), which also only catchesexcept Exception._run_actorcommand branches (list_tools,list_prompts,list_resources,list_resource_templates,read_resource,get_prompt) are left unchanged: none of them accept acancellation_tokeninMcpWorkbenchand none calllink_future, so they cannot trigger this bug.Validation
Added
python/packages/autogen-ext/tests/tools/test_mcp_cancellation.pywith two targeted regression tests, verified to FAIL with the exact reportedAttributeErrorbefore this fix (confirmed viagit stashon the two source files) and PASS after:test_call_tool_cancellation_stops_inflight_call: starts a mocked slowcall_tool, cancels the token mid-flight, assertstoken.cancel()does not raise and the call is cancelled promptly (< 1s) instead of running to completion (mocked 5s sleep).test_call_tool_pre_cancelled_token_reports_cancellation: cancels the token before callingcall_tool, assertsasyncio.CancelledErroris raised and the underlying session'scall_toolis never invoked.Commands run (all passing) from
python/, in theuv-managed virtualenv perpython/README.md:poe format(autogen-ext) — 1 file reformatted (the new test file), no other changes.poe lint(autogen-ext) — all checks passed.poe mypy(autogen-ext) — success, no issues in 198 source files.poe pyright(autogen-ext) — 0 errors, 0 warnings.pytest packages/autogen-ext/tests/tools/test_mcp_cancellation.py packages/autogen-ext/tests/tools/test_mcp_actor.py packages/autogen-ext/tests/tools/test_mcp_workbench_warnings_and_errors.py packages/autogen-ext/tests/tools/test_mcp_tools.py— 83 passed, 1 skipped (pre-existing GitHub-token-gated skip, unrelated).tests/tools -k mcp— 147 passed, 1 skipped.I also ran the full
autogen-extpackage test suite once; it showed unrelated failures (docker, playwright, llama_cpp model download, http_tool) that I traced to the sandbox's local disk being 99% full (166Mi free) — re-running the same MCP-adjacent tests without parallelism, on their own, passed cleanly, confirming those were environmental flakes and not caused by this change.Related issue number
Closes #8265
Checks