Skip to content

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
microsoft:mainfrom
pujitha24:auto/issue-8265
Open

Pujitha Paladugu (pujitha24) wants to merge 1 commit into
microsoft:mainfrom
pujitha24:auto/issue-8265

Conversation

@pujitha24

Copy link
Copy Markdown

Why are these changes needed?

McpWorkbench.call_tool links a CancellationToken to the object returned by McpSessionActor.call("call_tool", ...). That object was a bare coroutine (McpSessionActor._run_actor did result = session.call_tool(...) without ever wrapping it in a Task/Future), but CancellationToken.link_future calls future.cancel(), which a coroutine object does not implement. As a result:

  • Cancelling the token while a tool call is in flight raised AttributeError: 'coroutine' object has no attribute 'cancel' out of CancellationToken.cancel(), and the in-flight MCP tool call was never actually cancelled — it ran to completion anyway.
  • Cancelling the token before calling call_tool raised the same AttributeError inside call_tool's try block, which was swallowed by the generic except Exception handler and surfaced to the caller as a confusing is_error=True result containing the AttributeError string instead of a clean cancellation signal.

The sibling implementation McpToolAdapter._run (same package, _base.py) already does this correctly by wrapping the coroutine with asyncio.ensure_future(...) before linking the cancellation token, and by checking cancellation_token.is_cancelled() up front.

Approach

  • _actor.py: in McpSessionActor._run_actor's call_tool branch, wrap session.call_tool(...) in asyncio.ensure_future(...) before handing it back, so the object linked to the CancellationToken is an actual Task with a working .cancel().
  • _workbench.py: in McpWorkbench.call_tool, check cancellation_token.is_cancelled() before dispatching to the actor and raise asyncio.CancelledError immediately, mirroring McpToolAdapter._run's existing pattern, instead of starting a doomed tool call. asyncio.CancelledError is a BaseException (Python 3.8+), so it is not caught by the existing except Exception clause and propagates to the caller, matching the already-shipped behavior of McpToolAdapter._run (see existing tests test_mcp_tool_adapter_run_cancelled_before_call / _during_call) and of StaticWorkbench.call_tool (autogen_core), which also only catches except Exception.
  • 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 python/packages/autogen-ext/tests/tools/test_mcp_cancellation.py with two targeted regression tests, verified to FAIL with the exact reported AttributeError before this fix (confirmed via git stash on the two source files) and PASS after:

  • test_call_tool_cancellation_stops_inflight_call: starts a mocked slow call_tool, cancels the token mid-flight, asserts token.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 calling call_tool, asserts asyncio.CancelledError is raised and the underlying session's call_tool is never invoked.

Commands run (all passing) from python/, in the uv-managed virtualenv per python/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).
  • Full tests/tools -k mcp — 147 passed, 1 skipped.

I also ran the full autogen-ext package 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

…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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

McpWorkbench.call_tool passes a raw coroutine to CancellationToken.link_future: cancelling raises AttributeError and the tool call is never cancelled

1 participant