feat(mcp): add sampling callback support for MCP sessions

Merge #4718

### Link to Issue or Description of Change

**1. Link to an existing issue (if applicable):**

- Closes: N/A
- Related: N/A

**2. Or, if no issue exists, describe the change:**

**Problem**

ADK’s MCP integration currently does not expose the MCP sampling callback capability.
This prevents agent-side LLM sampling handlers from being used when interacting with MCP servers that support sampling.

The MCP Python SDK supports sampling callbacks, but these parameters are not propagated through the ADK MCP integration layers.

**Solution**

Add sampling callback support by propagating the parameters through the MCP stack:

- Add `sampling_callback` and `sampling_capabilities` parameters to `McpToolset`
- Forward them to `MCPSessionManager`
- Forward them to `SessionContext`
- Pass them into `ClientSession` initialization

This enables agent-side sampling handling when interacting with MCP servers.

---

### Testing Plan

**Unit Tests**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

Added `test_mcp_sampling_callback.py` to verify that the sampling callback is correctly invoked.

Example result:
pytest tests/unittests/tools/mcp_tool/test_mcp_sampling_callback.py
1 passed

**Manual End-to-End (E2E) Tests**

Manual testing was performed using a FastMCP sampling example server where the sampling callback was invoked from the agent side and returned the expected response.

---

### Checklist

- [x] I have read the CONTRIBUTING.md document.
- [x] I have performed a self-review of my own code.
- [x] I have commented my code where necessary.
- [x] I have added tests proving the feature works.
- [x] Unit tests pass locally.
- [x] I have manually tested the change end-to-end.

---

### Additional context

This change aligns ADK MCP support with the sampling capabilities available in the MCP Python SDK and enables agent implementations to handle sampling requests via a callback.

Co-authored-by: Kathy Wu <wukathy@google.com>
COPYBARA_INTEGRATE_REVIEW=#4718 from Piyushmrya:fix-mcp-sampling-callback 18f477f
PiperOrigin-RevId: 883401178

Author: 2 people

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

@@ -33,7 +33,9 @@
 from typing import Union
 
 from mcp import ClientSession
+from mcp import SamplingCapability
 from mcp import StdioServerParameters
+from mcp.client.session import SamplingFnT
 from mcp.client.sse import sse_client
 from mcp.client.stdio import stdio_client
 from mcp.client.streamable_http import create_mcp_http_client
@@ -195,6 +197,9 @@ def __init__(
           StreamableHTTPConnectionParams,
       ],
       errlog: TextIO = sys.stderr,
+      *,
+      sampling_callback: Optional[SamplingFnT] = None,
+      sampling_capabilities: Optional[SamplingCapability] = None,
   ):
     """Initializes the MCP session manager.
 
@@ -204,7 +209,13 @@ def __init__(
           parameters but it's not configurable for now.
         errlog: (Optional) TextIO stream for error logging. Use only for
           initializing a local stdio MCP session.
+        sampling_callback: Optional callback to handle sampling requests from the
+          MCP server.
+        sampling_capabilities: Optional capabilities for sampling.
     """
+    self._sampling_callback = sampling_callback
+    self._sampling_capabilities = sampling_capabilities
+
     if isinstance(connection_params, StdioServerParameters):
       # So far timeout is not configurable. Given MCP is still evolving, we
       # would expect stdio_client to evolve to accept timeout parameter like
@@ -475,6 +486,8 @@ async def create_session(
                     timeout=timeout_in_seconds,
                     sse_read_timeout=sse_read_timeout_in_seconds,
                     is_stdio=is_stdio,
+                    sampling_callback=self._sampling_callback,
+                    sampling_capabilities=self._sampling_capabilities,
                 )
             ),
             timeout=timeout_in_seconds,

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

@@ -29,7 +29,9 @@
 from typing import Union
 import warnings
 
+from mcp import SamplingCapability
 from mcp import StdioServerParameters
+from mcp.client.session import SamplingFnT
 from mcp.shared.session import ProgressFnT
 from mcp.types import ListResourcesResult
 from mcp.types import ListToolsResult
@@ -114,6 +116,8 @@ def __init__(
           Union[ProgressFnT, ProgressCallbackFactory]
       ] = None,
       use_mcp_resources: Optional[bool] = False,
+      sampling_callback: Optional[SamplingFnT] = None,
+      sampling_capabilities: Optional[SamplingCapability] = None,
   ):
     """Initializes the McpToolset.
 
@@ -150,10 +154,16 @@ def __init__(
       use_mcp_resources: Whether the agent should have access to MCP resources.
         This will add a `load_mcp_resource` tool to the toolset and include
         available resources in the agent context. Defaults to False.
+      sampling_callback: Optional callback to handle sampling requests from the
+        MCP server.
+      sampling_capabilities: Optional capabilities for sampling.
     """
 
     super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
 
+    self._sampling_callback = sampling_callback
+    self._sampling_capabilities = sampling_capabilities
+
     if not connection_params:
       raise ValueError("Missing connection params in McpToolset.")
 
@@ -166,6 +176,8 @@ def __init__(
     self._mcp_session_manager = MCPSessionManager(
         connection_params=self._connection_params,
         errlog=self._errlog,
+        sampling_callback=self._sampling_callback,
+        sampling_capabilities=self._sampling_capabilities,
     )
     self._auth_scheme = auth_scheme
     self._auth_credential = auth_credential

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

@@ -22,6 +22,8 @@
 from typing import Optional
 
 from mcp import ClientSession
+from mcp import SamplingCapability
+from mcp.client.session import SamplingFnT
 
 logger = logging.getLogger('google_adk.' + __name__)
 
@@ -54,6 +56,9 @@ def __init__(
       timeout: Optional[float],
       sse_read_timeout: Optional[float],
       is_stdio: bool = False,
+      *,
+      sampling_callback: Optional[SamplingFnT] = None,
+      sampling_capabilities: Optional[SamplingCapability] = None,
   ):
     """
     Args:
@@ -63,6 +68,9 @@ def __init__(
         sse_read_timeout: Timeout in seconds for reading data from the MCP SSE
             server.
         is_stdio: Whether this is a stdio connection (affects read timeout).
+        sampling_callback: Optional callback to handle sampling requests from the
+            MCP server.
+        sampling_capabilities: Optional capabilities for sampling.
     """
     self._client = client
     self._timeout = timeout
@@ -73,6 +81,8 @@ def __init__(
     self._close_event = asyncio.Event()
     self._task: Optional[asyncio.Task] = None
     self._task_lock = asyncio.Lock()
+    self._sampling_callback = sampling_callback
+    self._sampling_capabilities = sampling_capabilities
 
   @property
   def session(self) -> Optional[ClientSession]:
@@ -165,6 +175,8 @@ async def _run(self):
                   read_timeout_seconds=timedelta(seconds=self._timeout)
                   if self._timeout is not None
                   else None,
+                  sampling_callback=self._sampling_callback,
+                  sampling_capabilities=self._sampling_capabilities,
               )
           )
         else:
@@ -176,6 +188,8 @@ async def _run(self):
                   read_timeout_seconds=timedelta(seconds=self._sse_read_timeout)
                   if self._sse_read_timeout is not None
                   else None,
+                  sampling_callback=self._sampling_callback,
+                  sampling_capabilities=self._sampling_capabilities,
               )
           )
         await asyncio.wait_for(session.initialize(), timeout=self._timeout)

tests/unittests/tools/mcp_tool/test_mcp_toolset.py

@@ -612,3 +612,37 @@ async def test_read_resource(self, name, mime_type, content, encoding):
     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"