Implement sampling in Stdio #461
Conversation
WalkthroughThis change introduces full support for MCP sampling in both the client and server. It adds a bidirectional sampling handler to the client, enabling it to process sampling requests from servers. The server is updated to send sampling requests and handle asynchronous responses. Documentation and example implementations are included. Changes
Possibly related issues
Suggested labels
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (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 (
|
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (5)
client/transport/stdio_test.go (1)
151-151: Track and prioritize fixing the skipped test.While it's reasonable to skip this test if it's unrelated to the current sampling changes, the notification parsing issue should be tracked and fixed promptly to maintain test coverage.
Would you like me to help investigate the notification parsing issue or create a tracking issue for this TODO?
examples/sampling_client/README.md (1)
1-87: Comprehensive and well-structured documentation.The README provides excellent documentation covering all important aspects of the sampling client example, including features, usage instructions, and implementation details.
Fix fenced code block language specification.
The static analysis tool correctly identifies a missing language specification for the fenced code block.
-``` +```text Connected to server: sampling-example-server v1.0.0 Available tools: - ask_llm: Ask the LLM a question using sampling - greet: Greet the user --- Testing greet tool --- Greet result: Hello, Sampling Demo User! This server supports sampling - try using the ask_llm tool! --- Testing ask_llm tool (with sampling) --- Mock LLM received: What is the capital of France? System prompt: You are a helpful geography assistant. Max tokens: 1000 Temperature: 0.700000 Ask LLM result: LLM Response (model: mock-llm-v1): Mock LLM response to: 'What is the capital of France?'. This is a simulated response from a mock LLM handler. -``` +```client/transport/stdio.go (1)
354-366: Consider improving error handlingUsing
fmt.Printffor error logging is not ideal in a library context. Errors could be lost or interleaved with application output.Consider using a configurable logger or error callback:
+// Add to Stdio struct: +onError func(error) + // sendResponse sends a response back to the server. -func (c *Stdio) sendResponse(response JSONRPCResponse) { +func (c *Stdio) sendResponse(response JSONRPCResponse) error { responseBytes, err := json.Marshal(response) if err != nil { - fmt.Printf("Error marshaling response: %v\n", err) - return + return fmt.Errorf("failed to marshal response: %w", err) } responseBytes = append(responseBytes, '\n') if _, err := c.stdin.Write(responseBytes); err != nil { - fmt.Printf("Error writing response: %v\n", err) + return fmt.Errorf("failed to write response: %w", err) } + return nil }Then handle the error in
handleIncomingRequestwith the error callback.client/client.go (1)
433-479: Consider adding request validationThe implementation is solid, but could benefit from validating the request parameters before processing.
Add parameter validation after unmarshaling:
if err := json.Unmarshal(paramsBytes, ¶ms); err != nil { return nil, fmt.Errorf("failed to unmarshal params: %w", err) } + + // Validate required parameters + if len(params.Messages) == 0 { + return nil, fmt.Errorf("messages cannot be empty") + } + if params.MaxTokens <= 0 { + return nil, fmt.Errorf("maxTokens must be positive") + }This ensures the handler receives valid input and provides better error messages to the server.
server/stdio.go (1)
351-366: Consider tracking concurrent tool call errorsTool call errors are only logged and not tracked, which could hide failures from monitoring systems.
Consider tracking goroutine completion and errors:
+ // Add to stdioSession struct: + activeRequests sync.WaitGroup + if json.Unmarshal(rawMessage, &baseMessage) == nil && baseMessage.Method == "tools/call" { + s.activeRequests.Add(1) go func() { + defer s.activeRequests.Done() response := s.server.HandleMessage(ctx, rawMessage) if response != nil { if err := s.writeResponse(response, writer); err != nil { s.errLogger.Printf("Error writing tool response: %v", err) + // Consider incrementing an error metric here } } }() return nil }This would allow graceful shutdown by waiting for active requests to complete.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (15)
client/client.go(6 hunks)client/sampling.go(1 hunks)client/sampling_test.go(1 hunks)client/transport/interface.go(2 hunks)client/transport/stdio.go(4 hunks)client/transport/stdio_test.go(2 hunks)examples/sampling_client/README.md(1 hunks)examples/sampling_client/main.go(1 hunks)examples/sampling_server/README.md(1 hunks)examples/sampling_server/main.go(1 hunks)mcp/types.go(1 hunks)server/sampling.go(1 hunks)server/sampling_test.go(1 hunks)server/server.go(2 hunks)server/stdio.go(6 hunks)
🧰 Additional context used
🧠 Learnings (7)
client/transport/stdio_test.go (1)
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.
client/transport/interface.go (2)
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.
examples/sampling_server/README.md (1)
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.
examples/sampling_client/main.go (3)
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.
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.
examples/sampling_server/main.go (6)
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.
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: 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 mark3labs/mcp-go handles both InputSchema and RawInputSchema formats. When unmarshaling JSON, it first tries to parse into a structured ToolInputSchema format, and if that fails, it falls back to using the raw schema format, providing symmetry with the MarshalJSON method.
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: 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.
server/sampling_test.go (1)
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.
client/sampling_test.go (1)
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.
🧬 Code Graph Analysis (6)
client/transport/stdio_test.go (1)
mcp/types.go (1)
Params(162-162)
client/sampling.go (1)
mcp/types.go (2)
CreateMessageRequest(775-778)CreateMessageResult(795-802)
client/transport/interface.go (2)
mcp/types.go (3)
JSONRPCRequest(310-315)JSONRPCResponse(324-328)Result(231-235)testdata/mockstdio_server.go (2)
JSONRPCRequest(13-18)JSONRPCResponse(20-28)
server/stdio.go (3)
mcp/types.go (8)
JSONRPCNotification(318-321)CreateMessageResult(795-802)CreateMessageRequest(775-778)Params(162-162)CreateMessageParams(780-789)JSONRPC_VERSION(108-108)MethodSamplingCreateMessage(768-768)Result(231-235)server/session.go (2)
ClientSession(11-20)SessionWithClientInfo(43-49)server/sampling.go (1)
SessionWithSampling(53-56)
client/transport/stdio.go (2)
client/transport/interface.go (3)
RequestHandler(35-35)JSONRPCRequest(47-52)JSONRPCResponse(54-63)mcp/types.go (5)
JSONRPCRequest(310-315)JSONRPCResponse(324-328)JSONRPC_VERSION(108-108)METHOD_NOT_FOUND(350-350)INTERNAL_ERROR(352-352)
client/client.go (4)
client/sampling.go (1)
SamplingHandler(11-20)client/transport/interface.go (3)
BidirectionalInterface(39-45)JSONRPCRequest(47-52)JSONRPCResponse(54-63)mcp/types.go (9)
Params(162-162)JSONRPCRequest(310-315)JSONRPCResponse(324-328)MethodSamplingCreateMessage(768-768)CreateMessageParams(780-789)CreateMessageRequest(775-778)Request(153-156)JSONRPC_VERSION(108-108)Result(231-235)testdata/mockstdio_server.go (2)
JSONRPCRequest(13-18)JSONRPCResponse(20-28)
🪛 markdownlint-cli2 (0.17.2)
examples/sampling_client/README.md
54-54: Fenced code blocks should have a language specified
null
(MD040, fenced-code-language)
🪛 GitHub Check: lint
server/sampling_test.go
[failure] 152-152:
SA5011: possible nil pointer dereference (staticcheck)
[failure] 148-148:
SA5011(related information): this check suggests that the pointer can be nil (staticcheck)
🪛 GitHub Actions: golangci-lint
server/sampling_test.go
[error] 148-148: golangci-lint (staticcheck): SA5011 - this check suggests that the pointer can be nil
🔇 Additional comments (32)
mcp/types.go (1)
766-769: LGTM! Clean constant addition for sampling support.The new
MethodSamplingCreateMessageconstant follows the established naming convention and is properly documented. This constant will be used to identify sampling requests in the bidirectional MCP communication.server/server.go (2)
178-178: Good addition following established capability patterns.The new
samplingfield correctly follows the same pointer-based pattern used by other server capabilities (tools, resources, prompts), allowing it to be nil when sampling is not supported.
197-200: Well-structured sampling capabilities definition.The
samplingCapabilitiesstruct is appropriately simple with just anenabledfield, which is sufficient for the initial sampling implementation. This follows the established pattern of other capability structs in the codebase.client/transport/stdio_test.go (1)
185-200: Test logic correctly handles the mock server's notification format.The updated test logic properly validates that the mock server wraps the original notification in a "debug/test" method and extracts the original method from the params. The JSON marshaling/unmarshaling approach is appropriate for inspecting the notification structure.
client/sampling.go (1)
9-20: Excellent interface design with comprehensive documentation.The
SamplingHandlerinterface is well-crafted with:
- Proper use of context for cancellation/timeout handling
- Correct MCP request/result types from the types package
- Comprehensive documentation that clearly guides implementers through the expected steps
- Single responsibility principle with one focused method
The documentation particularly shines in outlining the implementation workflow (validate → approve → select → generate → return).
examples/sampling_server/README.md (1)
1-52: Comprehensive and well-structured documentation.This README provides excellent coverage of the sampling server example:
- Clear feature descriptions and tool documentation
- Practical usage instructions with build commands
- Helpful implementation details explaining key concepts like
EnableSampling()and response routing- Testing guidance with the companion client example
The documentation effectively guides users through understanding, building, and testing the sampling functionality.
client/transport/interface.go (3)
34-35: LGTM! Well-defined request handler interface.The
RequestHandlerfunction type follows Go conventions and provides a clean interface for handling incoming requests from servers.
37-45: LGTM! Clean extension for bidirectional communication.The
BidirectionalInterfaceproperly extends the baseInterfaceand provides the necessary method for registering request handlers to support server-initiated requests like sampling.
57-62: LGTM! Appropriate JSON field optimization.Adding
omitemptytags toResultandErrorfields prevents sending empty fields in JSON responses, which is correct for JSON-RPC where these fields should only be present when they have meaningful content.examples/sampling_server/main.go (4)
13-26: LGTM! Proper server initialization and capability setup.The server initialization correctly enables sampling capability and follows established patterns for MCP server setup.
21-89: LGTM! Well-implemented sampling tool with proper error handling.The
ask_llmtool demonstrates proper usage of the sampling API with:
- Correct parameter extraction using helper methods
- Proper sampling request construction
- Appropriate timeout handling (30 seconds)
- Good error handling and response formatting
The implementation serves as an excellent example for users implementing sampling functionality.
92-119: LGTM! Simple but informative greeting tool.The
greettool provides a good contrast to the sampling tool and properly informs users about available sampling functionality.
128-138: LGTM! Robust content extraction helper.The
getTextFromContentfunction properly handles different content types that might be received in sampling responses, including structured and unstructured formats.server/sampling_test.go (4)
10-25: LGTM! Good test coverage for enabling sampling.The test properly verifies the initial state and the effect of enabling sampling on server capabilities.
27-49: LGTM! Proper error handling test.The test correctly verifies that sampling requests fail when sampling is not enabled, with appropriate error message validation.
51-74: LGTM! Good test for no session scenario.The test properly validates error handling when no active session is available for sampling requests.
76-107: LGTM! Well-designed mock implementations.The mock session implementations provide good test doubles for the interfaces, enabling effective unit testing of the sampling functionality.
examples/sampling_client/main.go (4)
14-67: LGTM! Excellent mock sampling handler implementation.The
MockSamplingHandlerprovides a comprehensive demonstration of how to implement theSamplingHandlerinterface with:
- Proper error handling for edge cases
- Flexible content type handling (TextContent and generic maps)
- Good logging for debugging
- Realistic mock response generation
This serves as an excellent template for users implementing real LLM integrations.
69-113: LGTM! Proper client initialization and setup.The client setup demonstrates correct usage patterns:
- Proper command-line argument handling
- Correct transport and client configuration
- Proper use of
WithSamplingHandleroption- Standard MCP initialization flow
114-144: LGTM! Good tool discovery and testing pattern.The code properly demonstrates tool listing and basic tool testing, providing good examples for users learning the MCP client API.
146-190: LGTM! Comprehensive sampling functionality demonstration.The sampling tool tests effectively demonstrate:
- Different types of sampling requests
- Proper argument passing
- Response handling and logging
- Multiple use cases showing the flexibility of the sampling API
This provides excellent examples for developers implementing sampling clients.
server/sampling.go (2)
28-50: LGTM! Well-structured request handlingThe method properly validates preconditions, uses appropriate mutex locking for thread-safe capability checking, and returns descriptive error messages for each failure case.
52-56: LGTM! Clean interface extensionThe interface properly extends ClientSession and maintains consistency with the SamplingClient interface method signature.
client/transport/stdio.go (2)
163-169: LGTM! Thread-safe handler registrationThe method correctly uses mutex locking to ensure thread-safe assignment of the request handler.
188-194: LGTM! Clean request detection logicThe implementation correctly distinguishes requests from responses by checking for the presence of a method field, maintaining backward compatibility with existing response handling.
client/sampling_test.go (1)
24-87: LGTM! Comprehensive test coverageThe test properly covers both error and success cases using a table-driven approach. Error message validation ensures the correct error is returned.
client/client.go (3)
37-43: LGTM! Clean option pattern implementationThe function follows Go conventions for functional options and includes helpful documentation.
84-88: LGTM! Backward-compatible bidirectional supportThe type assertion pattern ensures backward compatibility with transports that don't implement bidirectional communication.
422-431: LGTM! Extensible request routingThe switch statement provides a clean, extensible pattern for handling different request types as they're added.
server/stdio.go (3)
55-70: LGTM! Well-designed concurrent request trackingThe struct properly uses atomic types and mutexes to ensure thread-safe access to shared state. The pending requests map with mutex protection enables safe concurrent request handling.
115-176: LGTM! Robust async request implementationExcellent implementation with:
- Proper resource cleanup using defer
- Context cancellation support
- Unique request ID generation
- Safe concurrent access patterns
346-349: LGTM! Clean response routingThe early return pattern cleanly separates sampling response handling from normal message processing.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
www/docs/pages/clients/advanced-sampling.mdx (1)
65-76: Add missinglog&contextimports to keep the snippet compile-readyThe snippet relies on
log.Fatalfandcontext.Background()but theimportblock is absent. Readers who copy / paste will hit a compile error.-func main() { +package main + +import ( + "context" + "log" + + "github.com/mark3labs/mcp-go/client" +) + +func main() {www/docs/pages/servers/advanced-sampling.mdx (1)
16-20: Addfmtto import list
fmt.Sprintfis used inside the tool handler (line 92 & 104) butfmtisn’t imported in this quick-start snippet, so it will not compile.import ( "context" + "fmt" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
server/sampling_test.go(1 hunks)www/docs/pages/clients/advanced-sampling.mdx(1 hunks)www/docs/pages/clients/operations.mdx(1 hunks)www/docs/pages/servers/advanced-sampling.mdx(1 hunks)www/docs/pages/servers/advanced.mdx(1 hunks)www/docs/pages/servers/index.mdx(1 hunks)www/vocs.config.ts(1 hunks)
✅ Files skipped from review due to trivial changes (4)
- www/vocs.config.ts
- www/docs/pages/servers/index.mdx
- www/docs/pages/servers/advanced.mdx
- www/docs/pages/clients/operations.mdx
🚧 Files skipped from review as they are similar to previous changes (1)
- server/sampling_test.go
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
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.
www/docs/pages/servers/advanced-sampling.mdx (2)
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: 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.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
| } | ||
|
|
||
| // Check if the session supports sampling requests | ||
| if samplingSession, ok := session.(SessionWithSampling); ok { |
There was a problem hiding this comment.
It would be better if we could find a way to ensure that session is an implementation of SessionWithSampling when EnableSampling is invoked.
Description
Fixes #<issue_number> (if applicable)
Type of Change
Checklist
MCP Spec Compliance
Additional Information
Summary by CodeRabbit
New Features
Documentation
Bug Fixes
Tests
Chores