Fix: Resolve CancelScope error in McpToolset session creation (issue #5729)

[9chait9]

Fixes #5729

This PR addresses an intermittent RuntimeError: Attempted to exit cancel scope in a different task that occurs during MCP session creation when LlmAgents are served via to_a2a() with multiple McpToolsets using StreamableHTTPServerParams. The root cause is asyncio.wait_for creating a nested task around the AnyIO CancelScope entry, leading to the cancellation happening from a different task.

The fix involves removing the problematic asyncio.wait_for wrapper around exit_stack.enter_async_context(self._client) in SessionContext._run. This ensures that the CancelScope is always entered and exited within the same task, satisfying AnyIO's requirements. The connection-establishment timeout is still enforced by MCPSessionManager.create_session via its outer asyncio.wait_for.

Test Plan

Objective: Verify that the RuntimeError: Attempted to exit cancel scope in a different task no longer occurs when using multiple McpToolsets with StreamableHTTPServerParams under concurrent load.

Scope:

• src/google/adk/tools/mcp_tool/session_context.py modifications.
• Interaction with MCPSessionManager.create_session.
• Stability of LlmAgent session creation with to_a2a() when configured with multiple McpToolsets and StreamableHTTPServerParams.

Test Cases:

1. Unit Test for SessionContext._run (New or Modified):

   • Description: Create a mock _client for SessionContext that simulates a delayed enter_async_context.
   • Expected Behavior:
     • The asyncio.wait_for around exit_stack.enter_async_context(self._client) should not be present.
     • The RuntimeError: Attempted to exit cancel scope in a different task should not be raised when the context entry eventually completes or is cancelled externally.
     • The MCPSessionManager.create_session should still enforce the overall connection establishment timeout.
   • Rationale: Directly verifies the removal of the problematic asyncio.wait_for and its impact on CancelScope behavior.

2. Integration Test: Multiple LlmAgents with McpToolset and StreamableHTTPServerParams:

   • Description:
     • Configure two or more LlmAgent instances, each served via to_a2a().
     • Each LlmAgent should use at least one McpToolset.
     • The StreamableHTTPServerParams should be configured for the McpToolsets.
     • Execute a series of concurrent requests to these agents, simulating high load and potential timeouts during session creation.
   • Expected Behavior:
     • All LlmAgent sessions should be created successfully without RuntimeError or ValueError: Tool '' not found.
     • Agents should operate with complete tool lists.
     • The system should remain stable under concurrent access.
   • Rationale: Replicates the original reported issue's conditions, confirming the fix at a higher integration level. This is critical as the problem scales with concurrent sessions.

3. Negative Test: StreamableHTTPServerParams with _MCP_GRACEFUL_ERROR_HANDLING (if applicable):

   • Description: If there's a way to explicitly enable/disable the _MCP_GRACEFUL_ERROR_HANDLING feature flag for testing, run the integration test (Test Case 2) with the flag both enabled and disabled (if the old code path can still be triggered).
   • Expected Behavior: The fix should be effective regardless of the feature flag's state, as the problematic asyncio.wait_for is being removed entirely.
   • Rationale: Ensures the fix is robust and doesn't depend on external feature flag configurations.

Testing Environment:

• A test environment capable of simulating concurrent LlmAgent and McpToolset interactions.
• The google/adk-python repository with the proposed changes.

Monitoring and Logging:

• Monitor logs for any RuntimeError: Attempted to exit cancel scope in a different task or ValueError: Tool '' not found messages.
• Verify that logger.debug('Session has been successfully initialized') appears as expected for all sessions.

Regression Testing:

• Ensure that existing test cases for McpToolset and LlmAgent functionality continue to pass without introducing new regressions.
• Specifically, verify that SSE transport (which was not affected by this bug) continues to function correctly.

[adk-bot]

Response from ADK Triaging Agent

Hello @9chait9, thank you for creating this PR!

We appreciate your contribution to resolving the CancelScope error in McpToolset session creation. To help us review and merge your changes more efficiently, please make sure your pull request follows our Contribution Guidelines.

Currently, the following sections are missing or incomplete in your PR description:

1. Testing Plan: Please include a testing plan section in your PR describing how you verified these changes (including unit tests or manual verification steps).
2. Logs or Screenshots: Since this is a bug fix, please provide logs or screenshots of the success/failure before and after the fix is applied so reviewers can better understand the resolution.

Thank you for your help in keeping the codebase robust!