feat: start and close ClientSession in a single task in McpSessionManager

[challenger71498]

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

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

• Closes:
  • Agent Engine Agent fails to use MCP. Error point at bad handling of session management in async context #3950
  • Critical Runtime Crash: RuntimeError: Attempted to exit cancel scope... when using MCPToolset on Vertex AI Agent Engine #3731
  • Bug Report: MCP Module Raises CancelledError and Prevents Custom Error Handling #3708

2. Or, if no issue exists, describe the change:

Problem:

• ClientSession of https://github.com/modelcontextprotocol/python-sdk uses AnyIO for async task management.
• AnyIO TaskGroup requires its start and close must happen in a same task.
• Since McpSessionManager does not create task per client, the client might be closed by different task, cause the error: Attempted to exit cancel scope in a different task than it was entered in.

Solution:

I Suggest 2 changes:

Handling the ClientSession in a single task

• To start and close ClientSession by the same task, we need to wrap the whole lifecycle of ClientSession to a single task.
• SessionContext wraps the initialization and disposal of ClientSession to a single task, ensures that the ClientSession will be handled only in a dedicated task.

Add timeout for ClientSession

• Since now we are using task per ClientSession, task should never be leaked.
• But McpSessionManager does not deliver timeout directly to ClientSession when the type is not STDIO.
  • There is only timeout for httpx client when MCP type is SSE or StreamableHTTP.
  • But the timeout applys only to httpx client, so if there is an issue in MCP client itself(e.g. cannot get response from await session.call_tool() modelcontextprotocol/python-sdk#262), a tool call waits the result FOREVER!
• To overcome this issue, I propagated the sse_read_timeout to ClientSession.
  • timeout is too short for timeout for tool call, since its default value is only 5s.
  • sse_read_timeout is originally made for read timeout of SSE(default value of 5m or 300s), but actually most of SSE implementations from server (e.g. FastAPI, etc.) sends ping periodically(about 15s I assume), so in a normal circumstances this timeout is quite useless.
  • If the server does not send ping, the timeout is equal to tool call timeout. Therefore, it would be appropriate to use sse_read_timeout as tool call timeout.
  • Most of tool calls should finish within 5 minutes, and sse timeout is adjustable if not.
• If this change is not acceptable, we could make a dedicate parameter for tool call timeout(e.g. tool_call_timeout).

Testing Plan

• Although this does not change the interface itself, it changes its own session management logics, some existing tests are no longer valid.
  • I made changes to those tests, especially those of which validate session states(e.g. checking whether initialize() called).
  • Since now session is encapsulated with SessionContext, we cannot validate the initialized state of the session in TestMcpSessionManager, should validate it at TestSessionContext.
• Added a simple test for reproducing the issue(test_create_and_close_session_in_different_tasks).
• Also made a test for the new component: SessionContext.

Unit Tests:

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

=================================================================================== 3689 passed, 1 skipped, 2205 warnings in 63.39s (0:01:03) ===================================================================================

Manual End-to-End (E2E) Tests:

Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix.

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. no deps has been changed

Additional context

This PR is related to modelcontextprotocol/python-sdk#1817 since it also fixes endless tool call awaiting.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Summary of Changes

Hello @challenger71498, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses critical issues in the McpSessionManager related to asynchronous session management and timeout handling. It introduces a new SessionContext class that encapsulates the entire lifecycle of a ClientSession within a dedicated AnyIO task, thereby resolving concurrency errors caused by improper task switching during session creation and closure. Additionally, it enhances timeout mechanisms by ensuring that sse_read_timeout is correctly applied to ClientSession instances, preventing indefinite waits for tool calls and improving the overall reliability of session interactions.

Highlights

• Structured Concurrency Fix: Introduced SessionContext to ensure ClientSession's lifecycle (start and close) occurs within a single AnyIO task, resolving Attempted to exit cancel scope in a different task than it was entered in errors.
• Improved Timeout Handling: Propagated sse_read_timeout to ClientSession for non-STDIO connections, preventing tool calls from waiting indefinitely and providing a more appropriate timeout for long-running operations.
• Encapsulated Session Management: Refactored McpSessionManager to delegate ClientSession creation and disposal to the new SessionContext, simplifying McpSessionManager's logic and enhancing robustness.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[adk-bot]

Response from ADK Triaging Agent

Hello @challenger71498, thank you for your contribution!

Before we can merge this PR, could you please sign the Contributor License Agreement (CLA)? You can find more information on how to do this in the "cla/google" check at the bottom of this page.

Also, could you please provide details on the manual end-to-end (E2E) tests you performed for these changes? This will help the reviewers to better understand and verify your work.

Thank you!

[gemini-code-assist]

Code Review

This pull request introduces a robust solution to a tricky asyncio task management issue with anyio by encapsulating the ClientSession lifecycle within a dedicated SessionContext class. This is a great architectural improvement that correctly handles structured concurrency. The addition of timeouts for tool calls is another valuable enhancement for the tool's reliability. The new SessionContext is thoroughly tested, and existing tests are updated accordingly, which is excellent. I have a couple of minor suggestions to further improve code clarity and cleanliness.