feat: Inprocess sampling support #487
Conversation
WalkthroughThis change introduces in-process client sampling support to the MCP framework. It adds new constructors, handler interfaces, and session management for in-process sampling, updates the transport and server logic to support sampling handlers, and provides comprehensive documentation and examples for using and integrating sampling in in-process scenarios. Changes
Possibly related PRs
Suggested labels
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
✨ Finishing Touches
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. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
client/transport/inprocess.go (1)
111-113: Fix GetSessionId to return the actual session ID.The method returns an empty string but should return the
sessionIDfield.Apply this fix:
func (c *InProcessTransport) GetSessionId() string { - return "" + return c.sessionID }
🧹 Nitpick comments (3)
examples/inprocess_sampling/README.md (1)
33-35: Fix markdown linting issue with fenced code block.The fenced code block is missing a language specification, which is flagged by markdownlint.
Apply this fix:
-``` +```text Tool result: LLM Response (model: mock-llm-v1): Mock LLM response to: 'What is the capital of France?'</blockquote></details> <details> <summary>server/inprocess_session.go (2)</summary><blockquote> `24-26`: **Consider removing the mutex as samplingHandler is immutable.** The `samplingHandler` field is only set during construction and never modified afterward. The mutex in `RequestSampling` (lines 79-81) may be unnecessary overhead. If the handler truly never changes after construction, you can remove the mutex: ```diff type InProcessSession struct { sessionID string notifications chan mcp.JSONRPCNotification initialized atomic.Bool loggingLevel atomic.Value clientInfo atomic.Value samplingHandler SamplingHandler - mu sync.RWMutex }And simplify
RequestSampling:func (s *InProcessSession) RequestSampling(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) { - s.mu.RLock() - handler := s.samplingHandler - s.mu.RUnlock() - - if handler == nil { + if s.samplingHandler == nil { return nil, fmt.Errorf("no sampling handler available") } - return handler.CreateMessage(ctx, request) + return s.samplingHandler.CreateMessage(ctx, request) }
91-93: Consider using a more collision-resistant ID generation method.While nanosecond timestamps are usually sufficient, consider using a UUID or adding a random component for better collision resistance in high-frequency scenarios.
Example using crypto/rand:
import ( "crypto/rand" "encoding/hex" ) func GenerateInProcessSessionID() string { b := make([]byte, 16) _, _ = rand.Read(b) return fmt.Sprintf("inprocess-%d-%s", time.Now().UnixNano(), hex.EncodeToString(b)) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (12)
client/inprocess.go(2 hunks)client/inprocess_sampling_test.go(1 hunks)client/transport/inprocess.go(3 hunks)examples/inprocess_sampling/README.md(1 hunks)examples/inprocess_sampling/main.go(1 hunks)server/inprocess_session.go(1 hunks)server/sampling.go(2 hunks)server/server.go(2 hunks)www/docs/pages/clients/advanced-sampling.mdx(5 hunks)www/docs/pages/servers/advanced-sampling.mdx(1 hunks)www/docs/pages/transports/index.mdx(2 hunks)www/docs/pages/transports/inprocess.mdx(1 hunks)
🧰 Additional context used
🧠 Learnings (11)
server/server.go (4)
Learnt from: octo
PR: mark3labs/mcp-go#149
File: mcptest/mcptest.go:0-0
Timestamp: 2025-04-21T21:26:32.945Z
Learning: In the mcptest package, prefer returning errors from helper functions rather than calling t.Fatalf() directly, giving callers flexibility in how to handle errors.
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
Learnt from: davidleitw
PR: mark3labs/mcp-go#451
File: mcp/tools.go:1192-1217
Timestamp: 2025-06-26T09:38:18.629Z
Learning: In mcp-go project, the maintainer prefers keeping builder pattern APIs simple without excessive validation for edge cases. The WithOutput* functions are designed to assume correct usage rather than defensive programming, following the principle of API simplicity over comprehensive validation.
www/docs/pages/transports/index.mdx (1)
Learnt from: leavez
PR: mark3labs/mcp-go#114
File: client/transport/sse.go:137-179
Timestamp: 2025-04-06T10:07:06.685Z
Learning: The SSE client implementation in the MCP-Go project uses a 30-second timeout for reading SSE events to match the behavior of the original implementation before the transport layer refactoring.
examples/inprocess_sampling/README.md (1)
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
server/sampling.go (1)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
www/docs/pages/clients/advanced-sampling.mdx (7)
Learnt from: leavez
PR: mark3labs/mcp-go#114
File: client/transport/sse.go:137-179
Timestamp: 2025-04-06T10:07:06.685Z
Learning: The SSE client implementation in the MCP-Go project uses a 30-second timeout for reading SSE events to match the behavior of the original implementation before the transport layer refactoring.
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
Learnt from: octo
PR: mark3labs/mcp-go#149
File: mcptest/mcptest.go:0-0
Timestamp: 2025-04-21T21:26:32.945Z
Learning: In the mcptest package, prefer returning errors from helper functions rather than calling t.Fatalf() directly, giving callers flexibility in how to handle errors.
Learnt from: davidleitw
PR: mark3labs/mcp-go#451
File: mcp/tools.go:1192-1217
Timestamp: 2025-06-26T09:38:18.629Z
Learning: In mcp-go project, the maintainer prefers keeping builder pattern APIs simple without excessive validation for edge cases. The WithOutput* functions are designed to assume correct usage rather than defensive programming, following the principle of API simplicity over comprehensive validation.
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
Learnt from: floatingIce91
PR: mark3labs/mcp-go#401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
client/inprocess_sampling_test.go (4)
Learnt from: octo
PR: mark3labs/mcp-go#149
File: mcptest/mcptest.go:0-0
Timestamp: 2025-04-21T21:26:32.945Z
Learning: In the mcptest package, prefer returning errors from helper functions rather than calling t.Fatalf() directly, giving callers flexibility in how to handle errors.
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:107-137
Timestamp: 2025-03-04T06:59:43.882Z
Learning: Tool responses from the MCP server shouldn't contain RawInputSchema, which is why the UnmarshalJSON method for the Tool struct is implemented to handle only the structured InputSchema format.
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
Learnt from: floatingIce91
PR: mark3labs/mcp-go#401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
www/docs/pages/transports/inprocess.mdx (1)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
examples/inprocess_sampling/main.go (3)
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:107-137
Timestamp: 2025-03-04T06:59:43.882Z
Learning: Tool responses from the MCP server shouldn't contain RawInputSchema, which is why the UnmarshalJSON method for the Tool struct is implemented to handle only the structured InputSchema format.
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
Learnt from: floatingIce91
PR: mark3labs/mcp-go#401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
server/inprocess_session.go (1)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
client/inprocess.go (4)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
Learnt from: octo
PR: mark3labs/mcp-go#149
File: mcptest/mcptest.go:0-0
Timestamp: 2025-04-21T21:26:32.945Z
Learning: In the mcptest package, prefer returning errors from helper functions rather than calling t.Fatalf() directly, giving callers flexibility in how to handle errors.
Learnt from: lariel-fernandes
PR: mark3labs/mcp-go#428
File: www/docs/pages/servers/prompts.mdx:218-234
Timestamp: 2025-06-20T20:39:51.870Z
Learning: In the mcp-go library, the GetPromptParams.Arguments field is of type map[string]string, not map[string]interface{}, so direct string access without type assertions is safe and correct.
Learnt from: floatingIce91
PR: mark3labs/mcp-go#401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
www/docs/pages/servers/advanced-sampling.mdx (1)
Learnt from: ezynda3
PR: mark3labs/mcp-go#461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.
🧬 Code Graph Analysis (4)
server/server.go (1)
server/inprocess_session.go (1)
GenerateInProcessSessionID(91-93)
server/sampling.go (2)
server/inprocess_session.go (1)
SamplingHandler(14-16)client/sampling.go (1)
SamplingHandler(11-20)
server/inprocess_session.go (3)
mcp/types.go (6)
CreateMessageRequest(775-778)CreateMessageResult(795-802)JSONRPCNotification(318-321)LoggingLevelError(758-758)Implementation(480-483)LoggingLevel(751-751)server/session.go (3)
ClientSession(11-20)SessionWithLogging(23-29)SessionWithClientInfo(43-49)server/sampling.go (1)
SessionWithSampling(39-42)
client/transport/inprocess.go (4)
server/server.go (1)
MCPServer(139-164)server/inprocess_session.go (4)
SamplingHandler(14-16)InProcessSession(18-26)GenerateInProcessSessionID(91-93)NewInProcessSession(28-34)mcp/types.go (1)
JSONRPCNotification(318-321)client/client.go (1)
WithSamplingHandler(39-43)
🪛 markdownlint-cli2 (0.17.2)
examples/inprocess_sampling/README.md
33-33: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🔇 Additional comments (28)
www/docs/pages/transports/index.mdx (2)
16-21: Excellent documentation update for sampling support.The addition of the "Sampling Support" column clearly communicates which transports support the new sampling feature. The technical accuracy aligns well with the implementation changes.
134-141: Good additions to In-Process transport use cases.The new use cases "Sampling-enabled applications" and "LLM-powered applications with bidirectional communication" effectively highlight the enhanced capabilities introduced by the sampling feature.
www/docs/pages/servers/advanced-sampling.mdx (2)
355-393: Comprehensive transport support documentation.The new "Transport Support" section effectively explains sampling capabilities across different transports. The technical explanations for why SSE and StreamableHTTP don't support sampling are accurate and help users understand the limitations.
398-399: Updated next steps section enhances discoverability.The addition of the in-process sampling documentation link provides good navigation to related content.
examples/inprocess_sampling/README.md (1)
1-39: Well-structured example documentation.The README provides clear explanations of the key components, running instructions, and integration guidance. The structure effectively guides users through understanding and using the in-process sampling feature.
client/inprocess.go (2)
17-29: Clean implementation of sampling-enabled client constructor.The new function properly creates an in-process client with sampling support. The wrapper pattern effectively adapts the client's
SamplingHandlerinterface to the server's interface, and the use of transport options maintains consistency with existing patterns.
31-38: Appropriate wrapper implementation for interface adaptation.The
inProcessSamplingHandlerWrappercorrectly implements the server'sSamplingHandlerinterface by forwarding calls to the client's handler. This design allows clean separation between client and server handler interfaces.server/sampling.go (2)
30-34: LGTM! Clean fallback mechanism for in-process sampling.This addition provides an elegant fallback that enables in-process sampling without modifying existing session-based sampling logic. The approach allows direct method calls instead of JSON-RPC serialization for better performance.
44-58: Well-implemented context helper functions.The context key type and helper functions follow Go best practices:
- Uses empty struct for the context key to avoid collisions
- Proper type assertion with ok check in the retrieval function
- Clear naming conventions
www/docs/pages/clients/advanced-sampling.mdx (3)
21-21: Good documentation update to reflect new API patterns.The documentation correctly updates the client instantiation pattern to use explicit transport creation, which aligns with the new in-process sampling support and provides better consistency across transports.
Also applies to: 71-71, 155-155, 379-379, 471-472
76-81: Excellent addition of client lifecycle management.The documentation now properly shows the client lifecycle with
Start()beforeConnect(), including error handling and proper cleanup order. This will help users avoid common setup mistakes.
462-497: Valuable Transport Support section.This new section provides clear guidance on which transports support sampling and how to use them. The distinction between STDIO requiring separate transport creation and in-process having a dedicated constructor is particularly helpful.
examples/inprocess_sampling/main.go (3)
13-42: Well-implemented MockSamplingHandler.The handler properly implements the
SamplingHandlerinterface:
- Correctly extracts user messages from the request
- Generates appropriate mock responses
- Returns properly structured
CreateMessageResultwith all required fields
67-117: Excellent tool implementation demonstrating sampling flow.The tool handler shows a complete sampling request flow:
- Proper parameter extraction with error handling
- Correct construction of
CreateMessageRequestwith all parameters- Appropriate error handling and response formatting
- Clear demonstration of how servers can leverage client-side LLM capabilities
119-166: Comprehensive client setup and usage example.The client setup demonstrates proper lifecycle management:
- Uses the new
NewInProcessClientWithSamplingHandlerfunction- Includes proper error handling and cleanup
- Shows complete initialization flow
- Demonstrates tool invocation with result handling
client/inprocess_sampling_test.go (2)
11-26: Simple and effective mock handler for testing.The test mock handler is well-designed:
- Returns a predictable response for verification
- Includes all required fields in the response
- Keeps the implementation simple to focus on testing the integration flow
28-148: Comprehensive integration test for in-process sampling.The test provides excellent coverage:
- Tests complete flow from server setup to client response
- Includes proper error handling at each step
- Validates the expected response content accurately
- Demonstrates proper lifecycle management with cleanup
This test effectively validates the end-to-end in-process sampling functionality.
www/docs/pages/transports/inprocess.mdx (3)
272-446: Comprehensive sampling support documentation.This new section provides excellent coverage of in-process sampling:
- Clear explanation of bidirectional communication capabilities
- Step-by-step example with proper error handling
- Complete code example showing the full flow
- Good balance of explanation and practical implementation
448-556: Excellent real LLM integration example.The OpenAI integration example is well-structured:
- Proper message format conversion between MCP and OpenAI
- Correct handling of system prompts and conversation flow
- Appropriate error handling for API calls
- Clear demonstration of production-ready implementation
558-616: Valuable guidance on parameters and error handling.The documentation effectively covers:
- Complete parameter reference with examples
- Best practices for error handling
- Practical patterns for graceful degradation
- Clear guidance on user experience considerations
client/transport/inprocess.go (5)
14-29: Clean implementation of sampling support fields and options.The addition of sampling-related fields and the functional option pattern are well-structured and follow Go best practices.
37-48: Well-designed constructor with options support.The constructor properly initializes the transport with a unique session ID and applies the provided options correctly.
51-58: Proper session lifecycle management in Start method.The conditional session creation and registration with appropriate error handling ensures clean session management.
68-72: Correct session context propagation.The session is properly injected into the context before forwarding the request to the server.
104-109: Proper session cleanup in Close method.The session is correctly unregistered during transport closure.
server/inprocess_session.go (3)
28-34: Verify that the notification channel buffer size is adequate.The channel has a fixed buffer of 100. If notifications aren't consumed quickly enough, sends could block.
Consider monitoring or making the buffer size configurable if blocking becomes an issue in production.
44-76: Well-implemented session interface methods.Proper use of atomic operations and nil checks for type safety.
96-101: Good practice with compile-time interface compliance checks.Ensures the struct properly implements all required interfaces.
Description
Fixes #<issue_number> (if applicable)
Type of Change
Checklist
MCP Spec Compliance
Additional Information
Summary by CodeRabbit
New Features
Documentation
Tests