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 50084f0767..87ef69d4bc 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( @@ -173,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. @@ -183,11 +239,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. @@ -195,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. @@ -213,11 +280,21 @@ 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. 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]: - return await _run_agent(agent_runner, request, ctx, sessions) + 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 + ) 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 6ffb8a21f6..f444c34547 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,117 @@ 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_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")