chore: import upstream snapshot with attribution

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