feat(mcp): add sampling callback support for MCP sessions#4718
Closed
Piyushmrya wants to merge 6 commits intogoogle:mainfrom
Closed
feat(mcp): add sampling callback support for MCP sessions#4718Piyushmrya wants to merge 6 commits intogoogle:mainfrom
Piyushmrya wants to merge 6 commits intogoogle:mainfrom
Conversation
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Collaborator
|
Hi @Piyushmrya ,Thank you for your contribution! We appreciate you taking the time to submit this pull request. |
copybara-service bot
pushed a commit
that referenced
this pull request
Mar 13, 2026
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
Collaborator
|
Thank you @Piyushmrya for your contribution! 🎉 Your changes have been successfully imported and merged via Copybara in commit 8f82697. Closing this PR as the changes are now in the main branch. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
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:
sampling_callbackandsampling_capabilitiesparameters toMcpToolsetMCPSessionManagerSessionContextClientSessioninitializationThis enables agent-side sampling handling when interacting with MCP servers.
Testing Plan
Unit Tests
Added
test_mcp_sampling_callback.pyto 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
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.