Implement MCP sampling protocol for dynamic tool discovery #2525
Conversation
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
Add comprehensive support for the Model Context Protocol (MCP) sampling feature, enabling MCP servers to request LLM completions through Amazon Q CLI. Key capabilities: - MCP servers can send sampling/createMessage requests to Amazon Q CLI - Amazon Q CLI processes requests using the configured LLM model - Responses follow MCP specification format with role, content, model fields - Supports text, image, and audio content types in sampling requests - Handles system prompts, model preferences, and token limits - Provides fallback responses when API client unavailable - Comprehensive error handling for malformed or invalid requests This enables powerful dynamic workflows where MCP servers can: - Analyze user tasks and intelligently enable/disable tools - Make context-aware decisions about which capabilities to expose - Provide adaptive user experiences based on LLM reasoning Implementation includes: - Full sampling request/response type definitions - Integration with Amazon Q's conversation API - Proper MCP protocol compliance and error handling - Extensive test coverage (16 sampling-specific tests) - End-to-end validation with real MCP server integration The sampling protocol is a key MCP feature that allows servers to leverage the client's LLM capabilities for intelligent decision-making.
Add end-to-end integration test that validates the complete MCP sampling workflow using the existing test server infrastructure. Key features: - Tests real MCP client-server communication over stdio transport - Validates sampling protocol request/response cycle - Includes discover_tools function that demonstrates dynamic tool discovery - Extends existing test server with tools/call and tools/list support - Follows established testing patterns from existing MCP tests - Provides comprehensive error handling and timeout management The test successfully demonstrates: - MCP server can request LLM completions via sampling/createMessage - Amazon Q CLI processes sampling requests correctly - Response format complies with MCP specification - End-to-end workflow from tool call to sampling completion This integration test provides confidence that the MCP sampling protocol implementation works correctly in real-world scenarios and can be used as a foundation for further MCP sampling development.
- Add 'sampling' boolean field to CustomToolConfig (defaults to false) - Add sampling_enabled field to ClientConfig and Client struct - Implement permission gating in handle_sampling_request method - Reject sampling requests with clear error when sampling not enabled - Add comprehensive test coverage for enabled/disabled scenarios - Update integration test to use sampling_enabled: true - Maintain backward compatibility with safe defaults This provides user control over MCP sampling while maintaining the simplicity of our direct MCP protocol implementation. Users must explicitly opt-in to sampling per server for security.
- Removed 60+ verbose debug and info logs with emojis - Kept only 3 essential error logs for legitimate error handling - Cleaned up DEBUG-prefixed logs that were added during development - Maintained all functional code and tests - Code still compiles and integration test passes This significantly reduces log noise while preserving error visibility.
- Break down massive 418-line handle_sampling_request method into focused helper methods: * validate_sampling_enabled() - validates sampling permission * parse_sampling_request() - parses and validates request format * create_fallback_response() - handles API client unavailable case * process_sampling_with_api() - processes request with API client * handle_successful_api_response() - converts API response to MCP format * create_error_response() - creates error responses * convert_sampling_response_to_json() - converts response to JSON - Fix production unwrap() call in pagination logic with proper error handling - Rename 'sampling' field to 'sampling_enabled' for better clarity - Remove unused imports and clean up code organization - All tests pass including integration test This addresses the 3 main blockers identified for PR readiness: 1. Method size (418 lines -> multiple focused methods ~20-50 lines each) 2. Error handling (replaced unwrap() with proper error handling) 3. Configuration clarity (sampling -> sampling_enabled)
- Remove empty debug blocks for tools/call and tools/list methods - Remove unnecessary allow(dead_code) annotation from set_api_client - Clean up leftover development debugging code
- Fix variable name from sampling to sampling_enabled in custom_tool.rs - Make API client optional in from_config method - Pass API client as Some() when sampling is enabled, None for tests - Remove unused re-export from transport/mod.rs - Fix import in test_server.rs to use direct path - Fix clippy warnings: let_and_return, unused self arguments - Convert methods to associated functions where self is unused - All tests pass and no clippy warnings remain
jsamuel1
reviewed
Aug 8, 2025
|
|
||
| // Sampling-related types for MCP sampling specification | ||
|
|
||
| /// Model preferences for sampling requests |
There was a problem hiding this comment.
Nitpick:
Can we have the links to the specific MCP protocol specification docs for these, so that we can rapidly validate/update these as the protocol evolves?
eg. add to the comment for this struct:
https://modelcontextprotocol.io/specification/2025-06-18/client/sampling#model-preferences
jsamuel1
reviewed
Aug 8, 2025
| @@ -418,6 +481,8 @@ where | |||
| "notifications/tools/list_changed" | "tools/list_changed" | |||
| if tools_list_changed_supported => | |||
| { | |||
| // Add a small delay to prevent rapid-fire loops | |||
| tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; | |||
There was a problem hiding this comment.
What is the rapid fire loop that can happen here? Is it repeatable?
Are we able to not wait on every list changed, but only on subsequent ones?
What is the impact if an MCP Gateway is in use and there are 100s of lists changed?
jsamuel1
reviewed
Aug 8, 2025
| for message in &sampling_request.messages { | ||
| let content = match &message.content { | ||
| SamplingContent::Text { text } => text.clone(), | ||
| SamplingContent::Image { .. } => "[Image content not supported in sampling]".to_string(), |
There was a problem hiding this comment.
Is this:
- A Q cli limitation
- A specification limitation
or - An v1 implementation of sampling limitation
We should ensure the wording of this message is reflective of which.
eg.
"[Image content sampling is not supported by this MCP client]"
"[Image content sampling is not supported by the MCP Protocol]"
or
"[Image content sampling is not supported by this MCP client]" // TODO - Implement.
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.
Implement MCP sampling protocol for dynamic tool discovery
Summary
This PR implements comprehensive support for the Model Context Protocol (MCP) sampling feature, enabling MCP servers to request LLM completions through Amazon Q CLI for intelligent, context-aware tool discovery and workflow management.
Key Features
MCP Sampling Protocol Support
sampling/createMessagerequests per MCP specificationDynamic Tool Discovery
discover_toolsfunctionality for iOS developmentConfiguration and Security
sampling: true/falsein agent configurationImplementation Details
Core Components
sampling/createMessagerequest handler in MCP clientTesting
Usage Example
{ "mcpServers": { "XcodeBuildMCP": { "command": "npx", "args": ["-y", "xcodebuildmcp@latest"], "sampling": true } } }When enabled, MCP servers can request LLM completions to make intelligent decisions about tool availability based on user context and development tasks.
Related Work
This builds upon the MCP foundation and enables the dynamic tool discovery patterns demonstrated in PR #2239, providing the sampling infrastructure needed for intelligent MCP server behavior.
Assisted by Amazon Q Developer