fix(mcp): cancel the in-flight tool call when the caller is cancelled

[AtulJoshi1206]

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

No existing issue; described below following the bug-report structure.

Describe the Bug:

SessionContext._run_guarded races the in-flight MCP call against the session's background task with asyncio.wait:

coro_task = asyncio.ensure_future(coro)

done, _ = await asyncio.wait(
    [coro_task, self._task],
    return_when=asyncio.FIRST_COMPLETED,
)

asyncio.wait does not own the futures it waits on. It re-raises CancelledError at the await and leaves coro_task running. So when the caller is cancelled, the live session.call_tool() is orphaned: nothing holds it, nothing awaits it, its exception is never retrieved, and the remote call is never cancelled.

Both branches that exit normally already cancel coro_task, and test_run_guarded_cancels_coro_when_task_dies_first pins that contract for the transport-crash case. The cancellation path is the one exit that misses it.

This path is the default: McpTool._run_async_impl calls _run_guarded whenever FeatureName._MCP_GRACEFUL_ERROR_HANDLING is on. Cancellation is routine here, arriving from an aborted agent run, a human-in-the-loop interrupt, an outer timeout, or a disconnecting client.

Impact:

McpTool._run_async_impl releases the session in a finally:

finally:
  self._mcp_session_manager._end_session_use(final_headers)

On cancellation that finally runs immediately, so the pool's in-flight count drops to zero and _session_last_used is stamped while the orphan is still reading the transport. _evict_idle_sessions is then free to close that transport underneath the orphaned call once the idle TTL passes. The pool's accounting is actively wrong for as long as the orphan lives.

Steps to Reproduce:

On current main (665ec9835), start a guarded call, cancel the caller, and check whether the inner call is still running:

caller = asyncio.create_task(session_context._run_guarded(slow_coro()))
await coro_started.wait()
caller.cancel()
with pytest.raises(asyncio.CancelledError):
    await caller
assert coro_was_cancelled is True

Observed Behavior:

assert cancelled is True, 'in-flight tool call was not cancelled'
AssertionError: in-flight tool call was not cancelled
assert False is True

The inner call runs to completion after the caller is gone.

Expected Behavior:

The in-flight call is cancelled and drained before _run_guarded propagates the cancellation, so the session is only released once nothing is still reading the transport.

Solution:

Cancel and drain coro_task when the wait itself is interrupted, then re-raise. The existing cancel-and-drain in the transport-crash branch is identical, so this change factors both into one _cancel_and_drain helper rather than duplicating it.

BaseException rather than CancelledError, so a KeyboardInterrupt or SystemExit through the same frame does not leak either.

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

Added test_run_guarded_cancels_coro_when_caller_is_cancelled next to the existing test_run_guarded_cancels_coro_when_task_dies_first, covering the mirror case. It fails on main with the assertion above and passes with this change.

$ pytest tests/unittests/tools/mcp_tool -q
391 passed

$ pytest tests/unittests -q
15226 passed, 86 skipped, 27 xfailed, 2 xpassed in 337.16s (0:05:37)
# 1 unrelated failure, test_import_loading.py::…[agent], is an artifact of the
# git-worktree checkout I built this on; it fails there on an unmodified main
# and passes in a normal clone.

$ pre-commit run --files src/google/adk/tools/mcp_tool/session_context.py tests/unittests/tools/mcp_tool/test_session_context.py
# all hooks pass

$ mypy src/google/adk/tools/mcp_tool/session_context.py
# 9 pre-existing errors, identical with and without this change; no new errors

Manual End-to-End (E2E) Tests:

Driven through SessionContext directly, since reaching this path end-to-end needs a live MCP server. With a started session and a call blocked on a long read, cancelling the caller task:

before (main):      events: ['caller cancelled', 'TOOL CALL RAN TO COMPLETION (orphaned)']
after (this PR):    events: ['tool call was cancelled (correct)', 'caller cancelled']

The ordering flip is the point: after the change the inner call has already unwound by the time the caller's cancellation propagates, so _end_session_use cannot run while the transport is still in use.

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules. (none)

Additional context

No public API change. _cancel_and_drain is module-private, and the normal-exit behaviour of _run_guarded is byte-for-byte what it was.

🤖 Generated with Claude Code