From 65f07e86e23da353f1aee9ff43f126e58004e505 Mon Sep 17 00:00:00 2001 From: harshal-96 Date: Thu, 17 Sep 2026 13:57:04 +0530 Subject: [PATCH 1/2] fix(tools): reap ADK sessions of dead MCP connections in to_mcp_server 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. --- .../adk/tools/mcp_tool/_agent_to_mcp.py | 72 +++++++++++- .../tools/mcp_tool/test_agent_to_mcp.py | 111 +++++++++++++++++- 2 files changed, 179 insertions(+), 4 deletions(-) diff --git a/src/google/adk/tools/mcp_tool/_agent_to_mcp.py b/src/google/adk/tools/mcp_tool/_agent_to_mcp.py index 50084f0767b..0e24fc93892 100644 --- a/src/google/adk/tools/mcp_tool/_agent_to_mcp.py +++ b/src/google/adk/tools/mcp_tool/_agent_to_mcp.py @@ -17,6 +17,7 @@ from __future__ import annotations import base64 +import logging from typing import Any from typing import MutableMapping from typing import Optional @@ -37,6 +38,8 @@ from ...runners import Runner from ...sessions.in_memory_session_service import InMemorySessionService +logger = logging.getLogger("google_adk." + __name__) + _MCP_USER_ID = "mcp_user" _INLINE_RESOURCE_URI = "resource://adk-agent/inline-data" @@ -109,11 +112,57 @@ class to tool functions yet. return getattr(session, "_connection", session) +async def _reap_orphaned_sessions( + runner: Runner, + sessions: MutableMapping[object, str], + created: set[str], +) -> None: + """Deletes ADK sessions whose MCP connection is gone. + + ``sessions`` holds its connections weakly, so an entry vanishes when its + connection is garbage-collected; the ADK session it pointed to would stay + in the session service forever. Under a stateless streamable HTTP transport + the connection lives for a single request, which turns that into one leaked + session per tool call. Reaping runs lazily from the next tool call because + a GC callback may fire without a running event loop. + + Args: + runner: The Runner whose session service owns the sessions. + sessions: Per-connection map from MCP connection to ADK session id. + created: Ids of every session ever entered into ``sessions``. Ids no + longer reachable through ``sessions`` are deleted and removed from it. + """ + live = set(sessions.values()) + for session_id in created - live: + if session_id not in created: + # A concurrent reap already took this one; the discard below and this + # check share one synchronous stretch, so each id is deleted once. + continue + created.discard(session_id) + try: + await runner.session_service.delete_session( + app_name=runner.app_name, + user_id=_MCP_USER_ID, + session_id=session_id, + ) + except Exception: # pylint: disable=broad-exception-caught + # Reaping is housekeeping; it must not fail the tool call that + # triggered it. Put the id back so a later call retries the delete. + created.add(session_id) + logger.warning( + "Failed to delete orphaned MCP agent session %s; will retry on a" + " later tool call.", + session_id, + exc_info=True, + ) + + async def _run_agent( runner: Runner, request: str, ctx: Optional[Context[ServerSession, Any]] = None, sessions: Optional[MutableMapping[object, str]] = None, + created: Optional[set[str]] = None, ) -> list[mcp_types.ContentBlock]: """Runs the agent for one request and returns its final response content. @@ -128,6 +177,8 @@ async def _run_agent( request: The user request text for this call. ctx: The MCP tool call context, used for progress and session reuse. sessions: Per-connection map from MCP connection to ADK session id. + created: Set recording the id of every session entered into ``sessions``, + so `_reap_orphaned_sessions` can delete the ones whose connection dies. Returns: The agent's final response as a list of MCP content blocks (text plus any @@ -144,7 +195,11 @@ async def _run_agent( ) session_id = session.id if sessions is not None and connection is not None: + # No await between the two writes: an id is either absent from both or + # present in both, so the reaper never sees a session it cannot delete. sessions[connection] = session_id + if created is not None: + created.add(session_id) new_message = types.Content(role="user", parts=[types.Part(text=request)]) final_content: list[mcp_types.ContentBlock] = [] async for event in runner.run_async( @@ -183,11 +238,16 @@ def to_mcp_server( lets harnesses that speak MCP drive an ADK agent. One ADK session is kept per MCP connection, so successive tool calls on the - same connection form a single multi-turn conversation. + same connection form a single multi-turn conversation. When a connection + goes away its ADK session is deleted from the session service on a later + tool call, so a long-running server does not accumulate dead conversations. The caller chooses the transport, e.g. ``server.run(transport="stdio")`` for a local host or ``server.run(transport="streamable-http")`` for a networked - one. + one. A stateless streamable HTTP deployment (``stateless_http=True``, e.g. + behind an autoscaler) gets a fresh connection per request, so every tool + call is its own single-turn conversation whose session is likewise + reclaimed. Args: agent: The ADK agent to serve. @@ -213,11 +273,17 @@ def to_mcp_server( # WeakKeyDictionary() instantiation below as abstract-class-instantiated. # pylint: disable-next=abstract-class-instantiated sessions: MutableMapping[object, str] = weakref.WeakKeyDictionary() + # Ids of every session in `sessions`, kept strongly so the sessions of + # collected connections can still be found and deleted. + created_session_ids: set[str] = set() async def call_agent( request: str, ctx: Context[ServerSession, Any] ) -> list[mcp_types.ContentBlock]: - return await _run_agent(agent_runner, request, ctx, sessions) + await _reap_orphaned_sessions(agent_runner, sessions, created_session_ids) + return await _run_agent( + agent_runner, request, ctx, sessions, created_session_ids + ) server.add_tool( call_agent, diff --git a/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py b/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py index 6ffb8a21f61..0ecfa7a6b7a 100644 --- a/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py +++ b/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py @@ -15,13 +15,16 @@ from __future__ import annotations import base64 +import gc from types import SimpleNamespace from typing import AsyncGenerator +import weakref from google.adk.agents.base_agent import BaseAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events.event import Event from google.adk.tools.mcp_tool._agent_to_mcp import _connection_key +from google.adk.tools.mcp_tool._agent_to_mcp import _reap_orphaned_sessions from google.adk.tools.mcp_tool._agent_to_mcp import _run_agent from google.adk.tools.mcp_tool._agent_to_mcp import to_mcp_server from google.genai import types @@ -76,12 +79,25 @@ def __init__(self, events: list[Event]): self._events = events self.create_session_calls = 0 self.session_ids: list[str] = [] - self.session_service = SimpleNamespace(create_session=self._create_session) + self.deleted_session_ids: list[str] = [] + self.failing_deletes = 0 + self.session_service = SimpleNamespace( + create_session=self._create_session, + delete_session=self._delete_session, + ) async def _create_session(self, *, app_name: str, user_id: str): self.create_session_calls += 1 return SimpleNamespace(id=f"session-{self.create_session_calls}") + async def _delete_session( + self, *, app_name: str, user_id: str, session_id: str + ): + if self.failing_deletes > 0: + self.failing_deletes -= 1 + raise ConnectionError("session service unavailable") + self.deleted_session_ids.append(session_id) + async def run_async( self, *, user_id: str, session_id: str, new_message: types.Content ) -> AsyncGenerator[Event, None]: @@ -271,6 +287,99 @@ async def test_run_agent_separates_connections_when_sessions_are_per_request(): assert runner.session_ids == ["session-1", "session-2"] +@pytest.mark.asyncio +async def test_reap_deletes_only_sessions_no_longer_reachable(): + runner = _FakeRunner([_text_event("ok")]) + connection = _Connection() + sessions: dict[object, str] = {connection: "session-live"} + created = {"session-live", "session-dead"} + + await _reap_orphaned_sessions(runner, sessions, created) + + assert runner.deleted_session_ids == ["session-dead"] + assert created == {"session-live"} + + +@pytest.mark.asyncio +async def test_session_of_a_collected_connection_is_reaped(): + """A conversation must not outlive its connection in the session service.""" + runner = _FakeRunner([_text_event("ok")]) + sessions: weakref.WeakKeyDictionary[object, str] = ( + # pylint: disable-next=abstract-class-instantiated + weakref.WeakKeyDictionary() + ) + created: set[str] = set() + ctx = _ConnCtx(_Connection()) + + await _run_agent(runner, "hi", ctx, sessions, created) + del ctx + gc.collect() + await _reap_orphaned_sessions(runner, sessions, created) + + assert runner.deleted_session_ids == ["session-1"] + assert not created + + +@pytest.mark.asyncio +async def test_per_request_connections_do_not_accumulate_sessions(): + """Stateless streamable HTTP builds a fresh connection per request; each + request's session must be reclaimed instead of leaking one per tool call.""" + runner = _FakeRunner([_text_event("ok")]) + sessions: weakref.WeakKeyDictionary[object, str] = ( + # pylint: disable-next=abstract-class-instantiated + weakref.WeakKeyDictionary() + ) + created: set[str] = set() + + for request in ("a", "b", "c"): + await _reap_orphaned_sessions(runner, sessions, created) + ctx = _RequestScopedCtx(_Connection()) + await _run_agent(runner, request, ctx, sessions, created) + del ctx + gc.collect() + await _reap_orphaned_sessions(runner, sessions, created) + + assert runner.deleted_session_ids == ["session-1", "session-2", "session-3"] + assert not created + + +@pytest.mark.asyncio +async def test_reap_failure_does_not_raise_and_is_retried(): + """A session service outage must not fail the live tool call, and the + orphaned session must be deleted once the service recovers.""" + runner = _FakeRunner([_text_event("ok")]) + runner.failing_deletes = 1 + connection = _Connection() + sessions: dict[object, str] = {connection: "session-live"} + created = {"session-live", "session-dead"} + + await _reap_orphaned_sessions(runner, sessions, created) + + assert runner.deleted_session_ids == [] + assert created == {"session-live", "session-dead"} + + await _reap_orphaned_sessions(runner, sessions, created) + + assert runner.deleted_session_ids == ["session-dead"] + assert created == {"session-live"} + + +@pytest.mark.asyncio +async def test_call_tool_reaps_conversation_of_closed_connection(): + agent = _EchoAgent(name="assistant") + runner = _FakeRunner([_text_event("ok")]) + server = to_mcp_server(agent, runner=runner) + + async with connected_client_session(server) as client: + await client.call_tool("assistant", {"request": "first"}) + gc.collect() + async with connected_client_session(server) as client: + await client.call_tool("assistant", {"request": "second"}) + + assert runner.session_ids == ["session-1", "session-2"] + assert runner.deleted_session_ids == ["session-1"] + + @pytest.mark.asyncio async def test_call_tool_reuses_session_across_calls_on_one_connection(): agent = _EchoAgent(name="assistant") From 4311d5cff829724857517c49cee74fe59ddf1429 Mon Sep 17 00:00:00 2001 From: harshal-96 Date: Fri, 18 Sep 2026 13:10:45 +0530 Subject: [PATCH 2/2] feat(tools): add delete_orphaned_sessions opt-out to to_mcp_server 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. --- src/google/adk/tools/mcp_tool/_agent_to_mcp.py | 17 ++++++++++++++--- .../tools/mcp_tool/test_agent_to_mcp.py | 18 ++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/src/google/adk/tools/mcp_tool/_agent_to_mcp.py b/src/google/adk/tools/mcp_tool/_agent_to_mcp.py index 0e24fc93892..87ef69d4bc7 100644 --- a/src/google/adk/tools/mcp_tool/_agent_to_mcp.py +++ b/src/google/adk/tools/mcp_tool/_agent_to_mcp.py @@ -228,6 +228,7 @@ def to_mcp_server( name: Optional[str] = None, instructions: Optional[str] = None, runner: Optional[Runner] = None, + delete_orphaned_sessions: bool = True, ) -> FastMCP: """Exposes an ADK agent as an MCP server. @@ -255,6 +256,12 @@ def to_mcp_server( instructions: Optional instructions the MCP host may show to its model. runner: A pre-built Runner. If omitted, one is created with in-memory services. + delete_orphaned_sessions: Whether to delete a connection's ADK session + from the session service once the connection is gone. Defaults to True, + which keeps a long-running server's memory bounded. Set to False to + retain finished conversations in the session service, e.g. when a + caller-supplied ``runner`` uses a persistent session service whose + records are read after the fact; the caller then owns their cleanup. Returns: A ``FastMCP`` server exposing the agent as a single tool. @@ -274,13 +281,17 @@ def to_mcp_server( # pylint: disable-next=abstract-class-instantiated sessions: MutableMapping[object, str] = weakref.WeakKeyDictionary() # Ids of every session in `sessions`, kept strongly so the sessions of - # collected connections can still be found and deleted. - created_session_ids: set[str] = set() + # collected connections can still be found and deleted. None disables the + # tracking and with it the reaping. + created_session_ids: Optional[set[str]] = ( + set() if delete_orphaned_sessions else None + ) async def call_agent( request: str, ctx: Context[ServerSession, Any] ) -> list[mcp_types.ContentBlock]: - await _reap_orphaned_sessions(agent_runner, sessions, created_session_ids) + if created_session_ids is not None: + await _reap_orphaned_sessions(agent_runner, sessions, created_session_ids) return await _run_agent( agent_runner, request, ctx, sessions, created_session_ids ) diff --git a/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py b/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py index 0ecfa7a6b7a..f444c345476 100644 --- a/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py +++ b/tests/unittests/tools/mcp_tool/test_agent_to_mcp.py @@ -380,6 +380,24 @@ async def test_call_tool_reaps_conversation_of_closed_connection(): assert runner.deleted_session_ids == ["session-1"] +@pytest.mark.asyncio +async def test_call_tool_retains_sessions_when_deletion_is_opted_out(): + """delete_orphaned_sessions=False keeps finished conversations in the + session service, for persistent services whose records are read later.""" + agent = _EchoAgent(name="assistant") + runner = _FakeRunner([_text_event("ok")]) + server = to_mcp_server(agent, runner=runner, delete_orphaned_sessions=False) + + async with connected_client_session(server) as client: + await client.call_tool("assistant", {"request": "first"}) + gc.collect() + async with connected_client_session(server) as client: + await client.call_tool("assistant", {"request": "second"}) + + assert runner.session_ids == ["session-1", "session-2"] + assert runner.deleted_session_ids == [] + + @pytest.mark.asyncio async def test_call_tool_reuses_session_across_calls_on_one_connection(): agent = _EchoAgent(name="assistant")