Fix streamable http sampling #531
Open
+614
−29
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
Implements sampling capability for HTTP transport, resolving issue mark3labs#419. Enables servers to send sampling requests to HTTP clients via SSE and receive LLM-generated responses. ## Key Changes ### Core Implementation - Add `BidirectionalInterface` support to `StreamableHTTP` - Implement `SetRequestHandler` for server-to-client requests - Enhance SSE parsing to handle requests alongside responses/notifications - Add `handleIncomingRequest` and `sendResponseToServer` methods ### HTTP-Specific Features - Leverage existing MCP headers (`Mcp-Session-Id`, `Mcp-Protocol-Version`) - Bidirectional communication via HTTP POST for responses - Proper JSON-RPC request/response handling over HTTP ### Error Handling - Add specific JSON-RPC error codes for different failure scenarios: - `-32601` (Method not found) when no handler configured - `-32603` (Internal error) for sampling failures - `-32800` (Request cancelled/timeout) for context errors - Enhanced error messages with sampling-specific context ### Testing & Examples - Comprehensive test suite in `streamable_http_sampling_test.go` - Complete working example in `examples/sampling_http_client/` - Tests cover success flows, error scenarios, and interface compliance ## Technical Details The implementation maintains full backward compatibility while adding bidirectional communication support. Server requests are processed asynchronously to avoid blocking the SSE stream reader. HTTP transport now supports the complete sampling flow that was previously only available in stdio and inprocess transports. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
This completes the server-side implementation of sampling support for HTTP transport, addressing the remaining requirements from issue mark3labs#419. Changes: - Enhanced streamableHttpSession to implement SessionWithSampling interface - Added bidirectional SSE communication for server-to-client requests - Implemented session registry for proper response correlation - Added comprehensive error handling with JSON-RPC error codes - Created extensive test suite covering all scenarios - Added working example server with sampling tools Key Features: - Server can send sampling requests to HTTP clients via SSE - Clients respond via HTTP POST with proper session correlation - Queue overflow protection and timeout handling - Compatible with existing HTTP transport architecture 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Replace flaky time.Sleep calls with proper synchronization using channels and sync.WaitGroup to make tests deterministic and avoid race conditions. Also improves error handling robustness in test servers with proper JSON decoding error checks. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Make request vs response detection more robust by checking for presence of "method" field instead of relying on nil Result/Error fields - Add nil pointer check in sendResponseToServer function to prevent panics These changes improve reliability against malformed messages and edge cases. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
The comment incorrectly stated that responses are broadcast to all sessions, but the implementation actually delivers responses to the specific session identified by sessionID using the activeSessions registry. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Previously, EnableSampling() was a no-op that didn't actually enable the sampling capability in the server's declared capabilities. Changes: - Add Sampling field to mcp.ServerCapabilities struct - Add sampling field to internal serverCapabilities struct - Update EnableSampling() to set the sampling capability flag - Update handleInitialize() to include sampling in capability response - Add test to verify sampling capability is properly declared Now when EnableSampling() is called, the server will properly declare sampling capability during initialization, allowing clients to know that the server supports sending sampling requests. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Replace unsafe type assertion result.Content.(mcp.TextContent).Text with safe type checking to handle cases where Content might not be a TextContent struct. Now gracefully handles different content types without panicking. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
The SamplingInterface test was missing the EnableSampling() call, which is necessary to activate sampling features for proper testing. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Replace single error test with comprehensive table-driven tests - Add test cases for invalid request IDs and malformed results - Replace t.Fatalf with t.Errorf to follow project conventions - Use proper session ID format for valid test scenarios 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Remove recursive call in RequestSampling that could cause stack overflow - Remove problematic response re-queuing to global channel - Update deliverSamplingResponse to route responses directly to dedicated request channels via samplingRequests map lookup - This prevents ordering issues and ensures responses reach the correct waiting request 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Modified deliverSamplingResponse to return error instead of just logging - Added proper error handling for disconnected sessions - Improved error messages for debugging - Updated test expectations to match new error behavior 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Add signal handling for SIGINT and SIGTERM - Move defer statement after error checking - Improve shutdown error handling 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Add timeout context for SSE response processing (30s default) - Add timeout for individual connection attempts in listenForever (10s) - Use context-aware sleep in retry logic - Ensure async goroutines properly respect context cancellation 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Make error message more descriptive and actionable - Provide clearer debugging information about why the channel is blocked 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Rename 'baseMessage' to 'jsonMessage' for more neutral naming - Improves code readability and follows consistent naming conventions 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Add test verifying that concurrent sampling requests are handled correctly when the second request completes faster than the first. The test ensures: - Responses are correctly associated with their request IDs - Server processes requests concurrently without blocking - Completion order follows actual processing time, not submission order 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Create new context with 30-second timeout for request handling to prevent long-running handlers from blocking indefinitely. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Replace all occurrences of interface{} with the modern Go any type alias
for improved readability and consistency with current Go best practices.
🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <[email protected]>
Create timeout context from parent context instead of context.Background() to ensure request handlers respect parent context cancellation. Addresses review comment about context handling in async goroutine. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
The samplingResponseChan field was declared but never used in the streamableHttpSession struct. Remove it and update tests accordingly. Addresses review comment about unused fields in session struct. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Add signal handling for SIGINT and SIGTERM to allow graceful shutdown of the sampling HTTP client example. This prevents indefinite blocking and provides better production-ready behavior. Addresses review comment about adding graceful shutdown handling. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Removes unused sync.RWMutex field that was flagged by golangci-lint. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
WalkthroughThe changes introduce fixes and enhancements to support persistent session reuse for streamable HTTP sampling, correct content parsing in the client, and add comprehensive end-to-end tests for HTTP-based sampling. Adjustments were made to both client and server logic to ensure proper request/response handling and session management for sampling scenarios. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Assessment against linked issues
Suggested labels
Suggested reviewers
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (4)
🧰 Additional context used🧠 Learnings (6)📚 Learning: in the mcptest package, prefer returning errors from helper functions rather than calling t.fatalf()...Applied to files:
📚 Learning: in go mcp server, servertool.tool field is only used for tool listing and indexing, not for tool exe...Applied to files:
📚 Learning: in the mark3labs/mcp-go project, the mcpserver.capabilities field is a struct value (servercapabilit...Applied to files:
📚 Learning: in the mcp-go library, the getpromptparams.arguments field is of type map[string]string, not map[str...Applied to files:
📚 Learning: tool responses from the mcp server shouldn't contain rawinputschema, which is why the unmarshaljson ...Applied to files:
📚 Learning: the tool struct in the mark3labs/mcp-go project should handle both inputschema and rawinputschema co...Applied to files:
🧬 Code Graph Analysis (2)e2e/sampling_http_test.go (1)
client/client.go (1)
🪛 ast-grep (0.38.6)e2e/sampling_http_test.go[warning] 83-83: "Detected a network listener listening on 0.0.0.0 or an empty string. (avoid-bind-to-all-interfaces-go) 🔇 Additional comments (6)
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
andig
changed the title
andig
commented
Aug 6, 2025
| @@ -478,6 +478,22 @@ func (c *Client) handleSamplingRequestTransport(ctx context.Context, request tra | |||
| return nil, fmt.Errorf("failed to unmarshal params: %w", err) | |||
| } | |||
| } | |||
|
|
|||
| // Fix content parsing - HTTP transport unmarshals TextContent as map[string]any | |||
| for i := range params.Messages { | |||
There was a problem hiding this comment.
This is actually
func getTextFromContent(content any) string
from the examples directory. Might be a good idea to make that a global helper function?
|
@ezynda3 how would you want to proceed? The fix here seems clear, not sure what to do about test, helper function and documentation?. |
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.
Fix #530
Add e2e test and fix sampling
TODO
getTextFromContenthelper functiontransport.WithContinuousListening()Summary by CodeRabbit
New Features
Bug Fixes
Refactor