Method _run

src/google/adk/tools/mcp_tool/session_context.py:248–324

Run the complete session context within a single task.

(self)

    await self.close()

  async def _run(self):
    """Run the complete session context within a single task."""
    try:
      async with AsyncExitStack() as exit_stack:
        if is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING):  # pylint: disable=protected-access
          # Post-fix: do NOT wrap in asyncio.wait_for. The MCP client uses
          # AnyIO TaskGroup/CancelScope internally, which must be entered
          # and exited in the same task. asyncio.wait_for runs its target
          # in a nested task and can cancel from a different task on
          # timeout, producing "Attempted to exit cancel scope in a
          # different task" errors. The connection-establishment timeout
          # is still enforced by MCPSessionManager.create_session via its
          # outer asyncio.wait_for around
          # exit_stack.enter_async_context(SessionContext(...)).
          transports = await exit_stack.enter_async_context(self._client)
        else:
          # Pre-fix behavior: wrap with asyncio.wait_for so the inner
          # context entry has its own timeout. Callers that depend on
          # this inner timeout firing rely on this path; without it,
          # mocks that delay `__aenter__` cause tests to time out at the
          # test framework limit instead of the configured per-step timeout.
          transports = await asyncio.wait_for(
              exit_stack.enter_async_context(self._client),
              timeout=self._timeout,
          )
        # The streamable http client returns a GetSessionCallback in addition
        # to the read/write MemoryObjectStreams needed to build the
        # ClientSession. We limit to the first two values to be compatible
        # with all clients.
        if self._is_stdio:
          session = await exit_stack.enter_async_context(
              ClientSession(
                  *transports[:2],
                  read_timeout_seconds=timedelta(seconds=self._timeout)
                  if self._timeout is not None
                  else None,
                  sampling_callback=self._sampling_callback,
                  sampling_capabilities=self._sampling_capabilities,
              )
          )
        else:
          # For SSE and Streamable HTTP clients, use the sse_read_timeout
          # instead of the connection timeout as the read_timeout for the session.
          session = await exit_stack.enter_async_context(
              ClientSession(
                  *transports[:2],
                  read_timeout_seconds=timedelta(seconds=self._sse_read_timeout)
                  if self._sse_read_timeout is not None
                  else None,
                  sampling_callback=self._sampling_callback,
                  sampling_capabilities=self._sampling_capabilities,
              )
          )
        # pylint: disable-next=protected-access
        if is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING):
          # Use anyio.fail_after to keep session.initialize within the AnyIO
          # cancel scope instead of asyncio.wait_for which runs in a nested
          # task.