fix(tools): reap ADK sessions of dead MCP connections in to_mcp_server - #7155
Open
harshal-96 wants to merge 2 commits into
Open
harshal-96 wants to merge 2 commits into
harshal-96 wants to merge 2 commits into
Conversation
to_mcp_server keeps one ADK session per MCP connection in a WeakKeyDictionary. When a connection is garbage-collected the map entry disappears, but the ADK session it pointed to stays in the session service forever, with its full event history. A long-running server therefore accumulates one dead conversation per closed connection, and a stateless streamable HTTP deployment, where the SDK builds a fresh transport for every request, leaks one session per tool call. Track the id of every session entered into the connection map and, at the start of each tool call, delete the sessions whose connection is no longer reachable. Reaping runs lazily from the tool call rather than a GC callback because finalizers may fire without a running event loop. Also document how a stateless streamable HTTP deployment behaves: each call is a fresh single-turn conversation whose session is reclaimed. Tested with mcp 1.26.0 and 2.2.0: 17 passed each.
|
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. |
Author
|
@googlebot I signed it! |
Callers who wire to_mcp_server to a persistent session service may want finished conversations to remain readable after their connection dies, e.g. for audit. delete_orphaned_sessions=False disables the reaping and leaves session lifecycle to the caller. The default stays True so a long-running server's memory is bounded out of the box.
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.
Fixes #7154
Description
to_mcp_serverkeeps one ADK session per MCP connection in aweakref.WeakKeyDictionary. When a connection is garbage collected the map entry disappears, but the ADK session it pointed to stays in the session service forever, with its full event history. A long running server therefore accumulates one dead conversation per closed connection, and a stateless streamable HTTP deployment, where the MCP SDK builds a fresh transport for every request, leaks one session per tool call.This PR tracks the id of every session entered into the connection map and, at the start of each tool call, deletes the sessions whose connection is no longer reachable through the weak map.
Design notes:
delete_sessionfailure is logged and retried on a later call, and never fails the live tool call. This matters for caller supplied runners with database or Vertex backed session services.mcp_userid, but I can add an opt-out flag if retention is wanted for some deployments.to_mcp_serverdocstring now documents stateless streamable HTTP behavior: each call is a fresh single turn conversation whose session is reclaimed.Testing plan
5 new unit tests in
tests/unittests/tools/mcp_tool/test_agent_to_mcp.py:test_reap_deletes_only_sessions_no_longer_reachable: the reap deletes exactly the unreachable ids.test_session_of_a_collected_connection_is_reaped: a conversation does not outlive its connection.test_per_request_connections_do_not_accumulate_sessions: the stateless per request pattern stays flat instead of leaking one session per call.test_reap_failure_does_not_raise_and_is_retried: a session service outage does not fail the live tool call, and the orphan is deleted once the service recovers.test_call_tool_reaps_conversation_of_closed_connection: end to end through a real in-memory MCP client and server; this test fails on main (nothing is ever deleted) and passes with the fix.Results:
pytest tests/unittests/tools/mcp_tool/test_agent_to_mcp.py: 18 passed on each of mcp 1.24.0, 1.26.0, and 2.2.0 (the pin admits 1.x and 2.x).pytest tests/unittests/tools/mcp_tool: 376 passed.run_streamable_http_async(stateless_http=True)on localhost, one freshstreamable_http_clientconnection per call. 9 calls leave 9 sessions in the service on main and 1 with this PR (the most recent call's session, reclaimed on the next call).sessions left in the service: 10on main,1with this PR.Formatting: pyink 25.12 and isort 8.0.1 (the pinned versions) produce no changes.