chore: import upstream snapshot with attribution

tests/unittests/tools/mcp_tool/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unittests/tools/mcp_tool/test_agent_to_mcp.py

@@ -0,0 +1,217 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import base64
+from types import SimpleNamespace
+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 _run_agent
+from google.adk.tools.mcp_tool._agent_to_mcp import to_mcp_server
+from google.genai import types
+from mcp.shared.memory import create_connected_server_and_client_session
+import pytest
+
+
+class _EchoAgent(BaseAgent):
+  """Minimal agent that emits a single final text event."""
+
+  reply: str = "hello from the agent"
+
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        author=self.name,
+        content=types.Content(
+            role="model", parts=[types.Part(text=self.reply)]
+        ),
+    )
+
+
+def _text_event(text: str, *, partial: bool = False) -> Event:
+  return Event(
+      author="a",
+      partial=partial,
+      content=types.Content(role="model", parts=[types.Part(text=text)]),
+  )
+
+
+def _image_event(data: bytes, mime_type: str) -> Event:
+  return Event(
+      author="a",
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part(inline_data=types.Blob(data=data, mime_type=mime_type))
+          ],
+      ),
+  )
+
+
+class _FakeRunner:
+  """Runner stub that yields a fixed event sequence."""
+
+  app_name = "fake"
+
+  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)
+
+  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 run_async(
+      self, *, user_id: str, session_id: str, new_message: types.Content
+  ) -> AsyncGenerator[Event, None]:
+    self.session_ids.append(session_id)
+    for event in self._events:
+      yield event
+
+
+class _ConnCtx:
+  """Fake MCP Context carrying a per-connection session object."""
+
+  def __init__(self, session: object):
+    self.session = session
+
+  async def report_progress(self, *, progress, total=None, message=None):
+    pass
+
+
+class _Connection:
+  """Stand-in for an MCP connection object (weak-referenceable)."""
+
+
+@pytest.mark.asyncio
+async def test_to_mcp_server_registers_agent_as_single_tool():
+  agent = _EchoAgent(name="my_agent", description="does useful things")
+
+  server = to_mcp_server(agent)
+  tools = await server.list_tools()
+
+  assert len(tools) == 1
+  assert tools[0].name == "my_agent"
+  assert tools[0].description == "does useful things"
+  assert "request" in tools[0].inputSchema["properties"]
+
+
+@pytest.mark.asyncio
+async def test_to_mcp_server_name_override():
+  agent = _EchoAgent(name="my_agent")
+
+  server = to_mcp_server(agent, name="custom")
+  tools = await server.list_tools()
+
+  assert tools[0].name == "custom"
+
+
+@pytest.mark.asyncio
+async def test_call_tool_runs_agent_end_to_end():
+  agent = _EchoAgent(name="assistant")
+  server = to_mcp_server(agent)
+
+  async with create_connected_server_and_client_session(server) as client:
+    result = await client.call_tool("assistant", {"request": "hi"})
+
+  assert not result.isError
+  assert "hello from the agent" in result.content[0].text
+
+
+@pytest.mark.asyncio
+async def test_run_agent_returns_only_final_text():
+  runner = _FakeRunner([_text_event("answer")])
+
+  result = await _run_agent(runner, "hi")
+
+  assert [block.type for block in result] == ["text"]
+  assert result[0].text == "answer"
+
+
+@pytest.mark.asyncio
+async def test_run_agent_reports_intermediate_events_as_progress():
+  reported: list[str] = []
+
+  class _Ctx:
+
+    async def report_progress(self, *, progress, total=None, message=None):
+      reported.append(message)
+
+  runner = _FakeRunner(
+      [_text_event("thinking", partial=True), _text_event("done")]
+  )
+
+  result = await _run_agent(runner, "hi", _Ctx())
+
+  assert result[0].text == "done"
+  assert reported == ["thinking"]
+
+
+@pytest.mark.asyncio
+async def test_run_agent_maps_image_output_to_image_content():
+  png = b"\x89PNG\r\n\x1a\n"
+  runner = _FakeRunner([_image_event(png, "image/png")])
+
+  result = await _run_agent(runner, "draw a logo")
+
+  assert len(result) == 1
+  assert result[0].type == "image"
+  assert result[0].mimeType == "image/png"
+  assert base64.b64decode(result[0].data) == png
+
+
+@pytest.mark.asyncio
+async def test_run_agent_reuses_one_session_per_connection():
+  runner = _FakeRunner([_text_event("ok")])
+  sessions: dict[object, str] = {}
+  ctx = _ConnCtx(_Connection())
+
+  await _run_agent(runner, "first", ctx, sessions)
+  await _run_agent(runner, "second", ctx, sessions)
+
+  assert runner.create_session_calls == 1
+  assert runner.session_ids == ["session-1", "session-1"]
+
+
+@pytest.mark.asyncio
+async def test_run_agent_uses_separate_sessions_across_connections():
+  runner = _FakeRunner([_text_event("ok")])
+  sessions: dict[object, str] = {}
+
+  await _run_agent(runner, "a", _ConnCtx(_Connection()), sessions)
+  await _run_agent(runner, "b", _ConnCtx(_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")
+  runner = _FakeRunner([_text_event("ok")])
+  server = to_mcp_server(agent, runner=runner)
+
+  async with create_connected_server_and_client_session(server) as client:
+    await client.call_tool("assistant", {"request": "first"})
+    await client.call_tool("assistant", {"request": "second"})
+
+  assert runner.create_session_calls == 1
+  assert runner.session_ids == ["session-1", "session-1"]

tests/unittests/tools/mcp_tool/test_conversion_utils.py

@@ -0,0 +1,209 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for MCP tool conversion utilities."""
+
+from __future__ import annotations
+
+from unittest import mock
+
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.mcp_tool.conversion_utils import adk_to_mcp_tool_type
+from google.genai import types
+import mcp.types as mcp_types
+
+
+class TestAdkToMcpToolType:
+  """Tests for adk_to_mcp_tool_type function."""
+
+  def test_tool_with_no_declaration(self):
+    """Test conversion when tool has no declaration."""
+    mock_tool = mock.Mock(spec=BaseTool)
+    mock_tool.name = "test_tool"
+    mock_tool.description = "Test tool"
+    mock_tool._get_declaration.return_value = None
+
+    result = adk_to_mcp_tool_type(mock_tool)
+
+    assert isinstance(result, mcp_types.Tool)
+    assert result.name == "test_tool"
+    assert result.description == "Test tool"
+    assert result.inputSchema == {}
+
+  def test_tool_with_parameters_schema(self):
+    """Test conversion when tool has parameters Schema object."""
+    mock_tool = mock.Mock(spec=BaseTool)
+    mock_tool.name = "get_weather"
+    mock_tool.description = "Gets weather information"
+
+    declaration = types.FunctionDeclaration(
+        name="get_weather",
+        description="Gets weather information",
+        parameters=types.Schema(
+            type=types.Type.OBJECT,
+            properties={
+                "location": types.Schema(
+                    type=types.Type.STRING,
+                    description="The location to get weather for",
+                ),
+                "units": types.Schema(
+                    type=types.Type.STRING,
+                    description="Temperature units",
+                ),
+            },
+            required=["location"],
+        ),
+    )
+    mock_tool._get_declaration.return_value = declaration
+
+    result = adk_to_mcp_tool_type(mock_tool)
+
+    assert isinstance(result, mcp_types.Tool)
+    assert result.name == "get_weather"
+    assert result.description == "Gets weather information"
+    assert "type" in result.inputSchema
+    assert result.inputSchema["type"] == "object"
+    assert "properties" in result.inputSchema
+    assert "location" in result.inputSchema["properties"]
+    assert "units" in result.inputSchema["properties"]
+    assert result.inputSchema["properties"]["location"]["type"] == "string"
+    assert "required" in result.inputSchema
+    assert "location" in result.inputSchema["required"]
+
+  def test_tool_with_parameters_json_schema(self):
+    """Test conversion when tool has parameters_json_schema."""
+    mock_tool = mock.Mock(spec=BaseTool)
+    mock_tool.name = "search_database"
+    mock_tool.description = "Searches a database"
+
+    json_schema = {
+        "type": "object",
+        "properties": {
+            "query": {
+                "type": "string",
+                "description": "The search query",
+            },
+            "limit": {
+                "type": "integer",
+                "description": "Maximum number of results",
+            },
+        },
+        "required": ["query"],
+    }
+
+    declaration = types.FunctionDeclaration(
+        name="search_database",
+        description="Searches a database",
+        parameters_json_schema=json_schema,
+    )
+    mock_tool._get_declaration.return_value = declaration
+
+    result = adk_to_mcp_tool_type(mock_tool)
+
+    assert isinstance(result, mcp_types.Tool)
+    assert result.name == "search_database"
+    assert result.description == "Searches a database"
+    # Should use the JSON schema directly
+    assert result.inputSchema == json_schema
+
+  def test_tool_with_no_parameters(self):
+    """Test conversion when tool has declaration but no parameters."""
+    mock_tool = mock.Mock(spec=BaseTool)
+    mock_tool.name = "get_current_time"
+    mock_tool.description = "Gets the current time"
+
+    declaration = types.FunctionDeclaration(
+        name="get_current_time",
+        description="Gets the current time",
+    )
+    mock_tool._get_declaration.return_value = declaration
+
+    result = adk_to_mcp_tool_type(mock_tool)
+
+    assert isinstance(result, mcp_types.Tool)
+    assert result.name == "get_current_time"
+    assert result.description == "Gets the current time"
+    assert not result.inputSchema
+
+  def test_tool_prefers_json_schema_over_parameters(self):
+    """Test that parameters_json_schema is preferred over parameters."""
+    mock_tool = mock.Mock(spec=BaseTool)
+    mock_tool.name = "test_tool"
+    mock_tool.description = "Test tool"
+
+    json_schema = {
+        "type": "object",
+        "properties": {
+            "json_param": {"type": "string"},
+        },
+    }
+
+    # Create a declaration with BOTH parameters and parameters_json_schema
+    declaration = types.FunctionDeclaration(
+        name="test_tool",
+        description="Test tool",
+        parameters=types.Schema(
+            type=types.Type.OBJECT,
+            properties={
+                "schema_param": types.Schema(type=types.Type.STRING),
+            },
+        ),
+        parameters_json_schema=json_schema,
+    )
+    mock_tool._get_declaration.return_value = declaration
+
+    result = adk_to_mcp_tool_type(mock_tool)
+
+    # Should use parameters_json_schema, not parameters
+    assert result.inputSchema == json_schema
+    assert "json_param" in result.inputSchema["properties"]
+    assert "schema_param" not in result.inputSchema["properties"]
+
+  def test_tool_with_complex_nested_schema(self):
+    """Test conversion with complex nested parameters_json_schema."""
+    mock_tool = mock.Mock(spec=BaseTool)
+    mock_tool.name = "create_user"
+    mock_tool.description = "Creates a new user"
+
+    json_schema = {
+        "type": "object",
+        "properties": {
+            "username": {"type": "string"},
+            "profile": {
+                "type": "object",
+                "properties": {
+                    "email": {"type": "string"},
+                    "age": {"type": "integer"},
+                    "tags": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                    },
+                },
+                "required": ["email"],
+            },
+        },
+        "required": ["username", "profile"],
+    }
+
+    declaration = types.FunctionDeclaration(
+        name="create_user",
+        description="Creates a new user",
+        parameters_json_schema=json_schema,
+    )
+    mock_tool._get_declaration.return_value = declaration
+
+    result = adk_to_mcp_tool_type(mock_tool)
+
+    assert isinstance(result, mcp_types.Tool)
+    assert result.inputSchema == json_schema

tests/unittests/tools/mcp_tool/test_mcp_toolset.py

@@ -0,0 +1,800 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import base64
+from io import StringIO
+import pickle
+import sys
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import Mock
+
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import HttpAuth
+from google.adk.auth.auth_credential import HttpCredentials
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.tools.load_mcp_resource_tool import LoadMcpResourceTool
+from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
+from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
+from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
+from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
+from google.adk.tools.mcp_tool.mcp_tool import MCPTool
+from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
+from google.adk.tools.tool_configs import ToolArgsConfig
+from mcp import StdioServerParameters
+from mcp.types import BlobResourceContents
+from mcp.types import ListResourcesResult
+from mcp.types import ReadResourceResult
+from mcp.types import Resource
+from mcp.types import TextResourceContents
+import pytest
+
+
+class MockMCPTool:
+  """Mock MCP Tool for testing."""
+
+  def __init__(self, name, description="Test tool description"):
+    self.name = name
+    self.description = description
+    self.inputSchema = {
+        "type": "object",
+        "properties": {"param": {"type": "string"}},
+    }
+
+
+class MockListToolsResult:
+  """Mock ListToolsResult for testing."""
+
+  def __init__(self, tools):
+    self.tools = tools
+
+
+class TestMcpToolset:
+  """Test suite for McpToolset class."""
+
+  def setup_method(self):
+    """Set up test fixtures."""
+    self.mock_stdio_params = StdioServerParameters(
+        command="test_command", args=[]
+    )
+    self.mock_session_manager = Mock(spec=MCPSessionManager)
+    self.mock_session = AsyncMock()
+    self.mock_session_manager.create_session = AsyncMock(
+        return_value=self.mock_session
+    )
+
+  def test_init_basic(self):
+    """Test basic initialization with StdioServerParameters."""
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+
+    # Note: StdioServerParameters gets converted to StdioConnectionParams internally
+    assert toolset._errlog == sys.stderr
+    assert toolset._auth_scheme is None
+    assert toolset._auth_credential is None
+    assert toolset._use_mcp_resources is False
+
+  def test_init_with_use_mcp_resources(self):
+    """Test initialization with use_mcp_resources."""
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params, use_mcp_resources=True
+    )
+    assert toolset._use_mcp_resources is True
+
+  def test_connection_params(self):
+    """Test getting connection params."""
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    assert toolset.connection_params == self.mock_stdio_params
+
+  def test_auth_scheme(self):
+    """Test getting auth scheme."""
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    assert toolset.auth_scheme is None
+
+  def test_auth_credential(self):
+    """Test getting auth credential."""
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    assert toolset.auth_credential is None
+
+  def test_error_log(self):
+    """Test getting error log."""
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    assert toolset.errlog == sys.stderr
+
+  def test_auth_scheme_with_value(self):
+    """Test getting auth scheme when provided at initialization."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={"read": "Read access"},
+            )
+        )
+    )
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        auth_scheme=auth_scheme,
+    )
+    assert toolset.auth_scheme == auth_scheme
+
+  def test_require_confirmation(self):
+    """Test getting require_confirmation flag."""
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        require_confirmation=True,
+    )
+    assert toolset.require_confirmation is True
+
+  def test_header_provider(self):
+    """Test getting header_provider."""
+    mock_header_provider = Mock()
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        header_provider=mock_header_provider,
+    )
+    assert toolset.header_provider == mock_header_provider
+
+  def test_auth_credential_with_value(self):
+    """Test getting auth credential when provided at initialization."""
+    mock_credential = Mock(spec=AuthCredential)
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        auth_credential=mock_credential,
+    )
+    assert toolset.auth_credential == mock_credential
+
+  def test_init_with_stdio_connection_params(self):
+    """Test initialization with StdioConnectionParams."""
+    stdio_params = StdioConnectionParams(
+        server_params=self.mock_stdio_params, timeout=10.0
+    )
+    toolset = McpToolset(connection_params=stdio_params)
+
+    assert toolset._connection_params == stdio_params
+
+  def test_init_with_sse_connection_params(self):
+    """Test initialization with SseConnectionParams."""
+    sse_params = SseConnectionParams(
+        url="https://example.com/mcp", headers={"Authorization": "Bearer token"}
+    )
+    toolset = McpToolset(connection_params=sse_params)
+
+    assert toolset._connection_params == sse_params
+
+  def test_init_with_streamable_http_params(self):
+    """Test initialization with StreamableHTTPConnectionParams."""
+    http_params = StreamableHTTPConnectionParams(
+        url="https://example.com/mcp",
+        headers={"Content-Type": "application/json"},
+    )
+    toolset = McpToolset(connection_params=http_params)
+
+    assert toolset._connection_params == http_params
+
+  def test_init_with_tool_filter_list(self):
+    """Test initialization with tool filter as list."""
+    tool_filter = ["tool1", "tool2"]
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params, tool_filter=tool_filter
+    )
+
+    # The tool filter is stored in the parent BaseToolset class
+    # We can verify it by checking the filtering behavior in get_tools
+    assert toolset._is_tool_selected is not None
+
+  def test_init_with_auth(self):
+    """Test initialization with authentication."""
+    # Create real auth scheme instances
+
+    auth_scheme = OAuth2(flows={})
+
+    auth_credential = AuthCredential(
+        auth_type="oauth2",
+        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
+    )
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
+    )
+
+    assert toolset._auth_scheme == auth_scheme
+    assert toolset._auth_credential == auth_credential
+
+  def test_init_with_auth_and_credential_key(self):
+    """Test initialization with authentication and a custom credential_key."""
+
+    auth_scheme = OAuth2(flows={})
+
+    auth_credential = AuthCredential(
+        auth_type="oauth2",
+        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
+    )
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
+        credential_key="my_custom_key",
+    )
+
+    assert toolset._auth_scheme == auth_scheme
+    assert toolset._auth_credential == auth_credential
+    assert toolset._auth_config.credential_key == "my_custom_key"
+
+  def test_from_config_with_credential_key(self):
+    """Test that from_config correctly parses credential_key."""
+
+    auth_scheme = OAuth2(flows={})
+
+    config = ToolArgsConfig(
+        stdio_server_params=self.mock_stdio_params,
+        auth_scheme=auth_scheme,
+        credential_key="my_custom_key",
+    )
+    toolset = McpToolset.from_config(config, "")
+
+    assert isinstance(toolset._auth_scheme, OAuth2)
+    assert toolset._auth_config.credential_key == "my_custom_key"
+
+  def test_init_missing_connection_params(self):
+    """Test initialization with missing connection params raises error."""
+    with pytest.raises(ValueError, match="Missing connection params"):
+      McpToolset(connection_params=None)
+
+  @pytest.mark.asyncio
+  async def test_get_tools_basic(self):
+    """Test getting tools without filtering."""
+    # Mock tools from MCP server
+    mock_tools = [
+        MockMCPTool("tool1"),
+        MockMCPTool("tool2"),
+        MockMCPTool("tool3"),
+    ]
+    self.mock_session.list_tools = AsyncMock(
+        return_value=MockListToolsResult(mock_tools)
+    )
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params, use_mcp_resources=True
+    )
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    tools = await toolset.get_tools()
+
+    assert len(tools) == 4
+    for tool in tools[:3]:
+      assert isinstance(tool, MCPTool)
+    assert isinstance(tools[3], LoadMcpResourceTool)
+    assert tools[0].name == "tool1"
+    assert tools[1].name == "tool2"
+    assert tools[2].name == "tool3"
+    assert tools[3].name == "load_mcp_resource"
+
+  @pytest.mark.asyncio
+  async def test_get_tools_with_list_filter(self):
+    """Test getting tools with list-based filtering."""
+    # Mock tools from MCP server
+    mock_tools = [
+        MockMCPTool("tool1"),
+        MockMCPTool("tool2"),
+        MockMCPTool("tool3"),
+    ]
+    self.mock_session.list_tools = AsyncMock(
+        return_value=MockListToolsResult(mock_tools)
+    )
+
+    tool_filter = ["tool1", "tool3"]
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params, tool_filter=tool_filter
+    )
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    tools = await toolset.get_tools()
+
+    assert len(tools) == 2
+    assert tools[0].name == "tool1"
+    assert tools[1].name == "tool3"
+
+  @pytest.mark.asyncio
+  async def test_get_tools_with_function_filter(self):
+    """Test getting tools with function-based filtering."""
+    # Mock tools from MCP server
+    mock_tools = [
+        MockMCPTool("read_file"),
+        MockMCPTool("write_file"),
+        MockMCPTool("list_directory"),
+    ]
+    self.mock_session.list_tools = AsyncMock(
+        return_value=MockListToolsResult(mock_tools)
+    )
+
+    def file_tools_filter(tool, context):
+      """Filter for file-related tools only."""
+      return "file" in tool.name
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params, tool_filter=file_tools_filter
+    )
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    tools = await toolset.get_tools()
+
+    assert len(tools) == 2
+    assert tools[0].name == "read_file"
+    assert tools[1].name == "write_file"
+
+  @pytest.mark.asyncio
+  async def test_get_tools_with_header_provider(self):
+    """Test get_tools with a header_provider."""
+    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
+    self.mock_session.list_tools = AsyncMock(
+        return_value=MockListToolsResult(mock_tools)
+    )
+    mock_readonly_context = Mock(spec=ReadonlyContext)
+    expected_headers = {"X-Tenant-ID": "test-tenant"}
+    header_provider = Mock(return_value=expected_headers)
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        header_provider=header_provider,
+    )
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    tools = await toolset.get_tools(readonly_context=mock_readonly_context)
+
+    assert len(tools) == 2
+    header_provider.assert_called_once_with(mock_readonly_context)
+    self.mock_session_manager.create_session.assert_called_once_with(
+        headers=expected_headers
+    )
+
+  @pytest.mark.asyncio
+  async def test_get_tools_with_async_header_provider(self):
+    """Test get_tools with an async header_provider."""
+    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
+    self.mock_session.list_tools = AsyncMock(
+        return_value=MockListToolsResult(mock_tools)
+    )
+    mock_readonly_context = Mock(spec=ReadonlyContext)
+    expected_headers = {"X-Tenant-ID": "test-tenant"}
+
+    async def header_provider(_context):
+      return expected_headers
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        header_provider=header_provider,
+    )
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    tools = await toolset.get_tools(readonly_context=mock_readonly_context)
+
+    assert len(tools) == 2
+    self.mock_session_manager.create_session.assert_called_once_with(
+        headers=expected_headers
+    )
+
+  @pytest.mark.asyncio
+  async def test_close_success(self):
+    """Test successful cleanup."""
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    await toolset.close()
+
+    self.mock_session_manager.close.assert_called_once()
+
+  @pytest.mark.asyncio
+  async def test_close_with_exception(self):
+    """Test cleanup when session manager raises exception."""
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    # Mock close to raise an exception
+    self.mock_session_manager.close = AsyncMock(
+        side_effect=Exception("Cleanup error")
+    )
+
+    # Should not raise exception, should log the warning
+    await toolset.close()
+
+  @pytest.mark.asyncio
+  async def test_get_tools_with_timeout(self):
+    """Test get_tools with timeout."""
+    stdio_params = StdioConnectionParams(
+        server_params=self.mock_stdio_params, timeout=0.01
+    )
+    toolset = McpToolset(connection_params=stdio_params)
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    async def long_running_list_tools():
+      await asyncio.sleep(0.1)
+      return MockListToolsResult([])
+
+    self.mock_session.list_tools = long_running_list_tools
+
+    with pytest.raises(
+        ConnectionError, match="Failed to get tools from MCP server."
+    ):
+      await toolset.get_tools()
+
+  @pytest.mark.asyncio
+  async def test_get_tools_retry_decorator(self):
+    """Test that get_tools has retry decorator applied."""
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+
+    # Check that the method has the retry decorator
+    assert hasattr(toolset.get_tools, "__wrapped__")
+
+  @pytest.mark.asyncio
+  async def test_mcp_toolset_with_prefix(self):
+    """Test that McpToolset correctly applies the tool_name_prefix."""
+    # Mock the connection parameters
+    mock_connection_params = MagicMock()
+    mock_connection_params.timeout = None
+
+    # Mock the MCPSessionManager and its create_session method
+    mock_session_manager = MagicMock()
+    mock_session = MagicMock()
+
+    # Mock the list_tools response from the MCP server
+    mock_tool1 = MagicMock()
+    mock_tool1.name = "tool1"
+    mock_tool1.description = "tool 1 desc"
+    mock_tool2 = MagicMock()
+    mock_tool2.name = "tool2"
+    mock_tool2.description = "tool 2 desc"
+    list_tools_result = MagicMock()
+    list_tools_result.tools = [mock_tool1, mock_tool2]
+    mock_session.list_tools = AsyncMock(return_value=list_tools_result)
+    mock_session_manager.create_session = AsyncMock(return_value=mock_session)
+
+    # Create an instance of McpToolset with a prefix
+    toolset = McpToolset(
+        connection_params=mock_connection_params,
+        tool_name_prefix="my_prefix",
+        use_mcp_resources=True,
+    )
+
+    # Replace the internal session manager with our mock
+    toolset._mcp_session_manager = mock_session_manager
+
+    # Get the tools from the toolset
+    tools = await toolset.get_tools()
+
+    # The get_tools method in McpToolset returns MCPTool objects, which are
+    # instances of BaseTool. The prefixing is handled by the BaseToolset,
+    # so we need to call get_tools_with_prefix to get the prefixed tools.
+    prefixed_tools = await toolset.get_tools_with_prefix()
+
+    # Assert that the tools are prefixed correctly
+    assert len(prefixed_tools) == 3
+    assert prefixed_tools[0].name == "my_prefix_tool1"
+    assert prefixed_tools[1].name == "my_prefix_tool2"
+    assert prefixed_tools[2].name == "my_prefix_load_mcp_resource"
+
+    # Assert that the original tools are not modified
+    assert tools[0].name == "tool1"
+    assert tools[1].name == "tool2"
+    assert tools[2].name == "load_mcp_resource"
+
+  def test_init_with_progress_callback(self):
+    """Test initialization with progress_callback."""
+
+    async def my_progress_callback(
+        progress: float, total: float | None, message: str | None
+    ) -> None:
+      pass
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        progress_callback=my_progress_callback,
+    )
+
+    assert toolset._progress_callback == my_progress_callback
+
+  @pytest.mark.asyncio
+  async def test_get_tools_passes_progress_callback_to_mcp_tools(self):
+    """Test that get_tools passes progress_callback to created MCPTool instances."""
+    progress_updates = []
+
+    async def my_progress_callback(
+        progress: float, total: float | None, message: str | None
+    ) -> None:
+      progress_updates.append((progress, total, message))
+
+    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
+    self.mock_session.list_tools = AsyncMock(
+        return_value=MockListToolsResult(mock_tools)
+    )
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        progress_callback=my_progress_callback,
+    )
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    tools = await toolset.get_tools()
+
+    assert len(tools) == 2
+    # Verify each tool has the progress_callback set
+    for tool in tools:
+      assert tool._progress_callback == my_progress_callback
+
+  def test_init_with_progress_callback_factory(self):
+    """Test initialization with a ProgressCallbackFactory."""
+
+    def my_callback_factory(tool_name: str, *, readonly_context=None, **kwargs):
+      async def callback(
+          progress: float, total: float | None, message: str | None
+      ) -> None:
+        pass
+
+      return callback
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        progress_callback=my_callback_factory,
+    )
+
+    assert toolset._progress_callback == my_callback_factory
+
+  @pytest.mark.asyncio
+  async def test_get_tools_passes_factory_to_mcp_tools(self):
+    """Test that get_tools passes factory directly to MCPTool instances.
+
+    The factory is resolved at runtime in McpTool._run_async_impl, not at
+    tool creation time. This allows the factory to receive ReadonlyContext.
+    """
+
+    def my_callback_factory(tool_name: str, *, readonly_context=None, **kwargs):
+      async def callback(
+          progress: float, total: float | None, message: str | None
+      ) -> None:
+        pass
+
+      return callback
+
+    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
+    self.mock_session.list_tools = AsyncMock(
+        return_value=MockListToolsResult(mock_tools)
+    )
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        progress_callback=my_callback_factory,
+    )
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    tools = await toolset.get_tools()
+
+    assert len(tools) == 2
+    # Factory is passed directly to each tool (resolved at runtime)
+    for tool in tools:
+      assert tool._progress_callback == my_callback_factory
+
+  @pytest.mark.asyncio
+  async def test_list_resources(self):
+    """Test listing resources."""
+    resources = [
+        Resource(
+            name="file1.txt", mime_type="text/plain", uri="file:///file1.txt"
+        ),
+        Resource(
+            name="data.json",
+            mime_type="application/json",
+            uri="file:///data.json",
+        ),
+    ]
+    list_resources_result = ListResourcesResult(resources=resources)
+    self.mock_session.list_resources = AsyncMock(
+        return_value=list_resources_result
+    )
+
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    result = await toolset.list_resources()
+
+    assert result == ["file1.txt", "data.json"]
+    self.mock_session.list_resources.assert_called_once()
+
+  @pytest.mark.asyncio
+  async def test_get_resource_info_success(self):
+    """Test getting resource info for an existing resource."""
+    resources = [
+        Resource(
+            name="file1.txt", mime_type="text/plain", uri="file:///file1.txt"
+        ),
+        Resource(
+            name="data.json",
+            mime_type="application/json",
+            uri="file:///data.json",
+        ),
+    ]
+    list_resources_result = ListResourcesResult(resources=resources)
+    self.mock_session.list_resources = AsyncMock(
+        return_value=list_resources_result
+    )
+
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    result = await toolset.get_resource_info("data.json")
+
+    assert result == {
+        "name": "data.json",
+        "mime_type": "application/json",
+        "uri": "file:///data.json",
+    }
+    self.mock_session.list_resources.assert_called_once()
+
+  @pytest.mark.asyncio
+  async def test_get_resource_info_not_found(self):
+    """Test getting resource info for a non-existent resource."""
+    resources = [
+        Resource(
+            name="file1.txt", mime_type="text/plain", uri="file:///file1.txt"
+        ),
+    ]
+    list_resources_result = ListResourcesResult(resources=resources)
+    self.mock_session.list_resources = AsyncMock(
+        return_value=list_resources_result
+    )
+
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    with pytest.raises(
+        ValueError, match="Resource with name 'other.json' not found."
+    ):
+      await toolset.get_resource_info("other.json")
+
+  @pytest.mark.parametrize(
+      "name,mime_type,content,encoding",
+      [
+          ("file1.txt", "text/plain", "hello world", None),
+          (
+              "data.json",
+              "application/json",
+              '{"key": "value"}',
+              None,
+          ),
+          (
+              "file1_b64.txt",
+              "text/plain",
+              base64.b64encode(b"hello world").decode("ascii"),
+              "base64",
+          ),
+          (
+              "data_b64.json",
+              "application/json",
+              base64.b64encode(b'{"key": "value"}').decode("ascii"),
+              "base64",
+          ),
+          (
+              "data.bin",
+              "application/octet-stream",
+              base64.b64encode(b"\x01\x02\x03").decode("ascii"),
+              "base64",
+          ),
+      ],
+  )
+  @pytest.mark.asyncio
+  async def test_read_resource(self, name, mime_type, content, encoding):
+    """Test reading various resource types."""
+    uri = f"file:///{name}"
+    # Mock list_resources for get_resource_info
+    resources = [Resource(name=name, mime_type=mime_type, uri=uri)]
+    list_resources_result = ListResourcesResult(resources=resources)
+    self.mock_session.list_resources = AsyncMock(
+        return_value=list_resources_result
+    )
+
+    # Mock read_resource
+    if encoding == "base64":
+      contents = [
+          BlobResourceContents(uri=uri, mimeType=mime_type, blob=content)
+      ]
+    else:
+      contents = [
+          TextResourceContents(uri=uri, mimeType=mime_type, text=content)
+      ]
+
+    read_resource_result = ReadResourceResult(contents=contents)
+    self.mock_session.read_resource = AsyncMock(
+        return_value=read_resource_result
+    )
+
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    result = await toolset.read_resource(name)
+
+    assert result == contents
+    self.mock_session.list_resources.assert_called_once()
+    self.mock_session.read_resource.assert_called_once_with(uri=uri)
+
+  @pytest.mark.asyncio
+  async def test_sampling_callback_invoked(self):
+
+    called = {"value": False}
+
+    async def mock_sampling_handler(messages, params=None, context=None):
+      called["value"] = True
+
+      assert isinstance(messages, list)
+      assert messages[0]["role"] == "user"
+
+      return {
+          "model": "test-model",
+          "role": "assistant",
+          "content": {"type": "text", "text": "sampling response"},
+          "stopReason": "endTurn",
+      }
+
+    toolset = McpToolset(
+        connection_params=StreamableHTTPConnectionParams(
+            url="http://localhost:9999",
+            timeout=10,
+        ),
+        sampling_callback=mock_sampling_handler,
+    )
+
+    messages = [{"role": "user", "content": {"type": "text", "text": "hello"}}]
+
+    result = await toolset._sampling_callback(messages)
+
+    assert called["value"] is True
+    assert result["role"] == "assistant"
+    assert result["content"]["text"] == "sampling response"
+
+  @pytest.mark.asyncio
+  async def test_get_auth_headers_includes_additional_headers(self):
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP,
+        http=HttpAuth(
+            scheme="bearer",
+            credentials=HttpCredentials(token="token"),
+            additional_headers={"X-API-Key": "secret"},
+        ),
+    )
+    auth_config = AuthConfig(
+        auth_scheme=OAuth2(flows={}),
+        raw_auth_credential=credential,
+    )
+    auth_config.exchanged_auth_credential = credential
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    toolset._auth_config = auth_config
+
+    headers = toolset._get_auth_headers()
+
+    assert headers["Authorization"] == "Bearer token"
+    assert headers["X-API-Key"] == "secret"
+
+  def test_pickle_mcp_toolset(self):
+    toolset = McpToolset(connection_params=self.mock_stdio_params)
+    pickled = pickle.dumps(toolset)
+    unpickled = pickle.loads(pickled)
+    assert unpickled._connection_params == self.mock_stdio_params
+    assert unpickled._errlog == sys.stderr

tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py

@@ -0,0 +1,298 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for MCPToolset authentication functionality."""
+
+import base64
+from unittest.mock import Mock
+
+from fastapi.openapi.models import APIKey as APIKeyScheme
+from fastapi.openapi.models import APIKeyIn
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import HttpAuth
+from google.adk.auth.auth_credential import HttpCredentials
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
+from mcp import StdioServerParameters
+import pytest
+
+
+class TestMcpToolsetGetAuthConfig:
+  """Tests for McpToolset.get_auth_config method."""
+
+  def test_get_auth_config_returns_none_without_auth_scheme(self):
+    """Test that get_auth_config returns None when no auth configured."""
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"])
+    )
+
+    assert toolset.get_auth_config() is None
+
+  def test_get_auth_config_returns_config_with_auth_scheme(self):
+    """Test that get_auth_config returns AuthConfig when auth configured."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={"read": "Read access"},
+            )
+        )
+    )
+    auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
+    )
+
+    auth_config = toolset.get_auth_config()
+    assert auth_config is not None
+    assert auth_config.auth_scheme == auth_scheme
+    assert auth_config.raw_auth_credential == auth_credential
+
+  def test_get_auth_config_returns_same_instance(self):
+    """Test that get_auth_config returns the same instance each time."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={},
+            )
+        )
+    )
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Should return the same instance
+    config1 = toolset.get_auth_config()
+    config2 = toolset.get_auth_config()
+    assert config1 is config2
+
+
+class TestMcpToolsetGetAuthHeaders:
+  """Tests for McpToolset._get_auth_headers method."""
+
+  @pytest.fixture
+  def toolset_with_oauth2(self):
+    """Create a toolset with OAuth2 auth configured."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={"read": "Read access"},
+            )
+        )
+    )
+    auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    return McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
+    )
+
+  def test_get_auth_headers_returns_none_without_auth_config(self):
+    """Test that _get_auth_headers returns None without auth config."""
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"])
+    )
+
+    assert toolset._get_auth_headers() is None
+
+  def test_get_auth_headers_returns_none_without_exchanged_credential(
+      self, toolset_with_oauth2
+  ):
+    """Test that _get_auth_headers returns None without exchanged credential."""
+    # No exchanged credential set yet
+    assert toolset_with_oauth2._get_auth_headers() is None
+
+  def test_get_auth_headers_oauth2_bearer_token(self, toolset_with_oauth2):
+    """Test that _get_auth_headers returns Bearer token for OAuth2."""
+    # Set exchanged credential with access token
+    toolset_with_oauth2._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="test-access-token"),
+    )
+
+    headers = toolset_with_oauth2._get_auth_headers()
+
+    assert headers is not None
+    assert headers["Authorization"] == "Bearer test-access-token"
+
+  def test_get_auth_headers_http_bearer_token(self):
+    """Test that _get_auth_headers returns Bearer token for HTTP bearer."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={},
+            )
+        )
+    )
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Set exchanged credential with HTTP bearer token
+    toolset._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP,
+        http=HttpAuth(
+            scheme="bearer",
+            credentials=HttpCredentials(token="test-bearer-token"),
+        ),
+    )
+
+    headers = toolset._get_auth_headers()
+
+    assert headers is not None
+    assert headers["Authorization"] == "Bearer test-bearer-token"
+
+  def test_get_auth_headers_http_basic_auth(self):
+    """Test that _get_auth_headers returns Basic auth for HTTP basic."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={},
+            )
+        )
+    )
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Set exchanged credential with HTTP basic auth
+    toolset._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP,
+        http=HttpAuth(
+            scheme="basic",
+            credentials=HttpCredentials(
+                username="testuser",
+                password="testpass",
+            ),
+        ),
+    )
+
+    headers = toolset._get_auth_headers()
+
+    assert headers is not None
+    expected_credentials = base64.b64encode(b"testuser:testpass").decode()
+    assert headers["Authorization"] == f"Basic {expected_credentials}"
+
+  def test_get_auth_headers_api_key_header(self):
+    """Test that _get_auth_headers returns API key in header."""
+    # Note: fastapi's APIKey model uses 'in' not 'in_', but accepts both
+    auth_scheme = APIKeyScheme(**{
+        "in": APIKeyIn.header,
+        "name": "X-API-Key",
+    })
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Set exchanged credential with API key
+    toolset._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.API_KEY,
+        api_key="test-api-key-12345",
+    )
+
+    headers = toolset._get_auth_headers()
+
+    assert headers is not None
+    assert headers["X-API-Key"] == "test-api-key-12345"
+
+  def test_get_auth_headers_api_key_non_header_logs_warning(self, caplog):
+    """Test that non-header API key logs a warning."""
+    # Note: fastapi's APIKey model uses 'in' not 'in_'
+    auth_scheme = APIKeyScheme(**{
+        "in": APIKeyIn.query,  # Query param, not header
+        "name": "api_key",
+    })
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Set exchanged credential with API key
+    toolset._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.API_KEY,
+        api_key="test-api-key",
+    )
+
+    headers = toolset._get_auth_headers()
+
+    # Should return None for non-header API key
+    assert headers is None
+
+  def test_get_auth_headers_reads_from_readonly_context(
+      self, toolset_with_oauth2
+  ):
+    """Test that _get_auth_headers reads from ReadonlyContext first."""
+    from google.adk.agents.readonly_context import ReadonlyContext
+
+    mock_readonly_context = Mock(spec=ReadonlyContext)
+    mock_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="token-from-context"),
+    )
+    mock_readonly_context.get_credential.return_value = mock_credential
+
+    # Even if exchanged_auth_credential has a different value
+    toolset_with_oauth2._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="token-from-config"),
+    )
+
+    headers = toolset_with_oauth2._get_auth_headers(
+        readonly_context=mock_readonly_context
+    )
+
+    assert headers is not None
+    assert headers["Authorization"] == "Bearer token-from-context"
+    mock_readonly_context.get_credential.assert_called_once_with(
+        toolset_with_oauth2._auth_config.credential_key
+    )

tests/unittests/tools/mcp_tool/test_session_context.py

@@ -0,0 +1,948 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import AsyncExitStack
+from datetime import timedelta
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from google.adk.features import FeatureName
+from google.adk.features._feature_registry import temporary_feature_override
+from google.adk.tools.mcp_tool.session_context import _format_exception
+from google.adk.tools.mcp_tool.session_context import SessionContext
+import httpx
+from mcp import ClientSession
+import pytest
+
+
+class MockClientSession:
+  """Mock ClientSession for testing."""
+
+  def __init__(self, *args, **kwargs):
+    self._initialized = False
+    self._args = args
+    self._kwargs = kwargs
+
+  async def initialize(self):
+    """Mock initialize method."""
+    self._initialized = True
+    return self
+
+  async def __aenter__(self):
+    return self
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb):
+    return False
+
+
+class MockClient:
+  """Mock MCP client."""
+
+  def __init__(
+      self,
+      transports=None,
+      raise_on_enter=None,
+      delay_on_enter=0,
+  ):
+    self._transports = transports or ('read_stream', 'write_stream')
+    self._raise_on_enter = raise_on_enter
+    self._delay_on_enter = delay_on_enter
+    self._entered = False
+    self._exited = False
+
+  async def __aenter__(self):
+    if self._delay_on_enter > 0:
+      await asyncio.sleep(self._delay_on_enter)
+    if self._raise_on_enter:
+      raise self._raise_on_enter
+    self._entered = True
+    return self._transports
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb):
+    self._exited = True
+    return False
+
+
+class TestSessionContext:
+  """Test suite for SessionContext class."""
+
+  @pytest.mark.asyncio
+  async def test_start_success_ready_event_set_and_session_returned(self):
+    """Test that start() sets _ready_event and returns session."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    # Mock ClientSession
+    mock_session = MockClientSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      session = await session_context.start()
+
+      # Verify ready_event was set
+      assert session_context._ready_event.is_set()
+
+      # Verify session was returned
+      assert session == mock_session
+      assert session_context.session == mock_session
+
+      # Verify initialize was called
+      assert mock_session._initialized
+
+      # Verify task was created and is still running (waiting for close)
+      assert session_context._task is not None
+      assert not session_context._task.done()
+
+      # Clean up
+      await session_context.close()
+
+  @pytest.mark.asyncio
+  async def test_start_raises_connection_error_on_exception(self):
+    """Test that start() raises ConnectionError when exception occurs."""
+    test_exception = ValueError('Connection failed')
+    mock_client = MockClient(raise_on_enter=test_exception)
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with pytest.raises(ConnectionError) as exc_info:
+      await session_context.start()
+
+    # Verify ConnectionError message contains original exception
+    assert 'Failed to create MCP session' in str(exc_info.value)
+    assert 'Connection failed' in str(exc_info.value)
+
+    # Verify ready_event was set (in finally block)
+    assert session_context._ready_event.is_set()
+
+  @pytest.mark.asyncio
+  async def test_start_raises_connection_error_on_cancelled_error(self):
+    """Test that start() raises ConnectionError when CancelledError occurs."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    # Mock session that will cause cancellation
+    mock_session = MockClientSession()
+
+    # Make initialize raise CancelledError
+    async def cancelled_initialize():
+      raise asyncio.CancelledError('Task cancelled')
+
+    mock_session.initialize = cancelled_initialize
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      # Should raise ConnectionError (not CancelledError directly)
+      with pytest.raises(ConnectionError) as exc_info:
+        await session_context.start()
+
+      # Verify it's a ConnectionError about cancellation
+      assert 'Failed to create MCP session' in str(exc_info.value)
+      assert 'task cancelled' in str(exc_info.value)
+
+      # Verify ready_event was set
+      assert session_context._ready_event.is_set()
+
+  @pytest.mark.asyncio
+  async def test_close_cleans_up_task(self):
+    """Test that close() properly cleans up the task."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    # Mock ClientSession
+    mock_session = MockClientSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      # Start the session context
+      await session_context.start()
+
+      # Verify task is running
+      assert session_context._task is not None
+      assert not session_context._task.done()
+
+      # Close the session context
+      await session_context.close()
+
+      # Wait a bit for cleanup
+      await asyncio.sleep(0.1)
+
+      # Verify close_event was set
+      assert session_context._close_event.is_set()
+
+      # Verify task completed (may take a moment)
+      # The task should finish after close_event is set
+      assert session_context._task.done()
+
+  @pytest.mark.asyncio
+  async def test_session_exception_does_not_break_event_loop(self):
+    """Test that session exceptions don't break the event loop."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    # Mock ClientSession that raises exception during use
+    mock_session = MockClientSession()
+
+    async def failing_operation():
+      raise RuntimeError('Session operation failed')
+
+    mock_session.failing_operation = failing_operation
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      # Start the session context
+      session = await session_context.start()
+
+      # Use session and trigger exception
+      with pytest.raises(RuntimeError, match='Session operation failed'):
+        await session.failing_operation()
+
+      # Close the session context - should not break event loop
+      await session_context.close()
+
+      # Verify event loop is still healthy by running another task
+      result = await asyncio.sleep(0.01)
+      assert result is None
+
+  @pytest.mark.asyncio
+  async def test_async_context_manager(self):
+    """Test using SessionContext as async context manager."""
+    mock_client = MockClient()
+    mock_session = MockClientSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      async with SessionContext(
+          mock_client, timeout=5.0, sse_read_timeout=None
+      ) as session:
+        assert session == mock_session
+        # Verify initialize was called by checking _initialized flag
+        assert session._initialized
+
+  @pytest.mark.asyncio
+  async def test_timeout_during_connection(self):
+    """Test timeout during client connection."""
+    # Client that takes longer than timeout
+    mock_client = MockClient(delay_on_enter=10.0)
+    session_context = SessionContext(
+        mock_client, timeout=0.1, sse_read_timeout=None
+    )
+
+    with pytest.raises(ConnectionError) as exc_info:
+      await session_context.start()
+
+    assert 'Failed to create MCP session' in str(exc_info.value)
+
+  @pytest.mark.asyncio
+  async def test_timeout_during_initialization(self):
+    """Test timeout during session initialization."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=0.1, sse_read_timeout=None
+    )
+
+    # Mock ClientSession with slow initialize
+    mock_session = MockClientSession()
+
+    async def slow_initialize():
+      await asyncio.sleep(1.0)
+      return mock_session
+
+    mock_session.initialize = slow_initialize
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      with pytest.raises(ConnectionError) as exc_info:
+        await session_context.start()
+
+      assert 'Failed to create MCP session' in str(exc_info.value)
+
+  @pytest.mark.asyncio
+  async def test_timeout_during_initialization_with_flag_on(self):
+    """Test timeout during session initialization with flag ON.
+
+    Verifies that session initialization uses `anyio.fail_after` under the
+    graceful error handling feature flag.
+    """
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=0.1, sse_read_timeout=None
+    )
+
+    # Mock ClientSession with slow initialize
+    mock_session = MockClientSession()
+
+    async def slow_initialize():
+      await asyncio.sleep(1.0)
+      return mock_session
+
+    mock_session.initialize = slow_initialize
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      with temporary_feature_override(
+          FeatureName._MCP_GRACEFUL_ERROR_HANDLING, True
+      ):
+        with pytest.raises(ConnectionError) as exc_info:
+          await session_context.start()
+
+      assert 'Failed to create MCP session' in str(exc_info.value)
+
+  @pytest.mark.asyncio
+  async def test_timeout_during_initialization_with_flag_off(self):
+    """Test timeout during session initialization with flag OFF.
+
+    Verifies that session initialization falls back to `asyncio.wait_for`
+    when graceful error handling is disabled.
+    """
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=0.1, sse_read_timeout=None
+    )
+
+    # Mock ClientSession with slow initialize
+    mock_session = MockClientSession()
+
+    async def slow_initialize():
+      await asyncio.sleep(1.0)
+      return mock_session
+
+    mock_session.initialize = slow_initialize
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      with temporary_feature_override(
+          FeatureName._MCP_GRACEFUL_ERROR_HANDLING, False
+      ):
+        with pytest.raises(ConnectionError) as exc_info:
+          await session_context.start()
+
+      assert 'Failed to create MCP session' in str(exc_info.value)
+
+  @pytest.mark.asyncio
+  async def test_uses_anyio_fail_after_when_flag_on(self):
+    """Test that session initialization structurally uses anyio.fail_after.
+
+    Asserts that the session runner enters `anyio.fail_after` context
+    with the timeout limit when graceful error handling is enabled.
+    """
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=2.5, sse_read_timeout=None
+    )
+    mock_session = MockClientSession()
+
+    with (
+        patch(
+            'google.adk.tools.mcp_tool.session_context.ClientSession'
+        ) as mock_session_class,
+        patch('anyio.fail_after') as mock_fail_after,
+    ):
+      mock_session_class.return_value = mock_session
+      # Configure mock_fail_after synchronous context manager to do nothing
+      mock_fail_after.return_value.__enter__ = Mock()
+      mock_fail_after.return_value.__exit__ = Mock(return_value=False)
+
+      with temporary_feature_override(
+          FeatureName._MCP_GRACEFUL_ERROR_HANDLING, True
+      ):
+        await session_context.start()
+
+        # Verify anyio.fail_after was called with the correct timeout
+        mock_fail_after.assert_called_once_with(2.5)
+
+      await session_context.close()
+
+  @pytest.mark.asyncio
+  async def test_stdio_client_with_read_timeout(self):
+    """Test stdio client includes read_timeout_seconds parameter."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None, is_stdio=True
+    )
+
+    mock_session = MockClientSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      await session_context.start()
+
+      # Verify ClientSession was called with read_timeout_seconds for stdio
+      call_args = mock_session_class.call_args
+      assert 'read_timeout_seconds' in call_args.kwargs
+      assert call_args.kwargs['read_timeout_seconds'] == timedelta(seconds=5.0)
+
+      await session_context.close()
+
+  @pytest.mark.asyncio
+  async def test_non_stdio_client_without_read_timeout(self):
+    """Test non-stdio client does not include read_timeout_seconds."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None, is_stdio=False
+    )
+
+    mock_session = MockClientSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      await session_context.start()
+
+      # Verify ClientSession was called with read_timeout_seconds=None for non-stdio
+      # when sse_read_timeout is None
+      call_args = mock_session_class.call_args
+      assert 'read_timeout_seconds' in call_args.kwargs
+      assert call_args.kwargs['read_timeout_seconds'] is None
+
+      await session_context.close()
+
+  @pytest.mark.asyncio
+  async def test_sse_read_timeout_passed_to_client_session(self):
+    """Test that sse_read_timeout is passed to ClientSession for non-stdio."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=300.0, is_stdio=False
+    )
+
+    mock_session = MockClientSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      await session_context.start()
+
+      # Verify ClientSession was called with sse_read_timeout
+      call_args = mock_session_class.call_args
+      assert 'read_timeout_seconds' in call_args.kwargs
+      assert call_args.kwargs['read_timeout_seconds'] == timedelta(
+          seconds=300.0
+      )
+
+      await session_context.close()
+
+  @pytest.mark.asyncio
+  async def test_close_multiple_times(self):
+    """Test that close() can be called multiple times safely."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    mock_session = MockClientSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      await session_context.start()
+
+      # Close multiple times
+      await session_context.close()
+      await session_context.close()
+      await session_context.close()
+
+      # Should not raise exception
+      assert session_context._close_event.is_set()
+
+  @pytest.mark.asyncio
+  async def test_close_before_start(self):
+    """Test that close() works even if start() was never called."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    # Close before starting should not raise
+    await session_context.close()
+
+    assert session_context._close_event.is_set()
+
+  @pytest.mark.asyncio
+  async def test_close_before_start_ends(self):
+    """Test that close() before start() ends the task."""
+    # Client has enough time to delay the start task
+    mock_client = MockClient(delay_on_enter=10.0)
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    start_task = asyncio.create_task(session_context.start())
+    await asyncio.sleep(0.1)
+    assert not start_task.done()
+
+    # Call close before start() ends the task
+    await session_context.close()
+    await asyncio.sleep(0.1)
+
+    assert start_task.done()
+    assert isinstance(
+        start_task.exception(), ConnectionError
+    ) and 'task cancelled' in str(start_task.exception())
+
+  @pytest.mark.asyncio
+  async def test_close_before_start_called(self):
+    """Test that close() before start() called sets the close event."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    # Call close() before start() called
+    await session_context.close()
+    await asyncio.sleep(0.1)
+
+    assert session_context._task is None
+    assert session_context._close_event.is_set()
+
+    with pytest.raises(ConnectionError) as exc_info:
+      await session_context.start()
+
+    assert 'session already closed' in str(exc_info.value)
+    assert session_context._task is None
+
+  @pytest.mark.asyncio
+  async def test_session_property(self):
+    """Test that session property returns the managed session."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    # Initially None
+    assert session_context.session is None
+
+    mock_session = MockClientSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = mock_session
+
+      await session_context.start()
+
+      # Should return the session
+      assert session_context.session == mock_session
+
+      await session_context.close()
+
+  @pytest.mark.asyncio
+  async def test_client_cleanup_on_exception(self):
+    """Test that client is properly cleaned up even when exception occurs."""
+    test_exception = RuntimeError('Test error')
+    mock_client = MockClient(raise_on_enter=test_exception)
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with pytest.raises(ConnectionError):
+      await session_context.start()
+
+    # Wait a bit for cleanup
+    await asyncio.sleep(0.1)
+
+    # Verify task completed
+    assert session_context._task.done()
+
+  @pytest.mark.asyncio
+  async def test_close_handles_cancelled_error(self):
+    """Test that close() handles CancelledError gracefully."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    mock_session = MockClientSession()
+
+    with (
+        patch(
+            'google.adk.tools.mcp_tool.session_context.ClientSession'
+        ) as mock_session_class,
+        patch(
+            'google.adk.tools.mcp_tool.session_context.logger'
+        ) as mock_logger,
+    ):
+      mock_session_class.return_value = mock_session
+
+      await session_context.start()
+
+      # Cancel the task
+      if session_context._task:
+        session_context._task.cancel()
+
+      # Close should handle CancelledError gracefully
+      await session_context.close()
+
+      # Should not raise exception
+      assert session_context._close_event.is_set()
+
+      # Verify no warning logs were generated
+      mock_logger.warning.assert_not_called()
+
+  @pytest.mark.asyncio
+  async def test_close_handles_exception_during_cleanup(self):
+    """Test that close() handles exceptions during cleanup gracefully."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    # Create a mock session that raises during exit
+    class FailingMockSession(MockClientSession):
+
+      async def __aexit__(self, exc_type, exc_val, exc_tb):
+        raise RuntimeError('Cleanup failed')
+
+    failing_session = FailingMockSession()
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = failing_session
+
+      await session_context.start()
+
+      # Close should handle the exception gracefully
+      await session_context.close()
+
+      # Should not raise exception
+      assert session_context._close_event.is_set()
+
+
+class TestSessionContextIsTaskAlive:
+  """Tests for the SessionContext._is_task_alive property."""
+
+  def test_is_task_alive_false_before_start(self):
+    """Before start(), there is no task and the property returns False."""
+    session_context = SessionContext(
+        MockClient(), timeout=5.0, sse_read_timeout=None
+    )
+    assert session_context._is_task_alive is False
+
+  @pytest.mark.asyncio
+  async def test_is_task_alive_true_while_session_running(self):
+    """After start(), the background task is alive until close()."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = MockClientSession()
+      await session_context.start()
+      try:
+        assert session_context._is_task_alive is True
+      finally:
+        await session_context.close()
+
+      assert session_context._is_task_alive is False
+
+
+class TestSessionContextRunGuarded:
+  """Tests for SessionContext._run_guarded.
+
+  This is the heart of the 5-minute-hang fix: the method races a
+  coroutine against the background session task and surfaces transport
+  crashes immediately.
+  """
+
+  @pytest.mark.asyncio
+  async def test_run_guarded_raises_when_task_not_started(self):
+    """If start() was never called, _run_guarded refuses to run the coro."""
+    session_context = SessionContext(
+        MockClient(), timeout=5.0, sse_read_timeout=None
+    )
+
+    async def coro():
+      return 'should never run'
+
+    with pytest.raises(ConnectionError, match='task has not been started'):
+      await session_context._run_guarded(coro())
+
+  @pytest.mark.asyncio
+  async def test_run_guarded_returns_result_on_success(self):
+    """When the coroutine completes first, its result is returned."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = MockClientSession()
+      await session_context.start()
+      try:
+
+        async def coro():
+          return 'expected_result'
+
+        result = await session_context._run_guarded(coro())
+        assert result == 'expected_result'
+      finally:
+        await session_context.close()
+
+  @pytest.mark.asyncio
+  async def test_run_guarded_propagates_coro_exception(self):
+    """A coroutine-level exception propagates as-is (not wrapped).
+
+    This is intentional: callers (McpTool) need to distinguish a
+    tool-level failure (McpError) from a transport-level failure
+    (ConnectionError).
+    """
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = MockClientSession()
+      await session_context.start()
+      try:
+
+        async def coro():
+          raise ValueError('tool error')
+
+        with pytest.raises(ValueError, match='tool error'):
+          await session_context._run_guarded(coro())
+      finally:
+        await session_context.close()
+
+  @pytest.mark.asyncio
+  async def test_run_guarded_raises_when_task_died_before_call(self):
+    """If the background task already died, surface ConnectionError immediately."""
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = MockClientSession()
+      await session_context.start()
+      # Simulate a transport crash by closing the session.
+      await session_context.close()
+
+      async def coro():
+        return 'should not run'
+
+      with pytest.raises(ConnectionError, match='already terminated'):
+        await session_context._run_guarded(coro())
+
+  @pytest.mark.asyncio
+  async def test_run_guarded_cancels_coro_when_task_dies_first(self):
+    """If the background task dies mid-flight, cancel the coro and raise.
+
+    This is the regression test for the 5-minute hang: when the MCP
+    transport crashes (e.g. AGW returns 403), the background task ends
+    quickly, and the in-flight call must be cancelled rather than
+    waiting for sse_read_timeout.
+    """
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = MockClientSession()
+      await session_context.start()
+
+      coro_started = asyncio.Event()
+      coro_was_cancelled = False
+
+      async def slow_coro():
+        nonlocal coro_was_cancelled
+        coro_started.set()
+        try:
+          # Pretend we're awaiting a 5-minute SSE read.
+          await asyncio.sleep(300)
+          return 'should never reach here'
+        except asyncio.CancelledError:
+          coro_was_cancelled = True
+          raise
+
+      async def kill_background_task():
+        await coro_started.wait()
+        # Simulate a transport crash by closing the session, which ends
+        # the background task quickly.
+        await session_context.close()
+
+      killer = asyncio.create_task(kill_background_task())
+
+      try:
+        with pytest.raises(ConnectionError, match='connection lost'):
+          await session_context._run_guarded(slow_coro())
+
+        assert coro_was_cancelled is True
+      finally:
+        await killer
+
+
+class TestSessionContextFlagOffPreservesPreFixBehavior:
+  """Pin down that flag=OFF reproduces pre-fix behavior exactly.
+
+  These tests guard against accidental changes leaking into the flag=OFF
+  path, which is the default. An earlier unconditional version of this
+  fix caused existing callers to hit a 3-minute hang because behavior
+  changes were applied to the default path. We must keep flag=OFF
+  byte-for-byte equivalent to pre-fix.
+  """
+
+  @pytest.mark.asyncio
+  async def test_inner_wait_for_is_used_when_flag_off(self):
+    """The inner asyncio.wait_for around enter_async_context must run.
+
+    Pre-fix code wrapped client entry in `asyncio.wait_for(..., timeout)`.
+    Callers that depend on that inner timeout firing for hanging mocks
+    rely on this behavior. With the flag OFF we must restore it.
+    """
+    delayed_client = MockClient(delay_on_enter=10.0)
+    session_context = SessionContext(
+        delayed_client, timeout=0.2, sse_read_timeout=None
+    )
+
+    with temporary_feature_override(
+        FeatureName._MCP_GRACEFUL_ERROR_HANDLING, False
+    ):
+      with pytest.raises(ConnectionError):
+        # The inner wait_for should fire at ~timeout=0.2s, surfacing as
+        # ConnectionError from start(). If the inner wait_for is missing
+        # (the AnyIO fix being applied unconditionally), this test would
+        # block until the OUTER timeout cancels - which doesn't exist
+        # here because we're calling start() directly.
+        await asyncio.wait_for(session_context.start(), timeout=2.0)
+
+    # And confirm: this would NOT raise quickly with the flag ON
+    # because the inner wait_for is removed. We don't actually run the
+    # flag-on case here because there's no outer timeout in this direct
+    # call - that's tested at the McpTool integration level.
+
+  @pytest.mark.asyncio
+  async def test_no_extra_none_check_when_flag_off(self):
+    """The 'session is None' raise must NOT happen when flag is off.
+
+    Pre-fix code returned `self._session` directly, even if it was
+    somehow None. Our new None check is gated to preserve that.
+    """
+    mock_client = MockClient()
+    session_context = SessionContext(
+        mock_client, timeout=5.0, sse_read_timeout=None
+    )
+
+    with patch(
+        'google.adk.tools.mcp_tool.session_context.ClientSession'
+    ) as mock_session_class:
+      mock_session_class.return_value = MockClientSession()
+      with temporary_feature_override(
+          FeatureName._MCP_GRACEFUL_ERROR_HANDLING, False
+      ):
+        # In normal flow, _session is set; the gated None check is moot.
+        # This test exists primarily to document and guard the flag-OFF
+        # code path.
+        result = await session_context.start()
+        try:
+          assert result is not None
+        finally:
+          await session_context.close()
+
+
+class TestFormatException:
+  """Test suite for _format_exception helper."""
+
+  def test_format_exception_normal(self):
+    exc = ValueError('normal error')
+    assert _format_exception(exc) == 'normal error'
+
+  def test_format_exception_http_status_error(self):
+    request = httpx.Request('GET', 'http://test')
+    response = httpx.Response(403, request=request, text='Forbidden access')
+    exc = httpx.HTTPStatusError(
+        '403 Forbidden', request=request, response=response
+    )
+
+    formatted = _format_exception(exc)
+    assert '403 Forbidden' in formatted
+    assert 'Forbidden access' in formatted
+
+  def test_format_exception_group(self):
+    class MockExceptionGroup(Exception):
+
+      def __init__(self, message, exceptions):
+        super().__init__(message)
+        self.exceptions = exceptions
+
+    request = httpx.Request('GET', 'http://test')
+    response = httpx.Response(403, request=request, text='Forbidden access')
+    exc1 = httpx.HTTPStatusError(
+        '403 Forbidden', request=request, response=response
+    )
+    exc2 = ValueError('another error')
+
+    eg = MockExceptionGroup('Group', [exc1, exc2])
+    formatted = _format_exception(eg)
+
+    assert '403 Forbidden' in formatted
+    assert 'Forbidden access' in formatted
+    assert 'another error' in formatted