fix: invalidate cached MCP session on server-side session loss - #7061
Open
sachiantany wants to merge 1 commit into
Open
fix: invalidate cached MCP session on server-side session loss#7061sachiantany wants to merge 1 commit into
sachiantany wants to merge 1 commit into
Conversation
When an MCP server scales to zero and back, its in-memory sessions are lost. The existing @retry_on_errors decorator retries the call, but create_session() returns the same cached (now dead) session because _is_session_disconnected() only checks transport-level stream closure — the transport is still alive, only the server-side session is gone. Add MCPSessionManager.invalidate_session() so callers can mark a session key as stale. _is_session_disconnected() now consults an _invalidated_sessions set before falling back to the stream check. McpToolset._execute_with_session() calls invalidate_session() on any exception, so the next retry from @retry_on_errors builds a fresh session instead of reusing the dead one.
|
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. |
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.
When an MCP server scales to zero and back, its in-memory sessions are lost. The existing @retry_on_errors decorator retries the call, but create_session() returns the same cached (now dead) session because _is_session_disconnected() only checks transport-level stream closure — the transport is still alive, only the server-side session is gone.
Add MCPSessionManager.invalidate_session() so callers can mark a session key as stale. _is_session_disconnected() now consults an _invalidated_sessions set before falling back to the stream check. McpToolset._execute_with_session() calls invalidate_session() on any exception, so the next retry from @retry_on_errors builds a fresh session instead of reusing the dead one.
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):
Problem:
When an MCP server (e.g. on Cloud Run) scales to zero and restarts, its in-memory sessions are lost. The agent's
@retry_on_errorsdecorator retries the failed call, butMCPSessionManager.create_session()returns the same cached dead session because_is_session_disconnected()only checks transport-level stream closure — the transport is still alive, only the server-side session is gone. Every retry reuses the stale session and fails, until all retries are exhausted.Solution:
MCPSessionManager.invalidate_session()method to mark a session key as stale_is_session_disconnected()to check an_invalidated_sessionsset before falling back to the transport stream checkinvalidate_session()fromMcpToolset._execute_with_session()on any exception, before re-raisingcreate_session()call, the invalidated session is replaced with a fresh one and the flag is clearedThis lets the existing
@retry_on_errorsmechanism work correctly for server-side session loss, not just transport failures.Testing Plan
Unit Tests:
105 tests passed, 0 failures in
tests/unittests/tools/mcp_tool/test_mcp_session_manager.py(5.30s). New tests added:test_is_session_disconnected_with_invalidated_key— invalidated key returns disconnected, other keys unaffectedtest_invalidate_session_marks_key— correct session key added to_invalidated_sessionstest_invalidated_session_is_replaced_on_next_create— full integration: inreturns new session → flag cleared_invalidated_sessionsround-trips as empty setManual End-to-End (E2E) Tests:
Tested with an ADK agent connected to an MCP server deployed on Cloud Run (Stith scale-to-zero enabled. After ~15 minutes idle, the server scaled to zero. On the next tool call, the server scaled back up — the agent detected the stale session, invalidated it, and successfully reconnected with a fresh session on retry instead of failing with repeated stale session errors.
Checklist
Additional context
3 files changed, 125 insertions(+), 3 deletions(-):
src/google/adk/tools/mcp_tool/mcp_session_manager.py— 45 additionssrc/google/adk/tools/mcp_tool/mcp_toolset.py— 6 additionstests/unittests/tools/mcp_tool/test_mcp_session_manager.py— 77 additions (3 new tests + 1 updated)