refactor: key the MCP agent-server session map on the connection

`to_mcp_server` keeps one ADK session per MCP connection, so successive tool
calls on that connection form a single conversation. It keyed that map on
`ctx.session`.

That key is correct on MCP SDK 1.x, where the server builds one session object
per connection. It is wrong on 2.x: the server builds a fresh `ServerSession`
for every inbound request and holds the connection on the session's private
`_connection`. The key would change on every call, so every tool call would
start a new conversation. Nothing raises. The agent just forgets.

Route the key through `_connection_key`, which reads `_connection` when the SDK
provides it and falls back to the session when it does not. That gives one key
per connection on both versions, and leaves 1.x behaviour unchanged.

The fallback degrades to one session per request on purpose. It must not fall
back to an object shared by all connections, because separate clients would
then share one conversation.

The private attribute is a stopgap. SDK 2.x already defines a public
`mcp.server.context.Context.connection`, but the server does not hand that
class to tool functions yet, so a tool's `Context` has no public route to its
connection.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 967462559

src/google/adk/tools/mcp_tool/_agent_to_mcp.py

@@ -82,6 +82,33 @@ def _part_to_content(part: types.Part) -> Optional[mcp_types.ContentBlock]:
   return None
 
 
+def _connection_key(ctx: Context[ServerSession, Any]) -> object:
+  """Returns the object that identifies the MCP connection behind ``ctx``.
+
+  The MCP SDK exposes no public per-connection handle. In SDK 1.x
+  ``ctx.session`` is itself one object per connection. In 2.x the server
+  builds a fresh ``ServerSession`` for every inbound message and keeps the
+  connection on its private ``_connection``, so ``ctx.session`` changes on
+  every call. Reading ``_connection`` when it is there gives one key per
+  connection on both versions.
+
+  TODO: Use a public accessor once the SDK adds one. SDK 2.x already defines
+  ``mcp.server.context.Context.connection``, but the server does not hand that
+  class to tool functions yet.
+
+  Args:
+    ctx: The MCP tool call context.
+
+  Returns:
+    The per-connection object, or the per-request session when the SDK gives
+    no connection. That fallback degrades to one agent session per request.
+    It must never fall back to an object shared by all connections, because
+    separate clients would then share one conversation.
+  """
+  session = ctx.session
+  return getattr(session, "_connection", session)
+
+
 async def _run_agent(
     runner: Runner,
     request: str,
@@ -107,15 +134,17 @@ async def _run_agent(
     images, audio, or other data the agent produced).
   """
   session_id: Optional[str] = None
+  connection: Optional[object] = None
   if ctx is not None and sessions is not None:
-    session_id = sessions.get(ctx.session)
+    connection = _connection_key(ctx)
+    session_id = sessions.get(connection)
   if session_id is None:
     session = await runner.session_service.create_session(
         app_name=runner.app_name, user_id=_MCP_USER_ID
     )
     session_id = session.id
-    if ctx is not None and sessions is not None:
-      sessions[ctx.session] = session_id
+    if sessions is not None and connection is not None:
+      sessions[connection] = 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(

tests/unittests/tools/mcp_tool/test_agent_to_mcp.py

@@ -21,6 +21,7 @@ from typing import AsyncGenerator
 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 _run_agent
 from google.adk.tools.mcp_tool._agent_to_mcp import to_mcp_server
 from google.genai import types
@@ -101,6 +102,29 @@ class _Connection:
   """Stand-in for an MCP connection object (weak-referenceable)."""
 
 
+class _RequestScopedSession:
+  """Stand-in for an MCP SDK 2.x ServerSession.
+
+  The 2.x server builds one of these per inbound request and holds the
+  connection on the private ``_connection``. It is hashable, like the real
+  class, so a stale per-request key silently starts a new conversation rather
+  than raising.
+  """
+
+  def __init__(self, connection: object):
+    self._connection = connection
+
+
+class _RequestScopedCtx:
+  """Fake MCP Context shaped like MCP SDK 2.x."""
+
+  def __init__(self, connection: object):
+    self.session = _RequestScopedSession(connection)
+
+  async def report_progress(self, *, progress, total=None, message=None):
+    pass
+
+
 @pytest.mark.asyncio
 async def test_to_mcp_server_registers_agent_as_single_tool():
   agent = _EchoAgent(name="my_agent", description="does useful things")
@@ -203,6 +227,48 @@ async def test_run_agent_uses_separate_sessions_across_connections():
   assert runner.session_ids == ["session-1", "session-2"]
 
 
+def test_connection_key_uses_the_session_when_it_has_no_connection():
+  """MCP SDK 1.x: the session is already one object per connection."""
+  session = _Connection()
+
+  assert _connection_key(_ConnCtx(session)) is session
+
+
+def test_connection_key_prefers_the_connection_behind_the_session():
+  """MCP SDK 2.x: the session is per request, the connection is not."""
+  connection = _Connection()
+  ctx = _RequestScopedCtx(connection)
+
+  assert _connection_key(ctx) is connection
+
+
+@pytest.mark.asyncio
+async def test_run_agent_reuses_one_session_when_sessions_are_per_request():
+  """Two requests on one connection stay in one conversation under SDK 2.x."""
+  runner = _FakeRunner([_text_event("ok")])
+  sessions: dict[object, str] = {}
+  connection = _Connection()
+
+  await _run_agent(runner, "first", _RequestScopedCtx(connection), sessions)
+  await _run_agent(runner, "second", _RequestScopedCtx(connection), sessions)
+
+  assert runner.create_session_calls == 1
+  assert runner.session_ids == ["session-1", "session-1"]
+
+
+@pytest.mark.asyncio
+async def test_run_agent_separates_connections_when_sessions_are_per_request():
+  """Separate clients never share a conversation under SDK 2.x."""
+  runner = _FakeRunner([_text_event("ok")])
+  sessions: dict[object, str] = {}
+
+  await _run_agent(runner, "a", _RequestScopedCtx(_Connection()), sessions)
+  await _run_agent(runner, "b", _RequestScopedCtx(_Connection()), sessions)
+
+  assert runner.create_session_calls == 2
+  assert runner.session_ids == ["session-1", "session-2"]
+
+
 @pytest.mark.asyncio
 async def test_call_tool_reuses_session_across_calls_on_one_connection():
   agent = _EchoAgent(name="assistant")